Update app.py

app.py

@@ -1,3 +1,4 @@
+# app.py
 import gradio as gr
 import os
 import torch
@@ -5,15 +6,23 @@ from langchain_community.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 transformers import AutoModelForCausalLM, AutoTokenizer, BitsAndBytesConfig
 
 # Configuration
 DOCS_DIR = "business_docs"
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 MODEL_NAME = "microsoft/phi-2"
 
+# Quantization config
+quant_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_quant_type="nf4",
+    bnb_4bit_compute_dtype=torch.float16,
+    bnb_4bit_use_double_quant=False
+)
+
 def initialize_system():
-    # Document verification
+    # Document processing
     if not os.path.exists(DOCS_DIR):
         raise FileNotFoundError(f"Missing {DOCS_DIR} folder")
     
@@ -21,10 +30,6 @@ def initialize_system():
                 for f in os.listdir(DOCS_DIR) 
                 if f.endswith(".pdf")]
     
-    if not pdf_files:
-        raise ValueError(f"No PDFs found in {DOCS_DIR}")
-
-    # Document processing
     text_splitter = RecursiveCharacterTextSplitter(
         chunk_size=800,
         chunk_overlap=100
@@ -56,7 +61,7 @@ def initialize_system():
         MODEL_NAME,
         trust_remote_code=True,
         device_map="auto",
-        load_in_4bit=True,
+        quantization_config=quant_config,
         torch_dtype=torch.float16
     )
 
@@ -64,47 +69,41 @@ def initialize_system():
 
 try:
     vector_store, model, tokenizer = initialize_system()
-    print("System initialized successfully ✅")
+    print("✅ System initialized successfully")
 except Exception as e:
-    print(f"Initialization failed ❌: {str(e)}")
+    print(f"❌ Initialization failed: {str(e)}")
     raise
 
 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-2 optimized prompt
         prompt = f"""<|system|>
-        You are a customer service bot. Answer only using:
-        {context}
-        - Max 3 sentences
-        - If unsure: "I'll check with the team"
-        </s>
-        <|user|>
-        {query}</s>
+        Answer using only this context: {context}
+        - Max 2 sentences
+        - If unsure: "I'll check with the team"</s>
+        <|user|>{query}</s>
         <|assistant|>"""
         
         inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
         outputs = model.generate(
             **inputs,
-            max_new_tokens=200,
+            max_new_tokens=150,
             temperature=0.1,
             pad_token_id=tokenizer.eos_token_id
         )
         
-        response = tokenizer.decode(outputs[0], skip_special_tokens=True)
-        return response.split("<|assistant|>")[-1].strip()
+        return tokenizer.decode(outputs[0], skip_special_tokens=True).split("<|assistant|>")[-1].strip()
     
     except Exception as e:
         return "Please try again later."
 
 # Gradio interface
-with gr.Blocks(theme=gr.themes.Soft()) as demo:
-    gr.Markdown("# Customer Support Chatbot")
+with gr.Blocks() as demo:
+    gr.Markdown("# Customer Service Chatbot")
     chatbot = gr.Chatbot()
-    msg = gr.Textbox(label="Ask about our services")
+    msg = gr.Textbox(label="Your question")
     clear = gr.ClearButton([msg, chatbot])
     
     def respond(message, history):