Update app.py

app.py

@@ -3,13 +3,12 @@ import gradio as gr
 import fitz  # PyMuPDF
 import faiss
 import numpy as np
-from io import BytesIO
 from sentence_transformers import SentenceTransformer
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from huggingface_hub import login
 
-# Authenticate with Hugging Face
+# Load Hugging Face Token from environment
 hf_token = os.environ.get("HUGGINGFACE_TOKEN")
 if not hf_token:
     raise ValueError("⚠️ Please set the HUGGINGFACE_TOKEN environment variable.")
@@ -18,42 +17,37 @@ login(token=hf_token)
 # Load embedding model
 embed_model = SentenceTransformer("BAAI/bge-base-en-v1.5")
 
-# ✅ 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)
+# Load small, fast LLM (great for CPU)
+model_id = "tiiuae/falcon-rw-1b"
+tokenizer = AutoTokenizer.from_pretrained(model_id, token=hf_token)
+model = AutoModelForCausalLM.from_pretrained(
+    model_id,
+    device_map={"": "cpu"},
+    torch_dtype="auto",
+    token=hf_token
+)
+llm = pipeline("text-generation", model=model, tokenizer=tokenizer)
 
 # Globals
 index = None
 doc_texts = []
 
-# Extract text from PDF or TXT
+# Extract text from PDF or TXT (handle Hugging Face Spaces file upload)
 def extract_text(file):
     text = ""
-    
-    # Handle if file is a NamedString (like on HF Spaces)
-    if isinstance(file, str):
-        with open(file, "rb") as f:
-            file_bytes = f.read()
-        filename = file
-    else:
-        file_bytes = file.read()
-        filename = file.name
-
-    if filename.endswith(".pdf"):
-        pdf_stream = BytesIO(file_bytes)
-        doc = fitz.open(stream=pdf_stream, filetype="pdf")
-        for page in doc:
-            text += page.get_text()
-    elif filename.endswith(".txt"):
-        text = file_bytes.decode("utf-8")
+    file_path = file.name if hasattr(file, 'name') else file
+    if file_path.endswith(".pdf"):
+        with fitz.open(file_path) as doc:
+            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()
     else:
         return "❌ Unsupported file type."
     return text
 
-
-# Process the file, build FAISS index
+# Process file and build FAISS index
 def process_file(file):
     global index, doc_texts
     text = extract_text(file)
@@ -68,42 +62,42 @@ def process_file(file):
     index = faiss.IndexFlatL2(dim)
     index.add(embeddings)
 
-    return "✅ File processed! You can now ask questions."
+    return "✅ File processed successfully! Ask your question below."
 
-# Generate answer using context + LLM
+# Generate answer
 def generate_answer(question):
     global index, doc_texts
     if index is None or not doc_texts:
-        return "⚠️ Please upload and process a file first."
+        return "⚠️ Please upload and process a document first."
 
     question_emb = embed_model.encode([question], convert_to_numpy=True)
     _, I = index.search(question_emb, k=3)
     context = "\n".join([doc_texts[i] for i in I[0]])
 
-    prompt = f"""Use the following context to answer the question.
+    prompt = f"""[System: You are a helpful assistant. Answer based on the context.]
 
 Context:
 {context}
 
 Question: {question}
-"""
+Answer:"""
 
-    response = llm(prompt, max_new_tokens=300)
-    return response[0]["generated_text"].strip()
+    result = llm(prompt, max_new_tokens=200, do_sample=True, temperature=0.7)
+    return result[0]["generated_text"].split("Answer:")[-1].strip()
 
 # Gradio UI
-with gr.Blocks(title="RAG Chatbot (Fast & CPU Compatible)") as demo:
-    gr.Markdown("## 📚 Upload PDF/TXT and Ask Questions using FLAN-T5")
+with gr.Blocks(title="RAG Chatbot (CPU-Optimized)") as demo:
+    gr.Markdown("## 📚 Upload PDF/TXT and Ask Questions (Fast CPU RAG Bot)")
 
     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_output = 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", placeholder="E.g. What is the document about?")
+        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)
+    file_input.change(fn=process_file, inputs=file_input, outputs=upload_output)
+    question_input.submit(fn=generate_answer, inputs=question_input, outputs=answer_output)
 
 demo.launch()