Update app.py

pagerank_score error fixed

app.py

@@ -9,21 +9,20 @@ import asyncio
 import networkx as nx
 from mistralai import Mistral
 from annoy import AnnoyIndex
-from sklearn.feature_extraction.text import TfidfVectorizer, CountVectorizer
-from sklearn.metrics.pairwise import cosine_similarity, euclidean_distances
-from sklearn.preprocessing import normalize
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
 from rank_bm25 import BM25Okapi
 from gensim.models import Word2Vec
 from typing import List, Optional, Tuple
 import gradio as gr
 
-
 logger = logging.getLogger(__name__)
+logging.basicConfig(level=logging.INFO)
+
 api_key = os.getenv("MISTRAL_API_KEY")
 client = Mistral(api_key=api_key)
 
-
-def get_text_embedding_with_rate_limit(text_list, initial_delay=2, max_retries=10):
+def get_text_embedding_with_rate_limit(text_list, initial_delay=2, max_retries=10, max_delay=60):
     embeddings = []
     for text in text_list:
         retries = 0
@@ -32,7 +31,7 @@ def get_text_embedding_with_rate_limit(text_list, initial_delay=2, max_retries=1
             try:
                 token_count = len(text.split())
                 if token_count > 16384:
-                    print("Warning: Text chunk exceeds the token limit. Truncating the text.")
+                    logger.warning("Text chunk exceeds the token limit. Truncating the text.")
                     text = " ".join(text.split()[:16384])
                 response = client.embeddings.create(model="mistral-embed", inputs=[text])
                 embeddings.extend([embedding.embedding for embedding in response.data])
@@ -40,15 +39,14 @@ def get_text_embedding_with_rate_limit(text_list, initial_delay=2, max_retries=1
                 break
             except Exception as e:
                 retries += 1
-                print(f"Rate limit exceeded, retrying in {delay} seconds... (Attempt {retries}/{max_retries})")
+                logger.warning(f"Rate limit exceeded, retrying in {delay} seconds... (Attempt {retries}/{max_retries})")
                 time.sleep(delay)
-                delay *= 2
+                delay = min(delay * 2, max_delay)
                 if retries == max_retries:
-                    print("Max retries reached. Skipping this chunk.")
+                    logger.error("Max retries reached. Skipping this chunk.")
                     break
     return embeddings
 
-
 def store_embeddings_in_vector_db(
     pdf_path: str,
     vector_db_path: str,
@@ -92,7 +90,6 @@ def store_embeddings_in_vector_db(
     annoy_index.save(annoy_index_path)
     logging.info(f"Annoy index built with {len(all_embeddings)} items and saved to {annoy_index_path}.")
 
-
 def split_text_into_chunks(text: str, chunk_size: int = 2048, overlap: int = 200) -> List[str]:
     tokens = text.split()
     chunks = []
@@ -112,10 +109,9 @@ class MistralRAGChatbot:
         self.word2vec_model = self.train_word2vec(self.texts)
         self.reranking_methods = {
             'advanced_fusion': self.advanced_fusion_retrieval
-            }
+        }
         logging.info("MistralRAGChatbot initialized successfully.")
 
-
     def load_vector_db(self, vector_db_path: str) -> Tuple[np.ndarray, List[str]]:
         with open(vector_db_path, "rb") as f:
             data = dill.load(f)
@@ -125,12 +121,13 @@ class MistralRAGChatbot:
         return embeddings, texts
 
     def load_annoy_index(self, annoy_index_path: str, embedding_dim: int) -> AnnoyIndex:
+        if not os.path.exists(annoy_index_path):
+            raise FileNotFoundError(f"Annoy index file {annoy_index_path} not found.")
         annoy_index = AnnoyIndex(embedding_dim, 'angular')
         annoy_index.load(annoy_index_path)
         logging.info(f"Loaded Annoy index from {annoy_index_path}.")
         return annoy_index
 
-
     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)
@@ -143,7 +140,7 @@ class MistralRAGChatbot:
             return np.array(response.data[0].embedding)
         except Exception as e:
             logging.error(f"Error fetching embedding: {e}")
-            return np.zeros((1024,))  
+            return np.zeros((1024,))
 
     def advanced_fusion_retrieval(self, user_query: str, docs: List[dict]) -> List[dict]:
         query_embedding = self.create_embeddings([user_query])[0]
@@ -166,22 +163,20 @@ class MistralRAGChatbot:
             combined_scores[idx] = (
                 0.5 * vector_scores.get(idx, 0) +
                 0.3 * bm25_scores.get(idx, 0) +
-                0.2 * pagerank_scores[idx]
+                0.2 * pagerank_scores[idx] if idx < len(pagerank_scores) else 0
             )
-        
+
         min_score = min(combined_scores.values())
         max_score = max(combined_scores.values())
-        
-        
+
         normalized_scores = {idx: (score - min_score) / (max_score - min_score) for idx, score in combined_scores.items()}
-        
-        
+
         sorted_indices = sorted(combined_scores, key=combined_scores.get, reverse=True)
 
         return [{'text': self.texts[i], 'method': 'advanced_fusion', 'score': normalized_scores[i], 'index': i} for i in sorted_indices[:5]]
 
     def create_embeddings(self, text_list: List[str]) -> np.ndarray:
-        expected_dim = 1024 
+        expected_dim = 1024
         embeddings = []
         for text in text_list:
             word_vectors = [self.word2vec_model.wv[token] for token in text.split() if token in self.word2vec_model.wv]
@@ -203,9 +198,9 @@ class MistralRAGChatbot:
         selected_reranking_methods: Optional[List[str]] = None
     ) -> Tuple[str, List[str], List[dict]]:
         if not selected_retrieval_methods:
-            selected_retrieval_methods = ['annoy', 'tfidf', 'bm25', 'word2vec', 'euclidean', 'jaccard']
+            selected_retrieval_methods = ['annoy', 'bm25']
         if not selected_reranking_methods:
-            selected_reranking_methods = ['reciprocal_rank_fusion', 'weighted_score_fusion', 'advanced_fusion']
+            selected_reranking_methods = ['advanced_fusion']
         query_embedding = await self.get_text_embedding(user_query)
         retrieved_docs = self.retrieve_documents(user_query, query_embedding, top_k, selected_retrieval_methods)
         reranked_docs = self.rerank_documents(user_query, retrieved_docs, selected_reranking_methods)
@@ -244,11 +239,10 @@ class MistralRAGChatbot:
     def retrieve_with_annoy(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:
         n_results = min(top_k, len(self.texts))
         indices, distances = self.annoy_index.get_nns_by_vector(query_embedding, n_results, include_distances=True)
-        scores = [1.0 - (dist / max(distances)) for dist in distances] 
+        scores = [1.0 - (dist / max(distances)) for dist in distances] if distances else []
         logging.debug(f"Annoy retrieval returned {len(indices)} documents.")
         return indices, scores
 
-
     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()
         scores = self.bm25.get_scores(tokenized_query)
@@ -256,21 +250,6 @@ 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 rerank_documents(
         self,
         user_query: str,
@@ -286,7 +265,6 @@ class MistralRAGChatbot:
 
         return reranked_docs
 
-
     def build_prompt(self, context: str, user_query: str, response_style: str) -> str:
         styles = {
             "detailed": "Provide a comprehensive and detailed answer based on the provided context.",
@@ -323,18 +301,16 @@ class MistralRAGChatbot:
         common_terms = query_terms.intersection(context_terms)
         return len(common_terms) > len(query_terms) * 0.2
 
-
 def create_vector_db_and_annoy_index(pdf_path, vector_db_path, annoy_index_path):
     store_embeddings_in_vector_db(pdf_path, vector_db_path, annoy_index_path)
     print("Vector database and Annoy index creation completed.")
 
-
 def chatbot_interface(file, user_query, response_style):
     vector_db_path = "vector_db.pkl"
     annoy_index_path = "vector_index.ann"
     chunk_size = 2048
     overlap = 200
-    store_embeddings_in_vector_db(file.name, 'vector_db.pkl', 'vector_index.ann', chunk_size, overlap)
+    store_embeddings_in_vector_db(file.name, vector_db_path, annoy_index_path, chunk_size, overlap)
 
     chatbot = MistralRAGChatbot(vector_db_path, annoy_index_path)
 
@@ -356,26 +332,22 @@ def chatbot_interface(file, user_query, response_style):
         formatted_response += f"Retrieval Method: {doc_info['method']}\n"
         if 'score' in doc_info:
             formatted_response += f"Precision Score: {doc_info['score']:.4f}\n"
-        for key, value in doc_info.items():
-            if key.endswith('_score') and key != 'score':
-                formatted_response += f"{key.replace('_', ' ').title()}: {value:.4f}\n"
 
     return formatted_response
-iface = gr.Blocks(theme="Rabbitt-AI/ChanceRAG")
+
+iface = gr.Blocks()
 with iface:
-    gr.Image("images/chanceRAG_logo.jpg", label="Image", show_label = False)
+    gr.Image("images/chanceRAG_logo.jpg", label="Image", show_label=False)
     gr.Interface(
-    fn=chatbot_interface,
-    theme="Rabbitt-AI/ChanceRAG",
-
-    inputs=[
-        gr.File(label="Upload a File"),
-        gr.Textbox(lines=5, label="User Query"),
-        gr.Dropdown(
-            ["Detailed", "Concise", "Creative", "Technical"], label="Retreival Style"
-        ),
-    ],
-    outputs=gr.Textbox(label="ChanceRAG Response"),
-)
+        fn=chatbot_interface,
+        inputs=[
+            gr.File(label="Upload a File"),
+            gr.Textbox(lines=5, label="User Query"),
+            gr.Dropdown([
+                "Detailed", "Concise", "Creative", "Technical"], label="Response Style"
+            ),
+        ],
+        outputs=gr.Textbox(label="ChanceRAG Response"),
+    )
 
 iface.launch(share=True)