Repository: keras-team/keras-nlp
Pull request #18 (2022-02-07 22:51:41)

Word piece
Adds a base class for tokenizers and a layerized WordPiece wrapping tf.text.

Will add tests shortly, opening to have a diff to comment on for API design questions.

fchollet: The main design question is probably this:

- Using the "raw" `tokenize`/`detokenize` methods is fine in tf.data, but when building a model with the functional API, you need to be using `__call__()`
- Hence, we can make `__call__()` support all possible modes (with routing to one mode or another via keyword arguments), and have `tokenize`/`detokenize` defer to `__call__()` with the appropriate mode arg
- However, this raises the question of what the layer implementer is supposed to override

Some options:

**Option A**: have new implementer-facing methods `tokenize_impl()` / `detokenize_impl()` or similar which will be called by the base class `call()` method. This mirrors a common pattern where this type of situation is handled by having the implementer override a method that's separate from the interface method, which wraps the overridden method. For instance, when writing layers, you implement `call()` but not `__call__()`.

**Option B**: have instantiation-time wrapping of `tokenize()` / `detokenize()`. It's like option A but where the overridden `tokenize` gets replaced with a wrapped version of itself upon instantiation. It works but is somewhat dirty. We do this on the `Metric` base class with `update_state()`.

**Option C**: have implementers implement `call(self, ..., mode="tokenize")` with the convention that they should do `if mode == "tokenize": ...` / `if mode == "detokenize": ...`

I think the simplest is probably option C. WDYT?
mattdangerw: Thanks! Agreed on the framing. I am not immediately sure on best solution.

We might have a few more things we need to cover as the base class evolves. Tf text has a `tokenizer_with_offsets` which gives byte offset into the original text (useful for things like named entity recognition I believe). If we start to have a few different flags we rout through call things might look a little more chaotic on the subclass. E.g.

```
class MyTokenizer:
  def call(self, inputs, detokenize=False, add_offsets=False):
    if detokenize:
      return some_thing
    if add_offsets:
      return other_thing
    if detokenize && add_offsets:
      raise "both not supported"
    return main_tokenization
```

So I might have a slight preference for **Option A**. It would reduce the need to replicate branching logic on the subclass, and allow us to default unhanded methods which an exception. That could be useful for a broader ecosystem. We add some new method down the line to the base class, and all existing tokenizers will correctly raise a `NotImplementError` until they update to support.
mattdangerw: I guess I was missing part of your proposal for **Option C**, we could follow a call block like:

```
  def call(self, inputs, mode="tokenize"):
    if mode == "tokenize":
      ...
    elif mode == "tokenize_with_offsets":
      ...
    elif mode == "detokenize":
      ...
    else:
      raise NotImplementedError
```

Given that, I'm probably not very opinionated between options A and C.
mattdangerw: Sounds good to me! Thanks for the input here. I'll finish up this PR with option C and we can see how it looks.
mattdangerw: Another thought, it would be good to be able to document the tokenize, detokenize methods for each class, with more details on how they operate per tokenizer. We did this for `adapt` on preprocessing layers by just redefining the method and chaining to super(), though that's another slightly awkward point for subclassers. Probably no need to worry about this until we start generating documentation, but something to think about.
fchollet: Perhaps we could simply require folks to add relevant sections in the class docstrings? Layer docs are typically generated purely from class docstrings.
mattdangerw: > Layer docs are typically generated purely from class docstrings.

