Update rag.py (#4)

- Update rag.py (730471acfb95df5fc45e201bec1d5d29d7130a32)


Co-authored-by: Prerna Aneja <[email protected]>

rag.py

@@ -1,13 +1,12 @@
 import faiss
+import numpy as np
 import pickle
 import threading
 import time
 import torch
-import numpy as np
+from sentence_transformers import SentenceTransformer
 from transformers import AutoModelForCausalLM, AutoTokenizer
 from rank_bm25 import BM25Okapi
-from sentence_transformers import SentenceTransformer
-from sklearn.metrics.pairwise import cosine_similarity
 
 class FinancialChatbot:
     def __init__(self):
@@ -15,10 +14,12 @@ class FinancialChatbot:
         self.faiss_index = faiss.read_index("financial_faiss.index")
         with open("index_map.pkl", "rb") as f:
             self.index_map = pickle.load(f)
-        
-        # Load BM25 keyword-based search
-        with open("bm25_corpus.pkl", "rb") as f:
-            self.bm25_corpus = pickle.load(f)
+
+        # Extract document texts for BM25 dynamically
+        self.documents = list(self.index_map.values())
+
+        # Build BM25 index dynamically
+        self.bm25_corpus = [doc.lower().split() for doc in self.documents]  # Tokenization
         self.bm25 = BM25Okapi(self.bm25_corpus)
 
         # Load SentenceTransformer for embedding-based retrieval
@@ -66,7 +67,7 @@ class FinancialChatbot:
         return results, confidence_scores
 
     def query_bm25(self, query, top_k=5):
-        """Retrieve relevant documents using BM25 keyword-based search."""
+        """Retrieve relevant documents using BM25 keyword-based search dynamically."""
         tokenized_query = query.lower().split()
         scores = self.bm25.get_scores(tokenized_query)
         top_indices = np.argsort(scores)[::-1][:top_k]
@@ -76,7 +77,7 @@ class FinancialChatbot:
 
         for idx in top_indices:
             if scores[idx] > 0:  # Ignore zero-score matches
-                results.append(self.bm25_corpus[idx])
+                results.append(self.documents[idx])
                 confidence_scores.append(scores[idx])
 
         return results, confidence_scores
@@ -98,7 +99,7 @@ class FinancialChatbot:
                 return
 
             if not self.moderate_query(query):
-                result[:] = ["I'm unable to process your request due to inappropriate language.", 1.0]
+                result[:] = ["I'm unable to process your request due to inappropriate language.", 0.0]
                 return
 
             faiss_results, faiss_conf = self.query_faiss(query)
@@ -107,25 +108,28 @@ class FinancialChatbot:
             all_results = faiss_results + bm25_results
             all_conf = faiss_conf + bm25_conf
 
-            # Check relevance
+            # Check if results are relevant
             if not all_results or max(all_conf, default=0) < self.min_similarity_threshold:
-                result[:] = ["No relevant information found", 1.0]
+                result[:] = ["No relevant information found", 0.0]
                 return
 
             context = " ".join(all_results)
             answer = self.generate_answer(context, query)
             
             last_index = answer.rfind("Answer")
-            if answer[last_index+9:11] == "--":
-                result[:] = ["No relevant information found", 1.0]
+            extracted_answer = answer[last_index:].strip() if last_index != -1 else ""
+
+            # Ensure the answer is grounded in the context
+            if not extracted_answer or "Answer" not in answer or extracted_answer.isnumeric():
+                result[:] = ["No relevant information found", 0.0]
             else:
-                result[:] = [answer[last_index:], max(all_conf, default=0.9)]
+                result[:] = [extracted_answer, max(all_conf, default=0.9)]
 
         thread = threading.Thread(target=task)
         thread.start()
         thread.join(timeout)
 
         if thread.is_alive():
-            return "No relevant information found", 1.0  # Timeout case
+            return "No relevant information found", 0.0  # Timeout case
 
-        return tuple(result)
+        return tuple(result)