Update app.py

app.py

@@ -1,23 +1,33 @@
 import gradio as gr
 import torch
 from sentence_transformers import SentenceTransformer, util
-from transformers import T5ForConditionalGeneration, T5Tokenizer
+# MUDANÇA: Usaremos AutoTokenizer e AutoModelForCausalLM para o novo modelo
+from transformers import AutoTokenizer, AutoModelForCausalLM
 from pypdf import PdfReader
 import os
 
-# --- 1. Carregamento dos Modelos (sem alterações) ---
+# --- 1. Carregamento dos Modelos ---
+# Modelo de recuperação não muda, ele é excelente para essa tarefa.
 print("Carregando o modelo de recuperação (Sentence Transformer)...")
 retriever_model = SentenceTransformer('all-MiniLM-L6-v2')
-print("Carregando o modelo de geração (Flan-T5)...")
-generator_tokenizer = T5Tokenizer.from_pretrained('google/flan-t5-base')
-generator_model = T5ForConditionalGeneration.from_pretrained('google/flan-t5-base')
+
+# MUDANÇA: Carregando o modelo DeepSeek
+print("Carregando o modelo de geração (DeepSeek)...")
+# Nota: "trust_remote_code=True" é necessário para carregar a arquitetura do DeepSeek
+generator_tokenizer = AutoTokenizer.from_pretrained(
+    'deepseek-ai/deepseek-coder-1.3b-instruct', 
+    trust_remote_code=True
+)
+generator_model = AutoModelForCausalLM.from_pretrained(
+    'deepseek-ai/deepseek-coder-1.3b-instruct', 
+    trust_remote_code=True
+)
 print("Modelos carregados com sucesso!")
 
-# --- 2. Função para Processar Arquivos Enviados ---
+# --- 2. Função para Processar Arquivos Enviados (sem alterações) ---
 def process_files(files):
     if not files:
         return None, "Por favor, envie um ou mais arquivos."
-
     knowledge_text = ""
     for file in files:
         file_path = file.name
@@ -40,11 +50,7 @@ def process_files(files):
     if not knowledge_text.strip():
         return None, "Não foi possível extrair texto dos arquivos fornecidos."
 
-    # MUDANÇA 1: Melhorando o "Chunking" (Divisão do Texto)
-    # Em vez de dividir por parágrafos (\n\n), dividimos por linha (\n).
-    # Isso cria chunks menores e mais focados, o que é melhor para documentos técnicos.
     text_chunks = [chunk.strip() for chunk in knowledge_text.split('\n') if chunk.strip() and len(chunk) > 10]
-    
     if not text_chunks:
         return None, "O texto extraído não continha blocos de texto válidos para processamento."
 
@@ -54,59 +60,56 @@ def process_files(files):
     
     return (text_chunks, knowledge_base_embeddings), f"✅ Sucesso! {len(files)} arquivo(s) processado(s), gerando {len(text_chunks)} blocos de texto."
 
-
-# --- 3. A Função Principal do RAG (com melhorias) ---
+# --- 3. A Função Principal do RAG (com prompt e decodificação ajustados) ---
 def answer_question(question, knowledge_state):
     if not question:
         return "Por favor, insira uma pergunta."
-        
     if not knowledge_state or not knowledge_state[0] or knowledge_state[1] is None:
         return "⚠️ A base de conhecimento está vazia. Por favor, processe alguns arquivos primeiro."
 
     knowledge_base, knowledge_base_embeddings = knowledge_state
 
-    # Etapa de Recuperação (Retrieval)
+    # Etapa de Recuperação (sem alterações)
     question_embedding = retriever_model.encode(question, convert_to_tensor=True)
-    
-    # MUDANÇA 2: Recuperando mais contexto (top_k=3)
-    # Em vez de pegar apenas o melhor chunk, pegamos os 3 melhores.
-    # Isso dá mais informação para a IA gerar uma resposta completa.
     cosine_scores = util.cos_sim(question_embedding, knowledge_base_embeddings)
-    top_k = min(3, len(knowledge_base)) # Garante que não tentamos pegar mais chunks do que existem
+    top_k = min(3, len(knowledge_base)) 
     top_results = torch.topk(cosine_scores, k=top_k, dim=-1)
-
-    # Junta os 3 melhores chunks em um único contexto
     retrieved_context = "\n---\n".join([knowledge_base[i] for i in top_results.indices[0]])
 
     print(f"\n--- Nova Pergunta de Auditoria ---")
     print(f"Pergunta: {question}")
     print(f"Contexto Recuperado (Top {top_k}):\n{retrieved_context}")
 
-    # Etapa de Geração (Generation)
-    
-    # MUDANÇA 3: Prompt de Geração Aprimorado
-    # Damos instruções mais claras para que a IA elabore a resposta.
-    prompt = f"""
-    Você é um assistente de auditoria especialista. Sua tarefa é responder à pergunta do usuário de forma clara e concisa, baseando-se exclusivamente nos trechos de documentos fornecidos abaixo. Elabore uma resposta completa em vez de simplesmente copiar o texto.
+    # MUDANÇA: Prompt ajustado para o formato de instrução do DeepSeek
+    prompt = f"""### Instruction:
+Você é um assistente de auditoria especialista. Sua tarefa é sintetizar as informações dos documentos fornecidos para responder à pergunta do auditor. Elabore uma resposta completa e concisa.
 
-    **Documentos Relevantes:**
-    {retrieved_context}
+**Documentos:**
+{retrieved_context}
 
-    **Pergunta do Auditor:**
-    {question}
+**Pergunta:**
+{question}
 
-    **Resposta Elaborada:**
-    """
+### Response:
+"""
 
     input_ids = generator_tokenizer(prompt, return_tensors="pt").input_ids
+    input_length = input_ids.shape[1]
+
+    # MUDANÇA: Ajuste nos parâmetros de geração
     outputs = generator_model.generate(
         input_ids,
-        max_length=256,
-        num_beams=5,
-        early_stopping=True
+        max_new_tokens=256,  # Controla o tamanho da *nova* resposta gerada
+        do_sample=False,     # Desativa a amostragem para respostas mais diretas
+        eos_token_id=generator_tokenizer.eos_token_id,
+        pad_token_id=generator_tokenizer.eos_token_id # Evita warnings
     )
 
-    answer = generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
+    # MUDANÇA: Decodificação correta para modelos Causal LM
+    # Precisamos remover o prompt inicial da saída gerada.
+    generated_tokens = outputs[0, input_length:]
+    answer = generator_tokenizer.decode(generated_tokens, skip_special_tokens=True)
+    
     return answer
 
 # --- 4. Interface Gráfica (sem alterações na estrutura) ---
@@ -114,7 +117,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as interface:
     knowledge_state = gr.State()
     gr.Markdown(
         """
-        # 🤖 RAG - Auditor de Documentos (v3)
+        # 🤖 RAG - Auditor de Documentos (v5 - DeepSeek)
         **1. Carregue seus arquivos**: Envie um ou mais certificados ou documentos nos formatos `.pdf` ou `.txt`.
         **2. Processe os arquivos**: Clique no botão para criar a base de conhecimento.
         **3. Faça perguntas**: Após o processamento, faça perguntas sobre o conteúdo dos documentos.