Update app.py

app.py

@@ -1,56 +1,73 @@
-import os
-import torch
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
+from PyPDF2 import PdfReader
 from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
-import PyPDF2
+import faiss
+import numpy as np
+from transformers import pipeline
 
-# Load LLM and Embedding model
-qa_model = "google/flan-t5-large"
-tokenizer = AutoTokenizer.from_pretrained(qa_model)
-model = AutoModelForSeq2SeqLM.from_pretrained(qa_model)
-qa_pipeline = pipeline("text2text-generation", model=model, tokenizer=tokenizer)
+# Load models once
+embedding_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+qa_model = pipeline("text2text-generation", model="google/flan-t5-base")
 
-embedder = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2")
+# Store docs and vectors
+doc_chunks = []
+doc_embeddings = []
+index = None
 
-# Global document store
-documents = []
-document_embeddings = []
+def read_pdf(file_path):
+    try:
+        reader = PdfReader(file_path)
+        text = ""
+        for page in reader.pages:
+            text += page.extract_text() or ""
+        return text
+    except Exception as e:
+        return f"Error reading PDF: {e}"
 
-def extract_text(file):
-    reader = PyPDF2.PdfReader(file)
-    return "\n".join(page.extract_text() for page in reader.pages if page.extract_text())
+def add_document(file_path):
+    global doc_chunks, doc_embeddings, index
+    text = read_pdf(file_path)
+    if not text.strip():
+        return "❌ Could not extract text from PDF."
+    
+    # Chunking the text (you can improve chunking logic)
+    chunks = [text[i:i+500] for i in range(0, len(text), 500)]
+    embeddings = embedding_model.encode(chunks)
 
-def add_document(file):
-    text = extract_text(file)
-    documents.append(text)
-    document_embeddings.append(embedder.encode(text))
-    return "Document uploaded and indexed successfully."
+    # Save to global
+    doc_chunks = chunks
+    doc_embeddings = embeddings
 
-def generate_answer(query):
-    if not documents:
-        return "Please upload a document first."
+    # Create FAISS index
+    dim = len(embeddings[0])
+    index = faiss.IndexFlatL2(dim)
+    index.add(np.array(embeddings))
 
-    query_embedding = embedder.encode(query)
-    similarities = cosine_similarity([query_embedding], document_embeddings)[0]
-    best_match_index = similarities.argmax()
-    relevant_text = documents[best_match_index][:3000]  # Truncate if too long
+    return f"✅ Uploaded & indexed {len(chunks)} chunks."
 
-    prompt = f"Answer this question based on the context:\n\nContext: {relevant_text}\n\nQuestion: {query}"
-    answer = qa_pipeline(prompt, max_new_tokens=300, temperature=0.3)[0]["generated_text"]
-    return answer.strip()
+def generate_answer(query):
+    if index is None:
+        return "⚠️ Please upload a document first."
+    
+    query_vec = embedding_model.encode([query])
+    D, I = index.search(np.array(query_vec), k=3)
+    context = " ".join([doc_chunks[i] for i in I[0]])
+    
+    # Use QA model
+    prompt = f"Context: {context}\n\nQuestion: {query}\nAnswer:"
+    result = qa_model(prompt, max_new_tokens=128)[0]["generated_text"]
+    return result.strip()
 
 # Gradio UI
 with gr.Blocks() as demo:
-    gr.Markdown("# 📄 Document Reader with RAG (Flan-T5)")
-    file_input = gr.File(label="Upload PDF", type="filepath")  # <-- FIXED HERE
+    gr.Markdown("## 📄 Document Q&A with PDF Upload")
+    file_input = gr.File(label="Upload PDF", type="filepath")
     upload_btn = gr.Button("Upload & Index")
-    query = gr.Textbox(label="Ask a question")
-    submit_btn = gr.Button("Get Answer")
-    answer_box = gr.Textbox(label="Answer")
+    query_input = gr.Textbox(label="Ask your question here")
+    submit_btn = gr.Button("Answer")
+    output_box = gr.Textbox(label="Answer")
 
-    upload_btn.click(fn=add_document, inputs=file_input, outputs=answer_box)
-    submit_btn.click(fn=generate_answer, inputs=query, outputs=answer_box)
+    upload_btn.click(fn=add_document, inputs=file_input, outputs=output_box)
+    submit_btn.click(fn=generate_answer, inputs=query_input, outputs=output_box)
 
 demo.launch()