Update app.py

app.py

@@ -1,70 +1,79 @@
 import pandas as pd
 from datasets import load_dataset
-from langchain_huggingface import HuggingFaceEmbeddings
-from langchain_chroma import Chroma
+from langchain.embeddings import HuggingFaceEmbeddings
+from langchain.vectorstores import Chroma
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
 
-# 1. Load Dataset
-ds = load_dataset("maxpro291/bankfaqs_dataset")
-data = ds['train'][:]
-Bank_Data = pd.DataFrame({
-    'question': [entry for entry in data['text'] if entry.startswith("Q:")],
-    'answer': [entry for entry in data['text'] if entry.startswith("A:")]
-})
+# 1. Load and prepare dataset
+def load_bank_data():
+    ds = load_dataset("maxpro291/bankfaqs_dataset")
+    data = ds['train'][:]
+    return pd.DataFrame({
+        'question': [entry for entry in data['text'] if entry.startswith("Q:")],
+        'answer': [entry for entry in data['text'] if entry.startswith("A:")]
+    })
 
-# 2. Setup Vector Store
-embed_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-vectorstore = Chroma.from_texts(
-    texts=[f"Q: {q}\nA: {a}" for q, a in zip(Bank_Data['question'], Bank_Data['answer'])],
-    embedding=embed_model,
-    persist_directory="./chroma_db"
-)
-retriever = vectorstore.as_retriever()
+# 2. Setup vector store
+def setup_retriever(data):
+    embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+    vectorstore = Chroma.from_texts(
+        texts=[f"Q: {q}\nA: {a}" for q, a in zip(data['question'], data['answer'])],
+        embedding=embeddings,
+        persist_directory="./chroma_db_bank"
+    )
+    return vectorstore.as_retriever()
 
-# 3. Load Model (using smaller, faster model)
-model_name = "microsoft/phi-2"
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    device_map="auto",
-    torch_dtype="auto",
-    trust_remote_code=True
-)
-llm = pipeline(
-    "text-generation",
-    model=model,
-    tokenizer=tokenizer,
-    max_new_tokens=150,
-    temperature=0.7
-)
+# 3. Load LLM
+def load_llm():
+    model_name = "microsoft/phi-2"
+    tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+    model = AutoModelForCausalLM.from_pretrained(
+        model_name,
+        device_map="auto",
+        torch_dtype="auto",
+        trust_remote_code=True
+    )
+    return pipeline(
+        "text-generation",
+        model=model,
+        tokenizer=tokenizer,
+        max_new_tokens=150,
+        temperature=0.7
+    )
 
-# 4. Response Generation
-def generate_response(question):
+# 4. Response generation
+def generate_response(question, retriever, llm):
     docs = retriever.get_relevant_documents(question)
     context = "\n".join([doc.page_content for doc in docs][:2]) if docs else ""
     
-    prompt = f"""|USER| You're a banking expert. {
-        f"Relevant info: {context}" if context else ""
-    }\nQuestion: {question}\n|ASSISTANT| """
+    prompt = f"""Instruct: You're a banking expert. {
+        f"Context: {context}" if context else ""
+    }\nQuestion: {question}\nAnswer: """
     
     response = llm(prompt)[0]['generated_text']
-    return response.split("|ASSISTANT|")[-1].strip()
+    return response.split("Answer: ")[-1].strip()
 
-# 5. Gradio Interface with explicit API setup
-with gr.Blocks(title="Banking Assistant") as app:
-    gr.Markdown("## 🏦 Free Banking Q&A Assistant")
-    chatbot = gr.Chatbot(height=400)
-    msg = gr.Textbox(label="Your banking question")
-    clear = gr.Button("Clear")
-    
-    def respond(message, chat_history):
-        response = generate_response(message)
-        chat_history.append((message, response))
-        return "", chat_history
-    
-    msg.submit(respond, [msg, chatbot], [msg, chatbot])
-    clear.click(lambda: None, None, chatbot, queue=False)
+# 5. Initialize components
+bank_data = load_bank_data()
+retriever = setup_retriever(bank_data)
+llm = load_llm()
+
+# 6. Gradio interface
+def chat_interface(question, history):
+    response = generate_response(question, retriever, llm)
+    return response
+
+demo = gr.ChatInterface(
+    fn=chat_interface,
+    title="Banking Assistant 🏦",
+    examples=[
+        "How do I open a savings account?",
+        "What's the difference between debit and credit cards?",
+        "How do I apply for a loan?"
+    ],
+    theme="soft"
+)
 
 if __name__ == "__main__":
-    app.launch(server_name="0.0.0.0", server_port=7860, share=False)
+    demo.launch(server_name="0.0.0.0", server_port=7860)