Update app.py

app.py

@@ -1,210 +1,210 @@
-# import gradio as gr
-# from langchain_mistralai.chat_models import ChatMistralAI
-# from langchain.prompts import ChatPromptTemplate
-# from langchain_deepseek import ChatDeepSeek
-# from langchain_google_genai import ChatGoogleGenerativeAI
-# import os
-# from pathlib import Path
-# import json
-# import faiss
-# import numpy as np
-# from langchain.schema import Document
-# import pickle
-# import re
-# import requests
-# from functools import lru_cache
-# import torch
-# from sentence_transformers import SentenceTransformer
-# from sentence_transformers.cross_encoder import CrossEncoder
-# import threading
-# from queue import Queue
-# import concurrent.futures
-# from typing import Generator, Tuple, Iterator
-# import time
-
-# class OptimizedRAGLoader:
-#     def __init__(self,
-#                  docs_folder: str = "./docs",
-#                  splits_folder: str = "./splits",
-#                  index_folder: str = "./index"):
+import gradio as gr
+from langchain_mistralai.chat_models import ChatMistralAI
+from langchain.prompts import ChatPromptTemplate
+from langchain_deepseek import ChatDeepSeek
+from langchain_google_genai import ChatGoogleGenerativeAI
+import os
+from pathlib import Path
+import json
+import faiss
+import numpy as np
+from langchain.schema import Document
+import pickle
+import re
+import requests
+from functools import lru_cache
+import torch
+from sentence_transformers import SentenceTransformer
+from sentence_transformers.cross_encoder import CrossEncoder
+import threading
+from queue import Queue
+import concurrent.futures
+from typing import Generator, Tuple, Iterator
+import time
+
+class OptimizedRAGLoader:
+    def __init__(self,
+                 docs_folder: str = "./docs",
+                 splits_folder: str = "./splits",
+                 index_folder: str = "./index"):
         
-#         self.docs_folder = Path(docs_folder)
-#         self.splits_folder = Path(splits_folder)
-#         self.index_folder = Path(index_folder)
+        self.docs_folder = Path(docs_folder)
+        self.splits_folder = Path(splits_folder)
+        self.index_folder = Path(index_folder)
         
-#         # Create folders if they don't exist
-#         for folder in [self.splits_folder, self.index_folder]:
-#             folder.mkdir(parents=True, exist_ok=True)
+        # Create folders if they don't exist
+        for folder in [self.splits_folder, self.index_folder]:
+            folder.mkdir(parents=True, exist_ok=True)
             
-#         # File paths
-#         self.splits_path = self.splits_folder / "splits.json"
-#         self.index_path = self.index_folder / "faiss.index"
-#         self.documents_path = self.index_folder / "documents.pkl"
+        # File paths
+        self.splits_path = self.splits_folder / "splits.json"
+        self.index_path = self.index_folder / "faiss.index"
+        self.documents_path = self.index_folder / "documents.pkl"
         
-#         # Initialize components
-#         self.index = None
-#         self.indexed_documents = None
+        # Initialize components
+        self.index = None
+        self.indexed_documents = None
         
-#         # Initialize encoder model
-#         self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-#         self.encoder = SentenceTransformer("intfloat/multilingual-e5-large")
-#         self.encoder.to(self.device)
-#         self.reranker = model = CrossEncoder("cross-encoder/mmarco-mMiniLMv2-L12-H384-v1",trust_remote_code=True)
+        # Initialize encoder model
+        self.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
+        self.encoder = SentenceTransformer("intfloat/multilingual-e5-large")
+        self.encoder.to(self.device)
+        self.reranker = model = CrossEncoder("cross-encoder/mmarco-mMiniLMv2-L12-H384-v1",trust_remote_code=True)
         
-#         # Initialize thread pool
-#         self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)
+        # Initialize thread pool
+        self.executor = concurrent.futures.ThreadPoolExecutor(max_workers=4)
         
-#         # Initialize response cache
-#         self.response_cache = {}
+        # Initialize response cache
+        self.response_cache = {}
         
