Update chatbot_rag.py

chatbot_rag.py

@@ -10,34 +10,39 @@ from langchain_huggingface import HuggingFacePipeline, HuggingFaceEmbeddings
 from langchain_chroma import Chroma
 
 def build_qa():
-    """Builds and returns the RAG QA pipeline."""
-
-    # 1. Load embeddings + DB
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
     vectorstore = Chroma(
         persist_directory="db",
         collection_name="rag-docs",
         embedding_function=embeddings,
     )
 
-    # 2. LLM (instruction-tuned preferred)
+    # 🔹 Use Phi-3 Mini (smaller, faster)
     model_id = "microsoft/phi-3-mini-4k-instruct"
-    tokenizer = AutoTokenizer.from_pretrained(model_id)
-    model = AutoModelForCausalLM.from_pretrained(model_id)
-    pipe = pipeline("text-generation", model=model, tokenizer=tokenizer, max_new_tokens=512)
-    llm = HuggingFacePipeline(pipeline=pipe)
 
-    # 3. QA Chain
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
-    qa = RetrievalQA.from_chain_type(llm=llm, retriever=retriever, return_source_documents=False)
-
-    return qa
+    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
+    )
 
+    pipe = pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=256,            # ✅ smaller output (faster)
+        temperature=0.2,               # ✅ more focused answers
+    )
 
-# Build once (so Hugging Face loads at startup)
-qa_pipeline = build_qa()
+    llm = HuggingFacePipeline(pipeline=pipe)
 
+    retriever = vectorstore.as_retriever(search_kwargs={"k": 3})
+    qa = RetrievalQA.from_chain_type(
+        llm=llm,
+        retriever=retriever,
+        return_source_documents=False
+    )
 
-def get_answer(query: str) -> str:
-    """Takes user query and returns chatbot response."""
-    return qa_pipeline.run(query)
+    return qa