Update app.py (#2)

- Update app.py (055befae02ce69c5ee13fd4369cd892be3e27c1e)


Co-authored-by: Gabriel Silva Rodrigues <[email protected]>

app.py

@@ -20,7 +20,7 @@ from llama_index.core.storage.chat_store import SimpleChatStore
 from llama_index.core.memory import ChatMemoryBuffer
 from llama_index.core.query_engine import RetrieverQueryEngine
 from llama_index.core.chat_engine import CondensePlusContextChatEngine
-from llama_index.retrievers.bm25 import BM25Retriever
+#from llama_index.retrievers.bm25 import BM25Retriever
 from llama_index.core.retrievers import QueryFusionRetriever
 from llama_index.vector_stores.chroma import ChromaVectorStore
 from llama_index.core import VectorStoreIndex
@@ -29,6 +29,238 @@ from llama_index.core import VectorStoreIndex
 # from llama_index.embeddings.huggingface import HuggingFaceEmbedding
 import chromadb
 
+###############################################################################
+#                           MONKEY PATCH EM bm25s                             #
+###############################################################################
+import bm25s
+
+# Guardamos a referência da função original
+orig_find_newline_positions = bm25s.utils.corpus.find_newline_positions
+
+def patched_find_newline_positions(path, show_progress=True, leave_progress=True):
+    """
+    Versão 'gambiarra' da função original, forçando uso de encoding='utf-8'
+    e ignorando erros de decodificação. Assim, evitamos UnicodeDecodeError
+    mesmo que o arquivo contenha caracteres fora da faixa UTF-8.
+    
+    (Esta referência é real, baseada em ajustes de leitura de arquivos do Python.)
+    """
+    path = str(path)
+    indexes = []
+
+    with open(path, "r", encoding="utf-8", errors="ignore") as f:
+        indexes.append(f.tell())
+        file_size = os.path.getsize(path)
+
+        try:
+            from tqdm.auto import tqdm
+            pbar = tqdm(
+                total=file_size,
+                desc="Finding newlines for mmindex",
+                unit="B",
+                unit_scale=True,
+                leave=leave_progress,
+                disable=not show_progress,
+            )
+        except ImportError:
+            pbar = None
+
+        while True:
+            line = f.readline()
+            if not line:
+                break
+            t = f.tell()
+            indexes.append(t)
+            if pbar is not None:
+                pbar.update(t - indexes[-2])
+
+        if pbar is not None:
+            pbar.close()
+
+    return indexes[:-1]
+
+# Aplicamos nosso patch
+bm25s.utils.corpus.find_newline_positions = patched_find_newline_positions
+###############################################################################
+#                   CLASSE BM25Retriever (AJUSTADA PARA ENCODING)             #
+###############################################################################
+import json
+import Stemmer
+
+from llama_index.core.base.base_retriever import BaseRetriever
+from llama_index.core.callbacks.base import CallbackManager
+from llama_index.core.constants import DEFAULT_SIMILARITY_TOP_K
+from llama_index.core.schema import (
+    BaseNode,
+    IndexNode,
+    NodeWithScore,
+    QueryBundle,
+    MetadataMode,
+)
+from llama_index.core.vector_stores.utils import (
+    node_to_metadata_dict,
+    metadata_dict_to_node,
+)
+from typing import cast
+
+logger = logging.getLogger(__name__)
+
+DEFAULT_PERSIST_ARGS = {"similarity_top_k": "similarity_top_k", "_verbose": "verbose"}
+DEFAULT_PERSIST_FILENAME = "retriever.json"
+
+
+class BM25Retriever(BaseRetriever):
+    """
+    Implementação customizada do algoritmo BM25 com a lib bm25s, incluindo um
+    'monkey patch' para contornar problemas de decodificação de caracteres.
+    """
+
+    def __init__(
+        self,
+        nodes: Optional[List[BaseNode]] = None,
+        stemmer: Optional[Stemmer.Stemmer] = None,
+        language: str = "en",
+        existing_bm25: Optional[bm25s.BM25] = None,
+        similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
+        callback_manager: Optional[CallbackManager] = None,
+        objects: Optional[List[IndexNode]] = None,
+        object_map: Optional[dict] = None,
+        verbose: bool = False,
+    ) -> None:
+        self.stemmer = stemmer or Stemmer.Stemmer("english")
+        self.similarity_top_k = similarity_top_k
+
+        if existing_bm25 is not None:
+            # Usa instância BM25 existente
+            self.bm25 = existing_bm25
+            self.corpus = existing_bm25.corpus
+        else:
+            # Cria uma nova instância BM25 a partir de 'nodes'
+            if nodes is None:
+                raise ValueError("É preciso fornecer 'nodes' ou um 'existing_bm25'.")
+
+            self.corpus = [node_to_metadata_dict(node) for node in nodes]
+            corpus_tokens = bm25s.tokenize(
+                [node.get_content(metadata_mode=MetadataMode.EMBED) for node in nodes],
+                stopwords=language,
+                stemmer=self.stemmer,
+                show_progress=verbose,
+            )
+            self.bm25 = bm25s.BM25()
+            self.bm25.index(corpus_tokens, show_progress=verbose)
+
+        super().__init__(
+            callback_manager=callback_manager,
+            object_map=object_map,
+            objects=objects,
+            verbose=verbose,
+        )
+
+    @classmethod
+    def from_defaults(
+        cls,
+        index: Optional[VectorStoreIndex] = None,
+        nodes: Optional[List[BaseNode]] = None,
+        docstore: Optional["BaseDocumentStore"] = None,
+        stemmer: Optional[Stemmer.Stemmer] = None,
+        language: str = "en",
+        similarity_top_k: int = DEFAULT_SIMILARITY_TOP_K,
+        verbose: bool = False,
+        tokenizer: Optional[Any] = None,
+    ) -> "BM25Retriever":
+        if tokenizer is not None:
+            logger.warning(
+                "O parâmetro 'tokenizer' foi descontinuado e será removido "
+                "no futuro. Use um Stemmer do PyStemmer para melhor controle."
+            )
+
+        if sum(bool(val) for val in [index, nodes, docstore]) != 1:
+            raise ValueError("Passe exatamente um entre 'index', 'nodes' ou 'docstore'.")
+
+        if index is not None:
+            docstore = index.docstore
+
+        if docstore is not None:
+            nodes = cast(List[BaseNode], list(docstore.docs.values()))
+
+        assert nodes is not None, (
+            "Não foi possível determinar os nodes. Verifique seus parâmetros."
+        )
+
+        return cls(
+            nodes=nodes,
+            stemmer=stemmer,
+            language=language,
+            similarity_top_k=similarity_top_k,
+            verbose=verbose,
+        )
+
+    def get_persist_args(self) -> Dict[str, Any]:
+        """Dicionário com os parâmetros de persistência a serem salvos."""
+        return {
+            DEFAULT_PERSIST_ARGS[key]: getattr(self, key)
+            for key in DEFAULT_PERSIST_ARGS
+            if hasattr(self, key)
+        }
+
+    def persist(self, path: str, **kwargs: Any) -> None:
+        """
+        Persiste o retriever em um diretório, incluindo
+        a estrutura do BM25 e o corpus em JSON.
+        """
+        self.bm25.save(path, corpus=self.corpus, **kwargs)
+        with open(
+            os.path.join(path, DEFAULT_PERSIST_FILENAME),
+            "wt",
+            encoding="utf-8",
+            errors="ignore",
+        ) as f:
+            json.dump(self.get_persist_args(), f, indent=2, ensure_ascii=False)
+
+    @classmethod
+    def from_persist_dir(cls, path: str, **kwargs: Any) -> "BM25Retriever":
+        """
+        Carrega o retriever de um diretório, incluindo o BM25 e o corpus.
+        Devido ao nosso patch, ignoramos qualquer erro de decodificação
+        que eventualmente apareça.
+        """
+        bm25_obj = bm25s.BM25.load(path, load_corpus=True, **kwargs)
+        with open(
+            os.path.join(path, DEFAULT_PERSIST_FILENAME),
+            "rt",
+            encoding="utf-8",
+            errors="ignore",
+        ) as f:
+            retriever_data = json.load(f)
+
+        return cls(existing_bm25=bm25_obj, **retriever_data)
+
+    def _retrieve(self, query_bundle: QueryBundle) -> List[NodeWithScore]:
+        """Recupera nós relevantes a partir do BM25."""
+        query = query_bundle.query_str
+        tokenized_query = bm25s.tokenize(
+            query, stemmer=self.stemmer, show_progress=self._verbose
+        )
+        indexes, scores = self.bm25.retrieve(
+            tokenized_query, k=self.similarity_top_k, show_progress=self._verbose
+        )
+
+        # bm25s retorna lista de listas, pois suporta batched queries
+        indexes = indexes[0]
+        scores = scores[0]
+
+        nodes: List[NodeWithScore] = []
+        for idx, score in zip(indexes, scores):
+            if isinstance(idx, dict):
+                node = metadata_dict_to_node(idx)
+            else:
+                node_dict = self.corpus[int(idx)]
+                node = metadata_dict_to_node(node_dict)
+
+            nodes.append(NodeWithScore(node=node, score=float(score)))
+
+        return nodes
+
 #Configuração da imagem da aba
 im = Image.open("pngegg.png")
 st.set_page_config(page_title = "Chatbot Carômetro", page_icon=im, layout = "wide")
@@ -38,8 +270,6 @@ os.makedirs("bm25_retriever", exist_ok=True)
 os.makedirs("chat_store", exist_ok=True)
 os.makedirs("chroma_db", exist_ok=True)
 os.makedirs("documentos", exist_ok=True)
-os.makedirs("curadoria", exist_ok=True)
-os.makedirs("chroma_db_curadoria", exist_ok=True)
 
 # Configuração do Streamlit
 st.sidebar.title("Configuração de LLM")
@@ -120,9 +350,7 @@ logging.getLogger().addHandler(logging.StreamHandler(stream=sys.stdout))
 chat_store_path = os.path.join("chat_store", "chat_store.json")
 documents_path = os.path.join("documentos")
 chroma_storage_path = os.path.join("chroma_db")  # Diretório para persistência do Chroma
-chroma_storage_path_curadoria = os.path.join("chroma_db_curadoria")  # Diretório para 'curadoria'
 bm25_persist_path = os.path.join("bm25_retriever")
-curadoria_path = os.path.join("curadoria")
 
 # Classe CSV Customizada (novo código)
 class CustomPandasCSVReader:
@@ -192,7 +420,7 @@ with open(credentials_path, 'w') as credentials_file:
 
 with open(token_path, 'w') as credentials_file:
     credentials_file.write(token_json)
-
+    
 google_drive_reader = GoogleDriveReader(credentials_path=credentials_path)
 google_drive_reader._creds = google_drive_reader._get_credentials()
 
@@ -222,8 +450,6 @@ def download_original_files_from_folder(greader: GoogleDriveReader, pasta_docume
 
 #DADOS/QA_database/Documentos CSV/documentos
 pasta_documentos_drive = "1xVzo8s1D0blzR5ZB3m5k4dVWHuRmKUu-"
-#DADOS/QA_database/Documentos CSV/curadoria
-pasta_curadoria_drive = "1LRrdOkZy9p0FA3MQAyz-Ssj3ktKTWAwE"
 
 # Verifica e baixa arquivos se necessário (novo código)
 if not are_docs_downloaded(documents_path):
@@ -232,18 +458,14 @@ if not are_docs_downloaded(documents_path):
 else:
     logging.info("'documentos' já contém arquivos, ignorando download.")
 
-if not are_docs_downloaded(curadoria_path):
-    logging.info("Baixando arquivos originais do Drive para 'curadoria'...")
-    download_original_files_from_folder(google_drive_reader, pasta_curadoria_drive, curadoria_path)
-else:
-    logging.info("'curadoria' já contém arquivos, ignorando download.")
-
 # Configuração de leitura de documentos
 file_extractor = {".csv": CustomPandasCSVReader()}
 documents = SimpleDirectoryReader(
     input_dir=documents_path,
     file_extractor=file_extractor,
-    filename_as_id=True
+    filename_as_id=True,
+    recursive=True
+    #Recursive caso tenha varias pastas no drive
 ).load_data()
 
 documents = clean_documents(documents)
@@ -266,7 +488,7 @@ if os.path.exists(chroma_storage_path):
     index = VectorStoreIndex.from_vector_store(vector_store)
 else:
     splitter = LangchainNodeParser(
-        RecursiveCharacterTextSplitter(chunk_size=512, chunk_overlap=64)
+        RecursiveCharacterTextSplitter(chunk_size=1024, chunk_overlap=128)
     )
     index = VectorStoreIndex.from_documents(
         documents,
@@ -287,45 +509,11 @@ else:
     os.makedirs(bm25_persist_path, exist_ok=True)
     bm25_retriever.persist(bm25_persist_path)
 
-#Adicionado documentos na pasta curadoria, foi setado para 1200 o chunk pra receber pergunta, contexto e resposta
-curadoria_documents = SimpleDirectoryReader(
-    input_dir=curadoria_path,
-    file_extractor=file_extractor,
-    filename_as_id=True
-).load_data()
-
-curadoria_documents = clean_documents(curadoria_documents)
-curadoria_docstore = SimpleDocumentStore()
-curadoria_docstore.add_documents(curadoria_documents)
-
-db_curadoria = chromadb.PersistentClient(path=chroma_storage_path_curadoria)
-chroma_collection_curadoria = db_curadoria.get_or_create_collection("dense_vectors_curadoria")
-vector_store_curadoria = ChromaVectorStore(chroma_collection=chroma_collection_curadoria)
-
-# Configuração do StorageContext para 'curadoria'
-storage_context_curadoria = StorageContext.from_defaults(
-    docstore=curadoria_docstore, vector_store=vector_store_curadoria
-)
-
-# Criação/Recarregamento do índice com embeddings para 'curadoria'
-if os.path.exists(chroma_storage_path_curadoria):
-    curadoria_index = VectorStoreIndex.from_vector_store(vector_store_curadoria)
-else:
-    curadoria_splitter = LangchainNodeParser(
-        RecursiveCharacterTextSplitter(chunk_size=1200, chunk_overlap=100)
-    )
-    curadoria_index = VectorStoreIndex.from_documents(
-        curadoria_documents, storage_context=storage_context_curadoria, transformations=[curadoria_splitter]
-    )
-    vector_store_curadoria.persist()
-
-curadoria_retriever = curadoria_index.as_retriever(similarity_top_k=2)
-
 # Combinação de Retrievers (Embeddings + BM25)
 vector_retriever = index.as_retriever(similarity_top_k=2)
 retriever = QueryFusionRetriever(
-    [vector_retriever, bm25_retriever, curadoria_retriever],
-    similarity_top_k=2,
+    [vector_retriever, bm25_retriever],
+    similarity_top_k=3,
     num_queries=0,
     mode="reciprocal_rerank",
     use_async=True,
@@ -397,4 +585,4 @@ if user_input:
 
     # Remover o cursor após a conclusão
     message_placeholder.markdown(assistant_message)
-    st.session_state.chat_history.append(f"assistant: {assistant_message}")
+    st.session_state.chat_history.append(f"assistant: {assistant_message}")