P8 / app.py
magomerob's picture
Upload 2 files
0e53abb verified
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()