Update app.py

app.py

@@ -1,217 +1,227 @@
 import gradio as gr
-from huggingface_hub import InferenceClient
 import PyPDF2
-from sentence_transformers import SentenceTransformer
+from transformers import AutoTokenizer, AutoModel
+from sentence_transformers import util
+import torch
 import numpy as np
-import faiss
-from typing import List, Tuple
-from rank_bm25 import BM25Okapi
+import os
+import re
+from typing import List, Dict, Any
+import requests
 
-# Inicialização do cliente e modelos
-client = InferenceClient("google/gemma-3-27b-it")
-embedder = SentenceTransformer('all-MiniLM-L6-v2')
+# Diretório para armazenar os PDFs
+PDF_DIR = "pdf_data"
+os.makedirs(PDF_DIR, exist_ok=True)
 
-# Classe para gerenciar a base de conhecimento
-class AdvancedPDFKnowledgeBase:
+# Modelo simples para embeddings
+tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+
+# Função para gerar embeddings
+def get_embeddings(texts):
+    inputs = tokenizer(texts, padding=True, truncation=True, return_tensors="pt", max_length=512)
+    with torch.no_grad():
+        outputs = model(**inputs)
+    embeddings = outputs.last_hidden_state.mean(dim=1)
+    return embeddings / embeddings.norm(dim=1, keepdim=True)
+
+# Classe RAG simplificada
+class SimpleRAG:
     def __init__(self):
-        self.chunks = []
-        self.chunk_embeddings = None
-        self.index = None
-        self.bm25 = None
+        self.documents = []
+        self.embeddings = None
         
-    def _split_into_chunks(self, text: str, chunk_size: int = 500) -> List[str]:
-        # Remove linhas duplicadas ou muito semelhantes
-        lines = text.split("\n")
-        unique_lines = []
-        seen = set()
-        for line in lines:
-            if line.strip() and line not in seen:
-                unique_lines.append(line)
-                seen.add(line)
+    def clear(self):
+        self.documents = []
+        self.embeddings = None
         
-        # Junta as linhas únicas em um único texto
-        cleaned_text = "\n".join(unique_lines)
+    def process_text(self, text):
+        # Limpa texto e divide em parágrafos
+        text = re.sub(r'\s+', ' ', text).strip()
+        paragraphs = [p for p in text.split('\n') if len(p) > 50]
+        return paragraphs
         
-        # Divide o texto limpo em chunks
-        words = cleaned_text.split()
-        return [' '.join(words[i:i + chunk_size]) 
-                for i in range(0, len(words), chunk_size)]
-    
-    def load_pdfs(self, pdf_files: List[gr.File]) -> str:
-        self.chunks = []
-        for file in pdf_files:
-            with open(file.name, 'rb') as pdf_file:
-                pdf_reader = PyPDF2.PdfReader(pdf_file)
+    def load_pdf(self, file_obj):
+        try:
+            # Salva o arquivo
+            file_path = os.path.join(PDF_DIR, file_obj.name)
+            with open(file_path, 'wb') as f:
+                f.write(file_obj.read())
+            
+            # Extrai o texto
+            with open(file_path, 'rb') as f:
+                pdf = PyPDF2.PdfReader(f)
                 text = ""
-                for page in pdf_reader.pages:
+                for page in pdf.pages:
                     text += page.extract_text() + "\n"
-                chunks = self._split_into_chunks(text)
-                for chunk in chunks:
-                    self.chunks.append({
-                        'filename': file.name.split('/')[-1],
-                        'content': chunk
-                    })
-        
-        if not self.chunks:
-            return "Nenhum PDF encontrado."
-        
-        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 = 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)
-        distances, indices = self.index.search(query_embedding, k)
-        candidates = [self.chunks[idx] for idx in indices[0]]
+            # Processa o texto em chunks
+            chunks = self.process_text(text)
+            
+            # Adiciona à base de conhecimento
+            filename = os.path.basename(file_path)
+            doc_chunks = [{"source": filename, "content": chunk} for chunk in chunks]
+            self.documents.extend(doc_chunks)
+            
+            # Recalcula embeddings
+            if self.documents:
+                contents = [doc["content"] for doc in self.documents]
+                self.embeddings = get_embeddings(contents)
+            
+            return f"Carregado: {filename} ({len(chunks)} segmentos)"
+        except Exception as e:
+            return f"Erro ao processar PDF: {str(e)}"
+            
+    def search(self, query, top_k=3):
+        if not self.documents or self.embeddings is None:
+            return []
+            
+        # Calcula embedding da query
+        query_embedding = get_embeddings([query])
+        
+        # Calcula similaridade
+        similarities = util.pytorch_cos_sim(query_embedding, self.embeddings)[0]
         
