Update app.py

app.py

@@ -1,66 +1,68 @@
 import gradio as gr
-from langchain.vectorstores import FAISS
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.chains import RetrievalQA
-from langchain.llms import HuggingFacePipeline
 from transformers import pipeline
 
-# ---------------- LOAD EMBEDDINGS ----------------
+from langchain.chains import RetrievalQA
+from langchain_community.vectorstores import FAISS
+from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.llms import HuggingFacePipeline
+
+# ------------------ LOAD EMBEDDINGS ------------------
 embeddings = HuggingFaceEmbeddings(
     model_name="sentence-transformers/all-MiniLM-L6-v2"
 )
 
-# ---------------- LOAD VECTOR STORE ----------------
+# ------------------ LOAD VECTOR STORE ------------------
 db = FAISS.load_local(
     "vectorstore/faiss_index",
     embeddings,
     allow_dangerous_deserialization=True
 )
 
-# ---------------- LOAD LLM ----------------
-generator = pipeline(
+# ------------------ LOAD LLM ------------------
+# NOTE: Use a LIGHT model for HF CPU
+text_gen_pipeline = pipeline(
     "text-generation",
-    model="mistralai/Mistral-7B-Instruct-v0.2",
+    model="microsoft/phi-2",
     max_new_tokens=512,
     temperature=0.2,
 )
 
-llm = HuggingFacePipeline(pipeline=generator)
+llm = HuggingFacePipeline(pipeline=text_gen_pipeline)
 
-# ---------------- RAG CHAIN ----------------
-qa = RetrievalQA.from_chain_type(
+# ------------------ RAG CHAIN ------------------
+qa_chain = RetrievalQA.from_chain_type(
     llm=llm,
     retriever=db.as_retriever(search_kwargs={"k": 3}),
     chain_type="stuff",
 )
 
-# ---------------- CHAT FUNCTION ----------------
-def chat(query, history):
-    if not query.strip():
+# ------------------ CHAT FUNCTION ------------------
+def chat(user_message, history):
+    if not user_message.strip():
         return history
 
-    answer = qa.run(query)
-
-    history.append((query, answer))
+    answer = qa_chain.run(user_message)
+    history.append((user_message, answer))
     return history
 
-# ---------------- GRADIO UI ----------------
-with gr.Blocks(title="RAG Document Chatbot") as demo:
+# ------------------ GRADIO UI ------------------
+with gr.Blocks(title="Document RAG Chatbot") as demo:
     gr.Markdown(
         """
-        # 📚 RAG Document Chatbot
-        Answers are **strictly based on the provided documents**.
+        # 📚 Document RAG Chatbot  
+        Answers are generated **strictly from the provided documents** using Retrieval-Augmented Generation.
         """
     )
 
-    chatbot = gr.Chatbot(height=400)
-    query_box = gr.Textbox(
-        placeholder="Ask a question from the documents...",
-        label="Your Question",
+    chatbot = gr.Chatbot(height=420)
+    query = gr.Textbox(
+        label="Ask a question",
+        placeholder="Ask something from the documents..."
     )
-    clear = gr.Button("Clear Chat")
 
-    query_box.submit(chat, [query_box, chatbot], chatbot)
-    clear.click(lambda: [], None, chatbot)
+    clear_btn = gr.Button("Clear Chat")
+
+    query.submit(chat, [query, chatbot], chatbot)
+    clear_btn.click(lambda: [], None, chatbot)
 
 demo.launch()