Testing with SMALL-100

app.py

@@ -1,19 +1,25 @@
 import streamlit as st
 import pandas as pd
 from pathlib import Path
-from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+#from transformers import MBartForConditionalGeneration, MBart50TokenizerFast
+from transformers import M2M100ForConditionalGeneration
+from tokenization_small100 import SMALL100Tokenizer
 
 st.set_page_config(page_title="Translation Demo", page_icon=":milky_way:", layout="wide")
 
 
 def get_translation(src_code, trg_code, src):
 
-    tokenizer.src_lang = src_code
+    #tokenizer.src_lang = src_code
+    #encoded = tokenizer(src, return_tensors="pt")
+    #generated_tokens = model.generate(
+        #**encoded,
+        #forced_bos_token_id=tokenizer.lang_code_to_id[trg_code]
+    #)
+    #trg = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
+    tokenizer.tgt_lang = trg_code
     encoded = tokenizer(src, return_tensors="pt")
-    generated_tokens = model.generate(
-        **encoded,
-        forced_bos_token_id=tokenizer.lang_code_to_id[trg_code]
-    )
+    generated_tokens = model.generate(**encoded)
     trg = tokenizer.batch_decode(generated_tokens, skip_special_tokens=True)
     
     return trg
@@ -29,16 +35,18 @@ def open_input(the_file):
     return parsed
 
 
-st.subheader("MBART-50 Translator")
+st.subheader("SMALL-100 Translator")
 
 source = "In the beginning the Universe was created. This has made a lot of people very angry and been widely regarded as a bad move."
 target = ""
-model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
-tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
-
+#model = MBartForConditionalGeneration.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
+#tokenizer = MBart50TokenizerFast.from_pretrained("facebook/mbart-large-50-many-to-many-mmt")
 
-valid_languages = ['de_DE', 'en_XX', 'it_IT']
+model = M2M100ForConditionalGeneration.from_pretrained("alirezamsh/small100")
+tokenizer = SMALL100Tokenizer.from_pretrained("alirezamsh/small100")
 
+#valid_languages = ['de_DE', 'en_XX', 'it_IT']
+valid_languages = ['de', 'it', 'en']
 
 
 valid_languages_tuple = (lang for lang in valid_languages)