Update app.py

app.py

@@ -14,11 +14,11 @@ from sentence_transformers import SentenceTransformer
 from transformers import pipeline
 from sklearn.metrics.pairwise import cosine_similarity
 
-# Load models once
+# Load models
 embed_model = SentenceTransformer('all-MiniLM-L6-v2')
 qa_pipeline = pipeline("question-answering", model="deepset/tinyroberta-squad2")
 
-# Globals to hold data
+# Globals
 all_chunks = []
 chunk_sources = []
 chunk_embeddings = None
@@ -82,11 +82,29 @@ def answer_question(question):
 
     q_emb = embed_model.encode([question], convert_to_numpy=True)
     sims = cosine_similarity(q_emb, chunk_embeddings)[0]
-    top_k_idx = sims.argsort()[::-1][:3]
 
-    context = "\n\n".join([all_chunks[i] for i in top_k_idx])
+    threshold = 0.5  # similarity threshold to filter relevant chunks
+    above_thresh_idx = [i for i, sim in enumerate(sims) if sim > threshold]
+
+    if not above_thresh_idx:
+        return "No relevant content found in the PDFs for your question."
+
+    # Sort by similarity descending
+    above_thresh_idx.sort(key=lambda i: sims[i], reverse=True)
+
+    max_context_chars = 2000
+    context_chunks = []
+    total_chars = 0
+    for i in above_thresh_idx:
+        chunk_len = len(all_chunks[i])
+        if total_chars + chunk_len > max_context_chars:
+            break
+        context_chunks.append(all_chunks[i])
+        total_chars += chunk_len
+
+    context = "\n\n".join(context_chunks)
     if not context.strip():
-        return "No relevant content found in PDFs."
+        return "No sufficient content to answer the question."
 
     try:
         result = qa_pipeline(question=question, context=context)
@@ -94,8 +112,8 @@ def answer_question(question):
     except Exception:
         return "Error generating answer from the model."
 
-    sources = "\n".join(set(chunk_sources[i] for i in top_k_idx))
-    confidence = np.mean([sims[i] for i in top_k_idx]) * 100
+    sources = "\n".join(set(chunk_sources[i] for i in above_thresh_idx[:len(context_chunks)]))
+    confidence = np.mean([sims[i] for i in above_thresh_idx[:len(context_chunks)]]) * 100
     return f"**Answer:** {answer}\n\n**Sources:**\n{sources}\n\n**Confidence:** {confidence:.2f}%"
 
 with gr.Blocks() as demo: