Update app.py

app.py

@@ -134,17 +134,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 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 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
 
     async def get_text_embedding(self, text: str, model: str = "mistral-embed") -> np.ndarray:
         try:
@@ -250,12 +250,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()
@@ -264,38 +264,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,
@@ -312,54 +312,54 @@ class MistralRAGChatbot:
 
         return reranked_docs
 
-    # def reciprocal_rank_fusion(self, user_query: str, docs: List[dict]) -> List[dict]:
-    #     k = 60  
-    #     method_ranks = {}
-    #     fused_scores = {} 
-    #     for doc in docs:
-    #         method = doc['method']
-    #         if method not in method_ranks:
-    #             method_ranks[method] = {doc['index']: 1}
-    #         else:
-    #             method_ranks[method][doc['index']] = len(method_ranks[method]) + 1
-    #     for doc in docs:
-    #         idx = doc['index']
-    #         if idx not in fused_scores:
-    #             fused_scores[idx] = sum(1 / (k + rank) for method_rank in method_ranks.values() for i, rank in method_rank.items() if i == idx)
-    #     reranked_docs = sorted(docs, key=lambda x: fused_scores.get(x['index'], 0), reverse=True)
-    #     for doc in reranked_docs:
-    #         doc['rrf_score'] = fused_scores.get(doc['index'], 0) 
-    #     return reranked_docs
-
-    # def weighted_score_fusion(self, user_query: str, docs: List[dict]) -> List[dict]:
-    #     method_weights = {
-    #         'annoy': 0.3,
-    #         'tfidf': 0.2,
-    #         'bm25': 0.2,
-    #         'word2vec': 0.1,
-    #         'euclidean': 0.1,
-    #         'jaccard': 0.1
-    #     }
-    #     fused_scores = {}
-    #     for doc in docs:
-    #         idx = doc['index']
-    #         if idx not in fused_scores:
-    #             fused_scores[idx] = doc['score'] * method_weights[doc['method']]
-    #         else:
-    #             fused_scores[idx] += doc['score'] * method_weights[doc['method']]
-
-    #     reranked_docs = sorted(docs, key=lambda x: fused_scores[x['index']], reverse=True)
-    #     for doc in reranked_docs:
-    #         doc['wsf_score'] = fused_scores[doc['index']]
-    #     return reranked_docs
-
-    # def semantic_similarity_reranking(self, user_query: str, docs: List[dict]) -> List[dict]:
-    #     query_embedding = np.mean([self.word2vec_model.wv[token] for token in user_query.split() if token in self.word2vec_model.wv], axis=0)
-    #     for doc in docs:
-    #         doc_embedding = np.mean([self.word2vec_model.wv[token] for token in doc['text'].split() if token in self.word2vec_model.wv], axis=0)
-    #         doc_embedding = doc_embedding if doc_embedding.shape == query_embedding.shape else np.zeros(query_embedding.shape)
-    #         doc['semantic_score'] = cosine_similarity([query_embedding], [doc_embedding])[0][0]
-    #     return sorted(docs, key=lambda x: x['semantic_score'], reverse=True)
+    def reciprocal_rank_fusion(self, user_query: str, docs: List[dict]) -> List[dict]:
+        k = 60  
+        method_ranks = {}
+        fused_scores = {} 
+        for doc in docs:
+            method = doc['method']
+            if method not in method_ranks:
+                method_ranks[method] = {doc['index']: 1}
+            else:
+                method_ranks[method][doc['index']] = len(method_ranks[method]) + 1
+        for doc in docs:
+            idx = doc['index']
+            if idx not in fused_scores:
+                fused_scores[idx] = sum(1 / (k + rank) for method_rank in method_ranks.values() for i, rank in method_rank.items() if i == idx)
+        reranked_docs = sorted(docs, key=lambda x: fused_scores.get(x['index'], 0), reverse=True)
+        for doc in reranked_docs:
+            doc['rrf_score'] = fused_scores.get(doc['index'], 0) 
+        return reranked_docs
+
+    def weighted_score_fusion(self, user_query: str, docs: List[dict]) -> List[dict]:
+        method_weights = {
+            'annoy': 0.3,
+            'tfidf': 0.2,
+            'bm25': 0.2,
+            'word2vec': 0.1,
+            'euclidean': 0.1,
+            'jaccard': 0.1
+        }
+        fused_scores = {}
+        for doc in docs:
+            idx = doc['index']
+            if idx not in fused_scores:
+                fused_scores[idx] = doc['score'] * method_weights[doc['method']]
+            else:
+                fused_scores[idx] += doc['score'] * method_weights[doc['method']]
+
+        reranked_docs = sorted(docs, key=lambda x: fused_scores[x['index']], reverse=True)
+        for doc in reranked_docs:
+            doc['wsf_score'] = fused_scores[doc['index']]
+        return reranked_docs
+
+    def semantic_similarity_reranking(self, user_query: str, docs: List[dict]) -> List[dict]:
+        query_embedding = np.mean([self.word2vec_model.wv[token] for token in user_query.split() if token in self.word2vec_model.wv], axis=0)
+        for doc in docs:
+            doc_embedding = np.mean([self.word2vec_model.wv[token] for token in doc['text'].split() if token in self.word2vec_model.wv], axis=0)
+            doc_embedding = doc_embedding if doc_embedding.shape == query_embedding.shape else np.zeros(query_embedding.shape)
+            doc['semantic_score'] = cosine_similarity([query_embedding], [doc_embedding])[0][0]
+        return sorted(docs, key=lambda x: x['semantic_score'], reverse=True)
 
     def build_prompt(self, context: str, user_query: str, response_style: str) -> str:
         styles = {