Update app.py

app.py

@@ -3,23 +3,30 @@ from huggingface_hub import InferenceClient
 import PyPDF2
 from sentence_transformers import SentenceTransformer
 import numpy as np
-from sklearn.metrics.pairwise import cosine_similarity
+import faiss
 import os
 from typing import List, Tuple
+from rank_bm25 import BM25Okapi
 
-# Inicialização do cliente de inferência e modelo de embeddings
+# Inicialização do cliente e modelos
 client = InferenceClient("google/gemma-3-27b-it")
 embedder = SentenceTransformer('all-MiniLM-L6-v2')
 
-# Classe para gerenciar o conhecimento dos PDFs
-class PDFKnowledgeBase:
+# Classe para gerenciar a base de conhecimento
+class AdvancedPDFKnowledgeBase:
     def __init__(self):
-        self.documents = []
-        self.embeddings = None
+        self.chunks = []
+        self.chunk_embeddings = None
+        self.index = None
+        self.bm25 = None
         
+    def _split_into_chunks(self, text: str, chunk_size: int = 500) -> List[str]:
+        words = text.split()
+        return [' '.join(words[i:i + chunk_size]) 
+                for i in range(0, len(words), chunk_size)]
+    
     def load_pdfs(self, pdf_directory: str):
-        """Carrega todos os PDFs de um diretório"""
-        self.documents = []
+        self.chunks = []
         for filename in os.listdir(pdf_directory):
             if filename.endswith('.pdf'):
                 pdf_path = os.path.join(pdf_directory, filename)
@@ -28,37 +35,51 @@ class PDFKnowledgeBase:
                     text = ""
                     for page in pdf_reader.pages:
                         text += page.extract_text() + "\n"
-                    self.documents.append({
-                        'filename': filename,
-                        'content': text
-                    })
+                    chunks = self._split_into_chunks(text)
+                    for chunk in chunks:
+                        self.chunks.append({
+                            'filename': filename,
+                            'content': chunk
+                        })
+        
+        if not self.chunks:
+            return "Nenhum PDF encontrado."
         
-        # Gera embeddings para todos os documentos
-        contents = [doc['content'] for doc in self.documents]
-        self.embeddings = embedder.encode(contents, convert_to_numpy=True)
+        contents = [chunk['content'] for chunk in self.chunks]
+        self.chunk_embeddings = embedder.encode(contents, convert_to_numpy=True)
+        dimension = self.chunk_embeddings.shape[1]
+        self.index = faiss.IndexFlatL2(dimension)
+        self.index.add(self.chunk_embeddings)
+        tokenized_chunks = [chunk['content'].split() for chunk in self.chunks]
+        self.bm25 = BM25Okapi(tokenized_chunks)
+        return f"Carregados {len(self.chunks)} chunks de {len(set(c['filename'] for c in self.chunks))} PDFs."
     
-    def get_relevant_context(self, query: str, k: int = 3) -> str:
-        """Recupera os k documentos mais relevantes para a query"""
-        if self.embeddings is None or len(self.documents) == 0:
+    def get_relevant_context(self, query: str, k: int = 5, rerank_k: int = 3) -> str:
+        if self.index is None or not self.chunks:
             return "Nenhum documento carregado ainda."
             
-        query_embedding = embedder.encode(query, convert_to_numpy=True)
-        similarities = cosine_similarity([query_embedding], self.embeddings)[0]
+        query_embedding = embedder.encode([query], convert_to_numpy=True)
+        distances, indices = self.index.search(query_embedding, k)
+        candidates = [self.chunks[idx] for idx in indices[0]]
         
-        # Obtém os índices dos k documentos mais similares
-        top_k_indices = np.argsort(similarities)[-k:][::-1]
+        tokenized_query = query.split()
+        bm25_scores = self.bm25.get_scores(tokenized_query)
+        candidate_scores = [(candidates[i], bm25_scores[indices[0][i]]) 
+                          for i in range(len(candidates))]
+        candidate_scores.sort(key=lambda x: x[1], reverse=True)
         
-        # Constrói o contexto relevante
+        top_chunks = candidate_scores[:rerank_k]
         context = ""
-        for idx in top_k_indices:
-            context += f"Documento: {self.documents[idx]['filename']}\n"
-            context += f"Trecho: {self.documents[idx]['content'][:500]}...\n\n"
-        
+        for chunk, score in top_chunks:
+            context += f"**Documento**: {chunk['filename']}\n"
+            context += f"**Trecho**: {chunk['content'][:500]}...\n"
+            context += f"**Score BM25**: {score:.2f}\n\n"
         return context
 
 # Inicializa a base de conhecimento
-knowledge_base = PDFKnowledgeBase()
+knowledge_base = AdvancedPDFKnowledgeBase()
 
