import gradio as gr
from transformers import pipeline

# Modelos geradores otimizados
model_a = pipeline("text-generation", model="distilgpt2")  # Modelo A
model_b = pipeline("text-generation", model="EleutherAI/gpt-neo-125M")  # Modelo B

# Modelo juiz
arbiter = pipeline("text-classification",model="google/flan-t5-small")  # Modelo Árbitro

# Julgamento por batch
def judge_response(response_a, response_b):
    results = arbiter([response_a, response_b])
    score_a = results[0]['score']
    score_b = results[1]['score']
    
    # Critério de julgamento: maior pontuação é melhor
    return ("Modelo A", response_a) if score_a > score_b else ("Modelo B", response_b)

# Função principal
def chatbot(prompt):
    # Gera respostas para ambos os modelos
    response_a = model_a(prompt, max_new_tokens=40)[0]['generated_text'].strip()
    response_b = model_b(prompt, max_new_tokens=40)[0]['generated_text'].strip()
    
    # Julga as respostas
    winner, final_response = judge_response(response_a, response_b)
    
    # Retorna as informações
    return {
        "prompt": prompt,
        "resposta_modelo_a": response_a,
        "resposta_modelo_b": response_b,
        "modelo_vencedor": winner,
        "resposta_escolhida": final_response,
        "informacoes_modelos": {
            "modelo_a": "distilgpt2 - Modelo de geração de texto leve.",
            "modelo_b": "EleutherAI/gpt-neo-125M - Modelo de geração de texto maior.",
            "modelo_juiz": "neuralmind/bert-base-portuguese-cased - Modelo de classificação para julgar as respostas."
        },
        "criterio_julgamento": "O modelo juiz classifica as respostas com base em uma pontuação de sentimento positivo."
    }

# Interface
iface = gr.Interface(
    fn=chatbot,
    inputs=gr.Textbox(label="Digite sua pergunta:"),
    outputs=[
        gr.JSON(label="Resultados"),
    ],
    title= "Atividade N2 Chatbot com Julgamento - Equipe: FDCH_FMU",
    description="Compara respostas de dois modelos e usa um árbitro para escolher a melhor com base no sentimento positivo."
)

iface.launch()