Spaces:
Sleeping
Sleeping
| import gradio as gr | |
| from transformers import MBartForConditionalGeneration, MBart50TokenizerFast | |
| import tiktoken | |
| # 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() |