Update app.py

app.py

@@ -2,11 +2,11 @@ import os
 import pandas as pd
 import logging
 from datasets import load_dataset
-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
+from langchain.embeddings import HuggingFaceEmbeddings  # Updated import
+from langchain.vectorstores import Chroma  # Updated import
+from langchain.prompts import PromptTemplate  # Updated import
+from langchain.schema.output_parser import StrOutputParser  # Updated import
+from langchain.schema.runnable import RunnablePassthrough  # Updated import
 import gradio as gr
 from transformers import AutoTokenizer, AutoModelForCausalLM, pipeline, BitsAndBytesConfig
 from fastapi import FastAPI, Header, HTTPException
@@ -14,15 +14,15 @@ import threading
 import uvicorn
 
 # ====================== CONFIGURATION ======================
-API_KEY = "Samson"  # Your hardcoded API key
-MODEL_NAME = "microsoft/phi-2"  # Using Phi-2 model
+API_KEY = "Samson"
+MODEL_NAME = "microsoft/phi-2" 
 # ===========================================================
 
 # Set up logging
 logging.basicConfig(level=logging.INFO, format='%(asctime)s - %(levelname)s - %(message)s')
 
 # ---------------------- RAG Setup --------------------------
-# 1. Load and prepare dataset
+# 1. Load dataset
 ds = load_dataset("maxpro291/bankfaqs_dataset")
 data = ds['train'][:]
 Bank_Data = pd.DataFrame({
@@ -39,7 +39,7 @@ vectorstore = Chroma.from_texts(
 )
 retriever = vectorstore.as_retriever()
 
-# 3. Initialize LLM with 4-bit quantization
+# 3. Initialize LLM
 quant_config = BitsAndBytesConfig(
     load_in_4bit=True,
     bnb_4bit_compute_dtype="float16",
@@ -52,27 +52,28 @@ model = AutoModelForCausalLM.from_pretrained(
     trust_remote_code=True,
     quantization_config=quant_config
 )
-pipe = pipeline(
+
+# Create LangChain pipeline
+llm_pipeline = pipeline(
     "text-generation",
     model=model,
     tokenizer=tokenizer,
     max_new_tokens=512,
     temperature=0.7,
-    top_p=0.95,
-    repetition_penalty=1.15
+    top_p=0.95
 )
-huggingface_model = HuggingFacePipeline(pipeline=pipe)
 
 # 4. Build RAG chain
 template = """You are a banking assistant. Use context if relevant:
 Context: {context}
 Question: {question}
 Answer:"""
-rag_prompt = PromptTemplate.from_template(template)
+prompt = PromptTemplate.from_template(template)
+
 rag_chain = (
     {"context": retriever, "question": RunnablePassthrough()}
-    | rag_prompt
-    | huggingface_model
+    | prompt
+    | llm_pipeline
     | StrOutputParser()
 )
 
@@ -85,47 +86,33 @@ def validate_api_key(api_key: str = Header(None)):
     return True
 
 @app.post("/chat")
-async def chat_endpoint(
-    question: str,
-    authorization: str = Header(None),
-):
+async def chat_endpoint(question: str, authorization: str = Header(None)):
     validate_api_key(authorization)
     response = ""
     for chunk in rag_chain.stream(question):
         response += chunk
     return {"response": response}
 
-@app.get("/health")
-async def health_check():
-    return {"status": "healthy"}
-
 # -------------------- Gradio Interface ---------------------
-def rag_memory_stream(message, history):
-    partial_text = ""
-    for new_text in rag_chain.stream(message):
-        partial_text += new_text
-        yield partial_text
+def respond(message, history):
+    return next(rag_chain.stream(message))
 
 demo = gr.ChatInterface(
-    fn=rag_memory_stream,
-    title="Banking Assistant 🔒 (API Key: Samson)",
-    description="Welcome! Use API key 'Samson' to access the /chat endpoint",
+    fn=respond,
+    title="Banking Assistant 🔒",
     examples=[
         "How do I open an account?",
-        "What's the interest rate for savings?",
+        "What's the interest rate?",
         "How do I apply for a loan?"
     ],
     theme="glass"
 )
 
 # --------------------- Launch Servers ----------------------
-def run_gradio():
-    demo.launch(server_name="0.0.0.0", server_port=7860)
-
 if __name__ == "__main__":
-    # Start Gradio in separate thread
-    gradio_thread = threading.Thread(target=run_gradio)
-    gradio_thread.start()
+    threading.Thread(
+        target=demo.launch,
+        kwargs={"server_name": "0.0.0.0", "server_port": 7860}
+    ).start()
     
-    # Start FastAPI
     uvicorn.run(app, host="0.0.0.0", port=8000)