import gradio as gr
from transformers import AutoTokenizer, AutoModelForSeq2SeqLM

# Selecciona el modelo preentrenado para traducción (inglés a español)
model_name = "Helsinki-NLP/opus-mt-en-es"
tokenizer = AutoTokenizer.from_pretrained(model_name)
model = AutoModelForSeq2SeqLM.from_pretrained(model_name)

# Función para traducir el texto
def translate_text(input_text):
    input_ids = tokenizer.encode(input_text, return_tensors="pt")
    output = model.generate(input_ids, max_length=500, num_return_sequences=1, do_sample=True)
    translated_text = tokenizer.decode(output[0], skip_special_tokens=True)
    return translated_text

# Configuración de la interfaz con Gradio
interface = gr.Interface(fn=translate_text,
                         inputs="text",
                         outputs="text",
                         title="Traductor Inglés a Español, que utiliza el modelo Helsinki-NLP/opus-mt-en-es-",
                         description="Introduce el texto en inglés y obtén la traducción en español. El máximo de palabras son 500.")

# Ejecuta la interfaz
interface.launch()