Update app.py

app.py

@@ -1,6 +1,7 @@
 import gradio as gr
-import fitz  # PyMuPDF
+import fitz
 import re
+import faiss
 import numpy as np
 
 from sentence_transformers import SentenceTransformer
@@ -8,14 +9,14 @@ from transformers import pipeline
 
 
 # =================================================
-# MODEL LOADING (ONCE AT STARTUP)
+# MODELS
 # =================================================
 
-# Embedding model (used for chunk relevance if needed later)
+# Embedding model (for retrieval)
 embedding_model = SentenceTransformer("multi-qa-MiniLM-L6-cos-v1")
 
-# Facebook BART summarization model (BEST accuracy on CPU)
-summarizer = pipeline(
+# BART summarization model (used as answer generator)
+bart = pipeline(
     "summarization",
     model="facebook/bart-large-cnn",
     tokenizer="facebook/bart-large-cnn"
@@ -41,99 +42,109 @@ def clean_text(text):
     return text.strip()
 
 
-def chunk_text(text, chunk_size=900, overlap=100):
-    """
-    Larger chunks are better for summarization
-    """
+def chunk_text(text, chunk_size=400, overlap=80):
     chunks = []
     start = 0
-    text_length = len(text)
-
-    while start < text_length:
+    while start < len(text):
         end = start + chunk_size
         chunks.append(text[start:end])
         start = end - overlap
-
     return chunks
 
 
 # =================================================
-# SUMMARIZATION LOGIC (ACCURATE & CPU SAFE)
+# VECTOR SEARCH
 # =================================================
 
-def summarize_chunks(chunks):
-    summaries = []
+def build_faiss_index(chunks):
+    embeddings = embedding_model.encode(chunks)
+    embeddings = np.array(embeddings).astype("float32")
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(embeddings)
+    return index, chunks
+
+
+def retrieve_chunks(question, index, chunks, top_k=3):
+    q_emb = embedding_model.encode([question]).astype("float32")
+    _, indices = index.search(q_emb, top_k)
+    return [chunks[i] for i in indices[0]]
+
+
+# =================================================
+# QUESTION–ANSWER USING BART
+# =================================================
 
-    for chunk in chunks:
-        summary = summarizer(
-            chunk,
-            max_length=150,
-            min_length=60,
-            do_sample=False
-        )[0]["summary_text"]
+def generate_answer(question, context_chunks):
+    context = " ".join(context_chunks)
 
-        summaries.append(summary)
+    prompt = f"""
+Answer the following question using ONLY the given context.
 
-    return " ".join(summaries)
+Context:
+{context}
+
+Question:
+{question}
+"""
+
+    result = bart(
+        prompt,
+        max_length=120,
+        min_length=30,
+        do_sample=False
+    )[0]["summary_text"]
+
+    return result
 
 
 # =================================================
 # MAIN PIPELINE
 # =================================================
 
-def pdf_summarizer(pdf_file):
-    if pdf_file is None:
-        return "Please upload a PDF document."
+def pdf_qa(pdf_file, question):
+    if pdf_file is None or question.strip() == "":
+        return "Please upload a PDF and ask a question."
 
-    # 1. Extract & clean text
-    raw_text = extract_text_from_pdf(pdf_file.name)
-    cleaned_text = clean_text(raw_text)
+    text = extract_text_from_pdf(pdf_file.name)
+    text = clean_text(text)
 
-    # 2. Chunk text
-    chunks = chunk_text(cleaned_text)
+    chunks = chunk_text(text)
+    index, chunks = build_faiss_index(chunks)
 
-    # 3. Summarize
-    final_summary = summarize_chunks(chunks)
+    relevant_chunks = retrieve_chunks(question, index, chunks)
+    answer = generate_answer(question, relevant_chunks)
 
-    return final_summary
+    return answer
 
 
 # =================================================
-# GRADIO UI (PRODUCTION READY)
+# GRADIO UI
 # =================================================
 
 with gr.Blocks() as demo:
 
     gr.Markdown("""
-    # 📄 PDF Summarizer (Open-Source AI)
-
-    Upload a **PDF document** to generate an **accurate, concise summary**.  
-    This system uses **Facebook BART**, a state-of-the-art open-source  
-    summarization model, optimized for **CPU execution**.
+    # 📄 PDF Question Answering System (BART Based)
 
-    ---
+    Upload a **PDF** and ask a **specific question**.  
+    The system retrieves relevant content and generates a **focused answer**,  
+    not a full summary.
     """)
 
     with gr.Row():
         with gr.Column(scale=1):
-            pdf_input = gr.File(
-                label="📤 Upload PDF",
-                file_types=[".pdf"]
+            pdf_input = gr.File(label="📤 Upload PDF", file_types=[".pdf"])
+            question_input = gr.Textbox(
+                label="❓ Ask your question",
+                placeholder="e.g. What is the objective of the project?",
+                lines=2
             )
-
-            summarize_btn = gr.Button("📝 Generate Summary")
+            btn = gr.Button("🔍 Get Answer")
 
         with gr.Column(scale=2):
-            summary_output = gr.Textbox(
-                label="📌 Summary",
-                lines=12
-            )
+            output = gr.Textbox(label="📌 Answer", lines=8)
 
-    summarize_btn.click(
-        fn=pdf_summarizer,
-        inputs=[pdf_input],
-        outputs=summary_output
-    )
+    btn.click(pdf_qa, [pdf_input, question_input], output)
 
     gr.Markdown("""
     ---