Allow tokenizer to customize stop_tokens (#84)

jetstream/engine/mock_utils.py

@@ -51,6 +51,7 @@ class TestVocab(Vocabulary):
   eos_id = 1
   bos_id = 2
   unk_id = 3
+  stop_tokens = {pad_id, eos_id}
   _base_vocab_size = 2**16
   tokenizer: TestTokenizer = TestTokenizer()
 

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:
@@ -395,6 +395,11 @@ def decode(self, token_ids: list[int]) -> str:
     """
     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:
     """ID of the pad token."""

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}