import gradio as gr
from transformers import pipeline

# Modelos pequenos e leves para rodar no ambiente gratuito
model_a = pipeline("text-generation", model="gpt2")
model_b = pipeline("text-generation", model="EleutherAI/gpt-neo-125M")

# Modelo juiz baseado em sentimento (positivo ou negativo)
arbiter = pipeline("text-classification", model="distilbert-base-uncased-finetuned-sst-2-english")

# Função para julgar a melhor resposta com base em sentimento positivo
def judge_response(prompt, response_a, response_b):
    combined_a = f"Question: {prompt} Answer: {response_a}"
    combined_b = f"Question: {prompt} Answer: {response_b}"

    score_a = arbiter(combined_a)[0]['score']
    score_b = arbiter(combined_b)[0]['score']

    if score_a > score_b:
        return "Model A", response_a
    else:
        return "Model B", response_b

# Função principal do chatbot
def chatbot(prompt):
    raw_response_a = model_a(prompt, max_new_tokens=60)[0]['generated_text']
    raw_response_b = model_b(prompt, max_new_tokens=60)[0]['generated_text']

    # Remove o prompt inicial da resposta
    response_a = raw_response_a[len(prompt):].strip()
    response_b = raw_response_b[len(prompt):].strip()

    winner, final_response = judge_response(prompt, response_a, response_b)

    return (
        prompt,
        response_a,
        response_b,
        winner,
        final_response
    )

# Interface Gradio
iface = gr.Interface(
    fn=chatbot,
    inputs=gr.Textbox(label="Digite sua pergunta:"),
    outputs=[
        gr.Textbox(label="Pergunta"),
        gr.Textbox(label="Resposta do Modelo A (GPT-2)"),
        gr.Textbox(label="Resposta do Modelo B (GPT-Neo-125M)"),
        gr.Textbox(label="Modelo Vencedor"),
        gr.Textbox(label="Resposta Escolhida"),
    ],
    title="Chatbot em Cascata com Julgamento",
    description="""
    Este chatbot utiliza dois modelos diferentes para responder à pergunta de um usuário.
    Um terceiro modelo avalia qual resposta possui maior sentimento positivo e a apresenta como a melhor.

    🧠 Modelos usados:
    - Modelo A: GPT-2 (`gpt2`)
    - Modelo B: GPT-Neo 125M (`EleutherAI/gpt-neo-125M`)
    - Árbitro: DistilBERT SST-2 (`distilbert-base-uncased-finetuned-sst-2-english`)

    Critério de julgamento: Resposta com maior sentimento positivo vence.
    """
)

iface.launch()