Update app.py

app.py

@@ -1,29 +1,60 @@
 import gradio as gr
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import FAISS
-from langchain_huggingface import HuggingFaceEmbeddings
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.llms import HuggingFacePipeline
 from langchain.chains import ConversationalRetrievalChain
-from langchain_huggingface import HuggingFaceEndpoint
+from langchain.memory import ConversationBufferMemory
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from pypdf import PdfReader
-import os
-from huggingface_hub import login
+import torch
 
 # Initialize global variables
 vectorstore = None
 qa_chain = None
-chat_history = []
+llm_pipeline = None
 
-def process_pdf(pdf_file, hf_token):
+def initialize_llm():
+    """Initialize the language model (done once at startup)"""
+    global llm_pipeline
+    
+    if llm_pipeline is not None:
+        return
+    
+    print("Loading language model...")
+    
+    # Use a smaller, efficient model that works without API
+    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",
+        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!")
+
+def process_pdf(pdf_file):
     """Process uploaded PDF and create vector store"""
-    global vectorstore, qa_chain, chat_history
+    global vectorstore, qa_chain
     
-    if not hf_token:
-        return "Please provide your Hugging Face API token!", None
+    if pdf_file is None:
+        return "Please upload a PDF file!", None, None
     
     try:
-        # Login to Hugging Face
-        login(token=hf_token)
-        
         # Extract text from PDF
         pdf_reader = PdfReader(pdf_file.name)
         text = ""
@@ -31,7 +62,7 @@ def process_pdf(pdf_file, hf_token):
             text += page.extract_text()
         
         if not text.strip():
-            return "Could not extract text from PDF. Please ensure it's a valid PDF with text content.", None
+            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(
@@ -41,69 +72,68 @@ def process_pdf(pdf_file, hf_token):
         )
         chunks = text_splitter.split_text(text)
         
-        # Create embeddings
+        # Create embeddings (using a lightweight model)
         embeddings = HuggingFaceEmbeddings(
-            model_name="sentence-transformers/all-MiniLM-L6-v2"
+            model_name="sentence-transformers/all-MiniLM-L6-v2",
+            model_kwargs={'device': 'cpu'}
         )
         
         # Create vector store
         vectorstore = FAISS.from_texts(chunks, embeddings)
         
-        # Initialize LLM
-        llm = HuggingFaceEndpoint(
-            repo_id="mistralai/Mistral-7B-Instruct-v0.2",
-            temperature=0.7,
-            max_new_tokens=512,
-            huggingfacehub_api_token=hf_token
+        # Initialize LLM if not already done
+        initialize_llm()
+        
+        # Create memory for conversation
+        memory = ConversationBufferMemory(
+            memory_key="chat_history",
+            return_messages=True,
+            output_key="answer"
         )
         
         # Create conversational chain
         qa_chain = ConversationalRetrievalChain.from_llm(
-            llm=llm,
+            llm=llm_pipeline,
             retriever=vectorstore.as_retriever(search_kwargs={"k": 3}),
+            memory=memory,
             return_source_documents=True,
             verbose=False
         )
         
-        # Reset chat history
-        chat_history = []
-        
-        return f"PDF processed successfully! Extracted {len(chunks)} text chunks. You can now ask questions!", None
+        return f"✅ PDF processed successfully! Extracted {len(chunks)} text chunks. You can now ask questions!", None, None
         
     except Exception as e:
-        return f"Error processing PDF: {str(e)}", None
+        return f"❌ Error processing PDF: {str(e)}", None, None
 
 def chat(message, history):
     """Handle chat interactions"""
-    global qa_chain, chat_history
+    global qa_chain
     
     if qa_chain is None:
-        return "Please upload and process a PDF first!"
+        return history + [[message, "⚠️ Please upload and process a PDF first!"]]
     
     if not message.strip():
-        return "Please enter a question!"
+        return history
     
     try:
         # Get response from chain
-        result = qa_chain({
-            "question": message,
-            "chat_history": chat_history
-        })
-        
+        result = qa_chain({"question": message})
         answer = result["answer"]
         
-        # Update chat history
-        chat_history.append((message, answer))
+        # Clean up the answer (remove any system prompts)
+        if "Answer:" in answer:
+            answer = answer.split("Answer:")[-1].strip()
         
-        return answer
+        return history + [[message, answer]]
         
     except Exception as e:
-        return f"Error: {str(e)}"
+        return history + [[message, f"❌ Error: {str(e)}"]]
 
 def clear_chat():
-    """Clear chat history"""
-    global chat_history
-    chat_history = []
+    """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
 
 # Create Gradio interface
@@ -111,82 +141,110 @@ 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!
+        Upload a PDF document and ask questions about its content - No API key required!
         
         **Instructions:**
-        1. Enter your Hugging Face API token (get one from [huggingface.co/settings/tokens](https://huggingface.co/settings/tokens))
-        2. Upload a PDF file
-        3. Click "Process PDF"
-        4. Start asking questions about your document!
+        1. Upload a PDF file
+        2. Click "Process PDF" and wait for confirmation
+        3. Start asking questions about your document!
         """
     )
     
     with gr.Row():
         with gr.Column(scale=1):
-            hf_token = gr.Textbox(
-                label="Hugging Face API Token",
-                type="password",
-                placeholder="hf_..."
-            )
             pdf_input = gr.File(
-                label="Upload PDF",
-                file_types=[".pdf"]
+                label="📎 Upload PDF",
+                file_types=[".pdf"],
+                type="filepath"
             )
-            process_btn = gr.Button("Process PDF", variant="primary")
+            process_btn = gr.Button("🔄 Process PDF", variant="primary", size="lg")
             status_output = gr.Textbox(
-                label="Status",
-                interactive=False
+                label="📊 Status",
+                interactive=False,
+                lines=3
+            )
+            
+            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=400
-            )
-            msg = gr.Textbox(
-                label="Your Question",
-                placeholder="Ask a question about your PDF...",
-                lines=2
+                label="💬 Chat History",
+                height=500,
+                bubble_full_width=False
             )
             with gr.Row():
-                submit_btn = gr.Button("Send", variant="primary")
-                clear_btn = gr.Button("Clear Chat")
+                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(
         """
-        ### Tips:
-        - Ask specific questions about the content
-        - You can ask follow-up questions
-        - The AI will use context from previous messages
-        - For best results, ensure your PDF has extractable text (not scanned images)
+        ---
+        ### 🔌 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?", []]
+            })
+        });
+        ```
         """
     )
     
     # Event handlers
     process_btn.click(
         fn=process_pdf,
-        inputs=[pdf_input, hf_token],
-        outputs=[status_output, chatbot]
+        inputs=[pdf_input],
+        outputs=[status_output, chatbot, msg]
     )
     
-    submit_btn.click(
+    msg.submit(
         fn=chat,
         inputs=[msg, chatbot],
-        outputs=[msg]
+        outputs=[chatbot]
     ).then(
-        fn=lambda m, h: (h + [[m, chat(m, h)]], ""),
-        inputs=[msg, chatbot],
-        outputs=[chatbot, msg]
+        fn=lambda: "",
+        outputs=[msg]
     )
     
-    msg.submit(
+    submit_btn.click(
         fn=chat,
         inputs=[msg, chatbot],
-        outputs=[msg]
+        outputs=[chatbot]
     ).then(
-        fn=lambda m, h: (h + [[m, chat(m, h)]], ""),
-        inputs=[msg, chatbot],
-        outputs=[chatbot, msg]
+        fn=lambda: "",
+        outputs=[msg]
     )
     
     clear_btn.click(
@@ -194,5 +252,8 @@ with gr.Blocks(theme=gr.themes.Soft(), title="Chat with PDF") as demo:
         outputs=[chatbot]
     )
 
+# Initialize model on startup
+initialize_llm()
+
 if __name__ == "__main__":
-    demo.launch()
+    demo.launch(share=False)