Update app.py

app.py

@@ -1,4 +1,4 @@
-# app.py
+# app.py (CPU-optimized)
 import gradio as gr
 import os
 import torch
@@ -6,21 +6,13 @@ 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, BitsAndBytesConfig
+from transformers import AutoModelForCausalLM, AutoTokenizer
 
 # 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 processing
     if not os.path.exists(DOCS_DIR):
@@ -31,8 +23,8 @@ def initialize_system():
                 if f.endswith(".pdf")]
     
     text_splitter = RecursiveCharacterTextSplitter(
-        chunk_size=800,
-        chunk_overlap=100
+        chunk_size=500,  # Smaller chunks for CPU
+        chunk_overlap=50
     )
     
     texts = []
@@ -50,7 +42,7 @@ def initialize_system():
     # Vector store
     vector_store = FAISS.from_documents(texts, embeddings)
 
-    # Model loading
+    # Load model without quantization
     tokenizer = AutoTokenizer.from_pretrained(
         MODEL_NAME,
         trust_remote_code=True,
@@ -60,9 +52,8 @@ def initialize_system():
     model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         trust_remote_code=True,
-        device_map="auto",
-        quantization_config=quant_config,
-        torch_dtype=torch.float16
+        torch_dtype=torch.float16,
+        device_map="cpu"  # Force CPU
     )
 
     return vector_store, model, tokenizer
@@ -76,22 +67,21 @@ except Exception as e:
 
 def generate_response(query):
     try:
-        docs = vector_store.similarity_search(query, k=2)
+        docs = vector_store.similarity_search(query, k=1)  # Less context
         context = "\n".join([d.page_content for d in docs])
         
         prompt = f"""<|system|>
-        Answer using only this context: {context}
-        - Max 2 sentences
+        Answer using: {context}
+        - Max 1 sentence
         - If unsure: "I'll check with the team"</s>
         <|user|>{query}</s>
         <|assistant|>"""
         
-        inputs = tokenizer(prompt, return_tensors="pt").to(model.device)
+        inputs = tokenizer(prompt, return_tensors="pt").to("cpu")
         outputs = model.generate(
             **inputs,
-            max_new_tokens=150,
-            temperature=0.1,
-            pad_token_id=tokenizer.eos_token_id
+            max_new_tokens=100,
+            temperature=0.1
         )
         
         return tokenizer.decode(outputs[0], skip_special_tokens=True).split("<|assistant|>")[-1].strip()
@@ -99,7 +89,7 @@ def generate_response(query):
     except Exception as e:
         return "Please try again later."
 
-# Gradio interface
+# Simplified interface
 with gr.Blocks() as demo:
     gr.Markdown("# Customer Service Chatbot")
     chatbot = gr.Chatbot()