Update chatbot_rag.py

chatbot_rag.py

@@ -1,18 +1,22 @@
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
 from langchain_community.llms import HuggingFacePipeline
-from transformers import AutoTokenizer, AutoModelForSeq2SeqLM, pipeline
-from langchain.chains import RetrievalQA
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 from langchain.prompts import PromptTemplate
+from langchain_core.runnables import RunnablePassthrough
+from langchain_core.output_parsers import StrOutputParser
 import traceback
 
+
 def build_qa():
-    """Builds and returns the RAG QA pipeline."""
+    """Builds and returns the RAG QA pipeline (rag_chain style)."""
     print("🚀 Starting QA pipeline...")
 
     # 1. Embeddings
     print("🔹 Loading embeddings...")
-    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    embeddings = HuggingFaceEmbeddings(
+        model_name="sentence-transformers/all-MiniLM-L6-v2"
+    )
 
     # 2. Load vector DB
     print("🔹 Loading Chroma DB...")
@@ -23,50 +27,66 @@ def build_qa():
     )
     print("📂 Docs in DB:", vectorstore._collection.count())
 
-    # 3. Load LLM (Flan-T5 small for lightweight QA)
+    # 3. Load LLM (Phi-3 mini)
     print("🔹 Loading LLM...")
     model_id = "microsoft/Phi-3-mini-4k-instruct"
     tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_id)
+    model = AutoModelForCausalLM.from_pretrained(model_id, device_map="auto")
 
     pipe = pipeline(
-    "text-generation", 
-    model=model, 
-    tokenizer=tokenizer, 
-    max_new_tokens=300,
-    do_sample=True, # Set to True to enable sampling and use temperature
-    temperature=0.2 # This is the temperature parameter
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=300,
+        do_sample=True,
+        temperature=0.2,
     )
     llm = HuggingFacePipeline(pipeline=pipe)
 
-    # 4. QA Chain with retrieval
-    print("🔹 Building RetrievalQA...")
-    retriever = vectorstore.as_retriever()
-    
-
-    template = """
-    Use the following context to answer the question at the end. 
-    If you don't know the answer, just say "I don't know" — do not make up an answer.
-    
-    Context:
-    {context}
-    
-    Question: {question}
-    Answer (one short sentence):
-    """
-    qa_prompt = PromptTemplate(template=template, input_variables=["context", "question"])
-    
-    qa = RetrievalQA.from_chain_type(
-        llm=llm,
-        retriever=retriever,
-        chain_type="stuff",
-        chain_type_kwargs={"prompt": qa_prompt},
-        return_source_documents=False,
+    # 4. Retriever
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+
+    # 5. Prompt
+    prompt = PromptTemplate(
+        input_variables=["context", "question"],
+        template="""
+        Use the following context to answer the question at the end. 
+        If you don't know the answer, just say "I don't know" — do not make up an answer.
+        
+        Context:
+        {context}
+        
+        Question: {question}
+        Answer (one short sentence):
+        """,
     )
 
+    # 6. Helper functions
+    def format_docs(docs):
+        return "\n".join(doc.page_content for doc in docs)
+
+    def hf_to_str(x):
+        """Convert Hugging Face pipeline output to plain string"""
+        if isinstance(x, list) and "generated_text" in x[0]:
+            return x[0]["generated_text"]
+        return str(x)
+
+    # 7. RAG chain
+    rag_chain = (
+        {
+            "context": retriever | format_docs,
+            "question": RunnablePassthrough(),
+        }
+        | prompt
+        | (lambda x: str(x))       # convert PromptTemplate value to str
+        | llm
+        | (lambda x: hf_to_str(x)) # clean HF output
+        | StrOutputParser()
+    )
 
     print("✅ QA pipeline ready.")
-    return qa
+    return rag_chain
+
 
 # Build once
 try:
@@ -81,17 +101,8 @@ def get_answer(query: str) -> str:
     """Takes user query and returns chatbot response."""
     if qa_pipeline is None:
         return "⚠️ QA pipeline not initialized."
-    
-    try:
-        retriever = qa_pipeline.retriever
-        docs = retriever.get_relevant_documents(query)
-        print("📂 Retrieved docs:", len(docs))
-        if not docs:
-            return "⚠️ No documents found in the DB. Check your `db/` folder."
-    except Exception as e:
-        return f"❌ Retriever error: {e}"
 
     try:
-        return qa_pipeline.run(query)
+        return qa_pipeline.invoke(query)
     except Exception as e:
         return f"❌ QA run failed: {e}"