+# Função principal de resposta
 def respond(
     message: str,
     history: List[Tuple[str, str]],
@@ -66,30 +87,27 @@ def respond(
     max_tokens: int,
     temperature: float,
     top_p: float,
-    pdf_directory: str
+    pdf_directory: str,
+    k_initial: int,
+    k_final: int
 ):
-    # Carrega os PDFs se ainda não foram carregados
-    if not knowledge_base.documents:
-        knowledge_base.load_pdfs(pdf_directory)
+    if not knowledge_base.chunks:
+        yield "Por favor, carregue os PDFs primeiro.", "", ""
+        return
     
-    # Obtém contexto relevante da base de conhecimento
-    context = knowledge_base.get_relevant_context(message)
+    context = knowledge_base.get_relevant_context(message, k_initial, k_final)
     
-    # Constrói o prompt com o contexto RAG
-    rag_prompt = f"""Você é Grok 3, criado por xAI. Use o seguinte contexto dos documentos para responder à pergunta:
+    rag_prompt = f"""Você é Grok 3, criado por xAI. Use o contexto dos documentos para responder:
 
 {context}
 
-Pergunta do usuário: {message}
-
-Responda de forma clara e precisa, utilizando o contexto quando relevante."""
+Pergunta: {message}
 
-    messages = [
-        {"role": "system", "content": system_message},
-        {"role": "user", "content": rag_prompt}
-    ]
+Responda com base no contexto quando relevante."""
+    
+    messages = [{"role": "system", "content": system_message},
+                {"role": "user", "content": rag_prompt}]
     
-    # Adiciona histórico se existir
     for user_msg, assistant_msg in history:
         if user_msg:
             messages.append({"role": "user", "content": user_msg})
@@ -106,27 +124,73 @@ Responda de forma clara e precisa, utilizando o contexto quando relevante."""
     ):
         token = message_chunk.choices[0].delta.content
         response += token
-        yield response
+        yield response, context, ""
+
+# Função para carregar PDFs
+def load_pdfs(pdf_directory: str):
+    status = knowledge_base.load_pdfs(pdf_directory)
+    return status
+
+# Interface Gradio personalizada
+with gr.Blocks(title="RAG Avançado com PDFs") as demo:
+    gr.Markdown("# Chatbot RAG com PDFs")
+    gr.Markdown("Carregue PDFs e faça perguntas baseadas nos documentos.")
+    
+    with gr.Row():
+        with gr.Column(scale=2):
+            chatbot = gr.Chatbot(label="Conversa")
+            msg = gr.Textbox(label="Sua pergunta", placeholder="Digite sua pergunta aqui...")
+            submit_btn = gr.Button("Enviar")
+        
+        with gr.Column(scale=1):
+            context_box = gr.Markdown(label="Contexto Recuperado", value="Contexto aparecerá aqui após a pergunta.")
+    
+    with gr.Accordion("Configurações", open=False):
+        with gr.Row():
+            with gr.Column():
+                pdf_dir = gr.Textbox(label="Diretório dos PDFs", value="./pdfs")
+                load_btn = gr.Button("Carregar PDFs")
+                load_status = gr.Textbox(label="Status do Carregamento", interactive=False)
+            
+            with gr.Column():
+                system_msg = gr.Textbox(
+                    label="Mensagem do Sistema",
+                    value="Você é um assistente útil que responde com base em documentos PDF."
+                )
+                max_tokens = gr.Slider(1, 2048, value=512, step=1, label="Max Tokens")
+                temperature = gr.Slider(0.1, 4.0, value=0.7, step=0.1, label="Temperature")
+                top_p = gr.Slider(0.1, 1.0, value=0.95, step=0.05, label="Top-p")
+        
+        with gr.Row():
+            k_initial = gr.Slider(1, 20, value=5, step=1, label="Candidatos Iniciais (FAISS)")
+            k_final = gr.Slider(1, 10, value=3, step=1, label="Resultados Finais (BM25)")
+
+    # Função para atualizar o chat
+    def submit_message(message, history, system_message, max_tokens, temperature, top_p, pdf_directory, k_initial, k_final):
+        history = history or []
+        for response, context, _ in respond(message, history, system_message, max_tokens, temperature, top_p, pdf_directory, k_initial, k_final):
+            history.append((message, response))
+            yield history, context, ""
+        yield history, context, ""
 
-# Interface do Gradio
-demo = gr.ChatInterface(
-    respond,
-    additional_inputs=[
-        gr.Textbox(value="Você é um assistente útil que responde com base em documentos PDF.", 
-                  label="System message"),
-        gr.Slider(minimum=1, maximum=2048, value=512, step=1, label="Max new tokens"),
-        gr.Slider(minimum=0.1, maximum=4.0, value=0.7, step=0.1, label="Temperature"),
-        gr.Slider(minimum=0.1, maximum=1.0, value=0.95, step=0.05, 
-                 label="Top-p (nucleus sampling)"),
-        gr.Textbox(value="./pdfs", label="Diretório dos PDFs"),
-    ],
-    title="RAG Chatbot com PDFs",
-    description="Faça perguntas e obtenha respostas baseadas em documentos PDF carregados."
-)
+    # Conexões de eventos
+    submit_btn.click(
+        submit_message,
+        inputs=[msg, chatbot, system_msg, max_tokens, temperature, top_p, pdf_dir, k_initial, k_final],
+        outputs=[chatbot, context_box, msg]
+    )
+    msg.submit(
+        submit_message,
+        inputs=[msg, chatbot, system_msg, max_tokens, temperature, top_p, pdf_dir, k_initial, k_final],
+        outputs=[chatbot, context_box, msg]
+    )
+    load_btn.click(
+        load_pdfs,
+        inputs=[pdf_dir],
+        outputs=[load_status]
+    )
 
 if __name__ == "__main__":
-    # Crie um diretório 'pdfs' e coloque seus PDFs lá
     if not os.path.exists("./pdfs"):
         os.makedirs("./pdfs")
-        
     demo.launch()