commit

translate.py

@@ -0,0 +1,59 @@
+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)