import gradio as gr
from transformers import pipeline

# Configuração dos modelos
modelo_rapido = pipeline("text-generation", model="facebook/bart-large-mnli")
modelo_detalhado = pipeline("text-generation", model="microsoft/DialoGPT-medium")
arbitro = pipeline("text2text-generation", model="google/flan-t5-large")


def avaliar_respostas(prompt, resposta1, resposta2):
    instrucao = f"""
    Avalie qual resposta é melhor para: '{prompt}':
    - Resposta 1: {resposta1}
    - Resposta 2: {resposta2}
    Critique com base em: clareza, precisão e relevância. Retorne apenas '1' ou '2'.
    """
    veredito = arbitro(instrucao, max_length=10)[0]['generated_text']
    return resposta1 if '1' in veredito else resposta2


def chatbot(prompt):
    # Gera respostas
    resposta_rapida = modelo_rapido(prompt, max_length=50)[0]['generated_text']
    resposta_detalhada = modelo_detalhado(prompt, max_length=100)[0]['generated_text']

    # Escolhe a melhor
    melhor_resposta = avaliar_respostas(prompt, resposta_rapida, resposta_detalhada)

    return melhor_resposta


# Interface do Gradio
with gr.Blocks() as interface:
    gr.Markdown("## 🤖 Chatbot Acadêmico (Funcionamento em Cascata)")
    with gr.Tab("Chat"):
        input_text = gr.Textbox(label="Digite sua pergunta")
        output_text = gr.Textbox(label="Resposta")
        btn = gr.Button("Enviar")
        btn.click(fn=chatbot, inputs=input_text, outputs=output_text)

    with gr.Tab("Sobre"):
        gr.Markdown("""
        **Modelos Utilizados**:
        - Resposta Rápida: `facebook/bart-large-mnli`
        - Resposta Detalhada: `microsoft/DialoGPT-medium`
        - Árbitro: `google/flan-t5-large`

        **Critérios de Avaliação**: Clareza, Precisão e Relevância.
        """)

interface.launch()