Spaces:
Sleeping
Sleeping
| mport streamlit as st | |
| import torch | |
| from transformers import MBartForConditionalGeneration, MBart50TokenizerFast | |
| ### Getting the Languages supported #### | |
| LanguageCovered = "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), Afrikaans (af_ZA), Azerbaijani (az_AZ), Bengali (bn_IN), Persian (fa_IR), Hebrew (he_IL), Croatian (hr_HR), Indonesian (id_ID), Georgian (ka_GE), Khmer (km_KH), Macedonian (mk_MK), Malayalam (ml_IN), Mongolian (mn_MN), Marathi (mr_IN), Polish (pl_PL), Pashto (ps_AF), Portuguese (pt_XX), Swedish (sv_SE), Swahili (sw_KE), Tamil (ta_IN), Telugu (te_IN), Thai (th_TH), Tagalog (tl_XX), Ukrainian (uk_UA), Urdu (ur_PK), Xhosa (xh_ZA), Galician (gl_ES), Slovene (sl_SI" | |
| LanguageCovered = LanguageCovered.split(",") | |
| languages_list = [a.strip() for a in LanguageCovered] | |
| languages_list = [a.split(" ") for a in languages_list] | |
| languages = [a[0] for a in languages_list] | |
| codes = [a[1] for a in languages_list] | |
| codes = [a.replace('(', '') for a in codes] | |
| codes = [a.replace(')', '') for a in codes] | |
| lang_dict = dict(zip(languages, codes)) | |
| model_name = "facebook/mbart-large-50-many-to-many-mmt" | |
| # tokenizer = AutoTokenizer.from_pretrained("facebook/mbart-large-50-many-to-many-mmt") | |
| # model = AutoModelForSeq2SeqLM.from_pretrained("facebook/mbart-large-50-many-to-many-mmt") | |
| tokenizer = MBart50TokenizerFast.from_pretrained(model_name) | |
| model = MBartForConditionalGeneration.from_pretrained(model_name) | |
| device = torch.device("cuda" if torch.cuda.is_available() else 'cpu') | |
| model = model.to(device) | |
| def translate_text(text, source_lang, target_lang): | |
| tokenizer.src_lang = source_lang | |
| encoded_text = tokenizer(text, return_tensors="pt").to(device) | |
| generated_tokens = model.generate(**encoded_text, forced_bos_token_id=tokenizer.lang_code_to_id[target_lang]) | |
| #Decode the output | |
| translated_text = tokenizer.decode(generated_tokens[0], skip_special_tokens=True) | |
| return translated_text | |
| st.markdown("### Language Translator") | |
| source_language = '' | |
| target_language = '' | |
| source = st.sidebar.selectbox('Source Language', languages) | |
| if source: | |
| source_language = lang_dict.get(source) | |
| st.write(source_language) | |
| target = st.sidebar.selectbox('Target Language', languages) | |
| if target: | |
| target_language = lang_dict.get(target) | |
| st.write(target_language) | |
| with st.form(key="myForm"): | |
| text = st.text_area("Enter your text") | |
| submit = st.form_submit_button("Submit", type='primary') | |
| if submit and text and source_language and target_language: | |
| with st.spinner(f"{source} to {target} translating"): | |
| translation = translate_text(text, source_language, target_language) | |
| st.write(translation) | |