model_run_utils: add STOP_SEQUENCE finish reason

Signed-off-by: Daniele Trifirò <dtrifiro@redhat.com>
caikit · Oct 2, 2023 · 4d15fda · 4d15fda
1 parent bb8ca31
commit 4d15fda
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
@@ -239,6 +239,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(