import gradio as gr
from transformers import pipeline, AutoTokenizer, AutoModelForSeq2SeqLM, AutoModelForCausalLM, AutoModelForSequenceClassification
import torch

# Modelo A - FLAN-T5
model_a_name = "google/flan-t5-small"
tokenizer_a = AutoTokenizer.from_pretrained(model_a_name)
model_a = AutoModelForSeq2SeqLM.from_pretrained(model_a_name)

# Modelo B - Falcon-RW
model_b_name = "tiiuae/falcon-rw-1b"
tokenizer_b = AutoTokenizer.from_pretrained(model_b_name)
model_b = AutoModelForCausalLM.from_pretrained(model_b_name)

# Árbitro - BART MNLI
arbiter_model = pipeline("zero-shot-classification", model="facebook/bart-large-mnli")

def generate_response(prompt):
    # Geração com Modelo A (FLAN-T5)
    input_ids_a = tokenizer_a(prompt, return_tensors="pt").input_ids
    output_a = model_a.generate(input_ids_a, max_new_tokens=100)
    response_a = tokenizer_a.decode(output_a[0], skip_special_tokens=True)

    # Geração com Modelo B (Falcon)
    input_ids_b = tokenizer_b(prompt, return_tensors="pt").input_ids
    output_b = model_b.generate(input_ids_b, max_new_tokens=100)
    response_b = tokenizer_b.decode(output_b[0], skip_special_tokens=True)

    # Julgamento com o Árbitro
    # A premissa é o prompt original, e as hipóteses são as duas respostas.
    choices = [response_a, response_b]
    judgment = arbiter_model(prompt, choices)

    best_response = judgment['labels'][0]  # Escolhe o que teve maior score

    return {
        "Resposta do Modelo A": response_a,
        "Resposta do Modelo B": response_b,
        "Resposta Escolhida pelo Árbitro": best_response
    }

# Interface Gradio
iface = gr.Interface(
    fn=generate_response,
    inputs=gr.Textbox(lines=2, placeholder="Digite sua pergunta..."),
    outputs=[
        gr.Textbox(label="Resposta do Modelo A"),
        gr.Textbox(label="Resposta do Modelo B"),
        gr.Textbox(label="Resposta Escolhida pelo Árbitro")
    ],
    title="Chatbot em Cascata com Árbitro",
    description="Dois modelos geram respostas. Um terceiro modelo julga qual é a melhor com base em clareza e coerência."
)

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