Update chatbot_rag.py

chatbot_rag.py

@@ -1,10 +1,9 @@
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFacePipeline
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
 from langchain.chains import RetrievalQA
-import traceback  # ✅ added
-
+import traceback
 
 def build_qa():
     """Builds and returns the RAG QA pipeline."""
@@ -23,46 +22,40 @@ def build_qa():
     )
     print("📂 Docs in DB:", vectorstore._collection.count())
 
-    # 3. LLM
+    # 3. Load LLM (Flan-T5 small for lightweight QA)
     print("🔹 Loading LLM...")
-    model_id = "sshleifer/tiny-gpt2"
-    tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
-    model = AutoModelForCausalLM.from_pretrained(
-        model_id,
-        device_map="auto",
-        torch_dtype="auto"
-    )
-    print("✅ LLM loaded.")
+    model_id = "google/flan-t5-small"
+    tokenizer = AutoTokenizer.from_pretrained(model_id)
+    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
 
     pipe = pipeline(
-        "question-answering",
+        "text2text-generation",
         model=model,
         tokenizer=tokenizer,
         max_new_tokens=256,
-        temperature=0.2,
     )
     llm = HuggingFacePipeline(pipeline=pipe)
 
-    # 4. QA Chain
+    # 4. QA Chain with retrieval
     print("🔹 Building RetrievalQA...")
     retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
     qa = RetrievalQA.from_chain_type(
         llm=llm,
         retriever=retriever,
-        return_source_documents=False
+        return_source_documents=False,
+        chain_type="stuff"  # simplest chain, passes context + question
     )
 
     print("✅ QA pipeline ready.")
     return qa
 
-
-# Build at import time (so it's ready when app runs)
+# Build once
 try:
     qa_pipeline = build_qa()
 except Exception as e:
     qa_pipeline = None
     print("❌ Failed to build QA pipeline:", e)
-    traceback.print_exc()   # ✅ added: full error details
+    traceback.print_exc()
 
 
 def get_answer(query: str) -> str: