Update app.py

app.py

@@ -1,79 +1,138 @@
+import os
 import pandas as pd
+import logging
 from datasets import load_dataset
-from langchain.embeddings import HuggingFaceEmbeddings
-from langchain.vectorstores import Chroma
+from langchain_huggingface import HuggingFaceEmbeddings, HuggingFacePipeline
+from langchain_chroma import Chroma
+from langchain_core.prompts import PromptTemplate
+from langchain_core.output_parsers import StrOutputParser
+from langchain_core.runnables import RunnablePassthrough
 import gradio as gr
-from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline
-
-# 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
-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 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, retriever, llm):
-    docs = retriever.get_relevant_documents(question)
-    context = "\n".join([doc.page_content for doc in docs][:2]) if docs else ""
-    
-    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("Answer: ")[-1].strip()
-
-# 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
+from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
+
+# Set up logging
+logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
+
+# ------------------------------------------------------------------
+# 1. Load and Prepare the Bank FAQ Dataset (UNCHANGED)
+# ------------------------------------------------------------------
+ds = load_dataset("maxpro291/bankfaqs_dataset")
+train_ds = ds['train']
+data = train_ds[:]  # load all examples
+
+questions = []
+answers = []
+for entry in data['text']:
+    if entry.startswith("Q:"):
+        questions.append(entry)
+    elif entry.startswith("A:"):
+        answers.append(entry)
+
+Bank_Data = pd.DataFrame({'question': questions, 'answer': answers})
+
+context_data = []
+for i in range(len(Bank_Data)):
+    context = f"Question: {Bank_Data.iloc[i]['question']} Answer: {Bank_Data.iloc[i]['answer']}"
+    context_data.append(context)
+
+# ------------------------------------------------------------------
+# 2. Create the Vector Store for Retrieval (UNCHANGED)
+# ------------------------------------------------------------------
+embed_model = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
+
+vectorstore = Chroma.from_texts(
+    texts=context_data,
+    embedding=embed_model,
+    persist_directory="./chroma_db_bank"
+)
+retriever = vectorstore.as_retriever()
+
+# ------------------------------------------------------------------
+# 3. Initialize PHI-2 Model (MODIFIED SECTION)
+# ------------------------------------------------------------------
+model_name = "microsoft/phi-2"
+
+# Configure 4-bit quantization for efficient loading
+quantization_config = BitsAndBytesConfig(
+    load_in_4bit=True,
+    bnb_4bit_compute_dtype="float16",
+    bnb_4bit_quant_type="nf4"
+)
+
+tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
+model = AutoModelForCausalLM.from_pretrained(
+    model_name,
+    device_map="auto",
+    trust_remote_code=True,
+    quantization_config=quantization_config
+)
+
+# Create text-generation pipeline with Phi-2 specific settings
+pipe = pipeline(
+    "text-generation",
+    model=model,
+    tokenizer=tokenizer,
+    max_new_tokens=512,
+    temperature=0.7,
+    top_p=0.95,
+    repetition_penalty=1.15,
+    do_sample=True
+)
+
+# Wrap the pipeline in LangChain's HuggingFacePipeline
+huggingface_model = HuggingFacePipeline(pipeline=pipe)
+
+# ------------------------------------------------------------------
+# 4. Build the RAG Chain (UNCHANGED)
+# ------------------------------------------------------------------
+template = (
+    "You are a helpful banking assistant. "
+    "Use the provided context if it is relevant to answer the question. "
+    "If not, answer using your general banking knowledge.\n"
+    "Question: {question}\n"
+    "Answer:"
+)
+rag_prompt = PromptTemplate.from_template(template)
+
+rag_chain = (
+    {"context": retriever, "question": RunnablePassthrough()}
+    | rag_prompt
+    | huggingface_model
+    | StrOutputParser()
+)
+
+# ------------------------------------------------------------------
+# 5. Gradio Chat Interface (UNCHANGED)
+# ------------------------------------------------------------------
+def rag_memory_stream(message, history):
+    partial_text = ""
+    for new_text in rag_chain.stream(message):
+        partial_text += new_text
+        yield partial_text
+
+examples = [
+    "I want to open an account", 
+    "What is a savings account?",
+    "How do I use an ATM?",
+    "How can I resolve a bank account issue?"
+]
+
+title = "Your Personal Banking Assistant 💬"
+description = (
+    "Welcome! I'm here to answer your questions about banking and related topics. "
+    "Ask me anything, and I'll do my best to assist you."
+)
 
 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"
+    fn=rag_memory_stream,
+    title=title,
+    description=description,
+    examples=examples,
+    theme="glass",
 )
 
+# ------------------------------------------------------------------
+# 6. Launch the App (UNCHANGED)
+# ------------------------------------------------------------------
 if __name__ == "__main__":
-    demo.launch(server_name="0.0.0.0", server_port=7860)
+    demo.launch(share=True)