Repository: keras-team/keras-nlp
Pull request #256 (2022-07-07 03:35:53)

Added WordPieceTokenizer training function

chenmoneygithub: /gcbrun

Patch:
diff --git a/keras_nlp/tokenizers/__init__.py b/keras_nlp/tokenizers/__init__.py
index c033b506..2124f1a8 100644
--- a/keras_nlp/tokenizers/__init__.py
+++ b/keras_nlp/tokenizers/__init__.py
@@ -19,3 +19,6 @@ from keras_nlp.tokenizers.unicode_character_tokenizer import (
     UnicodeCharacterTokenizer,
 )
 from keras_nlp.tokenizers.word_piece_tokenizer import WordPieceTokenizer
+from keras_nlp.tokenizers.word_piece_tokenizer_trainer import (
+    compute_word_piece_vocabulary,
+)
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index 75006ecb..0f84a299 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -63,8 +63,51 @@ WHITESPACE_AND_PUNCTUATION_REGEX = r"|".join(
 )
 
 
+def pretokenize(text, lowercase, strip_accents, split):
+    """Helper function that takes in a dataset element and pretokenizes it.
+
+    Args:
+        text: `tf.Tensor` or `tf.RaggedTensor`. Input to be pretokenized.
+        lowercase: bool, defaults to `True`. If true, the input text will be
+            lowercased before tokenization.
+        strip_accents: bool, defaults to `True`. If true, all accent marks will
+            be removed from text before tokenization.
+        split: bool, defaults to `True`. If true, input will be split on
+            whitespace and punctuation marks, and all punctuation marks will be
+            kept as tokens. If false, input should be split ("pre-tokenized")
+            before calling the tokenizer, and passed as a dense or ragged tensor
+            of whole words.
+
+    Returns:
+        A tensor containing the pre-processed and pre-tokenized `text`.
+    """
+    # Check for correct types.
+    if text.dtype != tf.string:
+        raise ValueError(
+            "The dataset elements in `data` must have string dtype. "
+            f"Recieved: {text.dtype}."
+        )
+    # Preprocess, lowercase, strip and split input data.
+    if text.shape.rank == 0:
+        text = tf.expand_dims(text, 0)
+    if lowercase:
+        text = tf_text.case_fold_utf8(text)
+    if strip_accents:
+        # Normalize unicode to NFD, which splits out accent mark characters.
+        text = tf_text.normalize_utf8(text, "NFD")
+        # Remove the accent marks.
+        text = tf.strings.regex_replace(text, r"\p{Mn}", "")
+    if split:
+        text = tf_text.regex_split(
+            text,
+            delim_regex_pattern=WHITESPACE_AND_PUNCTUATION_REGEX,
+            keep_delim_regex_pattern=PUNCTUATION_REGEX,
+        )
+    return text
+
+
 class WordPieceTokenizer(tokenizer.Tokenizer):
-    """A word piece tokenizer layer.
+    """A WordPiece tokenizer layer.
 
     This layer provides an efficient, in graph, implementation of the WordPiece
     algorithm used by BERT and other models.
@@ -100,23 +143,25 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
 
     Args:
         vocabulary: A list of strings or a 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
+            passing a list, each element of the list should be a single
+            WordPiece token string. If passing a filename, the file should be a
+            plain text file containing a single WordPiece token per line.
+        sequence_length: int. 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: If true, input will be split on whitespace and punctuation
-            marks, and all punctuation marks will be kept as tokens. If false,
-            input should be split ("pre-tokenized") before calling the
-            tokenizer, and passed as a dense or ragged tensor of whole words.
-        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.
+        lowercase: bool, defaults to `True`. If true, the input text will be
+            lowercased before tokenization.
+        strip_accents: bool, defaults to `True`. If true, all accent marks will
+            be removed from text before tokenization.
+        split: bool, defaults to `True`. If true, input will be split on
+            whitespace and punctuation marks, and all punctuation marks will be
+            kept as tokens. If false, input should be split ("pre-tokenized")
+            before calling the tokenizer, and passed as a dense or ragged tensor
+            of whole words.
+        suffix_indicator: str, defaults to "##". The characters prepended to a
+            WordPiece to indicate that it is a suffix to another subword.
+            E.g. "##ing".
+        oov_token: str, defaults to "[UNK]". 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/)
@@ -269,26 +314,13 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
     def tokenize(self, inputs):
         if not isinstance(inputs, (tf.Tensor, tf.RaggedTensor)):
             inputs = tf.convert_to_tensor(inputs)
+
         scalar_input = inputs.shape.rank == 0
