Update app.py

app.py

@@ -1,259 +1,185 @@
 import gradio as gr
-from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.vectorstores import FAISS
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.llms import HuggingFacePipeline
-from langchain.chains import ConversationalRetrievalChain
-from langchain.memory import ConversationBufferMemory
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from sentence_transformers import SentenceTransformer
+import numpy as np
 from pypdf import PdfReader
 import torch
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# Initialize global variables
-vectorstore = None
-qa_chain = None
-llm_pipeline = None
+# Global variables
+chunks = []
+embeddings = []
+model = None
+tokenizer = None
+embed_model = None
 
-def initialize_llm():
-    """Initialize the language model (done once at startup)"""
-    global llm_pipeline
+def initialize_models():
+    """Initialize models on startup"""
+    global model, tokenizer, embed_model
     
-    if llm_pipeline is not None:
-        return
+    print("Loading models...")
     
-    print("Loading language model...")
+    # Load embedding model
+    embed_model = SentenceTransformer('sentence-transformers/all-MiniLM-L6-v2')
     
-    # Use a smaller, efficient model that works without API
+    # Load language model
     model_name = "TinyLlama/TinyLlama-1.1B-Chat-v1.0"
-    
     tokenizer = AutoTokenizer.from_pretrained(model_name)
     model = AutoModelForCausalLM.from_pretrained(
         model_name,
-        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-        device_map="auto",
+        torch_dtype=torch.float32,
         low_cpu_mem_usage=True
     )
     
-    pipe = pipeline(
-        "text-generation",
-        model=model,
-        tokenizer=tokenizer,
-        max_new_tokens=512,
-        temperature=0.7,
-        top_p=0.95,
-        repetition_penalty=1.15
-    )
-    
-    llm_pipeline = HuggingFacePipeline(pipeline=pipe)
-    print("Model loaded successfully!")
+    print("Models loaded successfully!")
 
 def process_pdf(pdf_file):
-    """Process uploaded PDF and create vector store"""
-    global vectorstore, qa_chain
+    """Process PDF and create embeddings"""
+    global chunks, embeddings, embed_model
     
     if pdf_file is None:
-        return "Please upload a PDF file!", None, None
+        return "❌ Please upload a PDF file!", None
     
     try:
-        # Extract text from PDF
+        # Read PDF
         pdf_reader = PdfReader(pdf_file.name)
         text = ""
         for page in pdf_reader.pages:
-            text += page.extract_text()
+            text += page.extract_text() + "\n"
         
         if not text.strip():
-            return "Could not extract text from PDF. Please ensure it's a valid PDF with text content.", None, None
-        
-        # Split text into chunks
-        text_splitter = RecursiveCharacterTextSplitter(
-            chunk_size=1000,
-            chunk_overlap=200,
-            length_function=len
-        )
-        chunks = text_splitter.split_text(text)
-        
-        # Create embeddings (using a lightweight model)
-        embeddings = HuggingFaceEmbeddings(
-            model_name="sentence-transformers/all-MiniLM-L6-v2",
-            model_kwargs={'device': 'cpu'}
-        )
-        
-        # Create vector store
-        vectorstore = FAISS.from_texts(chunks, embeddings)
+            return "❌ Could not extract text from PDF!", None
         
-        # Initialize LLM if not already done
-        initialize_llm()
+        # Split into chunks
+        chunk_size = 1000
+        overlap = 200
+        chunks = []
         
-        # Create memory for conversation
-        memory = ConversationBufferMemory(
-            memory_key="chat_history",
-            return_messages=True,
-            output_key="answer"
-        )
+        for i in range(0, len(text), chunk_size - overlap):
+            chunk = text[i:i + chunk_size]
+            if chunk.strip():
+                chunks.append(chunk)
         
-        # Create conversational chain
-        qa_chain = ConversationalRetrievalChain.from_llm(
-            llm=llm_pipeline,
-            retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
-            memory=memory,
-            return_source_documents=True,
-            verbose=False
-        )
+        # Create embeddings
+        embeddings = embed_model.encode(chunks, show_progress_bar=False)
         
