model_run_utils: add STOP_SEQUENCE finish reason

Signed-off-by: Daniele Trifirò <[email protected]>
caikit · Sep 29, 2023 · bad7bc6 · bad7bc6
1 parent 7e31e6a
commit bad7bc6
Showing 1 changed file with 7 additions and 0 deletions.
diff --git a/caikit_nlp/toolkit/text_generation/model_run_utils.py b/caikit_nlp/toolkit/text_generation/model_run_utils.py
@@ -244,6 +244,13 @@ def generate_text_func(
         finish_reason = "EOS_TOKEN"
     elif generate_ids.size(1) - 1 == max_new_tokens:
         finish_reason = "MAX_TOKENS"
+    elif ("stopping_criteria" in gen_optional_params) and (
+        gen_optional_params["stopping_criteria"](
+            generate_ids,
+            None,  # scores, unused by SequenceStoppingCriteria
+        )
+    ):
+        finish_reason = "STOP_SEQUENCE"
     else:
         finish_reason = "OTHER"
     return GeneratedTextResult(