wip

app.py

@@ -14,12 +14,13 @@ def init():
     hf_token = os.environ.get("HF_TOKEN")
 
     model = MT5ForConditionalGeneration.from_pretrained("lm-spell/mt5-base-ft-ssc", token=hf_token)
+    model.eval()
+
     tokenizer = T5TokenizerFast.from_pretrained("google/mt5-base")
     tokenizer.add_special_tokens({'additional_special_tokens': ['<ZWJ>']})
-        
+
 
 def correct(text):
-    model.eval()
 
     text = re.sub(r'\u200d', '<ZWJ>', text)
     inputs = tokenizer(
@@ -48,10 +49,6 @@ def correct(text):
         token for token in tokens_list 
         if token == special_token_id_to_keep or token not in all_special_ids
     ]
-    # filtered_tokens = [
-    #     token for token in prediction 
-    #     if token == special_token_id_to_keep or token not in all_special_ids
-    # ]
 
     prediction_decoded = tokenizer.decode(filtered_tokens, skip_special_tokens=False).replace('\n', '').strip()