model_run_utils: fix MAX_TOKENS finish reason logic

Signed-off-by: Daniele Trifirò <[email protected]>
caikit · Oct 7, 2023 · c35fd62 · c35fd62
1 parent 9f83beb
commit c35fd62
Showing 1 changed file with 3 additions 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
@@ -237,7 +237,9 @@ def generate_text_func(
         generate_ids[0, -1] == tokenizer.eos_token_id
     ):
         finish_reason = "EOS_TOKEN"
-    elif generate_ids.size(1) - 1 == max_new_tokens:
+    elif (generate_ids.size(1) - 1 == max_new_tokens) or (
+        generate_ids.size(1) - inputs["input_ids"].size(1) == max_new_tokens
+    ):
         finish_reason = "MAX_TOKENS"
     elif ("stopping_criteria" in gen_optional_params) and (
         gen_optional_params["stopping_criteria"](