PR comment, update _tokenizer_type to global

Author: lessw2020

dist_run.py

@@ -50,6 +50,7 @@
 
 
 logger = SingletonLogger.get_logger()
+_tokenizer_type = None  # global variable to store the tokenizer type
 
 # Using model name to identify the model to load, for example "llama2-7b-chat".
 # You can change it to other values listed below.
@@ -85,8 +86,11 @@ def dict_to_args(dictionary: Dict[str, Any]) -> SimpleNamespace:
 def _build_chat_tokenizer(
     model_name: str,
     model_base_name: Optional[str] = None,
-) -> tuple[SentencePieceProcessor | TiktokenTokenizer, TokenizerType]:
-    """Builds a tokenizer for the given model name."""
+) -> SentencePieceProcessor | TiktokenTokenizer:
+    """Builds a tokenizer for the given model name, and sets the global tokenizer type variable"""
+
+    global _tokenizer_type
+
     # Try to infer the model base name from the model name:
     # e.g. "llama2-7b-chat" -> "llama2"
     if model_base_name is None:
@@ -113,15 +117,16 @@ def _build_chat_tokenizer(
     logger.info(
         f"using tokenizer = {tokenizer.__class__.__module__}.{tokenizer.__class__.__name__}"
     )
+    # set global variable _tokenizer_type
     if isinstance(tokenizer, TiktokenTokenizer):
-        tokenizer_type = TokenizerType.Tiktoken
+        _tokenizer_type = TokenizerType.Tiktoken
     elif isinstance(tokenizer, SentencePieceProcessor):
-        tokenizer_type = TokenizerType.SentencePiece
+        _tokenizer_type = TokenizerType.SentencePiece
     else:
         raise ValueError(f"Unknown tokenizer type: {tokenizer.__class__}")
 
-    logger.info(f"tokenizer type = {tokenizer_type}")
-    return tokenizer, tokenizer_type
+    logger.info(f"tokenizer type = {_tokenizer_type}")
+    return tokenizer
 
 
 def _load_model_weights(stage_module, distribution, device, model_config):
@@ -309,7 +314,7 @@ def main(args):
     config = TransformerArgs.from_params(model_config.transformer_args["text"])
     logger.info(f"Transformer Config: {config}")
 
-    tokenizer, tokenizer_type = _build_chat_tokenizer(model_name)
+    tokenizer = _build_chat_tokenizer(model_name)
 
     set_precision(model_dtype)
     logger.info(f"Using cache precision {model_dtype}")
@@ -554,15 +559,15 @@ def get_example_ins_outs(seqlen: int) -> Tuple[torch.Tensor, torch.Tensor]:
         # token ids. Thus cat'ing along dim 1.
         res = torch.cat(res, dim=1)
         res_list = res.tolist()
-        if tokenizer_type == TokenizerType.Tiktoken:
+        if _tokenizer_type == TokenizerType.Tiktoken:
             # For TiktokenTokenizer, we need to decode prompt by prompt.
             # TODO: is there a better way to do this?
             responses = [tokenizer.decode(sequence) for sequence in res_list]
-        elif tokenizer_type == TokenizerType.SentencePiece:  # SentencePieceProcessor
+        elif _tokenizer_type == TokenizerType.SentencePiece:  # SentencePieceProcessor
             # For SentencePieceProcessor, we can decode the entire 2D list at once.
             responses = tokenizer.decode(res_list)
         else:
-            raise ValueError(f"Unknown tokenizer type {tokenizer_type}")
+            raise ValueError(f"Unknown tokenizer type {_tokenizer_type}")
 
         # Show prompts and responses
         for prompt_text, response_text in zip(prompt, responses):