Repository: keras-team/keras-nlp

Issue:
Add a standalone tokenizer for GPT-2
We don't yet have a clear picture of all the preprocessing we would like to add for for GPT-2, but it should be relatively simpler to add a standalone tokenizer class based on https://github.com/keras-team/keras-nlp/pull/449

This will be a new class wrapping the `BytePairTokenizer`.

A few steps:

 - Add the `GPT2Tokenizer` class subclassing `BytePairTokenizer` following the example in https://github.com/keras-team/keras-nlp/pull/449
 - Check for the `
` token in the vocabulary and expose it as `self.end_token_id`.
 - Expose `GPT2Tokenizer` in `models/__init__.py`.
 - Add documentation and tests!


Discussion:
Working on this.

Resolving patch:
diff --git a/keras_nlp/models/__init__.py b/keras_nlp/models/__init__.py
--- a/keras_nlp/models/__init__.py
+++ b/keras_nlp/models/__init__.py
@@ -22,6 +22,7 @@
 )
 from keras_nlp.models.distilbert.distilbert_tasks import DistilBertClassifier
 from keras_nlp.models.gpt2.gpt2_models import GPT2
+from keras_nlp.models.gpt2.gpt2_preprocessing import GPT2Tokenizer
 from keras_nlp.models.roberta.roberta_models import Roberta
 from keras_nlp.models.roberta.roberta_preprocessing import RobertaPreprocessor
 from keras_nlp.models.roberta.roberta_tasks import RobertaClassifier
diff --git a/keras_nlp/models/gpt2/gpt2_preprocessing.py b/keras_nlp/models/gpt2/gpt2_preprocessing.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_preprocessing.py
@@ -0,0 +1,124 @@
+# 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.
+"""GPT-2 preprocessing layers."""
+
+from tensorflow import keras
+
+from keras_nlp.tokenizers.byte_pair_tokenizer import BytePairTokenizer
+from keras_nlp.utils.python_utils import classproperty
+
+
+@keras.utils.register_keras_serializable(package="keras_nlp")
+class GPT2Tokenizer(BytePairTokenizer):
+    """A GPT-2 tokenizer using Byte-Pair Encoding subword segmentation.
+
+    This tokenizer class will tokenize raw strings into integer sequences and
+    is based on `keras_nlp.tokenizers.BytePairTokenizer`. Unlike the
+    underlying tokenizer, it will check for all special tokens needed by GPT-2
+    models and provides a `from_preset()` method to automatically download
+    a matching vocabulary for a GPT-2 preset.
+
+    This tokenizer does not provide truncation or padding of inputs.
+
+    If input is a batch of strings (rank > 0), the layer will output a
+    `tf.RaggedTensor` where the last dimension of the output is ragged.
+
+    If input is a scalar string (rank == 0), the layer will output a dense
+    `tf.Tensor` with static shape `[None]`.
+
+    Args:
+        vocabulary: string or dict, maps token to integer ids. If it is a
+            string, it should be the file path to a json file.
+        merges: string or list, contains the merge rule. If it is a string,
+            it should be the file path to merge rules. The merge rule file
+            should have one merge rule per line. Every merge rule contains
+            merge entities separated by a space. Please refer to this example:
+            https://storage.googleapis.com/keras-nlp/models/roberta_base/merges.txt.
+
+    Examples:
+
+    Batched inputs.
+    >>> vocab = {"
": 0, "reful":1, "gent": 2, "Ġafter": 3}
+    >>> vocab = {**vocab, **{"noon": 4, "Ġsun": 5}}
+    >>> merges = ["Ġ a", "Ġ s", "r e", "f u", "g e", "n t", "e r", "n o", "o n"]
+    >>> merges += ["i g", "h t", "Ġs u", "Ġa f", "ge nt", "no on", "re fu"]
+    >>> merges += ["Ġsu n", "Ġaf t", "refu l", "Ġaft er"] # Ġ for whitespace
+    >>> inputs = [" afternoon sun", "refulgent sun"]
+    >>> tokenizer = keras_nlp.models.GPT2Tokenizer(
+    ...     vocabulary=vocab,
+    ...     merges=merges,
+    ... )
+    >>> tokenizer(inputs)
+    <tf.RaggedTensor [[3, 4, 5], [1, 2, 5]]>
+
+    Unbatched input.
+    >>> vocab = {"
": 0, "Ġafter": 1, "noon": 2, "Ġsun": 3}
+    >>> merges = ["Ġ a", "Ġ s", "e r", "n o", "o n", "i g", "h t", "Ġs u"]
+    >>> merges += ["Ġa f", "no on", "Ġsu n", "Ġaf t", "Ġaft er"]
+    >>> inputs = " afternoon sun"
+    >>> tokenizer = keras_nlp.models.GPT2Tokenizer(
+    ...     vocabulary=vocab,
+    ...     merges=merges,
+    ... )
+    >>> tokenizer(inputs)
+    <tf.Tensor: shape=(3,), dtype=int32, numpy=array([1, 2, 3], dtype=int32)>
+
+    Detokenization.
+    >>> vocab = {"
": 0, "Ġafter": 1, "noon": 2, "Ġsun": 3}
+    >>> merges = ["Ġ a", "Ġ s", "e r", "n o", "o n", "i g", "h t", "Ġs u"]
+    >>> merges += ["Ġa f", "no on", "Ġsu n", "Ġaf t", "Ġaft er"]
+    >>> inputs = " afternoon sun"
+    >>> tokenizer = keras_nlp.models.GPT2Tokenizer(
+    ...     vocabulary=vocab,
+    ...     merges=merges,
+    ... )
+    >>> tokenizer.detokenize(tokenizer.tokenize(inputs)).numpy().decode('utf-8')
+    ' afternoon sun'
+    """
+
+    def __init__(
+        self,
+        vocabulary,
+        merges,
+        **kwargs,
+    ):
+        super().__init__(
+            vocabulary=vocabulary,
+            merges=merges,
+            **kwargs,
+        )
+
+        # Check for necessary special tokens.
+        end_token = "
"
+
+        if end_token not in self.get_vocabulary():
+            raise ValueError(
+                f"Cannot find token `'{end_token}'` in the provided "
+                f"`vocabulary`. Please provide `'{end_token}'` in your "
+                "`vocabulary` or use a pretrained `vocabulary` name."
+            )
+
+        self.end_token_id = self.token_to_id(end_token)
+
+    @classproperty
+    def presets(cls):
+        raise NotImplementedError
+
+    @classmethod
+    def from_preset(
+        cls,
+        preset,
+        **kwargs,
+    ):
+        raise NotImplementedError


