Update app.py

app.py

@@ -1,119 +1,109 @@
-import os
 import gradio as gr
-import torch
-from langchain_community.document_loaders import PyPDFLoader
+import os
+from langchain.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
-from langchain_community.embeddings import HuggingFaceEmbeddings
-from langchain_community.vectorstores import FAISS
-from transformers import AutoModelForCausalLM, AutoTokenizer
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import FAISS
+from transformers import AutoModelForCausalLM, AutoTokenizer, TextStreamer
 
 # Configuration
 DOCS_DIR = "business_docs"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-MODEL_NAME = "microsoft/phi-3-mini-4k-instruct"  # CPU-optimized model
+MODEL_NAME = "microsoft/phi-2"
 
-# System Initialization
+# Initialize system components
 def initialize_system():
-    # Validate documents folder
+    # Load and process PDFs
     if not os.path.exists(DOCS_DIR):
-        raise FileNotFoundError(f"Missing documents folder: {DOCS_DIR}")
+        raise FileNotFoundError(f"'{DOCS_DIR}' folder not found")
+    
+    pdf_files = [os.path.join(DOCS_DIR, f) for f in os.listdir(DOCS_DIR) 
+                if f.endswith(".pdf")]
     
-    # Process PDFs
-    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}")
-
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=512,  # Optimized for CPU
-        chunk_overlap=50
+        chunk_size=1000,
+        chunk_overlap=200
     )
     
-    documents = []
-    for pdf_path in pdf_files:
-        try:
-            loader = PyPDFLoader(pdf_path)
-            documents.extend(loader.load_and_split(text_splitter))
-        except Exception as e:
-            print(f"Error processing {pdf_path}: {str(e)}")
+    texts = []
+    for pdf in pdf_files:
+        loader = PyPDFLoader(pdf)
+        pages = loader.load_and_split(text_splitter)
+        texts.extend(pages)
     
-    # Create embeddings
-    embeddings = HuggingFaceEmbeddings(
-        model_name=EMBEDDING_MODEL,
-        model_kwargs={'device': 'cpu'},
-        encode_kwargs={'normalize_embeddings': True}
-    )
-    
-    vector_store = FAISS.from_documents(documents, embeddings)
+    # Create vector store
+    embeddings = HuggingFaceEmbeddings(model_name=EMBEDDING_MODEL)
+    vector_store = FAISS.from_documents(texts, embeddings)
     
-    # Load CPU-optimized model
-    try:
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME)
-        model = AutoModelForCausalLM.from_pretrained(
-            MODEL_NAME,
-            trust_remote_code=True,
-            torch_dtype=torch.float32,
-            device_map="cpu"
-        )
-    except Exception as e:
-        raise RuntimeError(f"Model loading failed: {str(e)}")
+    # Load Phi-2 model 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
+    )
     
     return vector_store, model, tokenizer
 
-# Initialize system
 try:
     vector_store, model, tokenizer = initialize_system()
-    print("✅ System ready with business documents")
+    print("System ready with business documents loaded")
 except Exception as e:
-    print(f"❌ Initialization failed: {str(e)}")
-    raise
+    raise RuntimeError(f"Initialization failed: {str(e)}")
 
-# Response Generation
+# Response generation
 def generate_response(query):
-    try:
-        # Context retrieval
-        docs = vector_store.similarity_search(query, k=2)
-        context = "\n".join([d.page_content for d in docs])
-        
-        # Phi-3 prompt template
-        prompt = f"""<|system|>
-        Answer ONLY using the business documents. Respond to unknown queries with: "This information is not available in our current documentation."
-        
-        Context: {context}</s>
-        <|user|>
-        {query}</s>
-        <|assistant|>
-        """
-        
-        # Generate response
-        inputs = tokenizer(prompt, return_tensors="pt", return_attention_mask=False)
-        outputs = model.generate(
-            inputs.input_ids,
-            max_new_tokens=200,
-            temperature=0.1,
-            do_sample=True,
-            pad_token_id=tokenizer.eos_token_id
-        )
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response.split("<|assistant|>")[-1].strip()
+    # 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.'
+    
+    Context: {context}
+    
+    Question: {query}
+    
+    Answer:"""
     
-    except Exception as e:
-        return f"Error: Please try again. ({str(e)[:50]})"
+    # 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
+    )
+    
+    # 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
 
-# Gradio Interface
+# Chat interface
 with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# 📚 Business Documentation Assistant")
+    gr.Markdown("# Customer Care Assistant")
+    gr.Markdown("Ask questions about our products/services")
     
-    chatbot = gr.Chatbot(height=300)
-    msg = gr.Textbox(placeholder="Ask about our services...", label="")
+    chatbot = gr.Chatbot(height=400)
+    msg = gr.Textbox(label="Type your question here...")
     clear = gr.Button("Clear History")
     
-    def respond(message, history):
-        response = generate_response(message)
-        history.append((message, response))
-        return "", history
+    def respond(message, chat_history):
+        try:
+            response = generate_response(message)
+            if not response:
+                response = "I need to verify that information. Please contact [email protected]"
+        except Exception as e:
+            response = "Apologies, I'm experiencing technical difficulties. Please try again later."
+        
+        chat_history.append((message, response))
+        return "", chat_history
     
     msg.submit(respond, [msg, chatbot], [msg, chatbot])
     clear.click(lambda: None, None, chatbot, queue=False)
 
-if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+demo.launch(server_name="0.0.0.0", server_port=7860)