TextTranslation / app.py
Donlapark's picture
Update app.py
fb4f107 verified
import gradio as gr
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
# Load the mBART model and tokenizer for multilingual translation
model_name = "facebook/mbart-large-50-many-to-many-mmt"
model = MBartForConditionalGeneration.from_pretrained(model_name)
tokenizer = MBart50TokenizerFast.from_pretrained(model_name, src_lang="en_XX", tgt_lang="th_TH")
# Prediction function
def translate_text(input_text):
inputs = tokenizer.encode(input_text, return_tensors="pt")
outputs = model.generate(
inputs,
max_new_tokens=40,
do_sample=True,
top_k=30,
top_p=0.95
)
translated_text = tokenizer.decode(outputs[0], skip_special_tokens=True)
return translated_text
# Gradio interface
interface = gr.Interface(
fn=translate_text,
inputs="text",
outputs="text",
title="Language Translation",
description="Translate text using the my_awesome_opus_books_model."
)
# Launch the Gradio app
interface.launch()