Spaces:
Running
Running
File size: 2,049 Bytes
c239aa3 |
1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 |
import gradio as gr
from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
# Load the model and tokenizer
model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
# Define the language codes supported by the model
language_codes = {
"Arabic": "ar_AR",
"Czech": "cs_CZ",
"German": "de_DE",
"English": "en_XX",
"Spanish": "es_XX",
"Estonian": "et_EE",
"Finnish": "fi_FI",
"French": "fr_XX",
"Gujarati": "gu_IN",
"Hindi": "hi_IN",
"Italian": "it_IT",
"Japanese": "ja_XX",
"Kazakh": "kk_KZ",
"Korean": "ko_KR",
"Lithuanian": "lt_LT",
"Latvian": "lv_LV",
"Burmese": "my_MM",
"Nepali": "ne_NP",
"Dutch": "nl_XX",
"Romanian": "ro_RO",
"Russian": "ru_RU",
"Sinhala": "si_LK",
"Turkish": "tr_TR",
"Vietnamese": "vi_VN",
"Chinese": "zh_CN",
}
def translate(text, src_lang, tgt_lang):
# Set the source language
tokenizer.src_lang = language_codes[src_lang]
# Tokenize the input text
encoded = tokenizer(text, return_tensors="pt")
# Generate translation
generated_tokens = model.generate(
**encoded,
forced_bos_token_id=tokenizer.lang_code_to_id[language_codes[tgt_lang]]
)
# Decode the generated tokens
translation = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)[0]
return translation
# Create the Gradio interface
iface = gr.Interface(
fn=translate,
inputs=[
gr.Textbox(label="Input Text"),
gr.Dropdown(choices=list(language_codes.keys()), label="Source Language"),
gr.Dropdown(choices=list(language_codes.keys()), label="Target Language"),
],
outputs=gr.Textbox(label="Translated Text"),
title="Multilingual Translation with MBart",
description="Translate text between multiple languages using the MBart model.",
)
# Launch the interface
iface.launch() |