Update app.py

app.py

@@ -1,134 +1,124 @@
 import os
-import io
 import gradio as gr
 import numpy as np
 import faiss
-from huggingface_hub import InferenceClient
-from sentence_transformers import SentenceTransformer
 from PyPDF2 import PdfReader
+from sentence_transformers import SentenceTransformer
+from huggingface_hub import InferenceClient
 
-# =========================
-# Config
-# =========================
+# ============================================================
+# CONFIG
+# ============================================================
 
-# Set this in your HF Space secrets or local env
 HF_TOKEN = os.getenv("HF_TOKEN", None)
 
-# IBM Granite model on Hugging Face (free tier compatible)
+# Free IBM Granite model on Hugging Face
 GRANITE_MODEL = "ibm-granite/granite-8b-instruct"
 
-# Embedding model (runs locally in the Space CPU)
-EMBEDDING_MODEL_NAME = "sentence-transformers/all-MiniLM-L6-v2"
+# Local embedding model (fast + CPU friendly)
+EMBEDDING_MODEL = "sentence-transformers/all-MiniLM-L6-v2"
 
 # Load models once per process
-embedding_model = SentenceTransformer(EMBEDDING_MODEL_NAME)
-llm_client = InferenceClient(model=GRANITE_MODEL, token=HF_TOKEN)
+embedding_model = SentenceTransformer(EMBEDDING_MODEL)
+llm = InferenceClient(model=GRANITE_MODEL, token=HF_TOKEN)
 
 
-# =========================
-# PDF & RAG utilities
-# =========================
+# ============================================================
+# PDF PROCESSING
+# ============================================================
 
-def read_pdf(file_obj) -> str:
-    """Read all text from a PDF file object."""
-    if file_obj is None:
-        return ""
-
-    # Gradio gives a tempfile-like object
+def read_pdf(file_obj):
+    """Extract raw text from a PDF file."""
     reader = PdfReader(file_obj)
-    texts = []
+    pages = []
     for page in reader.pages:
         text = page.extract_text()
         if text:
-            texts.append(text)
-    return "\n".join(texts)
-
+            pages.append(text)
+    return "\n".join(pages)
 
-def chunk_text(text: str, chunk_size: int = 800, overlap: int = 200):
-    """Simple text splitter into overlapping chunks."""
-    if not text:
-        return []
 
+def chunk_text(text, chunk_size=800, overlap=200):
+    """Split text into overlapping chunks."""
     words = text.split()
     chunks = []
     start = 0
+
     while start < len(words):
         end = start + chunk_size
         chunk = " ".join(words[start:end])
         chunks.append(chunk)
         start += chunk_size - overlap
+
     return chunks
 
 
 def build_vectorstore(chunks):
-    """Create FAISS index and store embeddings."""
+    """Embed chunks and build FAISS index."""
     if not chunks:
         return None, None
 
     embeddings = embedding_model.encode(chunks, convert_to_numpy=True)
     dim = embeddings.shape[1]
+
     index = faiss.IndexFlatL2(dim)
     index.add(embeddings)
+
     return index, embeddings
 
 
-def retrieve_context(question: str, state: dict, k: int = 4) -> str:
-    """Retrieve top-k relevant chunks as context."""
+def retrieve_context(question, state, k=4):
+    """Retrieve top-k relevant chunks."""
     index = state.get("index")
     chunks = state.get("chunks")
 
-    if index is None or chunks is None or len(chunks) == 0:
+    if index is None or chunks is None:
         return ""
 
     q_emb = embedding_model.encode([question], convert_to_numpy=True)
     distances, indices = index.search(q_emb, k)
-    indices = indices[0]
 
-    selected_chunks = [chunks[i] for i in indices if 0 <= i < len(chunks)]
-    return "\n\n---\n\n".join(selected_chunks)
+    selected = [chunks[i] for i in indices[0] if i < len(chunks)]
+    return "\n\n---\n\n".join(selected)
+
 
+# ============================================================
+# LLM ANSWERING (IBM GRANITE)
+# ============================================================
 
-def generate_answer(question: str, context: str) -> str:
-    """Call Granite via HF Inference to generate an answer."""
-    if not context:
-        system_prompt = (
-            "You are a helpful assistant. The user asked a question, "
-            "but you have no document context. Answer briefly and honestly."
+def generate_answer(question, context):
+    """Generate an answer using IBM Granite."""
+    if context:
+        prompt = (
+            "You are a helpful assistant answering questions strictly from the provided context.\n"
+            "If the answer is not in the context, say you don't know.\n\n"
+            f"Context:\n{context}\n\n"
+            f"Question: {question}\nAnswer:"
         )
-        prompt = f"{system_prompt}\n\nQuestion: {question}\nAnswer:"
     else:
-        system_prompt = (
-            "You are a helpful assistant answering questions about a document. "
-            "Use ONLY the information in the context. If the answer is not in "
-            "the context, say you don't know."
-        )
         prompt = (
-            f"{system_prompt}\n\n"
-            f"Context:\n{context}\n\n"
+            "You are a helpful assistant. The user asked a question but no document context is available.\n\n"
             f"Question: {question}\nAnswer:"
         )
 
-    # Granite models on HF typically support text_generation-style calls
-    response = llm_client.text_generation(
+    response = llm.text_generation(
         prompt,
         max_new_tokens=512,
         temperature=0.2,
         top_p=0.9,
         repetition_penalty=1.05,
     )
+
     return response.strip()
 
 
-# =========================
-# Gradio callbacks
-# =========================
+# ============================================================
+# GRADIO CALLBACKS
+# ============================================================
 
 def load_pdf(file, state):
-    """
-    Called when a new PDF is uploaded.
-    We RESET all previous state here so no old PDF is kept in memory.
-    """
-    # Hard reset of PDF-related state
+    """Reset state and load a new PDF."""
+    # Reset everything to avoid memory persistence
     state["chunks"] = None
     state["index"] = None
 
@@ -140,7 +130,7 @@ def load_pdf(file, state):
         chunks = chunk_text(text)
 
         if not chunks:
-            return "Could not extract any text from the PDF.", state
+            return "Could not extract text from PDF.", state
 
         index, _ = build_vectorstore(chunks)
 
@@ -150,18 +140,15 @@ def load_pdf(file, state):
         return f"PDF loaded successfully. {len(chunks)} chunks created.", state
 
     except Exception as e:
-        return f"Error while processing PDF: {e}", state
+        return f"Error processing PDF: {e}", state
 
 
 def answer_question(question, state):
-    """
-    Called when the user submits a question.
-    Uses the current session's vectorstore only.
-    """
-    if not question or question.strip() == "":
+    """Answer a question using the current PDF context."""
+    if not question.strip():
         return "Please enter a question."
 
-    if state.get("index") is None or state.get("chunks") is None:
+    if state.get("index") is None:
         return "Please upload a PDF first."
 
     context = retrieve_context(question, state)
@@ -169,14 +156,13 @@ def answer_question(question, state):
     return answer
 
 
-# =========================
-# Gradio UI
-# =========================
+# ============================================================
+# GRADIO UI
+# ============================================================
 
 with gr.Blocks(title="Granite PDF QA") as demo:
-    gr.Markdown("# 📄🔍 PDF Q&A with IBM Granite\nUpload a PDF and ask questions about its content.")
+    gr.Markdown("## 📄🔍 PDF Q&A using IBM Granite\nUpload a PDF and ask questions about its content.")
 
-    # Per-session state: nothing leaks between users
     state = gr.State({"chunks": None, "index": None})
 
     with gr.Row():