app.py

import gradio as gr
import re
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Modelos utilizados
MODELOS = {
    "primario": "google/flan-t5-small",
    "secundario": "google/flan-t5-base",
    "arbitro": "google/flan-t5-large"
}

# Carregamento dos modelos e tokenizers
tokenizers = {nome: AutoTokenizer.from_pretrained(modelo)
              for nome, modelo in MODELOS.items()}
modelos = {nome: AutoModelForSeq2SeqLM.from_pretrained(modelo)
           for nome, modelo in MODELOS.items()}

# Base de capitais com erros comuns
BASE_CAPITAIS = {
    'brazil': {
        'correta': 'Brasília',
        'erros_comuns': ['sao paulo', 'rio de janeiro', 'brazil']
    },
    'germany': {
        'correta': 'Berlin',
        'erros_comuns': ['munich', 'frankfurt']
    },
    'france': {
        'correta': 'Paris',
        'erros_comuns': ['lyon', 'marseille']
    },
    # Adicione mais países conforme necessário
}

def gerar_resposta(nome_modelo, pergunta):
    """Gera resposta com o modelo especificado"""
    prompt = f"""Aja como um especialista em geografia. Responda APENAS com o nome da capital oficial.
Exemplos:
Q: Capital da França? A: Paris
Q: Capital do Brasil? A: Brasília
Q: Capital da Alemanha? A: Berlin
Q: {pergunta}
A:"""
    entrada = tokenizers[nome_modelo](prompt, return_tensors="pt")
    saida = modelos[nome_modelo].generate(**entrada, max_length=20)
    return tokenizers[nome_modelo].decode(saida[0], skip_special_tokens=True).strip()

def validar_corrigir(pergunta, resposta_bruta):
    """Valida e corrige a resposta com base na base de capitais"""
    pergunta = pergunta.lower()
    resposta = resposta_bruta.lower()

    for pais, dados in BASE_CAPITAIS.items():
        if pais in pergunta:
            if resposta in dados['erros_comuns']:
                return dados['correta']
            if resposta == dados['correta'].lower():
                return dados['correta']
            return dados['correta']
    return resposta_bruta.title()

def esta_confiante(resposta, pergunta):
    """Avalia se a resposta pode ser considerada confiável"""
    pergunta = pergunta.lower()
    resposta = resposta.lower()
    for pais, dados in BASE_CAPITAIS.items():
        if pais in pergunta:
            if resposta == dados['correta'].lower():
                return True
            if resposta in dados['erros_comuns']:
                return False
    return False

def arbitrar(pergunta, resposta1, resposta2):
    """Usa o modelo árbitro para escolher a melhor resposta"""
    corrigida1 = validar_corrigir(pergunta, resposta1)
    corrigida2 = validar_corrigir(pergunta, resposta2)

    for pais, dados in BASE_CAPITAIS.items():
        if pais in pergunta.lower():
            if corrigida1 == dados['correta']:
                return corrigida1, "Modelo 1 (validado)"
            if corrigida2 == dados['correta']:
                return corrigida2, "Modelo 2 (validado)"

    prompt = f"""Você é professor de geografia. Escolha a capital correta:
Pergunta: {pergunta}
Opção 1: {corrigida1}
Opção 2: {corrigida2}
Responda SOMENTE com "1" ou "2"."""
    
    entrada = tokenizers['arbitro'](prompt, return_tensors="pt")
    saida = modelos['arbitro'].generate(**entrada, max_length=3)
    escolha = tokenizers['arbitro'].decode(saida[0], skip_special_tokens=True).strip()

    if escolha == "1":
        return corrigida1, "Modelo 1 (árbitro)"
    else:
        return corrigida2, "Modelo 2 (árbitro)"

