Update src/qa.py

src/qa.py

@@ -1,8 +1,13 @@
-# ----------------------------
-# Hugging Face cache bootstrap
-# ----------------------------
 import os
+from sentence_transformers import SentenceTransformer
+from transformers import pipeline
+from vectorstore import search_faiss
 
+print("✅ qa.py loaded from:", __file__)
+
+# ----------------------------
+# Hugging Face cache setup
+# ----------------------------
 CACHE_DIR = "/tmp/hf_cache"
 os.makedirs(CACHE_DIR, exist_ok=True)
 
@@ -11,15 +16,6 @@ os.environ["TRANSFORMERS_CACHE"] = CACHE_DIR
 os.environ["HF_DATASETS_CACHE"] = CACHE_DIR
 os.environ["HF_MODULES_CACHE"] = CACHE_DIR
 
-print(f"✅ Using Hugging Face cache at {CACHE_DIR}")
-
-# ----------------------------
-# Imports AFTER cache bootstrap
-# ----------------------------
-from sentence_transformers import SentenceTransformer
-from transformers import pipeline
-from vectorstore import search_faiss
-
 # ----------------------------
 # Query embedding model
 # ----------------------------
@@ -32,7 +28,6 @@ _query_model = SentenceTransformer(
 # LLM for answers
 # ----------------------------
 MODEL_NAME = "google/flan-t5-small"
-
 _answer_model = pipeline(
     "text2text-generation",
     model=MODEL_NAME,
@@ -43,17 +38,27 @@ _answer_model = pipeline(
 # Functions
 # ----------------------------
 def retrieve_chunks(query, index, chunks, top_k=3):
+    """Embed the query and retrieve top-k chunks from FAISS."""
     q_emb = _query_model.encode([query], convert_to_numpy=True)[0]
     return search_faiss(q_emb, index, chunks, top_k)
 
 def generate_answer(query, retrieved_chunks):
+    """Generate an answer using retrieved chunks as context."""
     if not retrieved_chunks:
         return "Sorry, I could not find relevant information."
 
     context = " ".join(retrieved_chunks)
     prompt = (
-        "You are an assistant. Use the context to answer the question clearly.\n"
-        f"Context:\n{context}\n\nQuestion:\n{query}\n\nAnswer:"
+        "You are an assistant. Use the context below to answer the question clearly.\n\n"
+        f"Context:\n{context}\n\n"
+        f"Question:\n{query}\n\n"
+        "Answer:"
+    )
+
+    # ✅ Use max_new_tokens instead of max_length to avoid version mismatch errors
+    result = _answer_model(
+        prompt,
+        max_new_tokens=300,
+        do_sample=False
     )
-    result = _answer_model(prompt, max_length=300, do_sample=False)
     return result[0]["generated_text"].strip()