Build tokenizer from TokenizerArgs

Author: mreso

torchchat/distributed/dist_run.py

@@ -102,32 +102,11 @@ def decode(
 
 
 def _build_chat_tokenizer(
-    model_name: str,
-    model_base_name: Optional[str] = None,
+    tokenizer_args: TokenizerArgs,
 ) -> SentencePieceProcessor | TiktokenTokenizer:
     """Builds a tokenizer for the given model name"""
-
-    # Try to infer the model base name from the model name:
-    # e.g. "llama2-7b-chat" -> "llama2"
-    if model_base_name is None:
-        model_base_name = model_name.split("-")[0]
-        logger.info(
-            f"Using model base name '{model_base_name}' to build tokenizer. "
-            "If not found, please specify it using the `model_base_name` argument."
-        )
-
-    # Create base args for tokenizer
-    default_model_dir = Path(
-        os.getenv("TORCHCHAT_MODELDIR", "~/.torchchat/model-cache")
-    ).expanduser()
-
-    tokenconfig = {
-        "model_directory": default_model_dir,
-        "model": model_base_name,
-        "tokenizer_path": None,
-    }
-    args = dict_to_args(tokenconfig)
-    tokenizer_args = TokenizerArgs.from_args(args)
+    
+    tokenizer_args = TokenizerArgs.from_args(tokenizer_args)
     tokenizer = tokenizer_args.t
     assert tokenizer is not None, f"Failed to get tokenizer using {tokenconfig=}"
     logger.info(
@@ -313,9 +292,14 @@ def _cleanup():
 ]
 
 
-def main(args, pipe):
+def main(
+    builder_args,
+    tokenizer_args,
+    pipe,
+):
     model_name = "llama3"  # args.model_name
-    pp_degree = args.pp
+    # print(f"{builder_args.checkpoint_path=}")
+    pp_degree = builder_args.pp
 
     rank, world_size = _init_distributed()
     logger.info(f"Worker started: {rank=}, {world_size=}")
@@ -332,7 +316,7 @@ def main(args, pipe):
     config = TransformerArgs.from_params(model_config.transformer_args["text"])
     logger.info(f"Transformer Config: {config}")
 
-    tokenizer = _build_chat_tokenizer(model_name)
+    tokenizer = _build_chat_tokenizer(tokenizer_args)
 
     set_precision(model_dtype)
     logger.info(f"Using cache precision {model_dtype}")
@@ -385,7 +369,7 @@ def main(args, pipe):
     # Load weights
     logger.info(f"Loading weights for {pp_rank=} on {device=}")
     with CUDATrackTime() as timer:
-        _load_model_weights(model, distribution, device, config, args.chpt_from)
+        _load_model_weights(model, distribution, device, config, builder_args.chpt_from)
 
     logger.info(
         f"{color.green}Total weight loading time: {timer.get_time()} {timer.unit} for rank {rank}{color.reset}"

torchchat/distributed/generate.py

@@ -32,7 +32,9 @@ def _setup_env(world_size: int, rank: int, target: callable, *args, **kwargs):
     return target(*args, **kwargs)
 
 
-def _launch_distributed_inference(builder_args: BuilderArgs) -> None:
+def _launch_distributed_inference(
+    builder_args: BuilderArgs, tokenizer_args: TokenizerArgs
+) -> tuple[List]:
     # create programmatic elastic launch
     print("Launching distributed inference ...")
 
@@ -49,7 +51,7 @@ def _launch_distributed_inference(builder_args: BuilderArgs) -> None:
         pipes.append(server_pipe)
         proc = mp.Process(
             target=partial(_setup_env, num_processes_per_node, rank, main),
-            args=(builder_args, client_pipe),
+            args=(builder_args, tokenizer_args, client_pipe),
         )
         proc.start()
 
@@ -189,7 +191,9 @@ def __init__(
 
         self.check_args()
 
-        self.procs, self.pipes = _launch_distributed_inference(builder_args)
+        self.procs, self.pipes = _launch_distributed_inference(
+            builder_args, tokenizer_args
+        )
 
         self.loop = asyncio.new_event_loop()
         asyncio.set_event_loop(self.loop)

torchchat/generate.py

@@ -30,8 +30,8 @@
     BuilderArgs,
     TokenizerArgs,
 )
-from torchchat.model import Model, ModelType
 from torchchat.distributed.generate import DistributedGenerator
+from torchchat.model import Model, ModelType
 from torchchat.utils.build_utils import device_sync, set_precision
 from torchchat.utils.device_info import get_device_info
 
@@ -1228,7 +1228,6 @@ def main(args):
         )
         if torch.cuda.is_available():
             torch.cuda.reset_peak_memory_stats()
-        
 
         for _ in gen.chat(generator_args):
             pass
@@ -1248,5 +1247,3 @@ def main(args):
 
         print(f"Model output: {response}")
         dist_gen.shutdown()
-
-