Test patch:
diff --git a/keras_nlp/models/gpt2/gpt2_preprocessing_test.py b/keras_nlp/models/gpt2/gpt2_preprocessing_test.py
new file mode 100644
--- /dev/null
+++ b/keras_nlp/models/gpt2/gpt2_preprocessing_test.py
@@ -0,0 +1,83 @@
+# 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 GPT-2 preprocessing layers."""
+
+import os
+
+import tensorflow as tf
+from absl.testing import parameterized
+from tensorflow import keras
+
+from keras_nlp.models.gpt2.gpt2_preprocessing import GPT2Tokenizer
+
+
+class GPT2TokenizerTest(tf.test.TestCase, parameterized.TestCase):
+    def setUp(self):
+        vocab = {
+            "
": 0,
+            "Ġair": 1,
+            "plane": 2,
+            "Ġat": 3,
+            "port": 4,
+            "Ġkoh": 5,
+            "li": 6,
+            "Ġis": 7,
+            "Ġthe": 8,
+            "Ġbest": 9,
+        }
+
+        merges = ["Ġ a", "Ġ t", "Ġ k", "Ġ i", "Ġ b", "Ġa i", "p l", "n e"]
+        merges += ["Ġa t", "p o", "r t", "o h", "l i", "Ġi s", "Ġb e", "s t"]
+        merges += ["Ġt h", "Ġai r", "pl a", "Ġk oh", "Ġth e", "Ġbe st", "po rt"]
+        merges += ["pla ne"]
+
+        self.tokenizer = GPT2Tokenizer(vocabulary=vocab, merges=merges)
+
+    def test_tokenize(self):
+        input_data = " airplane at airport"
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [1, 2, 3, 1, 4])
+
+    def test_tokenize_batch(self):
+        input_data = tf.constant([" airplane at airport", " kohli is the best"])
+        output = self.tokenizer(input_data)
+        self.assertAllEqual(output, [[1, 2, 3, 1, 4], [5, 6, 7, 8, 9]])
+
+    def test_detokenize(self):
+        input_tokens = [1, 2, 3, 1, 4]
+        output = self.tokenizer.detokenize(input_tokens)
+        self.assertEqual(output, " airplane at airport")
+
+    def test_vocabulary_size(self):
+        self.assertEqual(self.tokenizer.vocabulary_size(), 10)
+
+    @parameterized.named_parameters(
+        ("save_format_tf", "tf"), ("save_format_h5", "h5")
+    )
+    def test_saving_model(self, save_format):
+        input_data = tf.constant([" airplane at airport"])
+
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = self.tokenizer(inputs)
+        model = keras.Model(inputs, outputs)
+
+        path = os.path.join(self.get_temp_dir(), "model")
+        model.save(path, save_format=save_format)
+
+        restored_model = keras.models.load_model(path)
+        self.assertAllEqual(
+            model(input_data),
+            restored_model(input_data),
+        )