-        if scalar_input:
-            inputs = tf.expand_dims(inputs, 0)
-
-        # 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:
-            inputs = tf_text.regex_split(
-                inputs,
-                delim_regex_pattern=WHITESPACE_AND_PUNCTUATION_REGEX,
-                keep_delim_regex_pattern=PUNCTUATION_REGEX,
-            )
+        inputs = pretokenize(
+            inputs, self.lowercase, self.strip_accents, self.split
+        )
 
-        # Apply word piece and coerce shape for outputs.
+        # Apply WordPiece 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
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py b/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py
new file mode 100644
index 00000000..7d455b43
--- /dev/null
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py
@@ -0,0 +1,151 @@
+# 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 WordPiece Tokenizer."""
+
+import tensorflow as tf
+from tensorflow_text.tools.wordpiece_vocab import (
+    wordpiece_tokenizer_learner_lib as learner,
+)
+
+from keras_nlp.tokenizers.word_piece_tokenizer import pretokenize
+
+
+def compute_word_piece_vocabulary(
+    data,
+    vocabulary_size,
+    vocabulary_output_file=None,
+    lowercase=True,
+    strip_accents=True,
+    split=True,
+    suffix_indicator="##",
+    reserved_tokens=["[PAD]", "[CLS]", "[SEP]", "[UNK]", "[MASK]"],
+):
+    r"""A utility to train a WordPiece vocabulary.
+
+    Trains a WordPiece vocabulary from an input dataset or a list of filenames.
+
+    For custom data loading and pretokenization (`split=False`), the input
+    `data` should be a `tf.data.Dataset`. If `data` is a list of filenames,
+    the file format is required to be plain text files, and the text would 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.
+        vocabulary_output_file: str, defaults to `None`. The location to write a
+            vocabulary file.
+        lowercase: bool, defaults to `True`. If true, the input text will be
+            lowercased before tokenization.
+        strip_accents: bool, defaults to `True`. If true, all accent marks will
+            be removed from text before tokenization.
+        split: bool, defaults to `True`. If true, input will be split on
+            whitespace and punctuation marks, and all punctuation marks will be
+            kept as tokens. If false, input should be split ("pre-tokenized")
+            before calling the tokenizer, and passed as a dense or ragged tensor
+            of whole words. `split` is required to be `True` when `data` is a
+            list of filenames.
+        suffix_indicator: str, defaults to "##". The characters prepended to a
+            WordPiece to indicate that it is a suffix to another subword.
+            E.g. "##ing".
+        reserved_tokens: list of strings. A list of tokens that must be included in the vocabulary.
+
+    Returns:
+        Returns a list of vocabulary terms.
+
+    Examples:
+
+    Basic Usage (from Dataset).
+    >>> inputs = tf.data.Dataset.from_tensor_slices(["bat sat pat mat rat"])
+    >>> vocab = compute_word_piece_vocabulary(inputs, 13)
+    >>> vocab
+    ['[PAD]', '[CLS]', '[SEP]', '[UNK]', '[MASK]', 'a', 'b', 'm', 'p', 'r', 's', 't', '##at']
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab, oov_token="[UNK]")
+    >>> outputs = inputs.map(tokenizer.tokenize)
+    >>> for x in outputs:
+    ...     print(x)
+    tf.Tensor([ 6 12 10 12  8 12  7 12  9 12], shape=(10,), dtype=int32)
+
+    Basic Usage (from filenames).
+    >>> with open("test.txt", "w+") as f: f.write("bat sat pat mat rat\n")
+    >>> inputs = ["test.txt"]
+    >>> vocab = compute_word_piece_vocabulary(inputs, 13)
+
+    Custom Split Usage (from Dataset).
+    >>> def normalize_and_split(x):
+    ...     "Strip punctuation and split on whitespace."
+    ...     x = tf.strings.regex_replace(x, r"\p{P}", "")
+    ...     return tf.strings.split(x)
+    >>> inputs = tf.data.Dataset.from_tensor_slices(["bat sat: pat mat rat.\n"])
+    >>> split_inputs = inputs.map(normalize_and_split)
+    >>> vocab = compute_word_piece_vocabulary(
+    ...     split_inputs, 13, split=False
+    ... )
+    >>> vocab
+    ['[PAD]', '[CLS]', '[SEP]', '[UNK]', '[MASK]', 'a', 'b', 'm', 'p', 'r', 's', 't', '##at']
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
+    >>> inputs.map(tokenizer.tokenize)
+
+    Custom Split Usage (from filenames).
+    >>> def normalize_and_split(x):
+    ...     "Strip punctuation and split on whitespace."
+    ...     x = tf.strings.regex_replace(x, r"\p{P}", "")
+    ...     return tf.strings.split(x)
+    >>> with open("test.txt", "w+") as f: f.write("bat sat: pat mat rat.\n")
+    >>> inputs = tf.data.TextLineDataset(["test.txt"])
+    >>> split_inputs = inputs.map(normalize_and_split)
+    >>> vocab = compute_word_piece_vocabulary(split_inputs, 13, split=False)
+    >>> vocab
+    ['[PAD]', '[CLS]', '[SEP]', '[UNK]', '[MASK]', 'a', 'b', 'm', 'p', 'r', 's', 't', '##at']
+    >>> tokenizer = keras_nlp.tokenizers.WordPieceTokenizer(vocabulary=vocab)
+    >>> inputs.map(tokenizer.tokenize)
+    """
+    # Read data files.
+
+    if not isinstance(data, (list, tf.data.Dataset)):
+        raise ValueError(
+            "The `data` argument must be either `tf.data.Dataset` or `list`. "
+            f"Recieved: {type(data)}."
+        )
+    if isinstance(data, list):
+        if not split:
+            raise ValueError(
+                "When learning a vocab from files, `split` must be `True`. "
+                "To compute a vocabulary with custom split rules, load your "
+                "data as a dataset, split it, and pass it to "
+                "`compute_word_piece_vocabulary()` with split=False."
+            )
+        data = tf.data.TextLineDataset(data)
+
+    words_data = data.map(
+        lambda text: pretokenize(text, lowercase, strip_accents, split),
+        num_parallel_calls=tf.data.AUTOTUNE,
+    )
+    word_counts = learner.count_words(words_data)
+    # Train tokenizer.
+    vocab = learner.learn(
+        word_counts,
+        vocab_size=vocabulary_size,
+        reserved_tokens=reserved_tokens,
+        include_joiner_token=True,
+        joiner=suffix_indicator,
+    )
+    if len(vocab) > vocabulary_size:
+        vocab = vocab[:vocabulary_size]
+    if vocabulary_output_file is not None:
+        vocab_text = "".join([line + "\n" for line in vocab])
+        # Write vocab to file.
+        with open(vocabulary_output_file, "w") as vocab_file:
+            vocab_file.write(vocab_text)
+    else:
+        return vocab


