Update src/qa.py

src/qa.py

@@ -1,11 +1,9 @@
 """
-qa.py — Fast Phi-2 Retrieval + Generation (Final Optimized Version)
+qa.py — Optimized Phi-2 Retrieval + Generation (Stable Fast Baseline)
 -------------------------------------------------------------------
-Uses:
-• intfloat/e5-small-v2 for embeddings
-• microsoft/phi-2 (quantized for CPU)
-• Reasoning toggle support (ON/OFF)
-Optimized for: speed + stability on Streamlit / Hugging Face Spaces
+✅ Best balance of speed + accuracy
+✅ Works perfectly on CPU (quantized)
+✅ Non-hallucinating (document-strict)
 """
 
 import os
@@ -13,9 +11,8 @@ import numpy as np
 import torch
 from sentence_transformers import SentenceTransformer
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-from sklearn.metrics.pairwise import cosine_similarity
 
-print("✅ qa.py (Final Fast Phi-2) loaded from:", __file__)
+print("✅ qa.py (FAST BASELINE) loaded from:", __file__)
 
 # ==========================================================
 # 1️⃣ Cache Setup
@@ -28,10 +25,9 @@ os.environ.update({
     "HF_DATASETS_CACHE": CACHE_DIR,
     "HF_MODULES_CACHE": CACHE_DIR
 })
-print(f"✅ Using cache dir: {CACHE_DIR}")
 
 # ==========================================================
-# 2️⃣ Embedding Model (fast + reliable)
+# 2️⃣ Embedding Model
 # ==========================================================
 try:
     _query_model = SentenceTransformer("intfloat/e5-small-v2", cache_folder=CACHE_DIR)
@@ -41,17 +37,18 @@ except Exception as e:
     _query_model = SentenceTransformer("sentence-transformers/all-MiniLM-L6-v2", cache_folder=CACHE_DIR)
 
 # ==========================================================
-# 3️⃣ Phi-2 LLM Setup (Quantized + CPU Optimized)
+# 3️⃣ Phi-2 LLM Setup (Quantized for CPU)
 # ==========================================================
 try:
     MODEL_NAME = "microsoft/phi-2"
     print(f"✅ Loading LLM: {MODEL_NAME} (quantized, CPU-optimized)")
 
     _tokenizer = AutoTokenizer.from_pretrained(MODEL_NAME, cache_dir=CACHE_DIR)
+
     _model = AutoModelForCausalLM.from_pretrained(
         MODEL_NAME,
         cache_dir=CACHE_DIR,
-        torch_dtype=torch.bfloat16 if not torch.cuda.is_available() else torch.float16,
+        torch_dtype=torch.float16 if torch.cuda.is_available() else torch.bfloat16,
         low_cpu_mem_usage=True,
     ).to("cpu")
 
@@ -63,34 +60,28 @@ try:
         model_kwargs={"torch_dtype": torch.bfloat16, "low_cpu_mem_usage": True},
     )
 
-    print("✅ Phi-2 pipeline ready (optimized).")
+    print("✅ Phi-2 text-generation pipeline ready (optimized).")
+
 except Exception as e:
     print(f"⚠️ Phi-2 load failed: {e}")
     _answer_model = None
 
 # ==========================================================
-# 4️⃣ Prompt Templates
+# 4️⃣ Prompt Template
 # ==========================================================
-STRICT_PROMPT = (
-    "You are an expert enterprise assistant.\n"
+PROMPT_TEMPLATE = (
+    "You are an expert assistant for enterprise document understanding.\n"
     "Use ONLY the context below to answer the question clearly and factually.\n"
-    "If the answer isn’t found in the context, reply exactly:\n"
+    "If the context doesn’t contain the answer, reply exactly:\n"
     "'I don't know based on the provided document.'\n\n"
     "Context:\n{context}\n\nQuestion: {query}\nAnswer:"
 )
 
-REASONING_PROMPT = (
-    "You are a reasoning-enabled enterprise assistant.\n"
-    "Use the CONTEXT below and your own reasoning ability to explain the answer clearly and logically.\n"
-    "If the answer isn’t explicit, infer based on context and domain understanding.\n\n"
-    "Context:\n{context}\n\nQuestion: {query}\nAnswer:"
-)
-
 # ==========================================================
-# 5️⃣ Retrieve Top-K Chunks (Balanced speed)
+# 5️⃣ Retrieve Top-K Chunks (Simple + Fast)
 # ==========================================================
-def retrieve_chunks(query: str, index, chunks: list, top_k: int = 3):
-    """Retrieve top-K relevant chunks and re-rank by cosine similarity for better precision."""
+def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5):
+    """Efficient FAISS retrieval using cosine similarity."""
     if not index or not chunks:
         return []
 
