Update app.py

app.py

@@ -108,9 +108,9 @@ class MistralRAGChatbot:
     def __init__(self, vector_db_path: str, annoy_index_path: str):
         self.embeddings, self.texts = self.load_vector_db(vector_db_path)
         self.annoy_index = self.load_annoy_index(annoy_index_path, self.embeddings.shape[1])
-        # self.tfidf_matrix, self.tfidf_vectorizer = self.calculate_tfidf(self.texts)
+        self.tfidf_matrix, self.tfidf_vectorizer = self.calculate_tfidf(self.texts)
         self.bm25 = BM25Okapi([text.split() for text in self.texts])
-        # self.word2vec_model = self.train_word2vec(self.texts)
+        self.word2vec_model = self.train_word2vec(self.texts)
         self.reranking_methods = {
             'advanced_fusion': self.advanced_fusion_retrieval
             }
@@ -131,17 +131,17 @@ class MistralRAGChatbot:
         logging.info(f"Loaded Annoy index from {annoy_index_path}.")
         return annoy_index
 
-    # def calculate_tfidf(self, texts: List[str]) -> Tuple[np.ndarray, TfidfVectorizer]:
-    #     vectorizer = TfidfVectorizer(stop_words='english')
-    #     tfidf_matrix = vectorizer.fit_transform(texts)
-    #     logging.info("TF-IDF matrix calculated.")
-    #     return tfidf_matrix, vectorizer
+    def calculate_tfidf(self, texts: List[str]) -> Tuple[np.ndarray, TfidfVectorizer]:
+        vectorizer = TfidfVectorizer(stop_words='english')
+        tfidf_matrix = vectorizer.fit_transform(texts)
+        logging.info("TF-IDF matrix calculated.")
+        return tfidf_matrix, vectorizer
 
-    # def train_word2vec(self, texts: List[str]) -> Word2Vec:
-    #     tokenized_texts = [text.split() for text in texts]
-    #     model = Word2Vec(sentences=tokenized_texts, vector_size=100, window=5, min_count=1, workers=4)
-    #     logging.info("Word2Vec model trained.")
-    #     return model
+    def train_word2vec(self, texts: List[str]) -> Word2Vec:
+        tokenized_texts = [text.split() for text in texts]
+        model = Word2Vec(sentences=tokenized_texts, vector_size=100, window=5, min_count=1, workers=4)
+        logging.info("Word2Vec model trained.")
+        return model
 
     async def get_text_embedding(self, text: str, model: str = "mistral-embed") -> np.ndarray:
         try:
@@ -247,12 +247,12 @@ class MistralRAGChatbot:
         logging.debug(f"Annoy retrieval returned {len(indices)} documents.")
         return indices, scores
 
-    # def retrieve_with_tfidf(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
-    #     query_vec = self.tfidf_vectorizer.transform([user_query])
-    #     similarities = cosine_similarity(query_vec, self.tfidf_matrix).flatten()
-    #     indices = np.argsort(-similarities)[:top_k]
-    #     logging.debug(f"TF-IDF retrieval returned {len(indices)} documents.")
-    #     return indices, similarities[indices].tolist()
+    def retrieve_with_tfidf(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
+        query_vec = self.tfidf_vectorizer.transform([user_query])
+        similarities = cosine_similarity(query_vec, self.tfidf_matrix).flatten()
+        indices = np.argsort(-similarities)[:top_k]
+        logging.debug(f"TF-IDF retrieval returned {len(indices)} documents.")
+        return indices, similarities[indices].tolist()
 
     def retrieve_with_bm25(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
         tokenized_query = user_query.split()
@@ -261,38 +261,38 @@ class MistralRAGChatbot:
         logging.debug(f"BM25 retrieval returned {len(indices)} documents.")
         return indices, scores[indices].tolist()
 
-    # def retrieve_with_word2vec(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
-    #     query_tokens = user_query.split()
-    #     query_vec = np.mean([self.word2vec_model.wv[token] for token in query_tokens if token in self.word2vec_model.wv], axis=0)
-    #     expected_dim = query_vec.shape[0]
-    #     doc_vectors = []
-    #     for doc in self.texts:
-    #         word_vectors = [self.word2vec_model.wv[token] for token in doc.split() if token in self.word2vec_model.wv]
-    #         avg_vector = np.mean(word_vectors, axis=0) if word_vectors else np.zeros(expected_dim)
-    #         doc_vectors.append(avg_vector)
-    #     doc_vectors = np.array(doc_vectors)
-    #     similarities = cosine_similarity([query_vec], doc_vectors).flatten()
-    #     indices = np.argsort(-similarities)[:top_k]
-    #     return indices, similarities[indices].tolist()
-
-    # def retrieve_with_euclidean(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
-    #     distances = euclidean_distances([query_embedding], self.embeddings).flatten()
-    #     indices = np.argsort(distances)[:top_k]
-    #     logging.debug(f"Euclidean retrieval returned {len(indices)} documents.")
-    #     return indices, distances[indices].tolist()
-
-    # def retrieve_with_jaccard(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
-    #     query_set = set(user_query.lower().split())
-    #     scores = []
-    #     for doc in self.texts:
-    #         doc_set = set(doc.lower().split())
-    #         intersection = query_set.intersection(doc_set)
-    #         union = query_set.union(doc_set)
-    #         score = float(len(intersection)) / len(union) if union else 0
-    #         scores.append(score)
-    #     indices = np.argsort(-np.array(scores))[:top_k]
-    #     logging.debug(f"Jaccard retrieval returned {len(indices)} documents.")
-    #     return indices.tolist(), [scores[i] for i in indices]
+    def retrieve_with_word2vec(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
+        query_tokens = user_query.split()
+        query_vec = np.mean([self.word2vec_model.wv[token] for token in query_tokens if token in self.word2vec_model.wv], axis=0)
+        expected_dim = query_vec.shape[0]
+        doc_vectors = []
+        for doc in self.texts:
+            word_vectors = [self.word2vec_model.wv[token] for token in doc.split() if token in self.word2vec_model.wv]
+            avg_vector = np.mean(word_vectors, axis=0) if word_vectors else np.zeros(expected_dim)
+            doc_vectors.append(avg_vector)
+        doc_vectors = np.array(doc_vectors)
+        similarities = cosine_similarity([query_vec], doc_vectors).flatten()
+        indices = np.argsort(-similarities)[:top_k]
+        return indices, similarities[indices].tolist()
+
+    def retrieve_with_euclidean(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
+        distances = euclidean_distances([query_embedding], self.embeddings).flatten()
+        indices = np.argsort(distances)[:top_k]
+        logging.debug(f"Euclidean retrieval returned {len(indices)} documents.")
+        return indices, distances[indices].tolist()
+
+    def retrieve_with_jaccard(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
+        query_set = set(user_query.lower().split())
+        scores = []
+        for doc in self.texts:
+            doc_set = set(doc.lower().split())
+            intersection = query_set.intersection(doc_set)
+            union = query_set.union(doc_set)
+            score = float(len(intersection)) / len(union) if union else 0
+            scores.append(score)
+        indices = np.argsort(-np.array(scores))[:top_k]
+        logging.debug(f"Jaccard retrieval returned {len(indices)} documents.")
+        return indices.tolist(), [scores[i] for i in indices]
 
     def rerank_documents(
         self,