import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import torch

model_name = "magomerob/opus-books-en-es"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)
device = "cuda" if torch.cuda.is_available() else "cpu"
model = model.to(device)

def translate(text):
    inputs = tokenizer(text, return_tensors="pt", max_length=128, truncation=True).to(device)
    outputs = model.generate(**inputs, max_new_tokens=128, num_beams=4)
    return tokenizer.decode(outputs[0], skip_special_tokens=True)

demo = gr.Interface(
    fn=translate,
    inputs=gr.Textbox(lines=4, placeholder="Escribe una frase en inglés...", label="Inglés"),
    outputs=gr.Textbox(label="Español"),
    title="Traductor Inglés → Español",
    description="Modelo fine-tuned con flan-t5-small sobre el dataset opus_books (en-es).",
    examples=[
        ["The student will choose the exam model."],
        ["She reads a book every night before sleeping."],
        ["The economy has grown significantly in the last decade."],
    ]
)

demo.launch()