@@ -98,46 +89,40 @@ def retrieve_chunks(query: str, index, chunks: list, top_k: int = 3):
         q_emb = _query_model.encode([f"query: {query.strip()}"], convert_to_numpy=True, normalize_embeddings=True)[0]
         distances, indices = index.search(np.array([q_emb]).astype("float32"), top_k * 2)
 
-        # Compute similarity scores for re-ranking
-        candidates = [chunks[i] for i in indices[0]]
-        cand_vecs = _query_model.encode(candidates, convert_to_numpy=True, normalize_embeddings=True)
-        sims = cosine_similarity([q_emb], cand_vecs)[0]
+        selected = set()
+        for idx in indices[0]:
+            for i in range(max(0, idx - 1), min(len(chunks), idx + 2)):
+                selected.add(i)
 
-        # Return top-K most semantically aligned
-        top_indices = np.argsort(sims)[::-1][:top_k]
-        return [candidates[i] for i in top_indices]
+        ordered_chunks = [chunks[i] for i in sorted(selected)]
+        return ordered_chunks
 
     except Exception as e:
         print(f"⚠️ Retrieval error: {e}")
         return []
 
-
 # ==========================================================
-# 6️⃣ Generate Answer (Reasoning or Strict Mode)
+# 6️⃣ Generate Answer (Fast)
 # ==========================================================
-def generate_answer(query: str, retrieved_chunks: list, reasoning_mode: bool = True):
-    """Generate concise, context-grounded answers using Phi-2."""
+def generate_answer(query: str, retrieved_chunks: list):
+    """Generate concise, grounded answers using Phi-2."""
     if not retrieved_chunks:
         return "Sorry, I couldn’t find relevant information in the document."
 
-    # ✅ Keep short context for faster inference
-    context = "\n".join(chunk.strip() for chunk in retrieved_chunks[:5])
-    prompt_template = REASONING_PROMPT if reasoning_mode else STRICT_PROMPT
-    prompt = prompt_template.format(context=context, query=query)
+    context = "\n".join(chunk.strip() for chunk in retrieved_chunks)
+    prompt = PROMPT_TEMPLATE.format(context=context, query=query)
 
     try:
         result = _answer_model(
             prompt,
-            max_new_tokens=140,   # fast but coherent answers
+            max_new_tokens=120,  # lower for faster completion
             do_sample=False,
             early_stopping=True,
             pad_token_id=_tokenizer.eos_token_id,
         )
-
         answer = result[0]["generated_text"].strip()
         if "Answer:" in answer:
             answer = answer.split("Answer:")[-1].strip()
-
         return answer
 
     except Exception as e:
@@ -145,11 +130,10 @@ def generate_answer(query: str, retrieved_chunks: list, reasoning_mode: bool = T
         return "⚠️ Error: Could not generate an answer at the moment."
 
 # ==========================================================
-# 7️⃣ Local Test (Optional)
+# 7️⃣ Local Test
 # ==========================================================
 if __name__ == "__main__":
     from vectorstore import build_faiss_index
-
     dummy_chunks = [
         "Step 1: Open the dashboard and navigate to reports.",
         "Step 2: Click 'Export' to download a CSV summary.",
@@ -164,4 +148,4 @@ if __name__ == "__main__":
     query = "What are the steps to export a report?"
     retrieved = retrieve_chunks(query, index, dummy_chunks)
     print("🔍 Retrieved:", retrieved)
-    print("💬 Answer:", generate_answer(query, retrieved, reasoning_mode=True))
+    print("💬 Answer:", generate_answer(query, retrieved))