test app

app.py

@@ -0,0 +1,104 @@
+import streamlit as st
+from langdetect import detect
+from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+
+
+@st.cache
+def load_data():
+  supported_languages = [
+    'ar_AR',
+    'cs_CZ',
+    'de_DE',
+    'en_XX',
+    'es_XX',
+    'et_EE',
+    'fi_FI',
+    'fr_XX',
+    'gu_IN',
+    'hi_IN',
+    'it_IT',
+    'ja_XX',
+    'kk_KZ',
+    'ko_KR',
+    'lt_LT',
+    'lv_LV',
+    'my_MM',
+    'ne_NP',
+    'nl_XX',
+    'ro_RO',
+    'ru_RU',
+    'si_LK',
+    'tr_TR',
+    'vi_VN',
+    'zh_CN',
+    'af_ZA',
+    'az_AZ',
+    'bn_IN',
+    'fa_IR',
+    'he_IL',
+    'hr_HR',
+    'id_ID',
+    'ka_GE',
+    'km_KH',
+    'mk_MK',
+    'ml_IN',
+    'mn_MN',
+    'mr_IN',
+    'pl_PL',
+    'ps_AF',
+    'pt_XX',
+    'sv_SE',
+    'sw_KE',
+    'ta_IN',
+    'te_IN',
+    'th_TH',
+    'tl_XX',
+    'uk_UA',
+    'ur_PK',
+    'xh_ZA',
+    'gl_ES',
+    'sl_SI'
+  ]
+  return {k.split('_')[0]:k for k in supported_languages}
+
+@st.cache(allow_output_mutation=True, suppress_st_warning=True)
+def load_model():
+    model_name = "facebook/mbart-large-50-many-to-many-mmt"
+    model = MBartForConditionalGeneration.from_pretrained(model_name)
+    tokenizer = MBart50TokenizerFast.from_pretrained(model_name)
+    return (model, tokenizer)
+
+data = load_data()
+
+def translate_to_english(model, tokenizer, text):
+    src_lang = detect(text)
+    if src_lang in data:
+        tokenizer.src_lang = src_lang
+        encoded_txt = tokenizer(text, return_tensors="pt")
+        generated_tokens = model.generate(
+            **encoded_txt,
+            forced_bos_token_id=tokenizer.lang_code_to_id["en_XX"]
+        )
+        return tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    else:
+        print(f"Language {src_lang} not found")
+        return
+
+st.title("Auto Translate (To English)")
+
+
+text = st.text_input(f"Write in any (1 of {len(data.keys())}) language")
+
+st.text("What you wrote: ")
+
+st.write(text)
+
+st.text("English Translation: ")
+
+if text:
+    model, tokenizer = load_model()
+    translated_text = translate_to_english(model, tokenizer, text)
+    st.write(translated_text[0] if translated_text else "No translation found"
+
+
+            )