Test patch:
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py
new file mode 100644
index 00000000..6bdf587d
--- /dev/null
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py
@@ -0,0 +1,176 @@
+# 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 WordPiece Tokenizer Trainer."""
+
+import os
+
+import tensorflow as tf
+
+from keras_nlp.tokenizers.word_piece_tokenizer_trainer import (
+    compute_word_piece_vocabulary,
+)
+
+
+class WordPieceTokenizerTrainerTest(tf.test.TestCase):
+    def test_dataset_input(self):
+        test_text = ["baa maa caa saa aaa"]
+        test_output = ["a", "b", "c", "m", "s", "##aa", "##a", "##b"]
+        data = tf.data.Dataset.from_tensor_slices(test_text)
+        vocab = compute_word_piece_vocabulary(data, 8, reserved_tokens=[])
+        self.assertAllEqual(vocab, test_output)
+
+    def test_filenames_input(self):
+        test_text = "baa maa caa saa aaa"
+        with open("test.txt", "w+") as f:
+            f.write(test_text + "\n")
+        test_output = ["a", "b", "c", "m", "s", "##aa", "##a", "##b"]
+        vocab = compute_word_piece_vocabulary(
+            ["test.txt"], 8, reserved_tokens=[]
+        )
+        self.assertAllEqual(vocab, test_output)
+        os.remove("test.txt")
+
+    def test_filenames_without_split(self):
+        test_text = "baa maa caa saa aaa"
+        with open("test.txt", "w+") as f:
+            f.write(test_text + "\n")
+
+        with self.assertRaisesRegex(
+            ValueError,
+            "When learning a vocab from files, `split` must be `True`. "
+            "To compute a vocabulary with custom split rules, load your "
+            "data as a dataset, split it, and pass it to "
+            r"`compute_word_piece_vocabulary\(\)` with split=False.",
+        ):
+            compute_word_piece_vocabulary(["test.txt"], 10, split=False)
+        os.remove("test.txt")
+
+    def test_invalid_input(self):
+        test_text_invalid = tf.data.Dataset.from_tensor_slices([1, 2, 3, 4])
+        with self.assertRaisesRegex(
+            ValueError,
+            "The dataset elements in `data` must have string dtype. "
+            "Recieved: <dtype: 'int32'>.",
+        ):
+            compute_word_piece_vocabulary(test_text_invalid, 10)
+        with self.assertRaisesRegex(
+            ValueError,
+            "The `data` argument must be either `tf.data.Dataset` or `list`. "
+            "Recieved: <class 'int'>.",
+        ):
+            compute_word_piece_vocabulary(4, 4)
+
+    def test_lowercase(self):
+        test_text = tf.data.Dataset.from_tensor_slices(["BaA Maa Caa Saa AAa"])
+        test_output = ["a", "b", "c", "m", "s", "##aa", "##a", "##b"]
+
+        vocab = compute_word_piece_vocabulary(
+            test_text, 8, lowercase=True, reserved_tokens=[]
+        )
+        self.assertAllEqual(vocab, test_output)
+
+    def test_skip_lowercase(self):
+        test_text = tf.data.Dataset.from_tensor_slices(["BAA MAA CAA SAA AAA"])
+        test_output = ["A", "B", "C", "M", "S", "##AA", "##A", "##B"]
+
+        vocab = compute_word_piece_vocabulary(
+            test_text, 8, lowercase=False, reserved_tokens=[]
+        )
+        self.assertAllEqual(vocab, test_output)
+
+    def test_split(self):
+        test_text = tf.data.Dataset.from_tensor_slices(
+            ["This string: would be split up."]
+        )
+        test_text_split = tf.data.Dataset.from_tensor_slices(
+            ["This", "string", ":", "would", "be", "split", "up", "."]
+        )
+        output_vocab_1 = compute_word_piece_vocabulary(
+            test_text, 20, split=True, lowercase=False, strip_accents=False
+        )
+        output_vocab_2 = compute_word_piece_vocabulary(
+            test_text_split,
+            20,
+            split=False,
+            lowercase=False,
+            strip_accents=False,
+        )
+        self.assertAllEqual(output_vocab_1, output_vocab_2)
+
+    def test_skip_split(self):
+        test_text = tf.data.Dataset.from_tensor_slices(
+            [
+                "This is a long line that isn't split up, and it exceeds maximum length."
+            ]
+        )
+        # The token would be removed for being too long.
+        vocab = compute_word_piece_vocabulary(
+            test_text, 20, split=False, reserved_tokens=[]
+        )
+        self.assertAllEqual(vocab, [])
+
+    def test_strip_accents(self):
+        test_text = tf.data.Dataset.from_tensor_slices(
+            ["áááá éááá íááá óááá úááá"]
+        )
+        output = ["a", "e", "i", "o", "u", "##aaa", "##a", "##e"]
+        vocab = compute_word_piece_vocabulary(
+            test_text, 8, strip_accents=True, reserved_tokens=[]
+        )
+        self.assertAllEqual(vocab, output)
+
+    def test_skip_strip_accents(self):
+        test_text = tf.data.Dataset.from_tensor_slices(
+            ["áááá éááá íááá óááá úááá"]
+        )
+        output = ["á", "é", "í", "ó", "ú", "##ááá", "##á", "##é"]
+        vocab = compute_word_piece_vocabulary(
+            test_text, 8, strip_accents=False, reserved_tokens=[]
+        )
+        self.assertAllEqual(vocab, output)
+
+    def test_output_file(self):
+        test_text = tf.data.Dataset.from_tensor_slices(["BaA Maa Caa Saa AAa"])
+        test_output = ["a", "b", "c", "m", "s", "##aa", "##a", "##b"]
+
+        compute_word_piece_vocabulary(
+            test_text, 8, vocabulary_output_file="test.txt", reserved_tokens=[]
+        )
+        vocab_from_file = []
+        with open("test.txt", "r") as f:
+            for line in f:
+                vocab_from_file.append(line.strip())
+        self.assertAllEqual(vocab_from_file, test_output)
+        os.remove("test.txt")
+
+    def test_reserved_tokens(self):
+        # This dummy text/token would be removed for being too long.
+        test_text = tf.data.Dataset.from_tensor_slices(
+            [
+                "The learner requires at least one input here, but this should be removed."
+            ]
+        )
+        output = ["token1", "token2", "token3", "token4"]
+        vocab = compute_word_piece_vocabulary(
+            test_text, 20, reserved_tokens=output, split=False
+        )
+        self.assertAllEqual(vocab, output)
+
+    def test_suffix_indicator(self):
+        test_text = tf.data.Dataset.from_tensor_slices(["baa maa caa saa aaa"])
+        test_output = ["a", "b", "c", "m", "s", "@@aa", "@@a", "@@b"]
+        vocab = compute_word_piece_vocabulary(
+            test_text, 8, suffix_indicator="@@", reserved_tokens=[]
+        )
+        self.assertAllEqual(vocab, test_output)