Repository: keras-team/keras-nlp

Issue:
Add `add_bos`, `add_eos` parameters to `SentencePieceTokenizer`.
**Is your feature request related to a problem? Please describe.**
[`SentencePieceTokenizer`](https://github.com/keras-team/keras-nlp/blob/v0.14.0/keras_nlp/src/tokenizers/sentence_piece_tokenizer.py#L47), which is a wrapper for [`tf_text.SentencepieceTokenizer`](https://github.com/tensorflow/text/blob/v2.16.1/tensorflow_text/python/ops/sentencepiece_tokenizer.py#L60), currently does not expose all the internal parameters that may be specified in [`tf_text.SentencepieceTokenizer.__init__()`](https://github.com/tensorflow/text/blob/v2.16.1/tensorflow_text/python/ops/sentencepiece_tokenizer.py#L75).
The parameters I am particularly interested in are `add_bos` and `add_eos`; in the current state, users must explicitly add the token ids for `'
'` and `'
'` (which default to 1 and 2) to the result of `SentencePieceTokenizer.tokenize()`.

**Describe the solution you'd like**
Add `add_bos=False, add_eos=False` to `SentencePieceTokenizer.__init__()`, save them in `self.add_bos` and `self.add_eos`, and use them in [`set_proto()`](https://github.com/keras-team/keras-nlp/blob/v0.14.0/keras_nlp/src/tokenizers/sentence_piece_tokenizer.py#L145) when initializing `tf_text.SentencepieceTokenizer`.

**Describe alternatives you've considered**
It's always possible to write a custom `Tokenizer` by subclassing, but the change seemed trivial enough.

Hints from issue comments:
@mattdangerw: The idea with KerasNLP tokenizers is to keep the basic functionality pretty pure. `string -> ragged int` and `ragged int -> string`. And then to add a layer like a `StartEndPacker` layer on top for bos, eos, padding, etc.

Is there a reason that doesn't work for your use case?@mattdangerw: It definitely might make sense to add the arguments here, as sentencepiece does have build in bos/eos tokens unlike other tokenizers. But would like to understand the overall use case first.@briango28: The primary use case I thought of was for porting preprocessing pipelines from other frameworks.
If a user wants to match the format of the original implementation, they must explicitly add a Packer layer after `SentencePieceTokenizer` unless the original implementation intentionally omitted the emission of bos & eos tokens.

While I feel this is perfectly fine in itself, I've found that having to manually load the sentencepiece model to copy bos/eos token ids or strings (in case of `dtype='string'`) to the Packer layer is rather unwieldy.
If adding too many options to `SentencePieceTokenizer.__init__()` is undesirable, perhaps adding attributes to expose this information would be acceptible?@mattdangerw: Yeah, I think this makes sense, just cause the beginning and end tokens are actual special in the sentencepiece library and proto. Let's just add them to init. Would you be interested in making the change?@briango28: Yes, I would. I'll follow up with a PR in the near future.@mattdangerw: This is fixed!

Solution patch:
diff --git a/keras_nlp/src/tokenizers/sentence_piece_tokenizer.py b/keras_nlp/src/tokenizers/sentence_piece_tokenizer.py
index 0d38998efe..37d609128b 100644
--- a/keras_nlp/src/tokenizers/sentence_piece_tokenizer.py
+++ b/keras_nlp/src/tokenizers/sentence_piece_tokenizer.py
@@ -67,6 +67,9 @@ class SentencePieceTokenizer(tokenizer.Tokenizer):
             for more details on the format.
         sequence_length: If set, the output will be converted to a dense
             tensor and padded/trimmed so all outputs are of `sequence_length`.
+        add_bos: Add beginning of sentence token to the result.
+        add_eos: Add end of sentence token to the result. Token is always
+            truncated if output is longer than specified `sequence_length`.
 
     References:
         - [Kudo and Richardson, 2018](https://arxiv.org/abs/1808.06226)
@@ -116,6 +119,8 @@ def __init__(
         proto=None,
         sequence_length=None,
         dtype="int32",
+        add_bos=False,
+        add_eos=False,
         **kwargs,
     ) -> None:
         if not is_int_dtype(dtype) and not is_string_dtype(dtype):
@@ -128,6 +133,8 @@ def __init__(
 
         self.proto = None
         self.sequence_length = sequence_length
+        self.add_bos = add_bos
+        self.add_eos = add_eos
         self.set_proto(proto)
         self.file_assets = [VOCAB_FILENAME]
 
@@ -172,6 +179,8 @@ def set_proto(self, proto):
         self._sentence_piece = tf_text.SentencepieceTokenizer(
             model=proto_bytes,
             out_type=self.compute_dtype,
+            add_bos=self.add_bos,
+            add_eos=self.add_eos,
         )
         # Keras cannot serialize a bytestring, so we base64 encode the model
         # byte array as a string for saving.
@@ -213,6 +222,8 @@ def get_config(self):
             {
                 "proto": None,  # Save vocabulary via an asset!
                 "sequence_length": self.sequence_length,
+                "add_bos": self.add_bos,
+                "add_eos": self.add_eos,
             }
         )
         return config


Test patch:
diff --git a/keras_nlp/src/tokenizers/sentence_piece_tokenizer_test.py b/keras_nlp/src/tokenizers/sentence_piece_tokenizer_test.py
index e4d09f44b6..16a76f810c 100644
--- a/keras_nlp/src/tokenizers/sentence_piece_tokenizer_test.py
+++ b/keras_nlp/src/tokenizers/sentence_piece_tokenizer_test.py
@@ -70,6 +70,29 @@ def test_string_tokenize(self):
             [["▁the", "▁quick", "▁brown", "▁fox."]],
         )
 
+    def test_scalar_bos_eos(self):
+        input_data = "the quick brown fox."
+        tokenizer = SentencePieceTokenizer(
+            proto=self.proto,
+            add_bos=True,
+            add_eos=True,
+        )
+        output_data = tokenizer(input_data)
+        self.assertAllEqual(output_data, [1, 6, 5, 3, 4, 2])
+
+    def test_string_bos_eos(self):
+        input_data = ["the quick brown fox."]
+        tokenizer = SentencePieceTokenizer(
+            proto=self.proto,
+            dtype="string",
+            add_bos=True,
+            add_eos=True,
+        )
+        output_data = tokenizer(input_data)
+        self.assertAllEqual(
+            output_data, [["
", "▁the", "▁quick", "▁brown", "▁fox.", "
"]]
+        )
+
     def test_detokenize(self):
         tokenizer = SentencePieceTokenizer(proto=self.proto)
         outputs = tokenizer.detokenize([6, 5, 3, 4])