app.py small100 fix (#6)

- app.py small100 fix (21a49a060dde151d2a200777c5699440cb29cf6a)

app.py

@@ -8,6 +8,7 @@ from transformers import (
     AutoModelForSeq2SeqLM
 )
 import torch
+from tokenization_small100 import SMALL100Tokenizer 
 
 # import your chunking helpers
 from chunking import get_max_word_length, chunk_text
@@ -50,7 +51,7 @@ MODEL_MAP = {
 # Cache loaded models/tokenizers
 MODEL_CACHE = {}
 
-def load_model(model_id: str):
+def load_model(model_id: str, target_lang: str):
     """
     Load & cache:
       - facebook/mbart-* via MBart50TokenizerFast & MBartForConditionalGeneration
@@ -62,7 +63,7 @@ def load_model(model_id: str):
             tokenizer = MBart50TokenizerFast.from_pretrained(model_id)
             model     = MBartForConditionalGeneration.from_pretrained(model_id)
         elif model_id == "alirezamsh/small100":
-            tokenizer = AutoTokenizer.from_pretrained(model_id)
+            tokenizer = SMALL100Tokenizer.from_pretrained(model_id, tgt_lang=target_lang)
             model     = AutoModelForSeq2SeqLM.from_pretrained(model_id)
         else:
             tokenizer = MarianTokenizer.from_pretrained(model_id)
@@ -91,7 +92,7 @@ async def translate(request: Request):
         safe_limit = get_max_word_length([target_lang])
         chunks     = chunk_text(text, safe_limit)
 
-        tokenizer, model = load_model(model_id)
+        tokenizer, model = load_model(model_id, target_lang)
         full_translation = []
 
         for chunk in chunks: