Update app.py

app.py

@@ -18,40 +18,42 @@ login(token=hf_token)
 # Load embedding model
 embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
 
-# Load Mistral without 4bit quantization (CPU-friendly)
+# Load Mistral LLM (CPU compatible)
 model_id = "mistralai/Mistral-7B-Instruct-v0.1"
 tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
 model = AutoModelForCausalLM.from_pretrained(
     model_id,
-    device_map={"": "cpu"},  # force CPU
-    torch_dtype="auto",      # safe for CPU
+    device_map={"": "cpu"},  # Force CPU
+    torch_dtype="auto",      # Safe for CPU
     token=hf_token
 )
 llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
-# Globals
+# Global state
 index = None
 doc_texts = []
 
-# Extract text from PDF or TXT
-def extract_text(file):
+# Extract text from uploaded file
+def extract_text(file_obj):
     text = ""
-    file_bytes = file.read()
-    if file.name.endswith(".pdf"):
-        pdf_stream = BytesIO(file_bytes)
+    file_path = file_obj.name
+    if file_path.endswith(".pdf"):
+        with open(file_path, "rb") as f:
+            pdf_stream = BytesIO(f.read())
         doc = fitz.open(stream=pdf_stream, filetype="pdf")
         for page in doc:
             text += page.get_text()
-    elif file.name.endswith(".txt"):
-        text = file_bytes.decode("utf-8")
+    elif file_path.endswith(".txt"):
+        with open(file_path, "r", encoding="utf-8") as f:
+            text = f.read()
     else:
         return "❌ Unsupported file type."
     return text
 
-# Process the file, build FAISS index
-def process_file(file):
+# Process file and build FAISS index
+def process_file(file_obj):
     global index, doc_texts
-    text = extract_text(file)
+    text = extract_text(file_obj)
     if text.startswith("❌"):
         return text
 
@@ -63,9 +65,9 @@ def process_file(file):
     index = faiss.IndexFlatL2(dim)
     index.add(embeddings)
 
-    return "✅ File processed! You can now ask questions."
+    return "✅ File processed successfully. You can now ask questions!"
 
-# Generate answer using context + LLM
+# Generate answer from FAISS context + LLM
 def generate_answer(question):
     global index, doc_texts
     if index is None or not doc_texts:
@@ -75,30 +77,30 @@ def generate_answer(question):
     _, I = index.search(question_emb, k=3)
     context = "\n".join([doc_texts[i] for i in I[0]])
 
-    prompt = f"""[System: You are a helpful assistant. Answer based on the context.]
+    prompt = f"""<s>[INST] You are a helpful assistant. Use the context below to answer the question.
 
 Context:
 {context}
 
 Question: {question}
-Answer:"""
+Answer: [/INST]</s>"""
 
     response = llm(prompt, max_new_tokens=300, do_sample=True, temperature=0.7)
     return response[0]["generated_text"].split("Answer:")[-1].strip()
 
 # Gradio UI
-with gr.Blocks(title="RAG Chatbot (CPU Compatible)") as demo:
-    gr.Markdown("## 📚 Upload PDF/TXT and Ask Questions using Mistral-7B")
+with gr.Blocks(title="RAG Chatbot with Mistral-7B (CPU-Friendly)") as demo:
+    gr.Markdown("## 🤖 Upload a PDF/TXT file and ask questions using Mistral-7B")
 
     with gr.Row():
-        file_input = gr.File(label="📁 Upload File (.pdf or .txt)", file_types=[".pdf", ".txt"])
-        upload_status = gr.Textbox(label="Upload Status", interactive=False)
+        file_input = gr.File(label="📁 Upload PDF or TXT", file_types=[".pdf", ".txt"])
+        upload_status = gr.Textbox(label="📥 Upload Status", interactive=False)
 
     with gr.Row():
-        question_box = gr.Textbox(label="❓ Ask a Question", placeholder="What would you like to know?")
-        answer_box = gr.Textbox(label="💬 Answer", interactive=False)
+        question_input = gr.Textbox(label="❓ Ask a Question")
+        answer_output = gr.Textbox(label="💬 Answer", interactive=False)
 
     file_input.change(fn=process_file, inputs=file_input, outputs=upload_status)
-    question_box.submit(fn=generate_answer, inputs=question_box, outputs=answer_box)
+    question_input.submit(fn=generate_answer, inputs=question_input, outputs=answer_output)
 
 demo.launch()