Repository: keras-team/keras-nlp

Issue:
Add a model proto training utility for sentence piece
Along with https://github.com/keras-team/keras-nlp/issues/248, we should add a utility to train a model proto for sentence piece. This can leverage the sentencepiece pip package to run their trainer.

We should not expose most of the option the sentencepiece trainer has, the goal will be for this to cover the majority of use cases, with advanced users till using the sentencepiece pip package directly.


Discussion:
@aflah02 here is the sentencepiece one!
Thanks @mattdangerw 😁

Resolving patch:
diff --git a/keras_nlp/tokenizers/__init__.py b/keras_nlp/tokenizers/__init__.py
--- a/keras_nlp/tokenizers/__init__.py
+++ b/keras_nlp/tokenizers/__init__.py
@@ -14,6 +14,9 @@
 
 from keras_nlp.tokenizers.byte_tokenizer import ByteTokenizer
 from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
+from keras_nlp.tokenizers.sentence_piece_tokenizer_trainer import (
+    compute_sentence_piece_proto,
+)
 from keras_nlp.tokenizers.tokenizer import Tokenizer
 from keras_nlp.tokenizers.unicode_codepoint_tokenizer import (
     UnicodeCodepointTokenizer,
diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer.py b/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer.py
@@ -0,0 +1,142 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Trainer for SentencePiece Tokenizer."""
+import io
+
+import tensorflow as tf
+
+try:
+    import sentencepiece as spm
+except ImportError:
+    spm = None
+
+
+def compute_sentence_piece_proto(
+    data,
+    vocabulary_size,
+    model_type="unigram",
+    proto_output_file=None,
+    lowercase=False,
+):
+    r"""A utility to train a SentencePiece vocabulary.
+
+    Trains a SentencePiece vocabulary from an input dataset or a list of
+    filenames.
+
+    If `data` is a list of filenames, the file format is required to be plain
+    text files, and the text will be read in line by line during training.
+
+    Args:
+        data: A `tf.data.Dataset`, or a list of filenames.
+        vocabulary_size: int. The maximum size of a vocabulary to be trained.
+        model_type: str, defaults to `"unigram"`. The model algorithm must be one
+            of `"unigram"`, `"bpe"`, `"word"` or `"char"`.
+        proto_output_file: str, defaults to `None`. If provided it will be used
+            as model_file which is passed to model_writer.
+            If `None`, the model_file will be `io.BytesIO` object.
+        lowercase: bool, defaults to `False`. If true, the input text will be
+            lowercased before tokenization.
+
+    Returns:
+        A `bytes` object with a serialized SentencePiece proto or
+        `None` if proto_output_file if provided.
+
+    Examples:
+
+    Basic Usage (from Dataset).
+    >>> inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+    >>> proto = keras_nlp.tokenizers.compute_sentence_piece_proto(inputs, vocabulary_size=15)
+    >>> tokenizer = keras_nlp.tokenizers.SentencePieceTokenizer(proto=proto)
+    >>> outputs = inputs.map(tokenizer)
+    >>> for output in outputs:
+    ...     print(output)
+    tf.Tensor([ 4  8 12  5  9 14  5  6 13  4  7 10 11  6 13],
+    shape=(15,), dtype=int32)
+
+    Basic Usage (with files).
+    ``` python
+    with open("test.txt", "w+") as f: f.write("Drifting Along\n")
+    inputs = ["test.txt"]
+    proto = keras_nlp.tokenizers.compute_sentence_piece_proto(
+         inputs, vocabulary_size=15, proto_output_file="model.spm")
+    tokenizer = keras_nlp.tokenizers.SentencePieceTokenizer(proto="model.spm")
+    ds = tf.data.Dataset.from_tensor_slices(["the quick brown fox."])
+    ds = ds.map(tokenizer)
+    for output in ds:
+         print(output)
+    tf.Tensor([ 4 14  1  4  1  5  1  4  1 12 11  1  6  4  9 11  1],
+    shape=(17,), dtype=int32)
+    ```
+
+    Usage with lowercase
+    >>> inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+    >>> proto = keras_nlp.tokenizers.compute_sentence_piece_proto(
+    ...     inputs, vocabulary_size=15, lowercase=True)
+    >>> tokenizer = keras_nlp.tokenizers.SentencePieceTokenizer(proto=proto)
+    >>> outputs = inputs.map(tokenizer)
+    >>> for output in outputs:
+    ...     print(output)
+    tf.Tensor([ 4  8 12  5  9 14  5  6 13  4  7 10 11  6 13],
+    shape=(15,), dtype=int32)
+    """
+
+    if spm is None:
+        raise ImportError(
+            "sentencepiece is not installed. Please install it via `pip install sentencepiece`."
+        )
+
+    if not isinstance(data, (list, tuple, tf.data.Dataset)):
+        raise ValueError(
+            "The `data` argument must be either `tf.data.Dataset` or `tuple` or `list`. "
+            f"Received: type(data)={type(data)}."
+        )
+
+    if model_type not in ["unigram", "bpe", "word", "char"]:
+        raise ValueError(
+            "The `model_type` argument must be one of `unigram`, `bpe`, `word`"
+            f"or `char`. Received: model_type={model_type}."
+        )
+
+    model_writer = (
+        open(proto_output_file, "wb") if proto_output_file else io.BytesIO()
+    )
+    is_dataset = isinstance(data, tf.data.Dataset)
+    if is_dataset:
+        spm.SentencePieceTrainer.train(
+            sentence_iterator=data.as_numpy_iterator(),
+            model_writer=model_writer,
+            vocab_size=vocabulary_size,
+            model_type=model_type,
+            normalization_rule_name="nmt_nfkc_cf" if lowercase else "nmt_nfkc",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+        )
+    else:
+        spm.SentencePieceTrainer.train(
+            input=data,
+            model_writer=model_writer,
+            vocab_size=vocabulary_size,
+            model_type=model_type,
+            normalization_rule_name="nmt_nfkc_cf" if lowercase else "nmt_nfkc",
+            pad_id=0,
+            unk_id=1,
+            bos_id=2,
+            eos_id=3,
+        )
+    if proto_output_file:
+        model_writer.close()
+    else:
+        return model_writer.getvalue()


Test patch:
diff --git a/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer_test.py b/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/tokenizers/sentence_piece_tokenizer_trainer_test.py
@@ -0,0 +1,98 @@
+# Copyright 2022 The KerasNLP Authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     https://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tests for SentencePiece Tokenizer Trainer."""
+
+import os
+import re
+
+import tensorflow as tf
+
+from keras_nlp.tokenizers.sentence_piece_tokenizer import SentencePieceTokenizer
+from keras_nlp.tokenizers.sentence_piece_tokenizer_trainer import (
+    compute_sentence_piece_proto,
+)
+
+
+class SentencePieceTokenizerTrainerTest(tf.test.TestCase):
+    def test_dataset_input(self):
+        test_text = ["Ninjas and Samurais"]
+        expected_output = [
+            [5, 9, 6, 7, 11, 4, 8, 5, 4, 7, 10, 5, 12, 4, 13, 15, 14, 4, 6, 8]
+        ]
+        data = tf.data.Dataset.from_tensor_slices(test_text)
+        proto = compute_sentence_piece_proto(data, 16)
+        tokenizer = SentencePieceTokenizer(proto=proto)
+        test_output = tokenizer(test_text)
+        self.assertAllEqual(expected_output, test_output)
+
+    def test_file_input(self):
+        test_text = "Ninja Land"
+        with open(os.path.join(self.get_temp_dir(), "test.txt"), "w+") as f:
+            f.write(test_text + "\n")
+        expected_output = [6, 8, 9, 5, 11, 4, 6, 7, 4, 5, 10]
+        proto = compute_sentence_piece_proto(
+            [os.path.join(self.get_temp_dir(), "test.txt")],
+            12,
+        )
+        tokenizer = SentencePieceTokenizer(proto=proto)
+        test_output = tokenizer(test_text)
+        self.assertAllEqual(expected_output, test_output)
+
+    def test_multiple_file_input(self):
+        with open(os.path.join(self.get_temp_dir(), "test1.txt"), "w+") as f:
+            f.write("Drifting Along\n")
+        with open(os.path.join(self.get_temp_dir(), "test2.txt"), "w+") as f:
+            f.write("Woah look there\n")
+        inputs = [
+            os.path.join(self.get_temp_dir(), "test1.txt"),
+            os.path.join(self.get_temp_dir(), "test2.txt"),
+        ]
+        proto = compute_sentence_piece_proto(inputs, 20)
+        tokenizer = SentencePieceTokenizer(proto=proto)
+        test_text = "Woah Along"
+        test_output = tokenizer(test_text)
+        expected_output = [4, 16, 5, 17, 9, 4, 15, 12, 5, 11, 18]
+        self.assertAllEqual(expected_output, test_output)
+
+    def test_invalid_input(self):
+        test_text_invalid = {"file1": "test.txt"}
+        with self.assertRaisesRegex(
+            ValueError,
+            re.escape(f"Received: type(data)={type(test_text_invalid)}."),
+        ):
+            compute_sentence_piece_proto(test_text_invalid, 10)
+
+    def test_lowercase(self):
+        inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+        proto = compute_sentence_piece_proto(
+            inputs, vocabulary_size=15, lowercase=True
+        )
+        tokenizer = SentencePieceTokenizer(proto=proto)
+        output = inputs.map(tokenizer).take(1).get_single_element()
+        expected_output = [4, 8, 12, 5, 9, 14, 5, 6, 13, 4, 7, 10, 11, 6, 13]
+        self.assertAllEqual(expected_output, output)
+
+    def test_proto_output_file(self):
+        inputs = tf.data.Dataset.from_tensor_slices(["Drifting Along"])
+        compute_sentence_piece_proto(
+            inputs,
+            vocabulary_size=15,
+            proto_output_file=os.path.join(self.get_temp_dir(), "model.spm"),
+        )
+        tokenizer = SentencePieceTokenizer(
+            proto=os.path.join(self.get_temp_dir(), "model.spm")
+        )
+        output = inputs.map(tokenizer).take(1).get_single_element()
+        expected_output = [4, 8, 12, 5, 9, 14, 5, 6, 13, 4, 7, 10, 11, 6, 13]
+        self.assertAllEqual(expected_output, output)