-        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)
+        # Encontra os top_k mais similares
+        top_results = torch.topk(similarities, min(top_k, len(self.documents)))
         
-        top_chunks = candidate_scores[:rerank_k]
-        context = ""
-        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
+        results = []
+        for score, idx in zip(top_results.values, top_results.indices):
+            results.append({
+                "score": score.item(),
+                "document": self.documents[idx]
+            })
+            
+        return results
 
-# Inicializa a base de conhecimento
-knowledge_base = AdvancedPDFKnowledgeBase()
+# Inicializa o RAG
+rag = SimpleRAG()
 
-# Função principal de resposta
-def respond(
-    message: str,
-    history: List[Tuple[str, str]],
-    system_message: str,
-    max_tokens: int,
-    temperature: float,
-    top_p: float,
-    k_initial: int,
-    k_final: int
-):
-    if not knowledge_base.chunks:
-        yield "Por favor, carregue os PDFs primeiro.", "", ""
-        return
-    
-    # Obtém o contexto relevante
-    context = knowledge_base.get_relevant_context(message, k_initial, k_final)
+# Configurações para LLM
+LLM_API_URL = "https://api-inference.huggingface.co/models/mistralai/Mistral-7B-Instruct-v0.2"
+headers = {"Authorization": "Bearer hf_XXXXXXXXXXXXXXXXXXXXXXX"}  # Substitua por sua API key
+
+def query_llm(prompt):
+    payload = {
+        "inputs": prompt,
+        "parameters": {
+            "max_new_tokens": 512,
+            "temperature": 0.7,
+            "top_p": 0.95
+        }
+    }
     
-    # Constrói o prompt RAG
-    rag_prompt = f"""Você é Grok 3, criado por xAI. Use o contexto dos documentos para responder:
-{context}
-Pergunta: {message}
-Responda com base no contexto quando relevante."""
+    try:
+        response = requests.post(LLM_API_URL, headers=headers, json=payload)
+        return response.json()[0]["generated_text"]
+    except Exception as e:
+        return f"Erro ao consultar o LLM: {str(e)}"
+
+# Função para processar a consulta
+def process_query(query, history):
+    # Busca documentos relevantes
+    results = rag.search(query)
     
-    # Inicializa a lista de mensagens
-    messages = [{"role": "system", "content": system_message}]
+    if not results:
+        return "Por favor, carregue alguns PDFs primeiro.", "Nenhum documento disponível."
     
-    # Adiciona mensagens do histórico, garantindo alternância entre "user" e "assistant"
-    for user_msg, assistant_msg in history:
-        messages.append({"role": "user", "content": user_msg})
-        messages.append({"role": "assistant", "content": assistant_msg})
+    # Formata o contexto
+    context = ""
+    for i, result in enumerate(results):
+        context += f"[{i+1}] Fonte: {result['document']['source']}\n"
+        context += f"Trecho: {result['document']['content'][:300]}...\n"
+        context += f"Relevância: {result['score']:.2f}\n\n"
     
-    # Adiciona a nova mensagem do usuário
-    messages.append({"role": "user", "content": rag_prompt})
+    # Constrói o prompt
+    prompt = f"""<s>[INST]Você é um assistente de IA especializado em responder perguntas usando apenas 
+o contexto fornecido. Considere apenas as informações nos documentos abaixo. Se a resposta não 
+puder ser derivada do contexto, diga que não possui informações suficientes.
 
-    response = ""
-    try:
-        for message_chunk in client.chat_completion(
-            messages=messages,
-            max_tokens=max_tokens,
-            stream=True,
-            temperature=temperature,
-            top_p=top_p,
-        ):
-            token = message_chunk.choices[0].delta.content
-            if token:
-                response += token
-                yield response, context, ""
-    except Exception as e:
-        yield f"Erro ao gerar resposta: {str(e)}", context, ""
+CONTEXTO:
+{context}
 
-# Função para carregar PDFs
-def load_pdfs(pdf_files: List[gr.File]):
-    if not pdf_files:
-        return "Nenhum arquivo selecionado."
-    status = knowledge_base.load_pdfs(pdf_files)
-    return status
+PERGUNTA: {query}[/INST]"""
+    
+    # Consulta o modelo
+    response = query_llm(prompt)
+    
+    # Extrai a resposta real (removendo o prompt)
+    actual_response = response.split("[/INST]")[-1].strip()
+    
+    return actual_response, context
 