-        return f"✅ PDF processed successfully! Extracted {len(chunks)} text chunks. You can now ask questions!", None, None
+        return f"✅ PDF processed! Created {len(chunks)} chunks. You can now ask questions!", None
         
     except Exception as e:
-        return f"❌ Error processing PDF: {str(e)}", None, None
+        return f"❌ Error: {str(e)}", None
+
+def find_relevant_chunks(query, top_k=3):
+    """Find most relevant chunks using cosine similarity"""
+    global chunks, embeddings, embed_model
+    
+    if not chunks:
+        return []
+    
+    query_embedding = embed_model.encode([query])[0]
+    
+    # Calculate cosine similarity
+    similarities = np.dot(embeddings, query_embedding) / (
+        np.linalg.norm(embeddings, axis=1) * np.linalg.norm(query_embedding)
+    )
+    
+    # Get top k indices
+    top_indices = np.argsort(similarities)[-top_k:][::-1]
+    
+    return [chunks[i] for i in top_indices]
+
+def generate_response(question, context):
+    """Generate response using the language model"""
+    global model, tokenizer
+    
+    prompt = f"""<|system|>
+You are a helpful assistant. Answer the question based on the provided context. Be concise and accurate.
+</s>
+<|user|>
+Context: {context}
+
+Question: {question}
+</s>
+<|assistant|>
+"""
+    
+    inputs = tokenizer(prompt, return_tensors="pt", truncation=True, max_length=2048)
+    
+    with torch.no_grad():
+        outputs = model.generate(
+            **inputs,
+            max_new_tokens=300,
+            temperature=0.7,
+            top_p=0.9,
+            do_sample=True,
+            pad_token_id=tokenizer.eos_token_id
+        )
+    
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    
+    # Extract only the assistant's response
+    if "<|assistant|>" in response:
+        response = response.split("<|assistant|>")[-1].strip()
+    
+    return response
 
 def chat(message, history):
-    """Handle chat interactions"""
-    global qa_chain
+    """Handle chat"""
+    global chunks
     
-    if qa_chain is None:
+    if not chunks:
         return history + [[message, "⚠️ Please upload and process a PDF first!"]]
     
     if not message.strip():
         return history
     
     try:
-        # Get response from chain
-        result = qa_chain({"question": message})
-        answer = result["answer"]
+        # Find relevant context
+        relevant_chunks = find_relevant_chunks(message)
+        context = "\n\n".join(relevant_chunks)
         
-        # Clean up the answer (remove any system prompts)
-        if "Answer:" in answer:
-            answer = answer.split("Answer:")[-1].strip()
+        # Generate response
+        response = generate_response(message, context)
         
-        return history + [[message, answer]]
+        return history + [[message, response]]
         
     except Exception as e:
         return history + [[message, f"❌ Error: {str(e)}"]]
 
-def clear_chat():
-    """Clear chat history and reset chain"""
-    global qa_chain
-    if qa_chain is not None and hasattr(qa_chain, 'memory'):
-        qa_chain.memory.clear()
-    return None
+def clear_all():
+    """Clear everything"""
+    global chunks, embeddings
+    chunks = []
+    embeddings = []
+    return None, "Ready to process a new PDF"
 
-# Create Gradio interface
-with gr.Blocks(theme=gr.themes.Soft(), title="Chat with PDF") as demo:
-    gr.Markdown(
-        """
-        # 📄 Chat with PDF using AI
-        Upload a PDF document and ask questions about its content - No API key required!
-        
-        **Instructions:**
-        1. Upload a PDF file
-        2. Click "Process PDF" and wait for confirmation
-        3. Start asking questions about your document!
-        """
-    )
+# Create UI
+with gr.Blocks(title="Chat with PDF") as demo:
+    gr.Markdown("# 📄 Chat with PDF - Simple Version")
     
     with gr.Row():
         with gr.Column(scale=1):
