Update app.py

app.py

@@ -1,15 +1,56 @@
 import streamlit as st
-from utils import ask_finance_bot
+import pandas as pd
+from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
+from sentence_transformers import SentenceTransformer
+import faiss
+import numpy as np
 
-st.set_page_config(page_title="FinanceBot", layout="centered")
-st.title("💼 FinanceBot")
+st.title("T5 QA Chatbot on CSV Content")
 
-user_input = st.text_input("Enter your finance-related question:")
+# Load models
+@st.cache_resource
+def load_models():
+    tokenizer = AutoTokenizer.from_pretrained("t5-small")
+    model = AutoModelForSeq2SeqLM.from_pretrained("t5-small")
+    embedder = SentenceTransformer("all-MiniLM-L6-v2")
+    return tokenizer, model, embedder
 
-if st.button("Ask"):
-    if user_input.strip():
-        with st.spinner("Generating answer..."):
-            answer = ask_finance_bot(user_input)
-            st.success(answer)
-    else:
-        st.warning("Please enter a question.")
+tokenizer, model, embedder = load_models()
+
+# Load data
+@st.cache_data
+def load_data():
+    df = pd.read_csv("train_data.csv").head(100)
+    df['content'] = df['answer']
+    return df
+
+data = load_data()
+
+# Build vector store
+@st.cache_resource
+def build_vector_store(texts):
+    embeddings = embedder.encode(texts)
+    dim = embeddings[0].shape[0]
+    index = faiss.IndexFlatL2(dim)
+    index.add(np.array(embeddings))
+    return index, embeddings
+
+texts = data['content'].tolist()
+index, embeddings = build_vector_store(texts)
+
+# Chat UI
+prompt = st.chat_input("Ask something about the content...")
+
+if prompt:
+    # Embed prompt and retrieve top 3
+    q_embed = embedder.encode([prompt])
+    _, I = index.search(np.array(q_embed), k=3)
+    context = " ".join([texts[i] for i in I[0]])
+
+    # Format prompt for T5
+    input_text = f"question: {prompt} context: {context}"
+    inputs = tokenizer(input_text, return_tensors="pt", truncation=True)
+    outputs = model.generate(**inputs, max_length=100)
+    answer = tokenizer.decode(outputs[0], skip_special_tokens=True)
+
+    st.markdown(f"**Answer:** {answer}")