Update app.py

app.py

@@ -5,7 +5,7 @@ import faiss
 import numpy as np
 from io import BytesIO
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from huggingface_hub import login
 
@@ -18,42 +18,35 @@ login(token=hf_token)
 # Load embedding model
 embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
 
-# 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
-    token=hf_token
-)
-llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
-
-# Global state
+# ✅ Load FLAN-T5 base (CPU-friendly)
+model_id = "google/flan-t5-base"
+tokenizer = AutoTokenizer.from_pretrained(model_id)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+llm = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+
+# Globals
 index = None
 doc_texts = []
 
-# Extract text from uploaded file
-def extract_text(file_obj):
+# Extract text from PDF or TXT
+def extract_text(file):
     text = ""
-    file_path = file_obj.name
-    if file_path.endswith(".pdf"):
-        with open(file_path, "rb") as f:
-            pdf_stream = BytesIO(f.read())
+    file_bytes = file.read()
+    if file.name.endswith(".pdf"):
+        pdf_stream = BytesIO(file_bytes)
         doc = fitz.open(stream=pdf_stream, filetype="pdf")
         for page in doc:
             text += page.get_text()
-    elif file_path.endswith(".txt"):
-        with open(file_path, "r", encoding="utf-8") as f:
-            text = f.read()
+    elif file.name.endswith(".txt"):
+        text = file_bytes.decode("utf-8")
     else:
         return "❌ Unsupported file type."
     return text
 
-# Process file and build FAISS index
-def process_file(file_obj):
+# Process the file, build FAISS index
+def process_file(file):
     global index, doc_texts
-    text = extract_text(file_obj)
+    text = extract_text(file)
     if text.startswith("❌"):
         return text
 
@@ -65,9 +58,9 @@ def process_file(file_obj):
     index = faiss.IndexFlatL2(dim)
     index.add(embeddings)
 
-    return "✅ File processed successfully. You can now ask questions!"
+    return "✅ File processed! You can now ask questions."
 
-# Generate answer from FAISS context + LLM
+# Generate answer using context + LLM
 def generate_answer(question):
     global index, doc_texts
     if index is None or not doc_texts:
@@ -77,30 +70,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"""<s>[INST] You are a helpful assistant. Use the context below to answer the question.
+    prompt = f"""Use the following context to answer the question.
 
 Context:
 {context}
 
 Question: {question}
-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()
+    response = llm(prompt, max_new_tokens=300)
+    return response[0]["generated_text"].strip()
 
 # Gradio UI
-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.Blocks(title="RAG Chatbot (Fast & CPU Compatible)") as demo:
+    gr.Markdown("## 📚 Upload PDF/TXT and Ask Questions using FLAN-T5")
 
     with gr.Row():
-        file_input = gr.File(label="📁 Upload PDF or TXT", file_types=[".pdf", ".txt"])
-        upload_status = gr.Textbox(label="📥 Upload Status", interactive=False)
+        file_input = gr.File(label="📁 Upload File (.pdf or .txt)", file_types=[".pdf", ".txt"])
+        upload_status = gr.Textbox(label="Upload Status", interactive=False)
 
     with gr.Row():
-        question_input = gr.Textbox(label="❓ Ask a Question")
-        answer_output = gr.Textbox(label="💬 Answer", interactive=False)
+        question_box = gr.Textbox(label="❓ Ask a Question", placeholder="What would you like to know?")
+        answer_box = gr.Textbox(label="💬 Answer", interactive=False)
 
     file_input.change(fn=process_file, inputs=file_input, outputs=upload_status)
-    question_input.submit(fn=generate_answer, inputs=question_input, outputs=answer_output)
+    question_box.submit(fn=generate_answer, inputs=question_box, outputs=answer_box)
 
 demo.launch()