Update app.py

app.py

@@ -1,26 +1,20 @@
 import gradio as gr
-import os
 from PyPDF2 import PdfReader
 from sentence_transformers import SentenceTransformer
-from transformers import pipeline, AutoTokenizer, AutoModelForCausalLM
-import torch
+from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM
 import numpy as np
 from sklearn.metrics.pairwise import cosine_similarity
 
-# Load sentence transformer for embeddings
+# Load embedding model (small and fast)
 embedder = SentenceTransformer("all-MiniLM-L6-v2")
 
-# Load Zephyr model (use accelerate for GPU/CPU device map)
-model_name = "HuggingFaceH4/zephyr-7b-beta"
+# Load FLAN-T5 model (CPU-friendly)
+model_name = "google/flan-t5-base"
 tokenizer = AutoTokenizer.from_pretrained(model_name)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16,
-    device_map="auto"  # ✅ Works if 'accelerate' is installed
-)
-rag_pipeline = pipeline("text-generation", model=model, tokenizer=tokenizer)
+model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
+rag_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
 
-# Function to extract text from PDF
+# Read PDF text
 def read_pdf(file_path):
     try:
         with open(file_path, "rb") as file:
@@ -34,12 +28,12 @@ def read_pdf(file_path):
     except Exception as e:
         return f"Error reading PDF: {str(e)}"
 
-# Function to split text into chunks (~500 words)
+# Split into ~500-word chunks
 def chunk_text(text, chunk_size=500):
     words = text.split()
     return [" ".join(words[i:i+chunk_size]) for i in range(0, len(words), chunk_size)]
 
-# Function to find top-k most relevant chunks
+# Find top-k relevant chunks via cosine similarity
 def retrieve_relevant_chunks(question, chunks, top_k=3):
     chunk_embeddings = embedder.encode(chunks)
     question_embedding = embedder.encode([question])
@@ -47,44 +41,38 @@ def retrieve_relevant_chunks(question, chunks, top_k=3):
     top_indices = np.argsort(scores)[-top_k:][::-1]
     return "\n\n".join([chunks[i] for i in top_indices])
 
-# Main function to process question and return answer
+# Main QA function
 def answer_question(uploaded_file, user_question):
     if uploaded_file is None:
         return "❌ Please upload a PDF file."
 
-    file_path = uploaded_file.name
-    document_text = read_pdf(file_path)
+    text = read_pdf(uploaded_file.name)
+    if not text or not isinstance(text, str):
+        return "❌ Could not extract text from the document."
 
-    if not document_text or not isinstance(document_text, str):
-        return "❌ Document is empty or could not be read."
-
-    chunks = chunk_text(document_text)
+    chunks = chunk_text(text)
     if not chunks:
-        return "❌ Document is too short to process."
-
-    relevant_context = retrieve_relevant_chunks(user_question, chunks)
+        return "❌ Document too short or empty."
 
-    prompt = f"""You are a helpful assistant. Use the context below to answer the user's question.\n\nContext:\n{relevant_context}\n\nQuestion: {user_question}\nAnswer:"""
+    context = retrieve_relevant_chunks(user_question, chunks)
 
+    prompt = f"Context: {context}\n\nQuestion: {user_question}\nAnswer:"
+    
     try:
-        result = rag_pipeline(prompt, max_new_tokens=300, do_sample=True, temperature=0.7)
-        answer = result[0]["generated_text"].split("Answer:")[-1].strip()
-        return str(answer)
+        result = rag_pipeline(prompt, max_new_tokens=256)
+        return result[0]["generated_text"].strip()
     except Exception as e:
-        return f"❌ Error generating answer: {str(e)}"
+        return f"❌ Error during generation: {str(e)}"
 
-# Gradio interface
+# Gradio Interface
 def create_interface():
     with gr.Blocks() as demo:
-        gr.Markdown("## 📄 Ask Questions from a PDF Document (RAG using Zephyr 7B)")
-
+        gr.Markdown("## 📄 Ask Questions from a PDF Document (RAG using FLAN-T5)")
         file_input = gr.File(label="Upload PDF", file_types=[".pdf"])
         question_input = gr.Textbox(label="Enter your question")
         answer_output = gr.Textbox(label="Answer", lines=10)
 
-        ask_button = gr.Button("Ask")
-        ask_button.click(fn=answer_question, inputs=[file_input, question_input], outputs=answer_output)
-
+        gr.Button("Ask").click(fn=answer_question, inputs=[file_input, question_input], outputs=[answer_output])
     return demo
 
 if __name__ == "__main__":