Update app.py

app.py

@@ -1,103 +1,35 @@
 import streamlit as st
-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}
+from transformers import T5ForConditionalGeneration, T5TokenizerFast
 
 @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)
+    model_name = "north/demo-nynorsk-base"
+    model = T5ForConditionalGeneration.from_pretrained(model_name)
+    tokenizer = T5TokenizerFast.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:
+def translate_to_nynorsk(model, tokenizer, text):
         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"]
+            **encoded_txt
         )
         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)")
+st.title("Translate To Nynorsk")
 
 
-text = st.text_input(f"Write in any (1 of {len(data.keys())}) language")
+text = st.text_input(f"Bokmål text: ")
 
-st.text("What you wrote: ")
+#st.text("What you wrote: ")
+#st.write(text)
 
-st.write(text)
-
-st.text("English Translation: ")
+st.text("Nynorsk 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"
-
+    translated_text = translate_to_nynorsk(model, tokenizer, text)
+    st.write(translated_text[0] if translated_text else "Unknown Error Translating Text"
 
-            )