Not for the tf docs (e.g. [adapt](https://www.tensorflow.org/api_docs/python/tf/keras/layers/StringLookup#adapt)), which I guess is why we did that workaround. But seems fine to do here.
mattdangerw: Okay, played around more with this and we now think **Option B** may be the best bet. **Option C** became cumbersome when subclassing a subclass of Tokenizer, when it was not easy to extend the tokenize and detokenize methods without redefining and chaining to the private helpers within call.

**Option B** is probably the nicest for end users. Who can just implement the same method that they call on instantiated objects, and not need to worry about functional models in the mental model at all.

A version of **Option B** base class is now on this PR.
mattdangerw: I have also tried moving these to a separate `keras_nlp.tokenzers.SomeTokenizer` namespace. Upon thinking about it more, I think there will be an expectation to have tokenizers grouped together from other similar APIs, and overall I think it will bring more clarity.

Since this is a custom base class, and users will not be implementing layers in the normal sense (the expectation is now to not implement call), I think this is a reasonable naming choice.

Patch:
diff --git a/keras_nlp/__init__.py b/keras_nlp/__init__.py
index bfe2d9fb..46b298c1 100644
--- a/keras_nlp/__init__.py
+++ b/keras_nlp/__init__.py
@@ -13,5 +13,6 @@
 # limitations under the License.
 
 from keras_nlp import layers
+from keras_nlp import tokenizers
 
 __version__ = "0.1.0"
diff --git a/keras_nlp/tokenizers/__init__.py b/keras_nlp/tokenizers/__init__.py
new file mode 100644
index 00000000..d6eb0b63
--- /dev/null
+++ b/keras_nlp/tokenizers/__init__.py
@@ -0,0 +1,16 @@
+# 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.
+
+from keras_nlp.tokenizers.tokenizer import Tokenizer
+from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
diff --git a/keras_nlp/tokenizers/tokenizer.py b/keras_nlp/tokenizers/tokenizer.py
new file mode 100644
index 00000000..62d98ecb
--- /dev/null
+++ b/keras_nlp/tokenizers/tokenizer.py
@@ -0,0 +1,96 @@
+# 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.
+
+from tensorflow import keras
+
+
+class Tokenizer(keras.layers.Layer):
+    """A base class for tokenizer layers.
+
+    The class is intended as a base class when implementing a tokenizer as a
+    `keras.layers.Layer`. It contains two new methods `tokenize()` and
+    `detokenize()`.
+
+    Subclassers should always implement the `tokenize` method, which will also
+    also be the default when invoking the layer directly on inputs
+
+    If a layer does not support detokenization (the tokenization step is not
+    reversible), the `detokenize()` method can be skipped.
+
+    Examples:
+
+    ```python
+    class WhitespaceTokenizer(keras_nlp.tokenizers.Tokenizer):
+        def tokenize(self, inputs):
+            return tf.strings.split(inputs).to_tensor()
+
+        def detokenizer(self, inputs):
+            return tf.strings.reduce_join([inputs], separator=" ", axis=-1)
+
+    tokenizer = WhitespaceTokenizer()
+    ```
+    """
+
+    def __new__(cls, *args, **kwargs):
+        # Wrap the `tokenize` and `detokenize` methods so they route through
+        # __call__. This is needed for functional model support.
+        obj = super().__new__(cls, *args, **kwargs)
+        obj._tokenize_without_call = obj.tokenize
+        obj._detokenize_without_call = obj.detokenize
+        obj.tokenize = obj._tokenize_with_call
+        obj.detokenize = obj._detokenize_with_call
+        return obj
+
+    def tokenize(self, inputs, *args, **kwargs):
+        """Transform input tensors of strings into output tokens.
+
+        Args:
+            inputs: Input tensor, or dict/list/tuple of input tensors.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        raise NotImplementedError(
+            "No implementation of `tokenize()` was found for "
+            f"{self.__class__.__name__}. All tokenizers should implement "
+            "`tokenize()`."
+        )
+
+    def detokenize(self, inputs, *args, **kwargs):
+        """Transform tokens back into strings.
+
+        Args:
+            inputs: Input tensor, or dict/list/tuple of input tensors.
+            *args: Additional positional arguments.
+            **kwargs: Additional keyword arguments.
+        """
+        raise NotImplementedError(
+            "No implementation of `detokenize()` was found for "
+            f"{self.__class__.__name__}."
+        )
+
+    def _tokenize_with_call(self, *args, **kwargs):
+        return self(*args, mode="tokenize", **kwargs)
+
+    def _detokenize_with_call(self, *args, **kwargs):
+        return self(*args, mode="detokenize", **kwargs)
+
+    def call(self, *args, mode="tokenize", training=None, **kwargs):
+        if mode == "tokenize":
+            return self._tokenize_without_call(*args, **kwargs)
+        elif mode == "detokenize":
+            return self._detokenize_without_call(*args, **kwargs)
+        else:
+            raise ValueError(
+                f"Unsupported tokenizer mode. Received: mode={mode}"
+            )
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
new file mode 100644
index 00000000..7e89ed49
--- /dev/null
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -0,0 +1,277 @@
+# 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.
+
+from typing import Any
+from typing import Dict
+from typing import Iterable
+from typing import List
+from typing import Union
+
+import tensorflow as tf
+import tensorflow_text as tf_text
+
+from keras_nlp.tokenizers import tokenizer
+
+# Matches whitespace and control characters.
+WHITESPACE_REGEX = r"|".join(
+    [
+        r"\s",
+        # Invisible control characters
+        r"\p{Cc}",
+        r"\p{Cf}",
+    ]
+)
+
+# Matches punctuation compatible with the original bert implementation.
+PUNCTUATION_REGEX = r"|".join(
+    [
+        # Treat all non-letter/number ASCII as punctuation.
+        # Characters such as "^", "$", and "`" are not in the Unicode
+        # Punctuation class but we treat them as punctuation anyways.
+        r"[!-/]",
+        r"[:-@]",
+        r"[\[-`]",
+        r"[{-~]",
+        # Unicode punctuation class.
+        r"[\p{P}]",
+        # More unicode ranges.
+        r"[\x{4E00}-\x{9FFF}]",
+        r"[\x{3400}-\x{4DBF}]",
+        r"[\x{20000}-\x{2A6DF}]",
+        r"[\x{2A700}-\x{2B73F}]",
+        r"[\x{2B740}-\x{2B81F}]",
+        r"[\x{2B820}-\x{2CEAF}]",
+        r"[\x{F900}-\x{FAFF}]",
+        r"[\x{2F800}-\x{2FA1F}]",
+    ]
+)
+
+# Matches both whitespace and punctuation.
+WHITESPACE_AND_PUNCTUATION_REGEX = r"|".join(
+    [
+        WHITESPACE_REGEX,
+        PUNCTUATION_REGEX,
+    ]
+)
+
+
+class WordPieceTokenizer(tokenizer.Tokenizer):
+    """A word piece tokenizer layer.
+
+    This layer provides an efficient, in graph, implementation of the WordPiece
+    algorithm used by BERT and other models.
+
+    To make this layer more useful out of the box, the layer will pre-tokenize
+    the input, which will optionally lower-case, strip accents, and split the
+    input on whitespace and punctuation. Each of these pre-tokenization steps is
+    not reversible. The `detokenize` method will join words with a space, and
+    will not invert `tokenize` exactly.
+
+    If a more custom pre-tokenization step is desired, the layer can be
+    configured to apply only the strict WordPiece algorithm by passing
+    `lowercase=False`, `strip_accents=False` and `split_pattern=None`. In
+    this case, inputs should be pre-split string tensors or ragged tensors.
+
+    By default, the layer will output a `tf.RaggedTensor` where the last
+    dimension of the output is ragged after whitespace splitting and sub-word
+    tokenizing. If `sequence_length` is set, the layer will output a dense
+    `tf.Tensor` where all inputs have been padded or truncated to
+    `sequence_length`. The output dtype can be controlled via the `dtype`
+    argument, which should be either an integer or string type.
+
+    Args:
+        vocabulary: A list of strings or a string string filename path. If
+            passing a list, each element of the list should be a single word
+            piece token string. If passing a filename, the file should be a
+            plain text file containing a single word piece token per line.
+        sequence_length: If set, the output will be converted to a dense
+            tensor and padded/trimmed so all outputs are of sequence_length.
+        lowercase: If true, the input text will be first lowered before
+            tokenization.
+        strip_accents: If true, all accent marks will be removed from text
+            before tokenization.
+        split_pattern: A regex pattern to match delimiters to split, or None
+            indicating that the input is pre-split and no splitting should be
+            performed. By default, all whitespace and punctuation marks will
+            be split on.
+        keep_pattern: A regex pattern of delimiters contained in the
+            `split_pattern` of delimeters that should be kept as independent
+            tokens. By default, all punctuation marks will be kept as tokens.
+        suffix_indicator: The characters prepended to a wordpiece to indicate
+            that it is a suffix to another subword.
+        oov_token: The string value to substitute for an unknown token. It
+            must be included in the vocab.
+
+    References:
+        - [Schuster and Nakajima, 2012](https://research.google/pubs/pub37842/)
+        - [Song et al., 2020](https://arxiv.org/abs/2012.15524)
+
+    Examples:
+
+    Ragged outputs.
+    >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
+    >>> inputs = ["The quick brown fox."]
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
+    >>> tokenizer(inputs)
+    <tf.RaggedTensor [[1, 2, 3, 4, 5, 6, 7]]>
+
+    Dense outputs.
+    >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
+    >>> inputs = ["The quick brown fox."]
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
+    ...     vocabulary=vocab, sequence_length=10)
+    >>> tokenizer(inputs)
+    <tf.Tensor: shape=(1, 10), dtype=int32,
+        numpy=array([[1, 2, 3, 4, 5, 6, 7, 0, 0, 0]]>
+
+    String output.
+    >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
+    >>> inputs = ["The quick brown fox."]
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(
+    ...     vocabulary=vocab, dtype="string")
+    >>> tokenizer(inputs)
+    <tf.RaggedTensor [[b'the', b'qu', b'##ick', b'br', b'##own', b'fox', b'.']]>
+
+    Detokenization.
+    >>> vocab = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
+    >>> inputs = ["The quick brown fox."]
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
+    >>> tokenizer.detokenize(tokenizer.tokenize(inputs))
+    <tf.Tensor: shape=(1,), dtype=string,
+        numpy=array([b"the quick brown fox ."], dtype=object)>
+
+    Custom splitting.
+    >>> vocab = ["[UNK]", "fox", ","]
+    >>> inputs = ["fox,,fox,fox"]
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,
+    ...     split_pattern=",", keep_pattern=",", dtype='string')(inputs)
+    <tf.RaggedTensor [[b'fox', b',', b',', b'fox', b',', b'fox']]>
+    >>> keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab,
+    ...     split_pattern=",", keep_pattern="", dtype='string')(inputs)
+    <tf.RaggedTensor [[b'fox', b'fox', b'fox']]>
+    """
+
+    def __init__(
+        self,
+        vocabulary: Union[Iterable[str], str] = None,
+        sequence_length: int = None,
+        lowercase: bool = True,
+        strip_accents: bool = True,
+        split_pattern: str = WHITESPACE_AND_PUNCTUATION_REGEX,
+        keep_pattern: str = PUNCTUATION_REGEX,
+        suffix_indicator: str = "##",
+        oov_token: str = "[UNK]",
+        **kwargs,
+    ) -> None:
+        # Check dtype and provide a default.
+        if "dtype" not in kwargs or kwargs["dtype"] is None:
+            kwargs["dtype"] = tf.int32
+        else:
+            dtype = tf.dtypes.as_dtype(kwargs["dtype"])
+            if not dtype.is_integer and dtype != tf.string:
+                raise ValueError(
+                    "Output dtype must be an integer type of a string. "
+                    f"Received: dtype={dtype}"
+                )
+
+        super().__init__(**kwargs)
+
+        if isinstance(vocabulary, str):
+            self._vocab = [
+                line.rstrip() for line in tf.io.gfile.GFile(vocabulary)
+            ]
+        elif isinstance(vocabulary, Iterable):
+            # Make a copy.
+            self._vocab = list(vocabulary)
+        else:
+            raise ValueError(
+                "Vocabulary must be an file path or list of terms. "
+                f"Received: vocabulary={vocabulary}"
+            )
+
+        self._sequence_length = sequence_length
+        self._lowercase = lowercase
+        self._strip_accents = strip_accents
+        self._split_pattern = split_pattern
+        self._keep_pattern = keep_pattern
+        self._suffix_indicator = suffix_indicator
+        self._oov_token = oov_token
+
+        self._fast_word_piece = tf_text.FastWordpieceTokenizer(
+            vocab=self._vocab,
+            token_out_type=self.compute_dtype,
+            suffix_indicator=suffix_indicator,
+            unknown_token=oov_token,
+            no_pretokenization=True,
+            support_detokenization=True,
+        )
+
+    def get_vocabulary(self) -> List[str]:
+        """Get the tokenizer vocabulary as a list of strings tokens."""
+        return self._vocab
+
+    def vocabulary_size(self) -> int:
+        """Get the size of the tokenizer vocabulary."""
+        return len(self._vocab)
+
+    def get_config(self) -> Dict[str, Any]:
+        config = super().get_config()
+        config.update(
+            {
+                # Ideally a vocabulary would be saved as a plain text asset in
+                # the saved model. We have no good way to support this
+                # currently, so we save the vocabulary in the config.
+                "vocabulary": self._vocab,
+                "sequence_length": self._sequence_length,
+                "lowercase": self._lowercase,
+                "strip_accents": self._strip_accents,
+                "split_pattern": self._split_pattern,
+                "keep_pattern": self._keep_pattern,
+                "suffix_indicator": self._suffix_indicator,
+                "oov_token": self._oov_token,
+            }
+        )
+        return config
+
+    def tokenize(self, inputs):
+        # Optionally normalize and split inputs.
+        if self._lowercase:
+            inputs = tf_text.case_fold_utf8(inputs)
+        if self._strip_accents:
+            # Normalize unicode to NFD, which splits out accent mark characters.
+            inputs = tf_text.normalize_utf8(inputs, "NFD")
+            # Remove the accent marks.
+            inputs = tf.strings.regex_replace(inputs, r"\p{Mn}", "")
+        if self._split_pattern:
+            inputs = tf_text.regex_split(
+                inputs,
+                delim_regex_pattern=self._split_pattern,
+                keep_delim_regex_pattern=self._keep_pattern,
+            )
+
+        # Apply word piece and coerce shape for outputs.
+        tokens = self._fast_word_piece.tokenize(inputs)
+        # By default tf.text tokenizes text with two ragged dimensions (one for
+        # split words and one for split subwords). We will collapse to a single
+        # ragged dimension which is a better out of box default.
+        tokens = tokens.merge_dims(-2, -1)
+        # Convert to a dense output if `sequence_length` is set.
+        if self._sequence_length:
+            output_shape = tokens.shape.as_list()
+            output_shape[-1] = self._sequence_length
+            tokens = tokens.to_tensor(shape=output_shape)
+        return tokens
+
+    def detokenize(self, inputs):
+        return self._fast_word_piece.detokenize(inputs)


Test patch:
diff --git a/keras_nlp/tokenizers/tokenizer_test.py b/keras_nlp/tokenizers/tokenizer_test.py
new file mode 100644
index 00000000..9df3b3e1
--- /dev/null
+++ b/keras_nlp/tokenizers/tokenizer_test.py
@@ -0,0 +1,59 @@
+# 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.
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.tokenizers.tokenizer import Tokenizer
+
+
+class SimpleTokenizer(Tokenizer):
+    __test__ = False  # for pytest
+
+    def tokenize(self, inputs):
+        return tf.strings.split(inputs).to_tensor()
+
+    def detokenize(self, inputs):
+        return tf.strings.reduce_join([inputs], separator=" ", axis=-1)
+
+
+class TokenizerTest(tf.test.TestCase):
+    def test_tokenize(self):
+        input_data = ["the quick brown fox"]
+        tokenizer = SimpleTokenizer()
+        tokenize_output = tokenizer.tokenize(input_data)
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(tokenize_output, [["the", "quick", "brown", "fox"]])
+        self.assertAllEqual(call_output, [["the", "quick", "brown", "fox"]])
+
+    def test_detokenize(self):
+        input_data = ["the", "quick", "brown", "fox"]
+        tokenizer = SimpleTokenizer()
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(detokenize_output, ["the quick brown fox"])
+
+    def test_functional_model(self):
+        input_data = tf.constant(["the   quick   brown   fox"])
+        tokenizer = SimpleTokenizer()
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = tokenizer.detokenize(tokenizer.tokenize(inputs))
+        model = keras.Model(inputs, outputs)
+        model_output = model(input_data)
+        # There appears to be a bug with shape inference for ragged reduce_join.
+        # The second dimension should be removed.
+        self.assertAllEqual(model_output, [["the quick brown fox"]])
+
+    def test_missing_tokenize_raises(self):
+        with self.assertRaises(NotImplementedError):
+            Tokenizer()(["the quick brown fox"])
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
new file mode 100644
index 00000000..d57c5427
--- /dev/null
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -0,0 +1,209 @@
+# 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.
+
+import os
+
+import tensorflow as tf
+from tensorflow import keras
+
+from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
+
+
+class WordPieceTokenizerTest(tf.test.TestCase):
+    def test_tokenize(self):
+        input_data = ["the quick brown fox."]
+        vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data)
+        call_output = tokenizer(input_data)
+        tokenize_output = tokenizer.tokenize(input_data)
+        self.assertIsInstance(call_output, tf.RaggedTensor)
+        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5, 6, 7]])
+        self.assertAllEqual(tokenize_output, [[1, 2, 3, 4, 5, 6, 7]])
+
+    def test_dense_output(self):
+        input_data = ["the quick brown fox."]
+        vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data, sequence_length=10
+        )
+        call_output = tokenizer(input_data)
+        self.assertIsInstance(call_output, tf.Tensor)
+        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5, 6, 7, 0, 0, 0]])
+
+    def test_string_tokenize(self):
+        input_data = ["the quick brown fox"]
+        vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data, dtype="string")
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(
+            call_output,
+            tf.ragged.constant([["the", "qu", "##ick", "br", "##own", "fox"]]),
+        )
+
+    def test_detokenize(self):
+        input_data = [[1, 2, 3, 4, 5, 6]]
+        vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data)
+        detokenize_output = tokenizer.detokenize(input_data)
+        self.assertAllEqual(detokenize_output, ["the quick brown fox"])
+
+    def test_accessors(self):
+        vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data)
+        self.assertEqual(tokenizer.vocabulary_size(), 7)
+        self.assertEqual(
+            tokenizer.get_vocabulary(),
+            ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"],
+        )
+
+    def test_special_tokens(self):
+        input_data = ["quick brown whale"]
+        vocab_data = ["@UNK@", "qu", "@@ick", "br", "@@own", "fox"]
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data,
+            oov_token="@UNK@",
+            suffix_indicator="@@",
+            dtype="string",
+        )
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(
+            call_output,
+            tf.ragged.constant([["qu", "@@ick", "br", "@@own", "@UNK@"]]),
+        )
+
+    def test_lowercase(self):
+        input_data = ["the QUicK brOWN FOX"]
+        vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data)
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5, 6]])
+
+    def test_skip_lowercase(self):
+        input_data = ["the QUicK brOWN FOX"]
+        vocab_data = ["[UNK]", "the", "QU", "##icK", "br", "##OWN", "fox"]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data, lowercase=False)
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5, 0]])
+
+    def test_strip_accents(self):
+        input_data = ["á é í ó ú"]
+        vocab_data = ["[UNK]", "a", "e", "i", "o", "u"]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data)
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5]])
+
+    def test_skip_strip_accents(self):
+        input_data = ["á é í ó ú"]
+        vocab_data = ["[UNK]", "á", "é", "í", "ó", "ú"]
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data, strip_accents=False
+        )
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5]])
+
+    def test_custom_spliting(self):
+        input_data = ["a,e,i,o,u"]
+        vocab_data = ["[UNK]", "a", "e", "i", "o", "u", ","]
+        # Don't keep commas.
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data, split_pattern=",", keep_pattern=""
+        )
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5]])
+        # Keep commas.
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data, split_pattern=",", keep_pattern=","
+        )
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [[1, 6, 2, 6, 3, 6, 4, 6, 5]])
+
+    def test_no_spliting(self):
+        input_data = ["t o k e n", "m i s s i n g", "t o k e n"]
+        vocab_data = ["[UNK]", "t o k e n"]
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data, split_pattern=None
+        )
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [1, 0, 1])
+
+    def test_word_piece_only(self):
+        input_data = ["the", "quíck", "Brówn", "Fóx"]
+        vocab_data = ["[UNK]", "the", "qu", "##íck", "Br", "##ówn", "Fóx"]
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data,
+            lowercase=False,
+            strip_accents=False,
+            split_pattern=None,
+        )
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [1, 2, 3, 4, 5, 6])
+
+    def test_functional_model(self):
+        input_data = tf.constant(["the quick brown fox"])
+        vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data)
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = tokenizer.detokenize(tokenizer.tokenize(inputs))
+        model = keras.Model(inputs, outputs)
+        model_output = model(input_data)
+        self.assertAllEqual(model_output, ["the quick brown fox"])
+
+    def test_from_file(self):
+        vocab_path = os.path.join(self.get_temp_dir(), "vocab.txt")
+        input_data = ["the quick brown fox."]
+        vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox", "."]
+        with tf.io.gfile.GFile(vocab_path, "w") as file:
+            for piece in vocab_data:
+                file.write(piece + "\n")
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_path)
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(call_output, [[1, 2, 3, 4, 5, 6, 7]])
+
+    def test_config(self):
+        input_data = ["quick brOWN whale"]
+        vocab_data = ["@UNK@", "qu", "@@ick", "br", "@@OWN", "fox"]
+        original_tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data,
+            lowercase=False,
+            oov_token="@UNK@",
+            suffix_indicator="@@",
+            dtype="string",
+        )
+        cloned_tokenizer = WordPieceTokenizer.from_config(
+            original_tokenizer.get_config()
+        )
+        self.assertAllEqual(
+            original_tokenizer(input_data),
+            cloned_tokenizer(input_data),
+        )
+
+    def test_saving(self):
+        input_data = tf.constant(["quick brOWN whale"])
+        vocab_data = ["@UNK@", "qu", "@@ick", "br", "@@OWN", "fox"]
+        tokenizer = WordPieceTokenizer(
+            vocabulary=vocab_data,
+            lowercase=False,
+            oov_token="@UNK@",
+            suffix_indicator="@@",
+            dtype="string",
+        )
+        inputs = keras.Input(dtype="string", shape=())
+        outputs = tokenizer(inputs)
+        model = keras.Model(inputs, outputs)
+        model.save(self.get_temp_dir())
+        restored_model = keras.models.load_model(self.get_temp_dir())
+        self.assertAllEqual(
+            model(input_data),
+            restored_model(input_data),
+        )