import gradio as gr
from transformers import pipeline

# Configura modelos com instruções em português
modelo_rapido = pipeline(
    "text-generation",
    model="pierreguillou/bert-base-cased-squad-v1.1-portuguese",
    max_length=50
)

modelo_detalhado = pipeline(
    "text-generation",
    model="pucpr/portuguese-gpt2-small",
    max_length=100
)

arbitro = pipeline(
    "text-classification",
    model="neuralmind/bert-base-portuguese-cased"
)

def avaliar_respostas(prompt, respostas):
    instrucao = f"""
    Avalie qual resposta é melhor para o prompt '{prompt}':
    1. {respostas[0]}
    2. {respostas[1]}
    
    Critérios:
    - Clareza (0-5)
    - Precisão científica (0-5)
    - Adequação ao português brasileiro (0-5)
    """
    
    # Força avaliação em PT-BR
    veredito = arbitro(instrucao[:512])[0]  # Limita ao tamanho máximo do BERT
    return respostas[0] if veredito['label'] == "LABEL_1" else respostas[1]

def chatbot(prompt):
    # Adiciona instrução explícita de idioma
    prompt_pt = f"Responda em português brasileiro: {prompt}"
    
    # Gera respostas com modelos específicos
    resposta_curta = modelo_rapido(prompt_pt, num_return_sequences=1)[0]['generated_text']
    resposta_long = modelo_detalhado(prompt_pt, num_return_sequences=1)[0]['generated_text']
    
    # Limpa respostas duplicadas
    resposta_curta = resposta_curta.split(".")[0] + "."
    resposta_long = resposta_long.split("\n")[0]
    
    return avaliar_respostas(prompt, [resposta_curta, resposta_long])

# Interface com exemplos em PT-BR
with gr.Blocks(title="Chatbot ABNT") as app:
    gr.Markdown("## 🇧🇷 Chatbot Acadêmico (Português Brasileiro)")
    with gr.Row():
        entrada = gr.Textbox(label="Digite sua pergunta", placeholder="Ex: Qual a fórmula da água?")
        saida = gr.Textbox(label="Resposta ABNT", lines=5)
    btn = gr.Button("Enviar")
    btn.click(chatbot, inputs=entrada, outputs=saida)
    gr.Examples(
        examples=["Qual a fórmula da água?", "Explique a teoria da relatividade"],
        inputs=entrada
    )

app.launch()