Repository: keras-team/keras-nlp
Pull request #318 (2022-08-27 06:43:35)

Add Support for CJK Char Splitting for WordPiece Tokenizer
Resolves https://github.com/keras-team/keras-nlp/issues/307

A few comments:

- [CJK regex](https://github.com/google-research/bert/blob/master/tokenization.py#L274-L281) was already present in [`PUNCTUATION_REGEX`](https://github.com/keras-team/keras-nlp/blob/master/keras_nlp/tokenizers/word_piece_tokenizer.py#L46-L54). Shifted it to a new variable, and added an arg for splitting on CJK (`split_on_cjk`).
- Made similar changes in WordPiece trainer.



Patch:
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index 7e2cdf0d..f261d109 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -43,7 +43,13 @@ PUNCTUATION_REGEX = r"|".join(
         r"[{-~]",
         # Unicode punctuation class.
         r"[\p{P}]",
-        # More unicode ranges.
+    ]
+)
+
+# Matches CJK characters. Obtained from
+# https://github.com/google-research/bert/blob/master/tokenization.py#L251.
+CJK_REGEX = r"|".join(
+    [
         r"[\x{4E00}-\x{9FFF}]",
         r"[\x{3400}-\x{4DBF}]",
         r"[\x{20000}-\x{2A6DF}]",
@@ -63,8 +69,30 @@ WHITESPACE_AND_PUNCTUATION_REGEX = r"|".join(
     ]
 )
 
+# Matches punctuation and CJK characters.
+PUNCTUATION_AND_CJK_REGEX = r"|".join(
+    [
+        PUNCTUATION_REGEX,
+        CJK_REGEX,
+    ]
+)
+
+# Matches whitespace, punctuation, and CJK characters.
+WHITESPACE_PUNCTUATION_AND_CJK_REGEX = r"|".join(
+    [
+        WHITESPACE_AND_PUNCTUATION_REGEX,
+        CJK_REGEX,
+    ]
+)
+
 
-def pretokenize(text, lowercase, strip_accents, split):
+def pretokenize(
+    text,
+    lowercase=True,
+    strip_accents=True,
+    split=True,
+    split_on_cjk=True,
+):
     """Helper function that takes in a dataset element and pretokenizes it.
 
     Args:
@@ -78,6 +106,10 @@ def pretokenize(text, lowercase, strip_accents, split):
             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_on_cjk: bool, defaults to `True`. If true, input will be split
+            on CJK characters, i.e., Chinese, Japanese, Korean and Vietnamese
+            characters (https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)).
+            Note that this is applicable only when `split` is true.
 
     Returns:
         A tensor containing the pre-processed and pre-tokenized `text`.
@@ -91,6 +123,8 @@ def pretokenize(text, lowercase, strip_accents, split):
     # Preprocess, lowercase, strip and split input data.
     if text.shape.rank == 0:
         text = tf.expand_dims(text, 0)
+    if split_on_cjk and split:
+        text = tf.strings.regex_replace(text, CJK_REGEX, r" \0 ")
     if lowercase:
         text = tf_text.case_fold_utf8(text)
     if strip_accents:
@@ -99,10 +133,16 @@ def pretokenize(text, lowercase, strip_accents, split):
         # Remove the accent marks.
         text = tf.strings.regex_replace(text, r"\p{Mn}", "")
     if split:
+        if split_on_cjk:
+            split_pattern = WHITESPACE_PUNCTUATION_AND_CJK_REGEX
+            keep_split_pattern = PUNCTUATION_AND_CJK_REGEX
+        else:
+            split_pattern = WHITESPACE_AND_PUNCTUATION_REGEX
+            keep_split_pattern = PUNCTUATION_REGEX
         text = tf_text.regex_split(
             text,
-            delim_regex_pattern=WHITESPACE_AND_PUNCTUATION_REGEX,
-            keep_delim_regex_pattern=PUNCTUATION_REGEX,
+            delim_regex_pattern=split_pattern,
+            keep_delim_regex_pattern=keep_split_pattern,
         )
     return text
 
@@ -159,6 +199,10 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
             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_on_cjk: bool, defaults to `True`. If true, input will be split
+            on CJK characters, i.e., Chinese, Japanese, Korean and Vietnamese
+            characters (https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)).
+            Note that this is applicable only when `split` is true.
         suffix_indicator: str, defaults to "##". The characters prepended to a
             WordPiece to indicate that it is a suffix to another subword.
             E.g. "##ing".
@@ -235,6 +279,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         lowercase: bool = False,
         strip_accents: bool = False,
         split: bool = True,
+        split_on_cjk: bool = True,
         suffix_indicator: str = "##",
         oov_token: str = "[UNK]",
         **kwargs,
@@ -271,6 +316,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         self.lowercase = lowercase
         self.strip_accents = strip_accents
         self.split = split
+        self.split_on_cjk = split_on_cjk
         self.suffix_indicator = suffix_indicator
         self.oov_token = oov_token
 
@@ -335,7 +381,11 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
 
         scalar_input = inputs.shape.rank == 0
         inputs = pretokenize(
-            inputs, self.lowercase, self.strip_accents, self.split
+            inputs,
+            self.lowercase,
+            self.strip_accents,
+            self.split,
+            self.split_on_cjk,
         )
 
         # Apply WordPiece and coerce shape for outputs.
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py b/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py
index bfdd258e..2d20ae49 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_trainer.py
@@ -28,6 +28,7 @@ def compute_word_piece_vocabulary(
     lowercase=False,
     strip_accents=False,
     split=True,
+    split_on_cjk=True,
     suffix_indicator="##",
     reserved_tokens=["[PAD]", "[CLS]", "[SEP]", "[UNK]", "[MASK]"],
 ):
@@ -55,6 +56,10 @@ def compute_word_piece_vocabulary(
             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.
+        split_on_cjk: bool, defaults to `True`. If true, input will be split
+            on CJK characters, i.e., Chinese, Japanese, Korean and Vietnamese
+            characters (https://en.wikipedia.org/wiki/CJK_Unified_Ideographs_(Unicode_block)).
+            Note that this is applicable only when `split` is true.
         suffix_indicator: str, defaults to "##". The characters prepended to a
             WordPiece to indicate that it is a suffix to another subword.
             E.g. "##ing".
@@ -138,7 +143,9 @@ def compute_word_piece_vocabulary(
         )
 
     words_data = data.map(
-        lambda text: pretokenize(text, lowercase, strip_accents, split),
+        lambda text: pretokenize(
+            text, lowercase, strip_accents, split, split_on_cjk
+        ),
         num_parallel_calls=tf.data.AUTOTUNE,
     )
     word_counts = learner.count_words(words_data)


Test patch:
diff --git a/keras_nlp/models/bert_test.py b/keras_nlp/models/bert_test.py
index e4fa6b14..f8aaca4f 100644
--- a/keras_nlp/models/bert_test.py
+++ b/keras_nlp/models/bert_test.py
@@ -22,79 +22,63 @@ from keras_nlp.models import bert
 
 
 class BertTest(tf.test.TestCase):
-    def test_valid_call_bert(self):
-        model = bert.BertCustom(
-            vocabulary_size=30522,
-            num_layers=12,
-            num_heads=12,
-            hidden_dim=768,
-            intermediate_dim=3072,
-            max_sequence_length=12,
+    def setUp(self):
+        self.model = bert.BertCustom(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
             name="encoder",
         )
-        input_data = {
-            "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+        self.batch_size = 8
+        self.input_data = {
+            "input_ids": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
-            "segment_ids": tf.constant(
-                [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            "segment_ids": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
-            "input_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            "input_mask": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
         }
-        model(input_data)
+
+    def test_valid_call_bert(self):
+        self.model(self.input_data)
 
     def test_variable_sequence_length_call_bert(self):
-        model = bert.BertCustom(
-            vocabulary_size=30522,
-            num_layers=12,
-            num_heads=12,
-            hidden_dim=768,
-            intermediate_dim=3072,
-            max_sequence_length=100,
-            name="encoder",
-        )
         for seq_length in (25, 50, 75):
             input_data = {
-                "input_ids": tf.ones((8, seq_length), dtype="int32"),
-                "segment_ids": tf.ones((8, seq_length), dtype="int32"),
-                "input_mask": tf.ones((8, seq_length), dtype="int32"),
+                "input_ids": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+                "segment_ids": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+                "input_mask": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
             }
-            model(input_data)
+            self.model(input_data)
 
     def test_valid_call_classifier(self):
-        model = bert.BertCustom(
-            vocabulary_size=30522,
-            num_layers=12,
-            num_heads=12,
-            hidden_dim=768,
-            intermediate_dim=3072,
-            max_sequence_length=12,
-            name="encoder",
-        )
-        input_data = {
-            "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=30522
-            ),
-            "segment_ids": tf.constant(
-                [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-            "input_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-        classifier = bert.BertClassifier(model, 4, name="classifier")
-        classifier(input_data)
+        classifier = bert.BertClassifier(self.model, 4, name="classifier")
+        classifier(self.input_data)
 
     def test_valid_call_bert_base(self):
-        model = bert.BertBase(vocabulary_size=10000, name="encoder")
+        model = bert.BertBase(vocabulary_size=1000, name="encoder")
         input_data = {
-            "input_ids": tf.random.uniform(
-                shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+            "input_ids": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
+            ),
+            "segment_ids": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
+            ),
+            "input_mask": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
-            "segment_ids": tf.constant([0] * 200 + [1] * 312, shape=(1, 512)),
-            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
         }
         model(input_data)
 
@@ -107,7 +91,7 @@ class BertTest(tf.test.TestCase):
         with self.assertRaises(ValueError):
             bert.BertBase(
                 weights="bert_base_uncased",
-                vocabulary_size=10000,
+                vocabulary_size=1000,
                 name="encoder",
             )
 
@@ -119,32 +103,12 @@ class BertTest(tf.test.TestCase):
             )
 
     def test_saving_model(self):
-        model = bert.BertCustom(
-            vocabulary_size=30522,
-            num_layers=12,
-            num_heads=12,
-            hidden_dim=768,
-            intermediate_dim=3072,
-            max_sequence_length=12,
-            name="encoder",
-        )
-        input_data = {
-            "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
-            ),
-            "segment_ids": tf.constant(
-                [0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-            "input_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-        model_output = model(input_data)
+        model_output = self.model(self.input_data)
         save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path)
+        self.model.save(save_path)
         restored_model = keras.models.load_model(save_path)
 
-        restored_output = restored_model(input_data)
+        restored_output = restored_model(self.input_data)
         self.assertAllClose(
             model_output["pooled_output"], restored_output["pooled_output"]
         )
diff --git a/keras_nlp/models/roberta_test.py b/keras_nlp/models/roberta_test.py
index 3efd6f14..6dee0fa8 100644
--- a/keras_nlp/models/roberta_test.py
+++ b/keras_nlp/models/roberta_test.py
@@ -22,102 +22,71 @@ from keras_nlp.models import roberta
 
 
 class RobertaTest(tf.test.TestCase):
-    def test_valid_call_roberta(self):
-        model = roberta.RobertaCustom(
-            vocabulary_size=50265,
-            num_layers=12,
-            num_heads=12,
-            hidden_dim=768,
-            intermediate_dim=3072,
-            max_sequence_length=12,
+    def setUp(self):
+        self.model = roberta.RobertaCustom(
+            vocabulary_size=1000,
+            num_layers=2,
+            num_heads=2,
+            hidden_dim=64,
+            intermediate_dim=128,
+            max_sequence_length=128,
             name="encoder",
         )
-        input_data = {
-            "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
+        self.batch_size = 8
+        self.input_data = {
+            "input_ids": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
-            "input_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
+            "input_mask": tf.ones(
+                (self.batch_size, self.model.max_sequence_length), dtype="int32"
             ),
         }
-        model(input_data)
+
+    def test_valid_call_roberta(self):
+        self.model(self.input_data)
 
     def test_valid_call_classifier(self):
-        model = roberta.RobertaCustom(
-            vocabulary_size=50265,
-            num_layers=12,
-            num_heads=12,
-            hidden_dim=768,
-            intermediate_dim=3072,
-            max_sequence_length=12,
-            name="encoder",
+        classifier = roberta.RobertaClassifier(
+            self.model, 4, 128, name="classifier"
         )
-        input_data = {
-            "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
-            ),
-            "input_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-        classifier = roberta.RobertaClassifier(model, 4, 768, name="classifier")
-        classifier(input_data)
+        classifier(self.input_data)
 
     def test_valid_call_roberta_base(self):
-        model = roberta.RobertaBase(vocabulary_size=10000, name="encoder")
+        model = roberta.RobertaBase(vocabulary_size=1000, name="encoder")
         input_data = {
-            "input_ids": tf.random.uniform(
-                shape=(1, 512), dtype=tf.int64, maxval=model.vocabulary_size
+            "input_ids": tf.ones(
+                (self.batch_size, model.max_sequence_length), dtype="int32"
+            ),
+            "input_mask": tf.ones(
+                (self.batch_size, model.max_sequence_length), dtype="int32"
             ),
-            "input_mask": tf.constant([1] * 512, shape=(1, 512)),
         }
         model(input_data)
 
     def test_variable_sequence_length_call_roberta(self):
-        model = roberta.RobertaCustom(
-            vocabulary_size=50265,
-            num_layers=12,
-            num_heads=12,
-            hidden_dim=768,
-            intermediate_dim=3072,
-            max_sequence_length=100,
-            name="encoder",
-        )
         for seq_length in (25, 50, 75):
             input_data = {
-                "input_ids": tf.ones((8, seq_length), dtype="int32"),
-                "input_mask": tf.ones((8, seq_length), dtype="int32"),
+                "input_ids": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
+                "input_mask": tf.ones(
+                    (self.batch_size, seq_length), dtype="int32"
+                ),
             }
-            output = model(input_data)
+            output = self.model(input_data)
             self.assertAllEqual(
-                tf.shape(output["sequence_output"]), [8, seq_length, 768]
+                tf.shape(output["sequence_output"]),
+                [self.batch_size, seq_length, self.model.hidden_dim],
             )
 
     def test_saving_model(self):
-        model = roberta.RobertaCustom(
-            vocabulary_size=50265,
-            num_layers=12,
-            num_heads=12,
-            hidden_dim=768,
-            intermediate_dim=3072,
-            max_sequence_length=12,
-            name="encoder",
-        )
-        input_data = {
-            "input_ids": tf.random.uniform(
-                shape=(1, 12), dtype=tf.int64, maxval=model.vocabulary_size
-            ),
-            "input_mask": tf.constant(
-                [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0], shape=(1, 12)
-            ),
-        }
-        model_output = model.predict(input_data)
+        model_output = self.model.predict(self.input_data)
 
         save_path = os.path.join(self.get_temp_dir(), "model")
-        model.save(save_path)
+        self.model.save(save_path)
         restored_model = keras.models.load_model(save_path)
 
-        restored_output = restored_model.predict(input_data)
+        restored_output = restored_model.predict(self.input_data)
         self.assertAllClose(
             model_output["sequence_output"],
             restored_output["sequence_output"],
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index cb7438bd..b10504e0 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -87,6 +87,16 @@ class WordPieceTokenizerTest(tf.test.TestCase, parameterized.TestCase):
             tf.ragged.constant([["qu", "@@ick", "br", "@@own", "@UNK@"]]),
         )
 
+    def test_cjk_tokens(self):
+        input_data = ["ah半推zz"]
+        vocab_data = ["[UNK]", "推", "敐", "乐", "半", "偷", "匕", "ah", "zz"]
+        tokenizer = WordPieceTokenizer(vocabulary=vocab_data, dtype="string")
+        call_output = tokenizer(input_data)
+        self.assertAllEqual(
+            call_output,
+            tf.ragged.constant([["ah", "半", "推", "zz"]]),
+        )
+
     def test_lowercase(self):
         input_data = ["the QUicK brOWN FOX"]
         vocab_data = ["[UNK]", "the", "qu", "##ick", "br", "##own", "fox"]
@@ -119,7 +129,7 @@ class WordPieceTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         call_output = tokenizer(input_data)
         self.assertAllEqual(call_output, [[1, 2, 3, 4, 5]])
 
-    def test_no_spliting(self):
+    def test_no_splitting(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=False)
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py
index 85b11be5..0a36d1e4 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_trainer_test.py
@@ -110,6 +110,29 @@ class WordPieceTokenizerTrainerTest(tf.test.TestCase):
         )
         self.assertAllEqual(output_vocab_1, output_vocab_2)
 
+    def test_split_on_cjk(self):
+        test_text = tf.data.Dataset.from_tensor_slices(["ah半推zz"])
+        test_text_split = tf.data.Dataset.from_tensor_slices(
+            ["ah", "半", "推", "zz"]
+        )
+        output_vocab_1 = compute_word_piece_vocabulary(
+            test_text,
+            4,
+            split=True,
+            split_on_cjk=True,
+            lowercase=False,
+            strip_accents=False,
+        )
+        output_vocab_2 = compute_word_piece_vocabulary(
+            test_text_split,
+            4,
+            split=False,
+            split_on_cjk=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(
             [
diff --git a/keras_nlp/utils/text_generation_test.py b/keras_nlp/utils/text_generation_test.py
index 9fb3f084..fab6ed8f 100644
--- a/keras_nlp/utils/text_generation_test.py
+++ b/keras_nlp/utils/text_generation_test.py
@@ -415,7 +415,7 @@ class RandomSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         outputs_count = np.array([0, 0, 0, 0])
         tf.random.set_seed(42)
-        for i in range(500):
+        for i in range(64):
             outputs = random_search(
                 token_probability_fn, inputs, max_length=max_length, seed=42
             )
@@ -629,7 +629,7 @@ class TopKSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         outputs_count = np.array([0, 0, 0, 0])
         tf.random.set_seed(42)
-        for i in range(500):
+        for _ in range(64):
             outputs = top_k_search(
                 token_probability_fn,
                 inputs,
@@ -859,7 +859,7 @@ class TopPSearchTextGenerationTest(tf.test.TestCase, parameterized.TestCase):
 
         outputs_count = np.array([0, 0, 0, 0])
         tf.random.set_seed(42)
-        for i in range(500):
+        for i in range(64):
             outputs = top_p_search(
                 token_probability_fn,
                 inputs,