Upload chatbot_rag.py

chatbot_rag.py

@@ -0,0 +1,47 @@
+# -*- coding: utf-8 -*-
+"""chatbot_rag.ipynb
+
+Automatically generated by Colab.
+
+Original file is located at
+    https://colab.research.google.com/drive/1Qb6RDSuj0-E-Jy6a7bze6Plz9l9e5D_q
+"""
+
+# rag_pipeline.py
+from langchain.vectorstores import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
+from langchain.llms import HuggingFacePipeline
+from langchain.chains import RetrievalQA
+
+
+def build_qa():
+    """Builds and returns the RAG QA pipeline."""
+
+    # 1. Embeddings
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+    # 2. Vector DB (already persisted in Kaggle, same folder name in HF Space)
+    vectorstore = Chroma(persist_directory="db", embedding_function=embeddings)
+
+    # 3. LLM
+    model_id = "microsoft/phi-2"
+    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)
+
+    # 4. 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
+
+
+# Build once (so Hugging Face loads at startup)
+qa_pipeline = build_qa()
+
+
+def get_answer(query: str) -> str:
+    """Takes user query and returns chatbot response."""
+    return qa_pipeline.run(query)