Tokenizer validation (#293)

* arg handling

* phase ordering issue resolved

* add args validation to eval, fix imports

* pass model, not model_args to validate_args

* read before write

* fixes

Author: mikekgfb

build/builder.py

@@ -124,13 +124,13 @@ def from_speculative_args(cls, args):  # -> BuilderArgs:
 @dataclass
 class TokenizerArgs:
     tokenizer_path: Optional[Union[Path, str]] = None
-    is_SentencePiece: bool = True
-    is_TikToken: bool = False
+    is_sentencepiece: bool = True
+    is_tiktoken: bool = False
 
     @classmethod
     def from_args(cls, args):  # -> TokenizerArgs:
-        is_SentencePiece = True
-        is_TikToken = False
+        is_sentencepiece = True
+        is_tiktoken = False
 
         if args.tokenizer_path:
             tokenizer_path = args.tokenizer_path
@@ -145,20 +145,20 @@ def from_args(cls, args):  # -> TokenizerArgs:
             raise RuntimeError(f"did not find tokenizer at {tokenizer_path}")
 
         if args.tiktoken:
-            is_SentencePiece = False
-            is_TikToken = True
+            is_sentencepiece = False
+            is_tiktoken = True
 
         return cls(
             tokenizer_path=tokenizer_path,
-            is_SentencePiece=is_SentencePiece,
-            is_TikToken=is_TikToken,
+            is_sentencepiece=is_sentencepiece,
+            is_tiktoken=is_tiktoken,
         )
 
 
 def _initialize_tokenizer(tokenizer_args: TokenizerArgs):
-    if tokenizer_args.is_SentencePiece:
+    if tokenizer_args.is_sentencepiece:
         return SentencePieceProcessor(model_file=str(tokenizer_args.tokenizer_path))
-    elif tokenizer_args.is_TikToken:
+    elif tokenizer_args.is_tiktoken:
         return TiktokenTokenizer(model_path=str(tokenizer_args.tokenizer_path))
     else:
         raise RuntimeError("must specify a valid tokenizer in TokenizerArgs")
@@ -347,3 +347,13 @@ def _initialize_model(
         model.to(dtype=builder_args.precision)
 
     return model
+
+def tokenizer_setting_to_name(tiktoken: bool = False) -> str:
+    return "TikToken" if tiktoken else "SentencePiece"
+
+def validate_args(model: Transformer, tokenizer_args: TokenizerArgs):
+    use_tiktoken = model.config.use_tiktoken
+    is_tiktoken = tokenizer_args.is_tiktoken
+    if use_tiktoken != is_tiktoken:
+        raise RuntimeError(f"model-specified tokenizer ({tokenizer_setting_to_name(use_tiktoken)} does not match provided tokenizer ({tokenizer_setting_to_name(is_tiktoken)}")
+    

build/model.py

@@ -37,7 +37,8 @@ class ModelArgs:
     norm_eps: float = 1e-5
     multiple_of: int = 256
     ffn_dim_multiplier: Optional[int] = None
-
+    use_tiktoken: bool = False
+    
     def __post_init__(self):
         if self.n_local_heads == -1:
             self.n_local_heads = self.n_heads
@@ -52,7 +53,10 @@ def __post_init__(self):
                 hidden_dim = int(self.ffn_dim_multiplier * hidden_dim)
             self.hidden_dim = find_multiple(hidden_dim, multiple_of)
         self.head_dim = self.dim // self.n_heads
+        if isinstance(self.use_tiktoken, str):
+            self.use_tiktoken = (self.use_tiktoken == "True")
 
+            
     @classmethod
     def from_params(cls, params_path):
         replace = [("rope_theta", "rope_base"), ("n_kv_heads", "n_local_heads")]
@@ -135,6 +139,7 @@ def from_name(cls, name: str):
         "n_layers": 32,
         "rope_base": 500000.0,  # rope_theta
         "vocab_size": 128256,
+        "use_tiktoken": True,
     },
     "Mistral-7B": {
         "n_layers": 32,

eval.py

@@ -14,6 +14,7 @@
 from build.builder import (
     _initialize_model,
     _initialize_tokenizer,
+    validate_args,
     BuilderArgs,
     TokenizerArgs,
 )
@@ -238,6 +239,7 @@ def main(args) -> None:
         builder_args,
         quantize,
     )
+    validate_args(model, tokenizer_args)
 
     if compile:
         assert not (

generate.py

@@ -21,6 +21,7 @@
     _initialize_model,
     _initialize_tokenizer,
     _load_model,
+    validate_args,
     BuilderArgs,
     TokenizerArgs,
 )
@@ -396,6 +397,7 @@ def _main(
 
     builder_args.setup_caches = False
     model = _initialize_model(builder_args, quantize)
+    validate_args(model, tokenizer_args)
 
     # will add a version of _initialize_model in future
     # (need additional args)