feat(tokenizers): Add and plumb the option to use the "tokenizers" tokenizer

Branch: GraniteCodeSupport

Signed-off-by: Gabe Goodhart <[email protected]>

Author: gabe-l-hart

torchchat/model.py

@@ -272,7 +272,9 @@ class TransformerArgs:
     norm_eps: float = 1e-5
     multiple_of: int = 256
     ffn_dim_multiplier: Optional[int] = None
+    # Select the desired tokenizer. Defaults to sentencepiece
     use_tiktoken: bool = False
+    use_tokenizers: bool = False
     max_seq_length: int = 8192
     rope_scaling: Optional[Dict[str, Any]] = None
     # For pipeline parallel
@@ -329,12 +331,14 @@ class ModelArgs:
     model_type: ModelType
     transformer_args: Dict[str, Dict[str, Any]]
     use_tiktoken: bool
+    use_tokenizers: bool
 
     def __init__(
         self,
         transformer_args: Dict[str, Dict[str, Any]],
         model_type: ModelType = ModelType.TextOnly,
         use_tiktoken: bool = False,
+        use_tokenizers: bool = False,
     ) -> None:
         self._sanity_check(transformer_args, model_type)
 
@@ -343,6 +347,7 @@ def __init__(
 
         # Model-level attributes
         self.use_tiktoken = use_tiktoken
+        self.use_tokenizers = use_tokenizers
 
     def _sanity_check(
         self,
@@ -369,7 +374,8 @@ def from_params(cls, params_path):
             }
 
         use_tiktoken = loaded_params.get("use_tiktoken", False)
-        return cls(transformer_args, model_type, use_tiktoken)
+        use_tokenizers = loaded_params.get("use_tokenizers", False)
+        return cls(transformer_args, model_type, use_tiktoken, use_tokenizers)
 
     @classmethod
     def from_table(cls, name: str):