-            pdf_input = gr.File(
-                label="📎 Upload PDF",
-                file_types=[".pdf"],
-                type="filepath"
-            )
-            process_btn = gr.Button("🔄 Process PDF", variant="primary", size="lg")
-            status_output = gr.Textbox(
-                label="📊 Status",
-                interactive=False,
-                lines=3
-            )
+            pdf_input = gr.File(label="📎 Upload PDF", file_types=[".pdf"])
+            process_btn = gr.Button("🔄 Process PDF", variant="primary")
+            status = gr.Textbox(label="Status", lines=3)
+            clear_all_btn = gr.Button("🗑️ Clear All")
             
-            gr.Markdown(
-                """
-                ### 💡 Tips:
-                - Processing may take 30-60 seconds
-                - Ask specific questions about the content
-                - You can ask follow-up questions
-                - Best with text-based PDFs (not scanned images)
-                """
-            )
-        
         with gr.Column(scale=2):
-            chatbot = gr.Chatbot(
-                label="💬 Chat History",
-                height=500,
-                bubble_full_width=False
-            )
+            chatbot = gr.Chatbot(label="💬 Chat", height=400)
+            msg = gr.Textbox(label="Question", placeholder="Ask about the PDF...")
             with gr.Row():
-                msg = gr.Textbox(
-                    label="Your Question",
-                    placeholder="Ask a question about your PDF...",
-                    lines=2,
-                    scale=4
-                )
-            with gr.Row():
-                submit_btn = gr.Button("📤 Send", variant="primary", scale=1)
-                clear_btn = gr.Button("🗑️ Clear Chat", scale=1)
-    
-    gr.Markdown(
-        """
-        ---
-        ### 🔌 API Access
-        Once deployed on Hugging Face Spaces, you can access this via API:
-        ```python
-        # Python example
-        from gradio_client import Client
-        
-        client = Client("YOUR_USERNAME/YOUR_SPACE_NAME")
-        
-        # Process PDF
-        result = client.predict("path/to/file.pdf", api_name="/process_pdf")
-        
-        # Ask questions
-        result = client.predict("What is this document about?", [], api_name="/chat")
-        ```
-        
-        ```javascript
-        // JavaScript example
-        const response = await fetch("https://YOUR_USERNAME-YOUR_SPACE_NAME.hf.space/api/predict", {
-            method: "POST",
-            headers: { "Content-Type": "application/json" },
-            body: JSON.stringify({
-                data: ["What is this document about?", []]
-            })
-        });
-        ```
-        """
-    )
+                send_btn = gr.Button("Send", variant="primary")
+                clear_btn = gr.Button("Clear Chat")
     
-    # Event handlers
-    process_btn.click(
-        fn=process_pdf,
-        inputs=[pdf_input],
-        outputs=[status_output, chatbot, msg]
-    )
+    # Events
+    process_btn.click(process_pdf, [pdf_input], [status, chatbot])
     
-    msg.submit(
-        fn=chat,
-        inputs=[msg, chatbot],
-        outputs=[chatbot]
-    ).then(
-        fn=lambda: "",
-        outputs=[msg]
-    )
+    msg.submit(chat, [msg, chatbot], [chatbot]).then(lambda: "", None, [msg])
+    send_btn.click(chat, [msg, chatbot], [chatbot]).then(lambda: "", None, [msg])
     
-    submit_btn.click(
-        fn=chat,
-        inputs=[msg, chatbot],
-        outputs=[chatbot]
-    ).then(
-        fn=lambda: "",
-        outputs=[msg]
-    )
-    
-    clear_btn.click(
-        fn=clear_chat,
-        outputs=[chatbot]
-    )
+    clear_btn.click(lambda: None, None, [chatbot])
+    clear_all_btn.click(clear_all, None, [chatbot, status])
 
-# Initialize model on startup
-initialize_llm()
+# Initialize on startup
+initialize_models()
 
 if __name__ == "__main__":
-    demo.launch(share=False)
+    demo.launch()