Update utils.py

utils.py

@@ -1,18 +1,33 @@
 import pandas as pd
-from sentence_transformers import SentenceTransformer
+from sentence_transformers import SentenceTransformer, util
 
-# Load embedding model
-embedding_model = SentenceTransformer("all-MiniLM-L6-v2")
-
-# Load CSV data
-def load_qa_data(file_path="qa_data.csv"):
+# Load and embed Q&A data
+def load_qa_data(file_path="train_data.csv"):
     df = pd.read_csv(file_path)
-    return list(zip(df["question"], df["answer"]))
+    questions = df["question"].tolist()
+    answers = df["answer"].tolist()
+    return questions, answers
+
+# Set up embeddings
+model = SentenceTransformer("all-MiniLM-L6-v2")
+questions, answers = load_qa_data()
+question_embeddings = model.encode(questions, convert_to_tensor=True)
+
+# Search top-k context
+def retrieve_context(query, top_k=3):
+    query_embedding = model.encode(query, convert_to_tensor=True)
+    hits = util.semantic_search(query_embedding, question_embeddings, top_k=top_k)[0]
+
+    context = ""
+    for hit in hits:
+        idx = hit["corpus_id"]
+        context += f"Q: {questions[idx]}\nA: {answers[idx]}\n\n"
+
+    return context.strip()
 
-# Basic semantic check (optional if you want to embed later)
-def is_financial_text(text: str) -> bool:
-    finance_keywords = [
-        "finance", "investment", "bank", "insurance", "credit", "budget", "economy",
-        "inflation", "debt", "interest", "mortgage", "pension", "retirement", "savings"
-    ]
-    return any(keyword in text.lower() for keyword in finance_keywords)
+# Use RAG to check if it's a finance question
+def is_financial_text_via_context(query, threshold=0.3):
+    query_embedding = model.encode(query, convert_to_tensor=True)
+    similarities = util.cos_sim(query_embedding, question_embeddings)[0]
+    max_sim = float(similarities.max())
+    return max_sim > threshold