from transformers import T5Tokenizer, T5ForConditionalGeneration
import gradio as gr

tokenizer = T5Tokenizer.from_pretrained("t5-large")
model = T5ForConditionalGeneration.from_pretrained("t5-large")

def translate_text(input_text, source_lang, target_lang):
  """Translates text using the T5 model."""
  prefix = f"translate {source_lang} to {target_lang}: "
  input_ids = tokenizer(prefix + input_text, return_tensors="pt").input_ids
  outputs = model.generate(input_ids, max_new_tokens=128) 
  translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
  return translated_text

iface = gr.Interface(
    fn=translate_text,
    inputs=[
        gr.Textbox(lines=2, placeholder="Digite o texto aqui...", label="Texto de Entrada"),
        gr.Dropdown(["English", "French", "German"], label="Idioma de Origem"),
        gr.Dropdown(["English", "French", "German"], label="Idioma de Destino"),
    ],
    outputs="text",
    title="Aplicativo de Tradução T5",
    description="Este aplicativo usa o modelo T5 para traduzir texto entre idiomas.",
)
iface.launch(share=True)