Merge pull request #353 from Zurnaz/llama_tpu_tokenizer_fix

fix: tpu tokenizers errors
henk717 · May 8, 2023 · 0f91298 · 0f91298
2 parents cb4af7e + d53726b
commit 0f91298
Show file tree

Hide file tree

Showing 2 changed files with 3 additions and 3 deletions.
diff --git a/modeling/inference_models/hf.py b/modeling/inference_models/hf.py
@@ -59,7 +59,7 @@ def decode_wrapper(self, token_ids, *args, **kwargs):
                         token_ids = [first]
                     elif len(token_ids) > 0:
                         first = int(token_ids[0])
-                elif token_ids:
+                elif token_ids is not None and len(token_ids) > 0:
                     first = token_ids[0]
                 result = original_decode(self, token_ids, *args, **kwargs)
                 if first is not None and first in has_prefix_space:

diff --git a/modeling/inference_models/hf_mtj.py b/modeling/inference_models/hf_mtj.py
@@ -17,6 +17,7 @@
     ModelCapabilities,
 )
 from modeling.inference_models.hf import HFInferenceModel
+from modeling.tokenizer import GenericTokenizer
 
 # This file shouldn't be imported unless using the TPU
 assert utils.koboldai_vars.use_colab_tpu
@@ -193,8 +194,7 @@ def _load(self, save_model: bool, initial_load: bool) -> None:
         utils.koboldai_vars.modeldim = int(
             tpu_mtj_backend.params.get("d_embed", tpu_mtj_backend.params["d_model"])
         )
-
-        self.tokenizer = tpu_mtj_backend.tokenizer
+        self.tokenizer = GenericTokenizer(tpu_mtj_backend.tokenizer)
 
         if (
             utils.koboldai_vars.badwordsids is koboldai_settings.badwordsids_default