fix: Added backwards compatibility for jamba tokenizer

AI21Labs · Jun 20, 2024 · a21d96d · a21d96d
1 parent 0df321c
commit a21d96d
Showing 1 changed file with 6 additions and 2 deletions.
diff --git a/ai21_tokenizer/tokenizer_factory.py b/ai21_tokenizer/tokenizer_factory.py
@@ -12,7 +12,8 @@
 
 class PreTrainedTokenizers:
     J2_TOKENIZER = "j2-tokenizer"
-    JAMBA_INSTRUCT_TOKENIZER = "jamba-tokenizer"
+    JAMBA_INSTRUCT_TOKENIZER = "jamba-instruct-tokenizer"
+    JAMBA_TOKENIZER = "jamba-tokenizer"
 
 
 class TokenizerFactory:
@@ -26,7 +27,10 @@ def get_tokenizer(
         cls,
         tokenizer_name: str = PreTrainedTokenizers.J2_TOKENIZER,
     ) -> BaseTokenizer:
-        if tokenizer_name == PreTrainedTokenizers.JAMBA_INSTRUCT_TOKENIZER:
+        if (
+            tokenizer_name == PreTrainedTokenizers.JAMBA_INSTRUCT_TOKENIZER
+            or tokenizer_name == PreTrainedTokenizers.JAMBA_TOKENIZER
+        ):
             return JambaInstructTokenizer(model_path=JAMBA_TOKENIZER_HF_PATH, cache_dir=os.getenv(_ENV_CACHE_DIR_KEY))
 
         if tokenizer_name == PreTrainedTokenizers.J2_TOKENIZER: