Update app.py

app.py

@@ -1,227 +1,287 @@
 import gradio as gr
+from transformers import pipeline
 import PyPDF2
-from transformers import AutoTokenizer, AutoModel
-from sentence_transformers import util
-import torch
-import numpy as np
 import os
 import re
-from typing import List, Dict, Any
-import requests
+import numpy as np
+from sklearn.feature_extraction.text import TfidfVectorizer
+from sklearn.metrics.pairwise import cosine_similarity
+import torch
 
-# Diretório para armazenar os PDFs
-PDF_DIR = "pdf_data"
+# Configuração de diretórios
+BASE_DIR = os.path.dirname(os.path.abspath(__file__))
+PDF_DIR = os.path.join(BASE_DIR, "pdf_data")
 os.makedirs(PDF_DIR, exist_ok=True)
 
-# Modelo simples para embeddings
-tokenizer = AutoTokenizer.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
-model = AutoModel.from_pretrained("sentence-transformers/all-MiniLM-L6-v2")
+# Inicialização de modelos gratuitos
+# Carregando o modelo de geração de texto
+device = 0 if torch.cuda.is_available() else -1
+summarizer = pipeline(
+    "summarization", 
+    model="facebook/bart-large-cnn",
+    device=device
+)
 
-# 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:
+# Classe RAG totalmente gratuita
+class FreeRAG:
     def __init__(self):
         self.documents = []
-        self.embeddings = None
+        self.vectorizer = TfidfVectorizer(stop_words='english')
+        self.vectors = None
         
     def clear(self):
         self.documents = []
-        self.embeddings = None
+        self.vectors = None
         
     def process_text(self, text):
-        # Limpa texto e divide em parágrafos
+        """Processa e divide o texto em chunks de tamanho razoável"""
+        # Limpa espaços extras
         text = re.sub(r'\s+', ' ', text).strip()
-        paragraphs = [p for p in text.split('\n') if len(p) > 50]
-        return paragraphs
         
+        # Divide por parágrafos
+        paragraphs = text.split('\n')
+        chunks = []
+        
+        current_chunk = ""
+        for p in paragraphs:
+            p = p.strip()
+            if not p:
+                continue
+                
+            # Se o parágrafo for muito grande, divide-o
+            if len(p) > 1000:
+                sentences = re.split(r'(?<=[.!?])\s+', p)
+                for sentence in sentences:
+                    if len(current_chunk) + len(sentence) < 1000:
+                        current_chunk += " " + sentence
+                    else:
+                        if current_chunk:
+                            chunks.append(current_chunk.strip())
+                        current_chunk = sentence
+            else:
+                if len(current_chunk) + len(p) < 1000:
+                    current_chunk += " " + p
+                else:
+                    if current_chunk:
+                        chunks.append(current_chunk.strip())
+                    current_chunk = p
+        
+        if current_chunk:
+            chunks.append(current_chunk.strip())
+            
+        return [c for c in chunks if len(c) > 100]  # Filtra chunks muito pequenos
+    
     def load_pdf(self, file_obj):
+        """Carrega um arquivo PDF"""
         try:
+            # Determina o caminho do arquivo
+            filename = os.path.basename(file_obj.name)
+            file_path = os.path.join(PDF_DIR, filename)
+            
             # 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
+            text = ""
             with open(file_path, 'rb') as f:
-                pdf = PyPDF2.PdfReader(f)
-                text = ""
-                for page in pdf.pages:
-                    text += page.extract_text() + "\n"
+                reader = PyPDF2.PdfReader(f)
+                for page in reader.pages:
+                    page_text = page.extract_text()
+                    if page_text:
+                        text += page_text + "\n"
             
-            # Processa o texto em chunks
+            if not text.strip():
+                return "Erro: Não foi possível extrair texto do PDF."
+            
+            # Processa o texto
             chunks = self.process_text(text)
+            if not chunks:
+                return "Erro: Conteúdo do PDF não pôde ser processado adequadamente."
             
             # 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)
+            for chunk in chunks:
+                self.documents.append({
+                    "source": filename,
+                    "content": chunk
+                })
+            
+            # Atualiza vetores TF-IDF
+            self._update_vectors()
             
-            # Recalcula embeddings
-            if self.documents:
-                contents = [doc["content"] for doc in self.documents]
-                self.embeddings = get_embeddings(contents)
+            return f"PDF carregado com sucesso: {filename} ({len(chunks)} segmentos)"
             
-            return f"Carregado: {filename} ({len(chunks)} segmentos)"
         except Exception as e:
             return f"Erro ao processar PDF: {str(e)}"
+    
+    def _update_vectors(self):
+        """Atualiza os vetores TF-IDF para todos os documentos"""
+        if not self.documents:
+            self.vectors = None
+            return
             
+        texts = [doc["content"] for doc in self.documents]
+        self.vectors = self.vectorizer.fit_transform(texts)
+    
     def search(self, query, top_k=3):
-        if not self.documents or self.embeddings is None:
+        """Busca documentos relevantes para a query"""
+        if not self.documents or self.vectors is None:
             return []
             
-        # Calcula embedding da query
-        query_embedding = get_embeddings([query])
+        # Vetoriza a query
+        query_vec = self.vectorizer.transform([query])
         
         # Calcula similaridade
-        similarities = util.pytorch_cos_sim(query_embedding, self.embeddings)[0]
+        similarity_scores = cosine_similarity(query_vec, self.vectors)[0]
         
-        # Encontra os top_k mais similares
-        top_results = torch.topk(similarities, min(top_k, len(self.documents)))
+        # Encontra os top-k resultados
+        top_indices = similarity_scores.argsort()[-top_k:][::-1]
         
         results = []
-        for score, idx in zip(top_results.values, top_results.indices):
+        for idx in top_indices:
             results.append({
-                "score": score.item(),
+                "score": similarity_scores[idx],
                 "document": self.documents[idx]
             })
             
         return results
 
-# Inicializa o RAG
-rag = SimpleRAG()
-
-# 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
+# Instância do RAG
+rag_engine = FreeRAG()
 
-def query_llm(prompt):
-    payload = {
-        "inputs": prompt,
-        "parameters": {
-            "max_new_tokens": 512,
-            "temperature": 0.7,
-            "top_p": 0.95
-        }
-    }
-    
+def generate_response(prompt, max_length=300):
+    """Gera uma resposta baseada no prompt usando o modelo carregado"""
     try:
-        response = requests.post(LLM_API_URL, headers=headers, json=payload)
-        return response.json()[0]["generated_text"]
+        # Limita o tamanho do prompt para evitar erros
+        if len(prompt) > 1024:
+            prompt = prompt[:1024]
+            
+        # Gera a resposta
+        response = summarizer(
+            prompt, 
+            max_length=max_length, 
+            min_length=50,
+            do_sample=False
+        )[0]['summary_text']
+        
+        return response
     except Exception as e:
-        return f"Erro ao consultar o LLM: {str(e)}"
+        return f"Erro ao gerar resposta: {str(e)}"
 
-# Função para processar a consulta
 def process_query(query, history):
-    # Busca documentos relevantes
-    results = rag.search(query)
+    """Processa uma consulta do usuário"""
+    # Verifica se há documentos carregados
+    if not rag_engine.documents:
+        return "Por favor, carregue alguns PDFs primeiro.", "Nenhum documento carregado."
     
-    if not results:
-        return "Por favor, carregue alguns PDFs primeiro.", "Nenhum documento disponível."
+    # Busca documentos relevantes
+    results = rag_engine.search(query, top_k=3)
     
     # Formata o contexto
     context = ""
     for i, result in enumerate(results):
-        context += f"[{i+1}] Fonte: {result['document']['source']}\n"
+        context += f"[{i+1}] Documento: {result['document']['source']}\n"
         context += f"Trecho: {result['document']['content'][:300]}...\n"
         context += f"Relevância: {result['score']:.2f}\n\n"
     
     # 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.
