Dynamic gen: Identify specific stop condition when ending generation

Author: turboderp

exllamav2/generator/dynamic.py

@@ -21,7 +21,7 @@
 import itertools
 from dataclasses import dataclass
 # import xxhash
-# from line_profiler import profile
+from line_profiler import profile
 
 # TODO:
 #  - ExLlamaV2StreamingGenerator wrapper
@@ -893,6 +893,11 @@ def iterate(self) -> list[dict]:
                         "stop_string"
                         "max_new_tokens"
                         "end_filter"
+                    optional, if "eos_reason" == "stop_token":
+                        "eos_triggering_token_id": int
+                        "eos_triggering_token_str": str
+                    optional, if "eos_reason" == "stop_string":
+                        "eos_triggering_string": str
                     "full_completion": str  - full text completion
                     "new_tokens": int  - number of tokens generated
                     "time_enqueued": float  - time from job was enqueued until it started, in seconds
@@ -1849,7 +1854,9 @@ def emit(
             eos_reason: str = None,
             emit_held = False,
             suppressed_text = None,
-            suppressed_tokens = None
+            suppressed_tokens = None,
+            stop_token: int = None,
+            stop_string: str = None
         ):
             r = {
                 "job": self,
@@ -1860,6 +1867,15 @@ def emit(
 
             if eos_reason is not None:
                 r.update({ "eos_reason": eos_reason })
+                if eos_reason == "stop_token":
+                    id_to_piece = self.generator.tokenizer.get_id_to_piece_list(True)
+                    r.update({
+                        "eos_triggering_token_id": stop_token,
+                        "eos_triggering_token_str": id_to_piece[stop_token]
+                    })
+                    pass
+                if eos_reason == "stop_string":
+                    r.update({ "eos_triggering_string": stop_string })
 
             if emit_held:
                 if self.held_text != "":
@@ -1913,7 +1929,7 @@ def emit(
         # End on stop tokens
 
         if next_token.item() in self.stop_tokens:
-            return emit(results, emit_eos = True, eos_reason = "stop_token")
+            return emit(results, emit_eos = True, eos_reason = "stop_token", stop_token = next_token.item())
 
         # Decode and buffer output
 
@@ -2032,8 +2048,12 @@ def rewind_checkpoint():
                 self.stop_strings_utf32_buffer
             )
             if match >= 0:
+                held = self.held_text[match:]
                 self.held_text = self.held_text[:match]
-                return emit(results, emit_eos = True, emit_held = True, eos_reason = "stop_string")
+                for s in self.stop_strings:
+                    if held.startswith(s):
+                        return emit(results, emit_eos = True, emit_held = True, eos_reason = "stop_string", stop_string = s)
+                assert False, "Detected stop string but couldn't identify it (logic error)"
             if match == -2:
                 return emit(results)