Update src/qa.py

src/qa.py

@@ -70,18 +70,14 @@ except Exception as e:
 # 4️⃣ Prompt Templates
 # ==========================================================
 STRICT_PROMPT = (
-    "You are an enterprise knowledge assistant.\n"
-    "Use ONLY the following context to answer the question clearly and factually.\n"
-    "If the context doesn’t contain the answer, reply exactly:\n"
-    "'I don't know based on the provided document.'\n\n"
+    "Answer based ONLY on the context below.\n"
+    "If the answer isn’t in the context, say: 'I don't know based on the provided document.'\n\n"
     "Context:\n{context}\n\nQuestion: {query}\nAnswer:"
 )
 
 REASONING_PROMPT = (
-    "You are a reasoning assistant with access to enterprise knowledge.\n"
-    "Use the context below to answer the question with logic, insight, and helpful reasoning.\n"
-    "If the context doesn’t provide enough detail, you may infer sensibly based on common knowledge, "
-    "but make it clear when you are reasoning beyond the text.\n\n"
+    "You are an expert assistant. Use the context and your reasoning ability to form a clear, step-by-step answer.\n"
+    "Be concise yet complete. If the context doesn’t contain the answer, say: 'I don't know based on the provided document.'\n\n"
     "Context:\n{context}\n\nQuestion: {query}\nAnswer:"
 )
 
@@ -116,47 +112,46 @@ def retrieve_chunks(query: str, index, chunks: list, top_k: int = 5):
 # ==========================================================
 def generate_answer(query: str, retrieved_chunks: list, reasoning_mode: bool = False):
     """
-    Generates answers with or without reasoning.
-    - reasoning_mode=True → uses world knowledge + logic (slower, more explanatory)
-    - reasoning_mode=False → sticks to chunks for fast factual accuracy
+    Generate answers with or without reasoning.
+    reasoning_mode=True → reasoning-rich, slightly slower
+    reasoning_mode=False → strict factual, fast
     """
     if not retrieved_chunks:
         return "Sorry, I couldn’t find relevant information in the document."
 
-    # Build the prompt with selected mode
     context = "\n".join(chunk.strip() for chunk in retrieved_chunks)
-    prompt = (REASONING_PROMPT if reasoning_mode else STRICT_PROMPT).format(
-        context=context,
-        query=query
-    )
+    prompt = (REASONING_PROMPT if reasoning_mode else STRICT_PROMPT).format(context=context, query=query)
 
     try:
-        # ⚙️ Tuned generation settings for balance
-        result = _answer_model(
-            prompt,
-            max_new_tokens=180 if reasoning_mode else 120,   # let reasoning finish sentences
-            temperature=0.6 if reasoning_mode else 0.2,      # more creative but controlled
-            top_p=0.9 if reasoning_mode else 0.8,            # smooth probability cutoff
-            do_sample=reasoning_mode,                        # only sample when reasoning
-            early_stopping=True,
-            pad_token_id=_tokenizer.eos_token_id,
-        )
-
-        # Clean answer text
-        answer = result[0]["generated_text"].strip()
-        if "Answer:" in answer:
-            answer = answer.split("Answer:")[-1].strip()
-
-        # ✅ Prevents mid-sentence cutoffs
-        if answer.endswith(("and", "or", ",")):
-            answer += " ..."
-
-        return answer
+        if reasoning_mode:
+            # 🧠 “Brainy mode” — restores the earlier expressive logic
+            result = _answer_model(
+                prompt,
+                max_new_tokens=180,
+                temperature=0.4,
+                do_sample=False,  # deterministic reasoning chain
+                pad_token_id=_tokenizer.eos_token_id,
+            )
+        else:
+            # ⚡ Strict factual fast mode
+            result = _answer_model(
+                prompt,
+                max_new_tokens=120,
+                temperature=0.2,
+                do_sample=False,
+                pad_token_id=_tokenizer.eos_token_id,
+            )
+
+        raw = result[0]["generated_text"]
+        if "Answer:" in raw:
+            raw = raw.split("Answer:")[-1].strip()
+        return raw.strip()
 
     except Exception as e:
         print(f"⚠️ Generation failed: {e}")
         return "⚠️ Error: Could not generate an answer."
 
+
 # ==========================================================
 # 7️⃣ Local Test
 # ==========================================================