-#     @lru_cache(maxsize=1000)
-#     def encode(self, text: str):
-#         """Cached encoding function"""
-#         with torch.no_grad():
-#             embeddings = self.encoder.encode(
-#                 text,
-#                 convert_to_numpy=True,
-#                 normalize_embeddings=True
-#             )
-#         return embeddings
+    @lru_cache(maxsize=1000)
+    def encode(self, text: str):
+        """Cached encoding function"""
+        with torch.no_grad():
+            embeddings = self.encoder.encode(
+                text,
+                convert_to_numpy=True,
+                normalize_embeddings=True
+            )
+        return embeddings
     
-#     def batch_encode(self, texts: list):
-#         """Batch encoding for multiple texts"""
-#         with torch.no_grad():
-#             embeddings = self.encoder.encode(
-#                 texts,
-#                 batch_size=32,
-#                 convert_to_numpy=True,
-#                 normalize_embeddings=True,
-#                 show_progress_bar=False
-#             )
-#         return embeddings
-
-#     def load_and_split_texts(self):
-#         if self._splits_exist():
-#             return self._load_existing_splits()
+    def batch_encode(self, texts: list):
+        """Batch encoding for multiple texts"""
+        with torch.no_grad():
+            embeddings = self.encoder.encode(
+                texts,
+                batch_size=32,
+                convert_to_numpy=True,
+                normalize_embeddings=True,
+                show_progress_bar=False
+            )
+        return embeddings
+
+    def load_and_split_texts(self):
+        if self._splits_exist():
+            return self._load_existing_splits()
             
-#         documents = []
-#         futures = []
+        documents = []
+        futures = []
         
-#         for file_path in self.docs_folder.glob("*.txt"):
-#             future = self.executor.submit(self._process_file, file_path)
-#             futures.append(future)
+        for file_path in self.docs_folder.glob("*.txt"):
+            future = self.executor.submit(self._process_file, file_path)
+            futures.append(future)
             
-#         for future in concurrent.futures.as_completed(futures):
-#             documents.extend(future.result())
+        for future in concurrent.futures.as_completed(futures):
+            documents.extend(future.result())
             
-#         self._save_splits(documents)
-#         return documents
+        self._save_splits(documents)
+        return documents
     
-#     def _process_file(self, file_path):
-#         with open(file_path, 'r', encoding='utf-8') as file:
-#             text = file.read()
-#             chunks = [s.strip() for s in re.split(r'(?<=[.!?])\s+', text) if s.strip()]
+    def _process_file(self, file_path):
+        with open(file_path, 'r', encoding='utf-8') as file:
+            text = file.read()
+            chunks = [s.strip() for s in re.split(r'(?<=[.!?])\s+', text) if s.strip()]
             
-#             return [
-#                 Document(
-#                     page_content=chunk,
-#                     metadata={
-#                         'source': file_path.name,
-#                         'chunk_id': i,
-#                         'total_chunks': len(chunks)
-#                     }
-#                 )
-#                 for i, chunk in enumerate(chunks)
-#             ]
-
-#     def load_index(self) -> bool:
-#         """
-#         Charge l'index FAISS et les documents associés s'ils existent
-
-#         Returns:
-#             bool: True si l'index a été chargé, False sinon
-#         """
-#         if not self._index_exists():
-#             print("Aucun index trouvé.")
-#             return False
-
-#         print("Chargement de l'index existant...")
-#         try:
-#             # Charger l'index FAISS
-#             self.index = faiss.read_index(str(self.index_path))
-
-#             # Charger les documents associés
-#             with open(self.documents_path, 'rb') as f:
-#                 self.indexed_documents = pickle.load(f)
-
-#             print(f"Index chargé avec {self.index.ntotal} vecteurs")
-#             return True
-
-#         except Exception as e:
-#             print(f"Erreur lors du chargement de l'index: {e}")
-#             return False
-
-#     def create_index(self, documents=None):
-#         if documents is None:
-#             documents = self.load_and_split_texts()
+            return [
+                Document(
+                    page_content=chunk,
+                    metadata={
+                        'source': file_path.name,
+                        'chunk_id': i,
+                        'total_chunks': len(chunks)
+                    }
+                )
+                for i, chunk in enumerate(chunks)
+            ]
+
+    def load_index(self) -> bool:
+        """
+        Charge l'index FAISS et les documents associés s'ils existent
+
+        Returns:
+            bool: True si l'index a été chargé, False sinon
+        """
+        if not self._index_exists():
+            print("Aucun index trouvé.")
+            return False
+
+        print("Chargement de l'index existant...")
+        try:
+            # Charger l'index FAISS
+            self.index = faiss.read_index(str(self.index_path))
+
+            # Charger les documents associés
+            with open(self.documents_path, 'rb') as f:
+                self.indexed_documents = pickle.load(f)
+
+            print(f"Index chargé avec {self.index.ntotal} vecteurs")
+            return True
+
+        except Exception as e:
+            print(f"Erreur lors du chargement de l'index: {e}")
+            return False
+
+    def create_index(self, documents=None):
+        if documents is None:
+            documents = self.load_and_split_texts()
             
