from transformers import AutoTokenizer, AutoModelForSeq2SeqLM
import gradio as gr

MODEL_ID = "adcelis/opus-books-en-es-flan-t5-small"

tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
model = AutoModelForSeq2SeqLM.from_pretrained(MODEL_ID)

def translate_text(text):
    if not text or not text.strip():
        return ""

    prompt = f"translate English to Spanish: {text}"

    inputs = tokenizer(prompt, return_tensors="pt", truncation=True)
    outputs = model.generate(
        **inputs,
        max_new_tokens=128
    )

    result = tokenizer.decode(outputs[0], skip_special_tokens=True)
    return result

demo = gr.Interface(
    fn=translate_text,
    inputs=gr.Textbox(
        lines=5,
        label="Texto en inglés",
        placeholder="Write here the English sentence you want to translate..."
    ),
    outputs=gr.Textbox(
        lines=5,
        label="Traducción al español"
    ),
    title="English to Spanish Translation",
    description="Demo del modelo fine-tuned con opus_books para traducir texto del inglés al español.",
    examples=[
        ["The house was quiet and the children were sleeping."],
        ["I have never seen such a beautiful garden."],
        ["She opened the window and looked at the sea."]
    ],
    flagging_mode="never"
)

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