Update smartbot.py

smartbot.py

@@ -1,41 +1,36 @@
 import streamlit as st
-import faiss
-import numpy as np
 import pandas as pd
-import torch
-import os
 from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
 from transformers import AutoTokenizer, AutoModelForCausalLM
-from huggingface_hub import login
-
-# --- HuggingFace login ---
-HF_TOKEN = os.getenv("HF_TOKEN")
-if HF_TOKEN:
-    login(token=HF_TOKEN)
 
 # --- Load data ---
-train = pd.read_csv('train_data.csv')  # Assumes it's in the same directory
-questions = train['question'].tolist()
-answers = train['answer'].tolist()
+@st.cache_data
+def load_data():
+    df = pd.read_csv("train_data.csv")
+    return df
 
-qa_pairs = [f"Q: {q} A: {a}" for q, a in zip(questions, answers)]
+df = load_data()
 
 # --- Embedding model ---
 embedding_model = SentenceTransformer('all-MiniLM-L6-v2')
-answer_embeddings = embedding_model.encode(answers)
 
 # --- FAISS index ---
-index = faiss.IndexFlatL2(answer_embeddings.shape[1])
-index.add(np.array(answer_embeddings))
-
-# --- LLaMA model setup ---
-model_name = "meta-llama/Llama-3-8B-Instruct"  # Update to a valid space-available model if needed
-tokenizer = AutoTokenizer.from_pretrained(model_name, trust_remote_code=True)
-model = AutoModelForCausalLM.from_pretrained(
-    model_name,
-    torch_dtype=torch.float16 if torch.cuda.is_available() else torch.float32,
-    device_map="auto"
-)
+@st.cache_resource
+def create_index(df):
+    embeddings = embedding_model.encode(df['question'].tolist(), show_progress_bar=True)
+    index = faiss.IndexFlatL2(embeddings.shape[1])
+    index.add(embeddings)
+    return index, embeddings
+
+index, embeddings = create_index(df)
+
+# --- Replace LLaMA with a smaller model ---
+# Using a smaller model for CPU inference (FLAN-T5)
+model_name = "google/flan-t5-base"
+tokenizer = AutoTokenizer.from_pretrained(model_name)
+model = AutoModelForCausalLM.from_pretrained(model_name)
 
 # --- Helper Functions ---
 def sanitize_answer(question, answer):
@@ -49,7 +44,7 @@ def is_finance_question(user_query):
         f"Question: {user_query}\n\n"
         f"Respond only with 'Yes' or 'No'."
     )
-    input_ids = tokenizer(check_prompt, return_tensors="pt").to(model.device)
+    input_ids = tokenizer(check_prompt, return_tensors="pt")
     output_ids = model.generate(
         **input_ids,
         max_new_tokens=10,
@@ -68,7 +63,7 @@ def ask_finance_bot(user_query, top_k=3):
     # Embed user query
     query_embedding = embedding_model.encode([user_query])
     D, I = index.search(np.array(query_embedding), top_k)
-    retrieved_answers = [answers[i] for i in I[0]]
+    retrieved_answers = [df.iloc[i]['answer'] for i in I[0]]
     context = "\n".join([f"- {text}" for text in retrieved_answers])
 
     temperature = min(0.7 + 0.1 * (count - 1), 1.0)
@@ -98,7 +93,7 @@ User question: {user_query}
 
 Answer:"""
 
-        input_ids = tokenizer(prompt, return_tensors="pt").to(model.device)
+        input_ids = tokenizer(prompt, return_tensors="pt")
         output_ids = model.generate(
             **input_ids,
             max_new_tokens=256,
@@ -116,8 +111,8 @@ Answer:"""
     return "I'm not confident in the response. Please consult a certified financial expert."
 
 # --- Streamlit App UI ---
-st.set_page_config(page_title="DiMowkayBot - Finance Assistant", layout="centered")
-st.title("💸 DiMowkayBot - Your Finance Q&A Assistant")
+st.set_page_config(page_title="Finsmart Bot", layout="centered")
+st.title(" Finsmart Bot - Your Finance Q&A Assistant")
 
 user_query = st.text_input("Enter your finance-related question:")