Update app.py

app.py

@@ -2,7 +2,6 @@ import os
 import time
 import gradio as gr
 import numpy as np
-import faiss
 from PyPDF2 import PdfReader
 from sentence_transformers import SentenceTransformer
 from huggingface_hub import InferenceClient
@@ -21,16 +20,14 @@ HF_TOKEN = os.getenv("HF_TOKEN", None)
 # LLM_MODEL ="HuggingFaceH4/zephyr-7b-beta"
 # LLM_MODEL = "NousResearch/Hermes-2-Pro-Mistral-7B" Does not work
 
-# Were Great
+# Were Great LLM
 # LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-7B"
 # LLM_MODEL = "deepseek-ai/DeepSeek-R1-Distill-Qwen-14B"
 
-# Great LLM
-##LLM_MODEL = 'meta-llama/Llama-3.1-8B-Instruct'
+# LLM model
 LLM_MODEL = "deepseek-ai/DeepSeek-V3"
 
 EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
-
 embedding_model = SentenceTransformer(EMBEDDING_MODEL)
 
 # IMPORTANT: do NOT pass model=... here
@@ -83,31 +80,42 @@ def chunk_text(text, chunk_size=800, overlap=200):
     return chunks
 
 
-def build_vectorstore(chunks):
-    if not chunks:
-        return None, None
+# ============================================================
+# PER-FILE RETRIEVAL
+# ============================================================
 
-    embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
-    dim = embeddings.shape[1]
+def retrieve_context_per_file(question, state, k=10):
+    """
+    For each file, retrieve top-k chunks separately,
+    then merge all selected chunks into a single context.
+    """
+    if not state["files"]:
+        return ""
 
-    index = faiss.IndexFlatL2(dim)
-    index.add(embeddings)
+    q_emb = embedding_model.encode([question], convert_to_numpy=True)[0]
+    all_selected = []
 
-    return index, embeddings
+    for filename, data in state["files"].items():
+        chunks = data["chunks"]
+        embeddings = data["embeddings"]
 
+        if len(chunks) == 0:
+            continue
 
-def retrieve_context(question, state, k=15):
-    index = state.get("index")
-    chunks = state.get("chunks")
+        # Compute L2 distances to the query embedding
+        distances = np.linalg.norm(embeddings - q_emb, axis=1)
 
-    if index is None or chunks is None:
-        return ""
+        # Get top-k indices for this file
+        top_k_idx = distances.argsort()[: min(k, len(chunks))]
 
-    q_emb = embedding_model.encode([question], convert_to_numpy=True)
-    distances, indices = index.search(q_emb, k)
+        # Collect chunks (optionally tag with filename)
+        selected = [f"[{filename}]\n{chunks[i]}" for i in top_k_idx]
+        all_selected.extend(selected)
 
-    selected = [chunks[i] for i in indices[0] if i < len(chunks)]
-    return "\n\n---\n\n".join(selected)
+    if not all_selected:
+        return ""
+
+    return "\n\n---\n\n".join(all_selected)
 
 
 # ============================================================
@@ -136,7 +144,7 @@ def generate_answer(question, context):
         temperature=0.25,
         top_p=0.9,
         frequency_penalty=0.1,
-        presence_penalty=0.0
+        presence_penalty=0.0,
     )
 
     return response.choices[0].message["content"].strip()
@@ -147,96 +155,113 @@ def generate_answer(question, context):
 # ============================================================
 
 def load_pdf(files, state):
-    # files is the *current* list from the UI (after add/remove)
-    # We rebuild everything from scratch to stay in sync.
-
+    """
+    files: current list from the UI (after add/remove).
+    We rebuild state from scratch to stay in sync.
+    State structure:
+
+    state = {
+        "files": {
+            "filename.pdf": {
+                "chunks": [...],
+                "embeddings": np.array([...])
+            },
+            ...
+        },
+        "pdf_count": int,
+        "total_chunks": int
+    }
+    """
     # If user cleared all files
     if not files:
-        state["chunks"] = []
         state["files"] = {}
         state["pdf_count"] = 0
-        state["index"] = None
+        state["total_chunks"] = 0
 
         return (
             "No file uploaded.",
             state,
-            gr.update(choices=[], value=None)  # empty dropdown
+            gr.update(choices=[], value=None),
         )
 
     try:
-        # Reset state and rebuild from current files only
-        state["chunks"] = []
+        # Reset and rebuild from current files only
         state["files"] = {}
         state["pdf_count"] = 0
+        state["total_chunks"] = 0
 
         # Deduplicate while preserving order
         unique_files = list(dict.fromkeys(files))
 
         for file in unique_files:
-            # In Gradio 6, gr.Files usually gives you temp paths (strings)
-            # PdfReader accepts a path or file-like object, so we can pass `file` directly.
             filename = os.path.basename(str(file))
 
             text = read_pdf(file)
-            new_chunks = chunk_text(text)
+            chunks = chunk_text(text)
 
-            if not new_chunks:
+            if not chunks:
                 continue
 
-            state["files"][filename] = new_chunks
-            state["chunks"].extend(new_chunks)
+            embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
+
+            state["files"][filename] = {
+                "chunks": chunks,
+                "embeddings": embeddings,
+            }
+
             state["pdf_count"] += 1
+            state["total_chunks"] += len(chunks)
 
