Update app.py

app.py

@@ -1,7 +1,7 @@
 import gradio as gr
 import torch
 from sentence_transformers import SentenceTransformer, util
-# MUDANÇA: Usaremos AutoTokenizer e AutoModelForCausalLM para o novo modelo
+# Usaremos AutoTokenizer e AutoModelForCausalLM para o novo modelo
 from transformers import AutoTokenizer, AutoModelForCausalLM
 from pypdf import PdfReader
 import os
@@ -11,7 +11,7 @@ import os
 print("Carregando o modelo de recuperação (Sentence Transformer)...")
 retriever_model = SentenceTransformer('all-MiniLM-L6-v2')
 
-# MUDANÇA: Carregando o modelo DeepSeek
+# Carregando o modelo de geração 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(
@@ -69,15 +69,14 @@ def answer_question(question, knowledge_state):
 
     knowledge_base, knowledge_base_embeddings = knowledge_state
 
-    # Etapa de Recuperação (sem alterações)
+    # Etapa de Recuperação
     question_embedding = retriever_model.encode(question, convert_to_tensor=True)
     cosine_scores = util.cos_sim(question_embedding, knowledge_base_embeddings)
-    # MUDANÇA: Aumentado para 7 para mais contexto
+    # Aumentado para 7 para mais contexto
     top_k = min(7, len(knowledge_base)) 
     top_results = torch.topk(cosine_scores, k=top_k, dim=-1)
     retrieved_context = "\n---\n".join([knowledge_base[i] for i in top_results.indices[0]])
 
-    # Verificação para garantir que o contexto não está vazio
     if not retrieved_context.strip():
         return "Não foi possível encontrar um contexto relevante nos documentos para responder a esta pergunta."
 
@@ -85,13 +84,15 @@ def answer_question(question, knowledge_state):
     print(f"Pergunta: {question}")
     print(f"Contexto Recuperado (Top {top_k}):\n{retrieved_context}")
 
-    # MUDANÇA PRINCIPAL: Prompt mais inteligente para interpretar documentos técnicos
+    # MUDANÇA PRINCIPAL: Prompt com regras explícitas de extração de entidades
     prompt = f"""### Instruction:
-Você é um assistente de IA especialista em análise de documentos técnicos. Sua tarefa é analisar o 'Contexto' para responder à 'Pergunta'.
-Regras importantes:
-1.  Use APENAS a informação do 'Contexto'. Não invente nada.
-2.  Não use cabeçalhos de tabelas ou definições de siglas (como 'VI' ou 'VRef') como respostas. Encontre os valores reais associados a eles no texto.
-3.  Responda de forma direta e concisa.
+Você é um assistente de IA especialista em extrair informações de documentos técnicos. Analise o 'Contexto' para responder à 'Pergunta' seguindo estas regras rigorosamente:
+
+**Regras de Extração:**
+1.  **Use APENAS a informação do 'Contexto'.** Não adicione informações externas.
+2.  **Para perguntas sobre 'cliente':** Procure por linhas que comecem com "Cliente:", "Contratante:", ou um nome de empresa claro. Ignore frases genéricas sobre escopo ou lotes.
+3.  **Para perguntas sobre 'instrumento':** Procure por linhas que comecem com "Instrumento:", "Descrição:", "Objeto:", "TAG:", ou um modelo específico. Ignore descrições de métodos de calibração.
+4.  **Para perguntas sobre 'título' ou 'resumo':** Resuma os dados principais, como o tipo de documento (ex: Certificado de Calibração), o nome do cliente e o instrumento calibrado.
 
 **Contexto:**
 {retrieved_context}
@@ -108,7 +109,7 @@ Regras importantes:
     # Ajuste nos parâmetros de geração
     outputs = generator_model.generate(
         input_ids,
-        # MUDANÇA: Aumentado para permitir respostas mais detalhadas
+        # Aumentado para permitir respostas mais detalhadas
         max_new_tokens=350,  
         do_sample=False,
         eos_token_id=generator_tokenizer.eos_token_id,
@@ -126,7 +127,7 @@ with gr.Blocks(theme=gr.themes.Soft()) as interface:
     knowledge_state = gr.State()
     gr.Markdown(
         """
-        # 🤖 RAG - Auditor de Documentos (v7 - Prompt Inteligente)
+        # 🤖 RAG - Auditor de Documentos (v8 - Extração de Entidades)
         **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.