+    prompt = f"""
+Com base nos seguintes documentos, responda à pergunta de forma concisa e informativa.
+Se a resposta não estiver nos documentos, diga que não há informações suficientes.
 
-CONTEXTO:
+DOCUMENTOS:
 {context}
 
-PERGUNTA: {query}[/INST]"""
-    
-    # Consulta o modelo
-    response = query_llm(prompt)
+PERGUNTA: {query}
+
+RESPOSTA:
+"""
     
-    # Extrai a resposta real (removendo o prompt)
-    actual_response = response.split("[/INST]")[-1].strip()
+    # Gera a resposta
+    response = generate_response(prompt)
     
-    return actual_response, context
+    return response, context
 
 # Interface Gradio
-with gr.Blocks(title="RAG PDF Simplificado") as demo:
-    with gr.Row():
-        with gr.Column(scale=2):
-            gr.Markdown("# RAG PDF Simplificado")
+with gr.Blocks(title="RAG PDF Gratuito") as demo:
+    gr.Markdown("# Sistema de RAG PDF (Retrieval Augmented Generation)")
+    gr.Markdown("Carregue PDFs e faça perguntas sobre eles.")
+    
+    with gr.Tab("Carregar PDFs"):
+        with gr.Row():
+            with gr.Column():
+                pdf_upload = gr.File(
+                    label="Selecionar PDF", 
+                    file_types=[".pdf"],
+                    file_count="single"
+                )
+                upload_btn = gr.Button("Carregar PDF")
+                clear_btn = gr.Button("Limpar Base de Conhecimento")
+                status = gr.Textbox(label="Status", interactive=False)
+    
+    with gr.Tab("Consultar"):
+        with gr.Row():
+            with gr.Column(scale=2):
+                chatbot = gr.Chatbot(label="Conversa")
+                query = gr.Textbox(
+                    label="Sua pergunta",
+                    placeholder="Digite sua pergunta sobre os documentos..."
+                )
+                submit_btn = gr.Button("Enviar")
             
-    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.Row():
-        with gr.Column(scale=2):
-            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_display = gr.Textbox(
-                label="Contexto Recuperado", 
-                interactive=False,
-                lines=10
-            )
-    
-    # Funções de callback
+            with gr.Column(scale=1):
+                context_box = gr.Textbox(
+                    label="Contexto Recuperado",
+                    interactive=False,
+                    lines=15
+                )
+    
+    # Callbacks
     def upload_pdf(file):
         if file is None:
             return "Nenhum arquivo selecionado."
-        return rag.load_pdf(file)
+        return rag_engine.load_pdf(file)
     
     def clear_knowledge_base():
-        rag.clear()
+        rag_engine.clear()
         return "Base de conhecimento limpa."
     
-    def submit_query(query, history):
-        history = history or []
-        response, context = process_query(query, history)
-        history.append((query, response))
-        return history, "", context
+    def handle_query(question, chat_history):
+        chat_history = chat_history or []
+        
+        # Processa a consulta
+        answer, context = process_query(question, chat_history)
+        
+        # Atualiza o histórico
+        chat_history.append((question, answer))
+        
+        return chat_history, "", context
     
     # Eventos
-    pdf_upload.upload(
+    upload_btn.click(
         upload_pdf,
         inputs=[pdf_upload],
-        outputs=[load_status]
+        outputs=[status]
     )
     
     clear_btn.click(
         clear_knowledge_base,
         inputs=[],
-        outputs=[load_status]
+        outputs=[status]
     )
     
-    query_btn.click(
-        submit_query,
-        inputs=[query_input, chatbot],
-        outputs=[chatbot, query_input, context_display]
+    submit_btn.click(
+        handle_query,
+        inputs=[query, chatbot],
+        outputs=[chatbot, query, context_box]
     )
     
-    query_input.submit(
-        submit_query,
-        inputs=[query_input, chatbot],
-        outputs=[chatbot, query_input, context_display]
+    query.submit(
+        handle_query,
+        inputs=[query, chatbot],
+        outputs=[chatbot, query, context_box]
     )
 
 if __name__ == "__main__":
+    # Inicializa a interface
     demo.launch()