-        # Rebuild FAISS
-        index, _ = build_vectorstore(state["chunks"])
-        state["index"] = index
+        status_msg = (
+            f"{state['pdf_count']} PDF(s) loaded. "
+            f"Total chunks: {state['total_chunks']}"
+        )
 
         return (
-            f"{state['pdf_count']} PDF(s) loaded. Total chunks: {len(state['chunks'])}",
+            status_msg,
             state,
-            gr.update(choices=list(state["files"].keys()), value=None)  # <-- key fix
+            gr.update(choices=list(state["files"].keys()), value=None),
         )
 
     except Exception as e:
         return (
             f"Error processing PDF: {e}",
             state,
-            gr.update(choices=list(state["files"].keys()), value=None)
+            gr.update(choices=list(state["files"].keys()), value=None),
         )
-        
+
+
 def remove_pdf(filename, state):
     if filename not in state["files"]:
         return (
             f"{filename} not found.",
             state,
-            gr.update(choices=list(state["files"].keys()), value=None)
+            gr.update(choices=list(state["files"].keys()), value=None),
         )
 
     # Remove the file entry
     state["files"].pop(filename, None)
 
-    # Rebuild chunks from remaining files
-    all_chunks = []
-    for chunks in state["files"].values():
-        all_chunks.extend(chunks)
+    # Recompute counts
+    pdf_count = len(state["files"])
+    total_chunks = sum(len(data["chunks"]) for data in state["files"].values())
 
-    state["chunks"] = all_chunks
-    state["pdf_count"] = len(state["files"])
+    state["pdf_count"] = pdf_count
+    state["total_chunks"] = total_chunks
 
-    if state["chunks"]:
-        index, _ = build_vectorstore(state["chunks"])
-        state["index"] = index
-    else:
-        state["index"] = None
+    status_msg = (
+        f"Removed {filename}. "
+        f"Total PDFs: {pdf_count}. Total chunks: {total_chunks}"
+    )
 
     return (
-        f"Removed {filename}. Total chunks: {len(state['chunks'])}",
+        status_msg,
         state,
-        gr.update(choices=list(state["files"].keys()), value=None)
+        gr.update(choices=list(state["files"].keys()), value=None),
     )
 
+
 def answer_question(question, state):
-    
     print("HF_TOKEN loaded:", HF_TOKEN is not None)
 
     if HF_TOKEN:
@@ -248,10 +273,10 @@ def answer_question(question, state):
     if not question.strip():
         return "Please enter a question."
 
-    if state.get("index") is None:
-        return "Please upload a PDF first."
+    if not state["files"]:
+        return "Please upload at least one PDF first."
 
-    context = retrieve_context(question, state)
+    context = retrieve_context_per_file(question, state, k=10)
     answer = generate_answer(question, context)
     return answer
 
@@ -261,26 +286,30 @@ def answer_question(question, state):
 # ============================================================
 
 with gr.Blocks(title="PDF Q&A Chatbot") as demo:
-
     gr.Markdown("## 📄 PDF Q&A Chatbot\nUpload PDFs and ask questions about their content.")
 
-    state = gr.State({
-        "chunks": [],
-        "index": None,
-        "pdf_count": 0,
-        "files": {}
-    })
+    state = gr.State(
+        {
+            "files": {},          # filename -> {chunks, embeddings}
+            "pdf_count": 0,
+            "total_chunks": 0,
+        }
+    )
 
     with gr.Row():
         # LEFT COLUMN — PDF upload + removal
         with gr.Column(scale=1):
-            pdf_input = gr.Files(label="Upload PDF(s)", file_types=[".pdf"])
+            pdf_input = gr.File(
+                label="Upload PDF(s)",
+                file_types=[".pdf"],
+                file_count="multiple",  # proper multi-upload in Gradio 6
+            )
             pdf_status = gr.Textbox(label="Status", interactive=False)
 
             file_selector = gr.Dropdown(
                 label="Remove a PDF",
                 choices=[],
-                interactive=True
+                interactive=True,
             )
             remove_button = gr.Button("Remove Selected PDF")
 
@@ -290,31 +319,35 @@ with gr.Blocks(title="PDF Q&A Chatbot") as demo:
                 question = gr.Textbox(
                     label="Your question",
                     lines=2,
-                    scale=3
+                    scale=3,
                 )
                 clear_button = gr.Button(
                     "Clear",
                     scale=1,
                     min_width=80,
-                    elem_classes="clear-btn"
+                    elem_classes="clear-btn",
                 )
 
             answer = gr.Textbox(label="Answer", lines=8)
             ask_button = gr.Button("Ask")
 
     # Bind events AFTER components exist
-    pdf_input.change(load_pdf, inputs=[pdf_input, state], outputs=[pdf_status, state, file_selector])
+    pdf_input.change(
+        load_pdf,
+        inputs=[pdf_input, state],
+        outputs=[pdf_status, state, file_selector],
+    )
 
     remove_button.click(
         remove_pdf,
         inputs=[file_selector, state],
-        outputs=[pdf_status, state, file_selector]
+        outputs=[pdf_status, state, file_selector],
     )
 
     ask_button.click(
         answer_question,
         inputs=[question, state],
-        outputs=answer
+        outputs=answer,
     )
 
     clear_button.click(lambda: "", outputs=question)
@@ -322,5 +355,5 @@ with gr.Blocks(title="PDF Q&A Chatbot") as demo:
 
 if __name__ == "__main__":
     demo.launch(
-    css=".clear-btn button { height: 100%; min-height: 52px; }"
-)
+        css=".clear-btn button { height: 100%; min-height: 52px; }"
+    )