Update chatbot_rag.py

chatbot_rag.py

@@ -1,43 +1,44 @@
 
-
-# rag_pipeline.py
 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 langchain.llms import HuggingFacePipeline
 from langchain.chains import RetrievalQA
-from langchain_huggingface import HuggingFacePipeline, HuggingFaceEmbeddings
-from langchain_chroma import Chroma
+
 
 def build_qa():
+    """Builds and returns the RAG QA pipeline."""
+
+    # 1. Embeddings
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
 
+    # 2. Load vector DB (must already exist in ./db)
     vectorstore = Chroma(
         persist_directory="db",
         collection_name="rag-docs",
         embedding_function=embeddings,
     )
 
-    # 🔹 Use Phi-3 Mini (smaller, faster)
+    # 3. LLM (lighter model = faster inference)
     model_id = "microsoft/phi-3-mini-4k-instruct"
-
     tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
     model = AutoModelForCausalLM.from_pretrained(
         model_id,
-        device_map="auto",             # ✅ auto place on GPU if available
-        torch_dtype="auto"             # ✅ better memory handling
+        device_map="auto",
+        torch_dtype="auto"
     )
 
     pipe = pipeline(
         "text-generation",
         model=model,
         tokenizer=tokenizer,
-        max_new_tokens=256,            # ✅ smaller output (faster)
-        temperature=0.2,               # ✅ more focused answers
+        max_new_tokens=256,
+        temperature=0.2,
     )
 
     llm = HuggingFacePipeline(pipeline=pipe)
 
+    # 4. RAG chain
     retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
     qa = RetrievalQA.from_chain_type(
         llm=llm,
@@ -46,3 +47,18 @@ def build_qa():
     )
 
     return qa
+
+
+# Build at import time (so it's ready when app runs)
+try:
+    qa_pipeline = build_qa()
+except Exception as e:
+    qa_pipeline = None
+    print("❌ Failed to build QA pipeline:", e)
+
+
+def get_answer(query: str) -> str:
+    """Takes user query and returns chatbot response."""
+    if qa_pipeline is None:
+        return "⚠️ QA pipeline not initialized."
+    return qa_pipeline.run(query)