def chatbot(pergunta):
    """Pipeline em cascata para determinar a capital"""
    resposta1 = gerar_resposta("primario", pergunta)
    corrigida1 = validar_corrigir(pergunta, resposta1)

    if corrigida1 == resposta1 and esta_confiante(corrigida1, pergunta):
        return [
            f"Resposta Selecionada: {corrigida1}\nModelo Escolhido: Modelo 1 (primário confiante)",
            f"Modelo 1 (primário): {resposta1}",
            f"Modelo 2 (secundário): Pulado"
        ]

    resposta2 = gerar_resposta("secundario", pergunta)
    corrigida2 = validar_corrigir(pergunta, resposta2)

    if corrigida2 == resposta2 and esta_confiante(corrigida2, pergunta):
        return [
            f"Resposta Selecionada: {corrigida2}\nModelo Escolhido: Modelo 2 (secundário confiante)",
            f"Modelo 1 (primário): {resposta1}",
            f"Modelo 2 (secundário): {resposta2}"
        ]

    resposta_final, modelo_escolhido = arbitrar(pergunta, resposta1, resposta2)
    return [
        f"Resposta Selecionada: {resposta_final}\nModelo Escolhido: {modelo_escolhido}",
        f"Modelo 1 (primário): {resposta1}",
        f"Modelo 2 (secundário): {resposta2}"
    ]

# Markdown content for the article section
article_markdown = """

Gustavo Henrique Xavier da Silva 1758309
Mateus Teixeira Maciel 1650601
Laís Aranda de Souza 1656007
Gustavo Queiroz Fonseca 1771365

### 1. Modelo Primário (google/flan-t5-small)
**Função**: Primeira linha de resposta  
**Características**:
- Modelo menor e mais rápido (80M parâmetros)
- Responsável por gerar a resposta inicial
- Ideal para casos simples e óbvios (como capitais bem conhecidas)

**Fluxo no código**:
- Chamado primeiro através da função `gerar_resposta("primario", pergunta)`
- Se sua resposta for validada como confiável (`esta_confiante()` retorna `True`), o sistema retorna imediatamente sem consultar outros modelos
- Caso contrário, o sistema prossegue para o modelo secundário

### 2. Modelo Secundário (google/flan-t5-base)
**Função**: Segunda opinião para casos mais complexos  
**Características**:
- Modelo médio (250M parâmetros)
- Maior capacidade que o primário
- Usado quando o primário não tem confiança suficiente

**Fluxo no código**:
- Chamado apenas se o primário não for confiável
- Se sua resposta for validada como confiável, o sistema retorna aqui
- Caso ambos (primário e secundário) não sejam confiáveis, o sistema chama o árbitro

### 3. Modelo Árbitro (google/flan-t5-large)
**Função**: Decisor final em casos de discordância  
**Características**:
- Modelo maior e mais capaz (800M parâmetros)
- Atua como "juiz" entre as respostas do primário e secundário
- Só é acionado quando os dois primeiros modelos dão respostas diferentes e não confiáveis

**Fluxo no código**:
- Chamado através da função `arbitrar()`
- Recebe a pergunta original e as duas respostas corrigidas
- Decide qual das duas respostas é mais provavelmente correta
- Pode também corrigir ambas as respostas se necessário

## Sistema de Validação e Correção

Além dos modelos, o sistema possui:
- **Base de Conhecimento**: `BASE_CAPITAIS` com respostas corretas e erros comuns
- **Função `validar_corrigir()`**: Compara respostas com a base de conhecimento
- **Função `esta_confiante()`**: Determina se uma resposta pode ser aceita sem arbitragem

## Fluxo de Decisão
1. Tenta primeiro com o modelo pequeno (rápido)
2. Se não confiável, tenta com o modelo médio
3. Se ainda não resolvido, usa o modelo grande como árbitro
4. Sempre aplica validação/correção baseada na base de conhecimento

Essa abordagem em cascata equilibra velocidade e precisão, usando modelos maiores apenas quando necessário.
"""

# Interface Gradio
interface = gr.Interface(
    fn=chatbot,
    inputs=gr.Textbox(label="Pergunte a capital de um país", placeholder="Qual é a capital do Brasil?"),
    outputs=[
        gr.Textbox(label="Resposta Final"),
        gr.Textbox(label="Resposta do Modelo 1"),
        gr.Textbox(label="Resposta do Modelo 2")
    ],
    title="🗺️ Especialista em Capitais (Cascata com Correção Automática)",
    description="Sistema com três modelos em cascata. Pergunte sobre a capital de qualquer país. Exemplos: Brasil, Alemanha, França...",
    article=article_markdown
)

if __name__ == "__main__":
    interface.launch()