Rabbitt-AI/ChanceRAG
6
1import time2import fitz3import numpy as np4import pickle5import os6import dill7import logging8import asyncio9import networkx as nx10from mistralai import Mistral11from annoy import AnnoyIndex12from sklearn.feature_extraction.text import TfidfVectorizer13from sklearn.metrics.pairwise import cosine_similarity14from rank_bm25 import BM25Okapi15from gensim.models import Word2Vec16from typing import List, Optional, Tuple17import gradio as gr18import moviepy.editor as mp19 20logger = logging.getLogger(__name__)21logging.basicConfig(level=logging.INFO)22 23api_key = os.getenv("MISTRAL_API_KEY")24client = Mistral(api_key=api_key)25 26from deepgram import Deepgram27 28dg_api_key = os.getenv("DEEPGRAM_API_KEY")29deepgram = Deepgram(dg_api_key)30 31def get_text_embedding_with_rate_limit(text_list, initial_delay=2, max_retries=10, max_delay=60):32 embeddings = []33 for text in text_list:34 retries = 035 delay = initial_delay36 while retries < max_retries:37 try:38 token_count = len(text.split())39 if token_count > 16384:40 logger.warning("Text chunk exceeds the token limit. Truncating the text.")41 text = " ".join(text.split()[:16384])42 response = client.embeddings.create(model="mistral-embed", inputs=[text])43 embeddings.extend([embedding.embedding for embedding in response.data])44 time.sleep(delay)45 break46 except Exception as e:47 retries += 148 logger.warning(f"Rate limit exceeded, retrying in {delay} seconds... (Attempt {retries}/{max_retries})")49 time.sleep(delay)50 delay = min(delay * 2, max_delay)51 if retries == max_retries:52 logger.error("Max retries reached. Skipping this chunk.")53 break54 return embeddings55 56def store_embeddings_in_vector_db(57 file_path: str,58 vector_db_path: str,59 annoy_index_path: str,60 chunk_size: int = 2048,61 overlap: int = 200,62 num_trees: int = 1063):64 all_texts = []65 if file_path.endswith(('.pdf', '.doc', '.docx' , '.pptx' , '.ppt' , '.xls', '.xlsx' , '.txt' )):66 doc = fitz.open(file_path)67 all_embeddings = []68 total_pages = doc.page_count69 logging.info(f"Processing PDF/DOC: {file_path} with {total_pages} pages.")70 71 for page_num in range(total_pages):72 page = doc.load_page(page_num)73 text = page.get_text()74 if text.strip():75 chunks = split_text_into_chunks(text, chunk_size, overlap)76 embeddings = get_text_embedding_with_rate_limit(chunks)77 all_embeddings.extend(embeddings)78 all_texts.extend(chunks)79 logging.info(f"Processed page {page_num + 1}/{total_pages}, extracted {len(chunks)} chunks.")80 else:81 logging.warning(f"No text found on page {page_num + 1}.")82 elif file_path.endswith(('.mp3', '.wav', '.m4a')):83 logging.info(f"Processing audio file: {file_path}")84 with open(file_path, 'rb') as audio_file:85 audio_content = audio_file.read()86 response = asyncio.run(deepgram.transcription.prerecorded({'buffer': audio_content, 'mimetype': 'audio/wav'}, {'punctuate': True}))87 text = response['results']['channels'][0]['alternatives'][0]['transcript']88 chunks = split_text_into_chunks(text, chunk_size, overlap)89 all_embeddings = get_text_embedding_with_rate_limit(chunks)90 all_texts.extend(chunks)91 elif file_path.endswith(('.mp4', '.avi', '.mov')):92 logging.info(f"Processing video file: {file_path}")93 video = mp.VideoFileClip(file_path)94 audio_path = "temp_audio.wav"95 video.audio.write_audiofile(audio_path)96 with open(audio_path, 'rb') as audio_file:97 audio_content = audio_file.read()98 response = asyncio.run(deepgram.transcription.prerecorded({'buffer': audio_content, 'mimetype': 'audio/wav'}, {'punctuate': True}))99 text = response['results']['channels'][0]['alternatives'][0]['transcript']100 os.remove(audio_path)101 chunks = split_text_into_chunks(text, chunk_size, overlap)102 all_embeddings = get_text_embedding_with_rate_limit(chunks)103 all_texts.extend(chunks)104 else:105 raise ValueError("Unsupported file format. Please upload a PDF, DOC, DOCX, MP3, WAV, M4A, MP4, AVI, or MOV file.")106 107 embeddings_np = np.array(all_embeddings).astype('float32')108 with open(vector_db_path, "wb") as f:109 dill.dump({'embeddings': embeddings_np, 'texts': all_texts}, f)110 logging.info(f"Stored embeddings and texts to {vector_db_path}.")111 112 if os.path.exists(annoy_index_path):113 os.remove(annoy_index_path)114 logging.info(f"Existing Annoy index at {annoy_index_path} removed.")115 116 embedding_dim = embeddings_np.shape[1]117 annoy_index = AnnoyIndex(embedding_dim, 'angular')118 for i, embedding in enumerate(embeddings_np):119 annoy_index.add_item(i, embedding)120 annoy_index.build(num_trees)121 annoy_index.save(annoy_index_path)122 logging.info(f"Annoy index built with {len(all_embeddings)} items and saved to {annoy_index_path}.")123 124def split_text_into_chunks(text: str, chunk_size: int = 2048, overlap: int = 200) -> List[str]:125 tokens = text.split()126 chunks = []127 start = 0128 while start < len(tokens):129 end = start + chunk_size130 chunk = " ".join(tokens[start:end])131 chunks.append(chunk)132 start += chunk_size - overlap133 return chunks134 135class MistralRAGChatbot:136 def __init__(self, vector_db_path: str, annoy_index_path: str):137 self.embeddings, self.texts = self.load_vector_db(vector_db_path)138 self.annoy_index = self.load_annoy_index(annoy_index_path, self.embeddings.shape[1])139 self.bm25 = BM25Okapi([text.split() for text in self.texts])140 self.word2vec_model = self.train_word2vec(self.texts)141 self.reranking_methods = {142 'advanced_fusion': self.advanced_fusion_retrieval143 }144 logging.info("MistralRAGChatbot initialized successfully.")145 146 def load_vector_db(self, vector_db_path: str) -> Tuple[np.ndarray, List[str]]:147 with open(vector_db_path, "rb") as f:148 data = dill.load(f)149 embeddings = np.array(data['embeddings'], dtype='float32')150 texts = data['texts']151 logging.info(f"Loaded vector database from {vector_db_path} with {len(texts)} entries.")152 return embeddings, texts153 154 def load_annoy_index(self, annoy_index_path: str, embedding_dim: int) -> AnnoyIndex:155 if not os.path.exists(annoy_index_path):156 raise FileNotFoundError(f"Annoy index file {annoy_index_path} not found.")157 annoy_index = AnnoyIndex(embedding_dim, 'angular')158 annoy_index.load(annoy_index_path)159 logging.info(f"Loaded Annoy index from {annoy_index_path}.")160 return annoy_index161 162 def train_word2vec(self, texts: List[str]) -> Word2Vec:163 tokenized_texts = [text.split() for text in texts]164 model = Word2Vec(sentences=tokenized_texts, vector_size=100, window=5, min_count=1, workers=4)165 logging.info("Word2Vec model trained.")166 return model167 168 async def get_text_embedding(self, text: str, model: str = "mistral-embed") -> np.ndarray:169 try:170 response = await client.embeddings.create_async(model=model, inputs=[text])171 return np.array(response.data[0].embedding)172 except Exception as e:173 logging.error(f"Error fetching embedding: {e}")174 return np.zeros((1024,))175 176 def advanced_fusion_retrieval(self, user_query: str, docs: List[dict]) -> List[dict]:177 query_embedding = self.create_embeddings([user_query])[0]178 179 vector_scores = {doc['index']: doc['score'] for doc in docs if doc['method'] == 'annoy'}180 bm25_scores = {doc['index']: doc['score'] for doc in docs if doc['method'] == 'bm25'}181 182 sim_graph = nx.Graph()183 sim_matrix = cosine_similarity(self.embeddings)184 for i in range(len(self.embeddings)):185 for j in range(i + 1, len(self.embeddings)):186 if sim_matrix[i, j] > 0.5:187 sim_graph.add_edge(i, j, weight=sim_matrix[i, j])188 189 pagerank_scores = np.array(list(nx.pagerank(sim_graph, weight='weight').values()))190 191 combined_scores = {}192 for doc in docs:193 idx = doc['index']194 combined_scores[idx] = (195 0.5 * vector_scores.get(idx, 0) +196 0.3 * bm25_scores.get(idx, 0) +197 0.2 * pagerank_scores[idx] if idx < len(pagerank_scores) else 0198 )199 200 min_score = min(combined_scores.values())201 max_score = max(combined_scores.values())202 203 if min_score == max_score:204 normalized_scores = {idx: 0.5 for idx in combined_scores}205 else:206 normalized_scores = {idx: (score - min_score) / (max_score - min_score) for idx, score in combined_scores.items()}207 208 sorted_indices = sorted(combined_scores, key=combined_scores.get, reverse=True)209 210 return [{'text': self.texts[i], 'method': 'advanced_fusion', 'score': normalized_scores[i], 'index': i} for i in sorted_indices[:5]]211 212 def create_embeddings(self, text_list: List[str]) -> np.ndarray:213 expected_dim = 1024214 embeddings = []215 for text in text_list:216 word_vectors = [self.word2vec_model.wv[token] for token in text.split() if token in self.word2vec_model.wv]217 avg_embedding = np.mean(word_vectors, axis=0, dtype=np.float32) if word_vectors else np.zeros(self.word2vec_model.vector_size, dtype=np.float32)218 if avg_embedding.shape[0] < expected_dim:219 avg_embedding = np.pad(avg_embedding, (0, expected_dim - avg_embedding.shape[0]), 'constant')220 elif avg_embedding.shape[0] > expected_dim:221 avg_embedding = avg_embedding[:expected_dim]222 embeddings.append(avg_embedding)223 return np.array(embeddings, dtype=np.float32)224 225 async def generate_response_with_rag(226 self,227 user_query: str,228 model: str = "mistral-small-latest",229 top_k: int = 10,230 response_style: str = "Detailed",231 selected_retrieval_methods: Optional[List[str]] = None,232 selected_reranking_methods: Optional[List[str]] = None233 ) -> Tuple[str, List[str], List[dict]]:234 if not selected_retrieval_methods:235 selected_retrieval_methods = ['annoy', 'bm25']236 if not selected_reranking_methods:237 selected_reranking_methods = ['advanced_fusion']238 query_embedding = await self.get_text_embedding(user_query)239 retrieved_docs = self.retrieve_documents(user_query, query_embedding, top_k, selected_retrieval_methods)240 reranked_docs = self.rerank_documents(user_query, retrieved_docs, selected_reranking_methods)241 context = "\n\n".join([doc['text'] for doc in reranked_docs[:5]])242 prompt = self.build_prompt(context, user_query, response_style)243 try:244 async_response = await client.chat.stream_async(model=model, messages=[{"role": "user", "content": prompt}])245 response = ""246 async for chunk in async_response:247 response += chunk.data.choices[0].delta.content248 logging.info("Response generated successfully.")249 except Exception as e:250 logging.error(f"Error generating response: {e}")251 response = "An error occurred while generating the response."252 return response, [doc['text'] for doc in reranked_docs[:5]], reranked_docs[:5]253 254 def retrieve_documents(255 self,256 user_query: str,257 query_embedding: np.ndarray,258 top_k: int,259 selected_methods: List[str]260 ) -> List[dict]:261 all_docs = []262 for method in selected_methods:263 indices, scores = getattr(self, f"retrieve_with_{method}")(user_query, query_embedding, top_k)264 for idx, score in zip(indices, scores):265 all_docs.append({266 'text': self.texts[idx],267 'method': method,268 'score': score,269 'index': idx270 })271 return all_docs272 273 def retrieve_with_annoy(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:274 n_results = min(top_k, len(self.texts))275 indices, distances = self.annoy_index.get_nns_by_vector(query_embedding, n_results, include_distances=True)276 scores = [1.0 - (dist / max(distances)) for dist in distances] if distances else []277 logging.debug(f"Annoy retrieval returned {len(indices)} documents.")278 return indices, scores279 280 def retrieve_with_bm25(self, user_query: str, query_embedding: np.ndarray, top_k: int) -> Tuple[List[int], List[float]]:281 tokenized_query = user_query.split()282 scores = self.bm25.get_scores(tokenized_query)283 indices = np.argsort(-scores)[:top_k]284 logging.debug(f"BM25 retrieval returned {len(indices)} documents.")285 return indices, scores[indices].tolist()286 287 def rerank_documents(288 self,289 user_query: str,290 retrieved_docs: List[dict],291 selected_methods: List[str]292 ) -> List[dict]:293 reranked_docs = retrieved_docs294 for method in selected_methods:295 if method == 'advanced_fusion':296 reranked_docs = self.advanced_fusion_retrieval(user_query, reranked_docs)297 else:298 reranked_docs = self.reranking_methods[method](user_query, reranked_docs)299 300 return reranked_docs301 302 def build_prompt(self, context: str, user_query: str, response_style: str) -> str:303 styles = {304 "detailed": "Provide a comprehensive and detailed answer based on the provided context.",305 "concise": "Provide a brief and concise answer based on the provided context.",306 "creative": "Provide a creative and engaging answer based on the provided context.",307 "technical": "Provide a technical and in-depth answer based on the provided context."308 }309 310 style_instruction = styles.get(response_style.lower(), styles["detailed"])311 312 if not context or not self.is_context_relevant(context, user_query):313 prompt = f"""You are an intelligent assistant.314 User Question:315 {user_query}316 Instruction:317 The document database does not contain relevant information to answer the question. Please inform the user that no relevant documents were found and refrain from generating an imaginative or unrelated response."""318 else:319 prompt = f"""You are an intelligent assistant.320 Context:321 {context}322 User Question:323 {user_query}324 Instruction:325 {style_instruction}"""326 327 logging.debug("Prompt constructed for response generation.")328 return prompt329 330 def is_context_relevant(self, context: str, user_query: str) -> bool:331 context_lower = context.lower()332 user_query_lower = user_query.lower()333 query_terms = set(user_query_lower.split())334 context_terms = set(context_lower.split())335 common_terms = query_terms.intersection(context_terms)336 return len(common_terms) > len(query_terms) * 0.2337 338def create_vector_db_and_annoy_index(pdf_path, vector_db_path, annoy_index_path):339 store_embeddings_in_vector_db(pdf_path, vector_db_path, annoy_index_path)340 print("Vector database and Annoy index creation completed.")341 342def chatbot_interface(file, user_query, response_style):343 vector_db_path = "vector_db.pkl"344 annoy_index_path = "vector_index.ann"345 chunk_size = 2048346 overlap = 200347 store_embeddings_in_vector_db(file.name, vector_db_path, annoy_index_path, chunk_size, overlap)348 349 chatbot = MistralRAGChatbot(vector_db_path, annoy_index_path)350 351 selected_retrieval_methods_list = ['annoy', 'bm25']352 selected_reranking_methods_list = ["advanced_fusion"]353 354 response, retrieved_docs, source_info = asyncio.run(chatbot.generate_response_with_rag(355 user_query=user_query,356 response_style=response_style,357 selected_retrieval_methods=selected_retrieval_methods_list,358 selected_reranking_methods=selected_reranking_methods_list359 ))360 361 formatted_response = f"# **ChanceRAG Response:**\n\n{response}\n\n"362 formatted_response += "Retrieved and Reranked Documents:\n"363 for idx, doc_info in enumerate(source_info, start=1):364 formatted_response += f"\nDocument {idx}:\n"365 formatted_response += f"Content Preview: {doc_info['text'][:200]}...\n"366 formatted_response += f"Retrieval Method: {doc_info['method']}\n"367 if 'score' in doc_info:368 formatted_response += f"Precision Score: {doc_info['score']:.4f}\n"369 return formatted_response370 371iface = gr.Blocks(theme="Rabbitt-AI/ChanceRAG")372with iface:373 gr.Image("images/chanceRAG_logo.jpg", label="Image", show_label=False)374 gr.Interface(375 fn=chatbot_interface,376 theme="Rabbitt-AI/ChanceRAG",377 inputs=[378 gr.File(label="Upload a File"),379 gr.Textbox(lines=5, label="User Query"),380 gr.Dropdown([381 "Detailed", "Concise", "Creative", "Technical"], label="Response Style"382 ),383 ],384 outputs= gr.Markdown(value="# **ChanceRAG Response**"),385 )386 387iface.launch(share=True)