Fix tokenizer path if it is not correct

nemo_deploy/llm/inference/inference_base.py

@@ -245,7 +245,22 @@ def setup_megatron_model_and_tokenizer_for_inference(
     """
     dist_config = DistributedInitConfig(distributed_backend="nccl")
     torch_distributed_init(dist_config)
+
     model_config, mlm_args = load_model_config(checkpoint_path)
+
+    new_tokenizer_path = None
+    if hasattr(mlm_args, "tokenizer_model") and mlm_args.tokenizer_model:
+        tokenizer_model_path = Path(mlm_args.tokenizer_model)
+        if not tokenizer_model_path.exists():
+            # Attempt to reconstruct tokenizer path from checkpoint_path
+            checkpoint_dir = Path(checkpoint_path)
+            if checkpoint_dir.is_file():
+                checkpoint_dir = checkpoint_dir.parent
+            # Use the filename of the original tokenizer_model (if possible)
+            tokenizer_filename = tokenizer_model_path.name
+            new_tokenizer_path = checkpoint_dir / tokenizer_filename
+            mlm_args.tokenizer_model = str(new_tokenizer_path)
+
     if tensor_model_parallel_size is not None:
         model_config.tensor_model_parallel_size = tensor_model_parallel_size
     if pipeline_model_parallel_size is not None:
@@ -254,6 +269,7 @@ def setup_megatron_model_and_tokenizer_for_inference(
         model_config.context_parallel_size = context_parallel_size
     if expert_model_parallel_size is not None:
         model_config.expert_model_parallel_size = expert_model_parallel_size
+
     # Initialize Megatron for inference
     rng_config = RNGConfig(inference_rng_tracker=True)
     initialize_megatron_for_inference(model_config, dist_config, rng_config, micro_batch_size)
@@ -264,7 +280,10 @@ def setup_megatron_model_and_tokenizer_for_inference(
         megatron_args=mlm_args,
         use_cpu_init=False,
     )
-    tokenizer = load_tokenizer(checkpoint_path)
+    if new_tokenizer_path:
+        tokenizer = load_tokenizer(checkpoint_path, tokenizer_model=str(new_tokenizer_path))
+    else:
+        tokenizer = load_tokenizer(checkpoint_path)
     return model, tokenizer, mlm_args