Update

data_utils.py

@@ -21,7 +21,7 @@ from tensorflow.python.framework import ops
 from tensorflow.python.ops import array_ops
 from transformers import PreTrainedTokenizerFast
 
-from . import seqio_tokenizer as vocab
+from .seqio_tokenizer import SentencePieceVocabulary, HfTokenizerWrapper, OLMoTokenizerWrapper
 from .constants import *
 from .utils import pop_metadata
 from .util import is_url
@@ -43,21 +43,21 @@ def build_tokenizer(
     return cache[cache_key]
 
   if tokenizer_type == 'llama':
-    tok = vocab.SentencePieceVocabulary(
+    tok = SentencePieceVocabulary(
       os.path.join(tokenizer_dir, "llama_tokenizer.model"),
       extra_ids=DEFAULT_EXTRA_IDS,
       reverse_extra_ids=True,
       extra_tokens=EXTRA_TOKENS if has_extra_token else None,
     )
   elif tokenizer_type == 'yi':
-      tok = vocab.SentencePieceVocabulary(
+      tok = SentencePieceVocabulary(
           os.path.join(tokenizer_dir, "yi_tokenizer.model"),
           extra_ids=DEFAULT_EXTRA_IDS,
           reverse_extra_ids=True,
           extra_tokens=EXTRA_TOKENS if has_extra_token else None,
       )
   elif tokenizer_type == 'mistral':
-      tok = vocab.SentencePieceVocabulary(
+      tok = SentencePieceVocabulary(
       os.path.join(tokenizer_dir, "mistral_tokenizer.model"),
       extra_ids=DEFAULT_EXTRA_IDS,
       reverse_extra_ids=True,
@@ -65,14 +65,14 @@ def build_tokenizer(
     )
 
   elif tokenizer_type == "mistral0.3":
-      tok = vocab.SentencePieceVocabulary(
+      tok = SentencePieceVocabulary(
           os.path.join(tokenizer_dir, "mistral0.3_tokenizer.model.v3"),
           extra_ids=DEFAULT_EXTRA_IDS,
           reverse_extra_ids=True,
           extra_tokens=EXTRA_TOKENS if has_extra_token else None,
       )
   elif tokenizer_type == 'gemma':
-      tok = vocab.SentencePieceVocabulary(
+      tok = SentencePieceVocabulary(
       os.path.join(tokenizer_dir, "gemma_tokenizer.model"),
       extra_ids=DEFAULT_EXTRA_IDS,
       reverse_extra_ids=True,
@@ -114,7 +114,7 @@ def build_tokenizer(
               ids = tokenizer.encode(tok, add_special_tokens=False)
               assert ids == [pad_tokenizer_to + ix]
 
-      tok = vocab.HfTokenizerWrapper(tokenizer, bos_token_id=bos_token_id, adds_space=adds_space)
+      tok = HfTokenizerWrapper(tokenizer, bos_token_id=bos_token_id, adds_space=adds_space)
   elif tokenizer_type.startswith("olmo-"):
       from olmo.tokenizer import Tokenizer
       assert Path(tokenizer_type[5:]).is_file()
@@ -123,7 +123,7 @@ def build_tokenizer(
           eos_token_id=olmo_eos_token_id,
           pad_token_id=-1,
       )
-      tok = vocab.OLMoTokenizerWrapper(tokenizer, bos_token_id=olmo_bos_token_id, adds_space=adds_space)
+      tok = OLMoTokenizerWrapper(tokenizer, bos_token_id=olmo_bos_token_id, adds_space=adds_space)
   else:
     raise NotImplementedError(tokenizer_type)
   cache[cache_key] = tok
@@ -131,7 +131,7 @@ def build_tokenizer(
 
 
 def get_special_token_ids(tokenizer):
-  if isinstance(tokenizer, (vocab.HfTokenizerWrapper, vocab.OLMoTokenizerWrapper)):
+  if isinstance(tokenizer, (HfTokenizerWrapper, OLMoTokenizerWrapper)):
       ids = tokenizer.encode("".join(EXTRA_TOKENS))
       if len(ids) == len(EXTRA_TOKENS) + 1:
           ids = ids[1:]