Update app.py

app.py

@@ -1,28 +1,30 @@
 import gradio as gr
 import os
-from langchain.document_loaders import PyPDFLoader
+from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import FAISS
-from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.vectorstores import FAISS
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # Configuration
 DOCS_DIR = "business_docs"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 MODEL_NAME = "microsoft/phi-2"
 
-# Initialize system components
 def initialize_system():
-    # Load and process PDFs
+    # Verify documents
     if not os.path.exists(DOCS_DIR):
-        raise FileNotFoundError(f"'{DOCS_DIR}' folder not found")
-    
+        raise FileNotFoundError(f"Missing {DOCS_DIR} folder")
+        
     pdf_files = [os.path.join(DOCS_DIR, f) for f in os.listdir(DOCS_DIR) 
                 if f.endswith(".pdf")]
+    if not pdf_files:
+        raise ValueError(f"No PDFs found in {DOCS_DIR}")
     
+    # Process documents
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=1000,
-        chunk_overlap=200
+        chunk_size=800,  # Reduced for Phi-2's context window
+        chunk_overlap=100
     )
     
     texts = []
@@ -31,79 +33,76 @@ def initialize_system():
         pages = loader.load_and_split(text_splitter)
         texts.extend(pages)
     
+    # Create embeddings
+    embeddings = HuggingFaceEmbeddings(
+        model_name=EMBEDDING_MODEL,
+        model_kwargs={'device': 'cpu'},  # Force CPU for compatibility
+        encode_kwargs={'normalize_embeddings': False}
+    )
+    
     # Create vector store
-    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
     vector_store = FAISS.from_documents(texts, embeddings)
     
-    # Load Phi-2 model with 4-bit quantization
+    # Load Phi-2 with 4-bit quantization
     tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         trust_remote_code=True,
         device_map="auto",
-        load_in_4bit=True
+        load_in_4bit=True,
+        torch_dtype=torch.float16
     )
     
     return vector_store, model, tokenizer
 
 try:
     vector_store, model, tokenizer = initialize_system()
-    print("System ready with business documents loaded")
+    print("System initialized successfully")
 except Exception as e:
-    raise RuntimeError(f"Initialization failed: {str(e)}")
+    raise RuntimeError(f"Initialization error: {str(e)}")
 
-# Response generation
 def generate_response(query):
-    # Retrieve relevant context
-    docs = vector_store.similarity_search(query, k=3)
-    context = "\n".join([doc.page_content for doc in docs])
-    
-    # Create custom prompt template
-    prompt = f"""Instruct: Answer the customer's question using only the provided context.
-    If you don't know the answer, say 'I need to check with our team about that.'
+    # Retrieve context
+    docs = vector_store.similarity_search(query, k=2)  # Fewer docs for Phi-2
+    context = "\n".join([d.page_content for d in docs])
     
+    # Phi-2 specific prompt format
+    prompt = f"""Question: {query}
     Context: {context}
-    
-    Question: {query}
+    Instructions:
+    - Answer only using the context
+    - Keep responses under 3 sentences
+    - If unsure, say "I'll need to check with the team"
     
     Answer:"""
     
-    # Generate response
     inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False).to(model.device)
     outputs = model.generate(
         **inputs,
-        max_new_tokens=300,
-        temperature=0.2,
-        repetition_penalty=1.2,
-        do_sample=True
+        max_new_tokens=200,
+        temperature=0.1,
+        do_sample=True,
+        pad_token_id=tokenizer.eos_token_id
     )
     
-    # Decode and clean response
-    full_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
-    answer = full_text.split("Answer:")[-1].strip()
-    return answer.split("\n\n")[0]  # Return first paragraph
+    response = tokenizer.decode(outputs[0], skip_special_tokens=True)
+    return response.split("Answer:")[-1].strip()
 
-# Chat interface
+# Simplified Gradio interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# Customer Care Assistant")
-    gr.Markdown("Ask questions about our products/services")
+    gr.Markdown("# Customer Service Chatbot")
+    chatbot = gr.Chatbot()
+    msg = gr.Textbox(label="Your question")
+    clear = gr.ClearButton([msg, chatbot])
     
-    chatbot = gr.Chatbot(height=400)
-    msg = gr.Textbox(label="Type your question here...")
-    clear = gr.Button("Clear History")
-    
-    def respond(message, chat_history):
+    def respond(message, history):
         try:
             response = generate_response(message)
-            if not response:
-                response = "I need to verify that information. Please contact [email protected]"
+            return response
         except Exception as e:
-            response = "Apologies, I'm experiencing technical difficulties. Please try again later."
-        
-        chat_history.append((message, response))
-        return "", chat_history
+            return "I'm having trouble answering that right now. Please try again later."
     
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: None, None, chatbot, queue=False)
+    msg.submit(respond, [msg, chatbot], chatbot)
+    msg.submit(lambda: "", None, msg)
 
-demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch(server_port=7860)