Allow tokenizer to customize stop_tokens

Author: qihqi

jetstream/engine/token_utils.py

@@ -187,7 +187,7 @@ def process_result_tokens(
   slot_valid = slot_data.valid
   slot_lengths = slot_data.lengths
   samples, speculations = slot_tokens.shape
-  stop_tokens = [tokenizer.eos_id, tokenizer.pad_id]
+  stop_tokens = tokenizer.stop_tokens
   # Stop anything which has reached it's max length.
   complete = complete | (slot_lengths > slot_max_length)
   if debug:
@@ -349,6 +349,11 @@ def bos_id(self) -> int:
     """ID of the BOS token."""
     return self.vocab.bos_id
 
+  @property
+  def stop_tokens(self) -> set[int]:
+    """ID of the stop token."""
+    return {self.eos_id, self.pad_id}
+
 
 class TikToken(tokenizer_api.Tokenizer):
   """Tokenizer to convert strings to token ids and vice-versa."""
@@ -394,6 +399,11 @@ def decode(self, token_ids: list[int]) -> str:
       str: String generated from the token ids.
     """
     return self.tokenizer.decode(token_ids)
+  
+  @property
+  def stop_tokens(self) -> set[int]:
+    """ID of the stop token."""
+    return self.tokenizer.stop_tokens
 
   @property
   def pad_id(self) -> int:

jetstream/engine/tokenizer_api.py

@@ -65,3 +65,8 @@ def eos_id(self) -> int:
   @abc.abstractmethod
   def bos_id(self) -> int:
     """ID of BOS token."""
+
+  @property
+  def stop_tokens(self) -> set[int]:
+    """ID of the stop token."""
+    return {self.eos_id, self.pad_id}