Update app.py

app.py

@@ -4,7 +4,11 @@ from langchain_community.document_loaders import PyPDFLoader
 from langchain.text_splitter import RecursiveCharacterTextSplitter
 from langchain_community.vectorstores import Chroma
 from langchain_community.embeddings import HuggingFaceEmbeddings
+from langchain_community.retrievers import BM25Retriever
+from langchain.retrievers import EnsembleRetriever
 import os
+import re
+from unidecode import unidecode
 
 # CSS para estilização
 css = '''
@@ -16,28 +20,62 @@ footer {visibility: hidden}
 # Inicializar o cliente de inferência
 client = InferenceClient("mistralai/Mistral-7B-Instruct-v0.3")
 
-# Configurar o retriever globalmente
-def initialize_retriever(file_objs):
-    """Carrega documentos PDFs e cria um retriever."""
+# Função de pré-processamento de texto
+def preprocess_text(text):
+    """Pré-processa o texto removendo ruídos e normalizando."""
+    # Remover números de página (ex.: "Página 1", "Page 1 of 10")
+    text = re.sub(r'(Página|Page)\s+\d+(?:\s+of\s+\d+)?', '', text, flags=re.IGNORECASE)
+    # Remover múltiplos espaços e quebras de linha
+    text = re.sub(r'\s+', ' ', text).strip()
+    # Normalizar texto (remover acentos e converter para minúsculas)
+    text = unidecode(text.lower())
+    return text
+
+# Configurar o retriever com pré-processamento e indexação avançada
+def initialize_retriever(file_objs, persist_directory="chroma_db"):
+    """Carrega documentos PDFs, pré-processa e cria um retriever híbrido."""
     if not file_objs:
         return None, "Nenhum documento carregado."
     
-    # Carregar e dividir documentos
+    # Carregar e pré-processar documentos
     documents = []
     for file_obj in file_objs:
         loader = PyPDFLoader(file_obj.name)
-        documents.extend(loader.load())
+        raw_docs = loader.load()
+        for doc in raw_docs:
+            doc.page_content = preprocess_text(doc.page_content)
+            # Adicionar metadados (exemplo: página e origem)
+            doc.metadata.update({"source": os.path.basename(file_obj.name)})
+        documents.extend(raw_docs)
     
     # Dividir em pedaços menores
     text_splitter = RecursiveCharacterTextSplitter(chunk_size=2048, chunk_overlap=128)
     splits = text_splitter.split_documents(documents)
     
-    # Criar embeddings e banco de vetores
+    # Criar embeddings e banco de vetores (Chroma)
     embeddings = HuggingFaceEmbeddings(model_name="sentence-transformers/all-MiniLM-L6-v2")
-    vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings)
-    retriever = vectorstore.as_retriever(search_kwargs={"k": 2})  # Retorna 2 documentos mais relevantes
+    try:
+        # Tentar carregar um banco existente
+        vectorstore = Chroma(persist_directory=persist_directory, embedding_function=embeddings)
+        vectorstore.add_documents(splits)  # Adicionar novos documentos
+    except:
+        # Criar um novo banco se não existir
+        vectorstore = Chroma.from_documents(documents=splits, embedding=embeddings, persist_directory=persist_directory)
     
-    return retriever, "Documentos processados com sucesso!"
+    # Configurar retriever semântico
+    semantic_retriever = vectorstore.as_retriever(search_kwargs={"k": 2})
+    
+    # Configurar retriever lexical (BM25)
+    bm25_retriever = BM25Retriever.from_documents(splits)
+    bm25_retriever.k = 2
+    
+    # Combinar em um retriever híbrido
+    ensemble_retriever = EnsembleRetriever(
+        retrievers=[semantic_retriever, bm25_retriever],
+        weights=[0.6, 0.4]  # Mais peso para busca semântica
+    )
+    
+    return ensemble_retriever, "Documentos processados com sucesso!"
 
 # Formatar o prompt para RAG
 def format_prompt(message, history, retriever=None, system_prompt=None):
@@ -54,9 +92,8 @@ def format_prompt(message, history, retriever=None, system_prompt=None):
     
     # Adicionar contexto recuperado, se houver retriever
     if retriever:
-        # Buscar documentos relevantes
         docs = retriever.get_relevant_documents(message)
-        context = "\n".join([doc.page_content for doc in docs])
+        context = "\n".join([f"[{doc.metadata.get('source', 'Unknown')}, Page {doc.metadata.get('page', 'N/A')}] {doc.page_content}" for doc in docs])
         prompt += f"[CONTEXT] {context} [/CONTEXT]"
     
     # Adicionar a mensagem do usuário
@@ -81,7 +118,7 @@ def generate(
         seed=42,
     )
 
-    # Formatar o prompt com contexto RAG, se disponível
+    # Formatar o prompt com contexto RAG
     formatted_prompt = format_prompt(prompt, history, retriever, system_prompt)
 
     # Gerar resposta em streaming
@@ -113,7 +150,7 @@ def create_demo():
         chat_interface = gr.ChatInterface(
             fn=generate,
             additional_inputs=[
-                gr.State(value=retriever_state),  # Passa o retriever como entrada adicional
+                gr.State(value=retriever_state),
                 gr.Textbox(label="System Prompt", placeholder="Digite um prompt de sistema (opcional)", value=None)
             ],
             title="",
@@ -130,4 +167,4 @@ def create_demo():
 
 # Lançar a aplicação
 demo = create_demo()
-demo.queue().launch(share=False)
+demo.queue().launch(share=False)