Merge pull request caikit#228 from opendatahub-io/fix-input-token-count

model_run_utils: return correct input_token_count
gkumbhat · Oct 10, 2023 · a1b5674 · a1b5674
2 parents f6d04f0 + 0955ffc
commit a1b5674
Showing 1 changed file with 1 addition and 1 deletion.
diff --git a/caikit_nlp/toolkit/text_generation/model_run_utils.py b/caikit_nlp/toolkit/text_generation/model_run_utils.py
@@ -190,7 +190,7 @@ def generate_text_func(
     )
     inputs = {k: v.to(model.device) for k, v in tok_tensors.items()}
 
-    input_token_count = len(tok_tensors)
+    input_token_count = tok_tensors["input_ids"].size(1)
 
     gen_optional_params = __process_gen_args(
         tokenizer,