-#         if not documents:
-#             return False
+        if not documents:
+            return False
             
-#         texts = [doc.page_content for doc in documents]
-#         embeddings = self.batch_encode(texts)
+        texts = [doc.page_content for doc in documents]
+        embeddings = self.batch_encode(texts)
         
-#         dimension = embeddings.shape[1]
-#         self.index = faiss.IndexFlatL2(dimension)
+        dimension = embeddings.shape[1]
+        self.index = faiss.IndexFlatL2(dimension)
         
-#         if torch.cuda.is_available():
-#             # Use GPU for FAISS if available
-#             res = faiss.StandardGpuResources()
-#             self.index = faiss.index_cpu_to_gpu(res, 0, self.index)
+        if torch.cuda.is_available():
+            # Use GPU for FAISS if available
+            res = faiss.StandardGpuResources()
+            self.index = faiss.index_cpu_to_gpu(res, 0, self.index)
             
-#         self.index.add(np.array(embeddings).astype('float32'))
-#         self.indexed_documents = documents
+        self.index.add(np.array(embeddings).astype('float32'))
+        self.indexed_documents = documents
         
-#         # Save index and documents
-#         cpu_index = faiss.index_gpu_to_cpu(self.index) if torch.cuda.is_available() else self.index
-#         faiss.write_index(cpu_index, str(self.index_path))
+        # Save index and documents
+        cpu_index = faiss.index_gpu_to_cpu(self.index) if torch.cuda.is_available() else self.index
+        faiss.write_index(cpu_index, str(self.index_path))
         
-#         with open(self.documents_path, 'wb') as f:
-#             pickle.dump(documents, f)
+        with open(self.documents_path, 'wb') as f:
+            pickle.dump(documents, f)
             
-#         return True
+        return True
 
-#     def _index_exists(self) -> bool:
-#         """Vérifie si l'index et les documents associés existent"""
-#         return self.index_path.exists() and self.documents_path.exists()
+    def _index_exists(self) -> bool:
+        """Vérifie si l'index et les documents associés existent"""
+        return self.index_path.exists() and self.documents_path.exists()
 
-#     def get_retriever(self, k: int = 10):
-#         if self.index is None:
-#             if not self.load_index():
-#                 if not self.create_index():
-#                     raise ValueError("Unable to load or create index")
+    def get_retriever(self, k: int = 10):
+        if self.index is None:
+            if not self.load_index():
+                if not self.create_index():
+                    raise ValueError("Unable to load or create index")
 
-#         def retriever_function(query: str) -> list:
-#             # Check cache first
-#             cache_key = f"{query}_{k}"
-#             if cache_key in self.response_cache:
-#                 return self.response_cache[cache_key]
+        def retriever_function(query: str) -> list:
+            # Check cache first
+            cache_key = f"{query}_{k}"
+            if cache_key in self.response_cache:
+                return self.response_cache[cache_key]
 
-#             query_embedding = self.encode(query)
+            query_embedding = self.encode(query)
             
-#             distances, indices = self.index.search(
-#                 np.array([query_embedding]).astype('float32'),
-#                 k
-#             )
+            distances, indices = self.index.search(
+                np.array([query_embedding]).astype('float32'),
+                k
+            )
             
-#             results = [
-#                 self.indexed_documents[idx]
-#                 for idx in indices[0]
-#                 if idx != -1
-#             ]
+            results = [
+                self.indexed_documents[idx]
+                for idx in indices[0]
+                if idx != -1
+            ]
             
-#             # Cache the results
-#             self.response_cache[cache_key] = results
-#             return results
+            # Cache the results
+            self.response_cache[cache_key] = results
+            return results
             
-#         return retriever_function
+        return retriever_function
 
 # # # Initialize components
 # # mistral_api_key = os.getenv("mistral_api_key")