app.py

import gradio as gr
import torch
from sentence_transformers import SentenceTransformer, util
from transformers import T5ForConditionalGeneration, T5Tokenizer

# --- 1. Carregamento dos Modelos (faça isso apenas uma vez) ---
# Esta parte não muda. Usaremos os mesmos modelos eficientes.

# Modelo para criar embeddings (vetores) a partir do texto
print("Carregando o modelo de recuperação (Sentence Transformer)...")
retriever_model = SentenceTransformer('all-MiniLM-L6-v2')

# Modelo para gerar as respostas (um T5 do Hugging Face)
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')
print("Modelos carregados com sucesso!")


# --- 2. Base de Conhecimento: DADOS DOS CERTIFICADOS DE CALIBRAÇÃO ---
# Esta é a principal modificação.
# Cada string representa os dados essenciais de um certificado de calibração.
# Em um sistema real, isso viria de um banco de dados ou da leitura de PDFs.

knowledge_base = [
    "Certificado ID: CAL-2023-001. Instrumento TAG: PI-101 (Manômetro de Pressão). Data da Calibração: 15/01/2023. Próxima Calibração: 15/01/2024. Resultado: APROVADO. Incerteza da medição: 0.05 bar. Padrão utilizado: Fluke 754 (ID: P-05). Técnico: Ana Oliveira.",
    "Certificado ID: CAL-2023-002. Instrumento TAG: TT-205 (Transmissor de Temperatura). Data da Calibração: 20/02/2023. Próxima Calibração: 20/08/2023. Resultado: REPROVADO 'como encontrado', APROVADO 'como deixado' após ajuste. Tolerância: ±0.1°C. Observação: Encontrado desvio de +0.3°C no ponto de 100°C.",
    "Certificado ID: CAL-2023-003. Instrumento TAG: FV-300 (Válvula de Controle). Data do Teste: 05/03/2023. Próximo Teste: 05/03/2025. Procedimento: POP-123-Rev02. Resultado: APROVADO. Observação: Histerese de 1.5% encontrada, dentro do limite de 2.0%.",
    "Certificado ID: CAL-2023-004. Instrumento TAG: BAL-01 (Balança Analítica). Data da Calibração: 10/04/2023. Próxima Calibração: 10/04/2024. Resultado: APROVADO. Incerteza da medição: ±0.001g. Padrões de massa com rastreabilidade à RBC/Inmetro.",
    "Certificado ID: CAL-2023-005. Instrumento TAG: PI-102 (Manômetro de Pressão de reserva). Data da Calibração: 18/01/2023. Próxima Calibração: 18/01/2024. Resultado: APROVADO. Técnico: Carlos Pereira. O instrumento estava armazenado e foi calibrado antes do uso potencial."
]

# --- 3. Pré-processamento da Base de Conhecimento ---
# Esta parte não muda. Convertemos nossa nova base de conhecimento em vetores.
print("Processando a base de conhecimento (certificados)...")
knowledge_base_embeddings = retriever_model.encode(knowledge_base, convert_to_tensor=True)
print("Base de conhecimento pronta!")


# --- 4. A Função Principal do RAG ---
# Esta função é agnóstica ao domínio, então não precisa de alterações.
def answer_question(question):
    """
    Esta função recebe uma pergunta, encontra o certificado mais relevante na base de conhecimento
    e gera uma resposta baseada nos dados daquele certificado.
    """
    # Etapa de Recuperação (Retrieval)
    question_embedding = retriever_model.encode(question, convert_to_tensor=True)
    cosine_scores = util.cos_sim(question_embedding, knowledge_base_embeddings)
    best_doc_index = torch.argmax(cosine_scores)
    retrieved_context = knowledge_base[best_doc_index]

    # Log para depuração, para vermos qual certificado foi recuperado.
    print(f"\n--- Nova Pergunta de Auditoria ---")
    print(f"Pergunta: {question}")
    print(f"Certificado Recuperado (Contexto): {retrieved_context}")

    # Etapa de Geração (Generation)
    prompt = f"""
    Contexto: {retrieved_context}

    Pergunta: {question}

    Com base estritamente no contexto do certificado fornecido, responda à pergunta do auditor.
    Resposta:
    """

    input_ids = generator_tokenizer(prompt, return_tensors="pt").input_ids

    outputs = generator_model.generate(
        input_ids,
        max_length=150,
        num_beams=5,
        early_stopping=True
    )

    answer = generator_tokenizer.decode(outputs[0], skip_special_tokens=True)
    return answer

# --- 5. Criação da Interface com Gradio ---
# Modificamos o título e a descrição para a nossa nova aplicação.
interface = gr.Interface(
    fn=answer_question,
    inputs=gr.Textbox(
        lines=3,
        placeholder="Faça uma pergunta sobre um certificado ou instrumento...\nEx: Qual o resultado do manômetro PI-101?\nQuando vence a calibração do TT-205?\nQuem calibrou a balança BAL-01?",
        label="Pergunta do Auditor"
    ),
    outputs=gr.Textbox(label="Resposta Baseada no Certificado"),
    title="🤖 RAG - Auditor de Certificados de Calibração",
    description="Este sistema usa RAG para responder perguntas sobre certificados de calibração. Pergunte sobre o status, datas, técnicos ou resultados de um instrumento específico (PI-101, TT-205, FV-300, BAL-01).",
    examples=[
        ["Qual foi o resultado da calibração do PI-101?"],
        ["Quando é a próxima calibração do transmissor TT-205?"],
        ["Qual a incerteza da balança BAL-01?"],
        ["Houve alguma observação no certificado da válvula FV-300?"]
    ]
)

# --- 6. Lançamento do App ---
if __name__ == "__main__":
    interface.launch()