Update gen_model_answer.py

Signed-off-by: Chenjie Luo <[email protected]>

Author: cjluo-nv

examples/llm_eval/gen_model_answer.py

@@ -178,21 +178,18 @@ def get_model_answers(
     nim_model=None,
 ):
     # Model Optimizer modification
+    tokenizer = get_tokenizer(model_path, trust_remote_code=args.trust_remote_code)
     if checkpoint_dir:
-        tokenizer = get_tokenizer(model_path, trust_remote_code=args.trust_remote_code)
-        if checkpoint_dir:
-            # get model type
-            last_part = os.path.basename(checkpoint_dir)
-            model_type = last_part.split("_")[0]
-            # Some models require to set pad_token and eos_token based on external config (e.g., qwen)
-            if model_type == "qwen":
-                tokenizer.pad_token = tokenizer.convert_ids_to_tokens(151643)
-                tokenizer.eos_token = tokenizer.convert_ids_to_tokens(151643)
-
-            assert LLM is not None, "tensorrt_llm APIs could not be imported."
-            model = LLM(checkpoint_dir, tokenizer=tokenizer)
-        else:
-            raise ValueError("checkpoint_dir is required for TensorRT LLM inference.")
+        # get model type
+        last_part = os.path.basename(checkpoint_dir)
+        model_type = last_part.split("_")[0]
+        # Some models require to set pad_token and eos_token based on external config (e.g., qwen)
+        if model_type == "qwen":
+            tokenizer.pad_token = tokenizer.convert_ids_to_tokens(151643)
+            tokenizer.eos_token = tokenizer.convert_ids_to_tokens(151643)
+
+        assert LLM is not None, "tensorrt_llm APIs could not be imported."
+        model = LLM(checkpoint_dir, tokenizer=tokenizer)
     elif not nim_model:
         model, _ = load_model(
             model_path,
@@ -205,7 +202,6 @@ def get_model_answers(
             cpu_offloading=False,
             debug=False,
         )
-        tokenizer = get_tokenizer(model_path, trust_remote_code=args.trust_remote_code)
         if args.quant_cfg:
             quantize_model(
                 model,