-# Interface Gradio personalizada
-with gr.Blocks(title="RAG Avançado com PDFs", theme=gr.themes.Soft()) as demo:
+# Interface Gradio
+with gr.Blocks(title="RAG PDF Simplificado") as demo:
     with gr.Row():
         with gr.Column(scale=2):
-            gr.Markdown("# Chatbot RAG com PDFs")
-            gr.Markdown("Arraste e solte seus PDFs abaixo ou clique para selecionar.")
+            gr.Markdown("# RAG PDF Simplificado")
+            
+    with gr.Row():
+        with gr.Column():
+            pdf_upload = gr.File(
+                label="Carregar PDF", 
+                file_types=[".pdf"], 
+                file_count="single"
+            )
+            load_status = gr.Textbox(label="Status", interactive=False)
+            clear_btn = gr.Button("Limpar Base de Conhecimento")
         
-        with gr.Column(scale=1):
-            load_status = gr.Textbox(label="Status do Carregamento", interactive=False)
-
     with gr.Row():
         with gr.Column(scale=2):
-            chatbot = gr.Chatbot(label="Conversa", height=400)
-            msg = gr.Textbox(label="Sua pergunta", placeholder="Digite sua pergunta aqui...")
-            submit_btn = gr.Button("Enviar")
+            chatbot = gr.Chatbot(label="Conversa")
+            query_input = gr.Textbox(
+                label="Sua pergunta", 
+                placeholder="Digite sua pergunta sobre os documentos..."
+            )
+            query_btn = gr.Button("Enviar")
         
         with gr.Column(scale=1):
-            context_box = gr.Markdown(label="Contexto Recuperado")
-
-    with gr.Accordion("Configurações", open=False):
-        with gr.Row():
-            with gr.Column():
-                pdf_upload = gr.File(label="Carregar PDFs", file_types=[".pdf", ".txt"], file_count="multiple", interactive=True)
-                load_btn = gr.Button("Carregar PDFs")
-            
-            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, k_initial, k_final):
+            context_display = gr.Textbox(
+                label="Contexto Recuperado", 
+                interactive=False,
+                lines=10
+            )
+    
+    # Funções de callback
+    def upload_pdf(file):
+        if file is None:
+            return "Nenhum arquivo selecionado."
+        return rag.load_pdf(file)
+    
+    def clear_knowledge_base():
+        rag.clear()
+        return "Base de conhecimento limpa."
+    
+    def submit_query(query, history):
         history = history or []
-        return_history = history.copy()
-        return_history.append((message, ""))
-        yield return_history, "", ""
-        
-        for response, context, _ in respond(message, history, system_message, max_tokens, temperature, top_p, k_initial, k_final):
-            return_history[-1] = (message, response)
-            yield return_history, context, ""
-
-    # Conexões de eventos
-    submit_btn.click(
-        submit_message,
-        inputs=[msg, chatbot, system_msg, max_tokens, temperature, top_p, k_initial, k_final],
-        outputs=[chatbot, context_box, msg]
-    )
-    msg.submit(
-        submit_message,
-        inputs=[msg, chatbot, system_msg, max_tokens, temperature, top_p, k_initial, k_final],
-        outputs=[chatbot, context_box, msg]
-    )
-    load_btn.click(
-        load_pdfs,
+        response, context = process_query(query, history)
+        history.append((query, response))
+        return history, "", context
+    
+    # Eventos
+    pdf_upload.upload(
+        upload_pdf,
         inputs=[pdf_upload],
         outputs=[load_status]
     )
+    
+    clear_btn.click(
+        clear_knowledge_base,
+        inputs=[],
+        outputs=[load_status]
+    )
+    
+    query_btn.click(
+        submit_query,
+        inputs=[query_input, chatbot],
+        outputs=[chatbot, query_input, context_display]
+    )
+    
+    query_input.submit(
+        submit_query,
+        inputs=[query_input, chatbot],
+        outputs=[chatbot, query_input, context_display]
+    )
 
 if __name__ == "__main__":
     demo.launch()