Repository: keras-team/keras-nlp
Pull request #743 (2023-02-10 22:22:17)

Improve test coverage for special tokens in model tokenizers


Patch:
diff --git a/keras_nlp/models/roberta/roberta_tokenizer.py b/keras_nlp/models/roberta/roberta_tokenizer.py
index c5ba13b5..0e3284fa 100644
--- a/keras_nlp/models/roberta/roberta_tokenizer.py
+++ b/keras_nlp/models/roberta/roberta_tokenizer.py
@@ -105,7 +105,7 @@ class RobertaTokenizer(BytePairTokenizer):
         pad_token = "<pad>"
         end_token = "
"
         mask_token = "<mask>"
-        for token in [start_token, pad_token, end_token]:
+        for token in [start_token, pad_token, end_token, mask_token]:
             if token not in self.get_vocabulary():
                 raise ValueError(
                     f"Cannot find token `'{token}'` in the provided "
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer.py b/keras_nlp/tokenizers/word_piece_tokenizer.py
index ed593a32..c5e0e578 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer.py
@@ -331,7 +331,7 @@ class WordPieceTokenizer(tokenizer.Tokenizer):
         self.oov_token = oov_token
 
         if oov_token not in self.vocabulary:
-            raise RuntimeError(
+            raise ValueError(
                 f'Cannot find `oov_token="{self.oov_token}"` in the '
                 "vocabulary.\n"
                 "You can either update the vocabulary to include "


Test patch:
diff --git a/keras_nlp/models/albert/albert_tokenizer_test.py b/keras_nlp/models/albert/albert_tokenizer_test.py
index b7099b8d..d80657db 100644
--- a/keras_nlp/models/albert/albert_tokenizer_test.py
+++ b/keras_nlp/models/albert/albert_tokenizer_test.py
@@ -65,8 +65,20 @@ class AlbertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEqual(output, tf.constant(["the quick brown fox"]))
 
     def test_vocabulary_size(self):
-        tokenizer = AlbertTokenizer(proto=self.proto)
-        self.assertEqual(tokenizer.vocabulary_size(), 10)
+        self.assertEqual(self.tokenizer.vocabulary_size(), 10)
+
+    def test_errors_missing_special_tokens(self):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=iter(["abc"]),
+            model_writer=bytes_io,
+            vocab_size=5,
+            pad_id=-1,
+            eos_id=-1,
+            bos_id=-1,
+        )
+        with self.assertRaises(ValueError):
+            AlbertTokenizer(proto=bytes_io.getvalue())
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/models/bart/bart_tokenizer_test.py b/keras_nlp/models/bart/bart_tokenizer_test.py
index cd8fd624..1ec896eb 100644
--- a/keras_nlp/models/bart/bart_tokenizer_test.py
+++ b/keras_nlp/models/bart/bart_tokenizer_test.py
@@ -65,6 +65,10 @@ class BartTokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 12)
 
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            BartTokenizer(vocabulary=["a", "b", "c"], merges=[])
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/bert/bert_tokenizer_test.py b/keras_nlp/models/bert/bert_tokenizer_test.py
index 5f72a59f..1d17d9aa 100644
--- a/keras_nlp/models/bert/bert_tokenizer_test.py
+++ b/keras_nlp/models/bert/bert_tokenizer_test.py
@@ -27,17 +27,16 @@ class BertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
         self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
         self.vocab += ["the", "quick", "brown", "fox"]
+        self.tokenizer = BertTokenizer(vocabulary=self.vocab)
 
     def test_tokenize(self):
         input_data = "THE QUICK BROWN FOX."
-        tokenizer = BertTokenizer(vocabulary=self.vocab)
-        output = tokenizer(input_data)
+        output = self.tokenizer(input_data)
         self.assertAllEqual(output, [5, 6, 7, 8, 1])
 
     def test_tokenize_batch(self):
         input_data = tf.constant(["THE QUICK BROWN FOX.", "THE FOX."])
-        tokenizer = BertTokenizer(vocabulary=self.vocab)
-        output = tokenizer(input_data)
+        output = self.tokenizer(input_data)
         self.assertAllEqual(output, [[5, 6, 7, 8, 1], [5, 8, 1]])
 
     def test_lowercase(self):
@@ -48,13 +47,15 @@ class BertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_detokenize(self):
         input_tokens = [[5, 6, 7, 8]]
-        tokenizer = BertTokenizer(vocabulary=self.vocab)
-        output = tokenizer.detokenize(input_tokens)
+        output = self.tokenizer.detokenize(input_tokens)
         self.assertAllEqual(output, ["THE QUICK BROWN FOX"])
 
     def test_vocabulary_size(self):
-        tokenizer = BertTokenizer(vocabulary=self.vocab)
-        self.assertEqual(tokenizer.vocabulary_size(), 13)
+        self.assertEqual(self.tokenizer.vocabulary_size(), 13)
+
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            BertTokenizer(vocabulary=["a", "b", "c"])
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
index 1be123c9..b85f3491 100644
--- a/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
+++ b/keras_nlp/models/deberta_v3/deberta_v3_tokenizer_test.py
@@ -65,8 +65,20 @@ class DebertaV3TokenizerTest(tf.test.TestCase, parameterized.TestCase):
         self.assertEqual(output, tf.constant(["the quick brown fox"]))
 
     def test_vocabulary_size(self):
-        tokenizer = DebertaV3Tokenizer(proto=self.proto)
-        self.assertEqual(tokenizer.vocabulary_size(), 10)
+        self.assertEqual(self.tokenizer.vocabulary_size(), 10)
+
+    def test_errors_missing_special_tokens(self):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=iter(["abc"]),
+            model_writer=bytes_io,
+            vocab_size=5,
+            pad_id=-1,
+            eos_id=-1,
+            bos_id=-1,
+        )
+        with self.assertRaises(ValueError):
+            DebertaV3Tokenizer(proto=bytes_io.getvalue())
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py b/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
index c3fb810a..b51c2247 100644
--- a/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
+++ b/keras_nlp/models/distil_bert/distil_bert_tokenizer_test.py
@@ -29,17 +29,16 @@ class DistilBertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         self.vocab = ["[PAD]", "[UNK]", "[CLS]", "[SEP]", "[MASK]"]
         self.vocab += ["THE", "QUICK", "BROWN", "FOX"]
         self.vocab += ["the", "quick", "brown", "fox"]
+        self.tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
 
     def test_tokenize(self):
         input_data = "THE QUICK BROWN FOX."
-        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
-        output = tokenizer(input_data)
+        output = self.tokenizer(input_data)
         self.assertAllEqual(output, [5, 6, 7, 8, 1])
 
     def test_tokenize_batch(self):
         input_data = tf.constant(["THE QUICK BROWN FOX.", "THE FOX."])
-        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
-        output = tokenizer(input_data)
+        output = self.tokenizer(input_data)
         self.assertAllEqual(output, [[5, 6, 7, 8, 1], [5, 8, 1]])
 
     def test_lowercase(self):
@@ -50,13 +49,15 @@ class DistilBertTokenizerTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_detokenize(self):
         input_tokens = [[5, 6, 7, 8]]
-        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
-        output = tokenizer.detokenize(input_tokens)
+        output = self.tokenizer.detokenize(input_tokens)
         self.assertAllEqual(output, ["THE QUICK BROWN FOX"])
 
     def test_vocabulary_size(self):
-        tokenizer = DistilBertTokenizer(vocabulary=self.vocab)
-        self.assertEqual(tokenizer.vocabulary_size(), 13)
+        self.assertEqual(self.tokenizer.vocabulary_size(), 13)
+
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            DistilBertTokenizer(vocabulary=["a", "b", "c"])
 
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
diff --git a/keras_nlp/models/f_net/f_net_tokenizer_test.py b/keras_nlp/models/f_net/f_net_tokenizer_test.py
index a4257ffc..b5c186a2 100644
--- a/keras_nlp/models/f_net/f_net_tokenizer_test.py
+++ b/keras_nlp/models/f_net/f_net_tokenizer_test.py
@@ -68,6 +68,19 @@ class FNetTokenizerTest(tf.test.TestCase, parameterized.TestCase):
         tokenizer = FNetTokenizer(proto=self.proto)
         self.assertEqual(tokenizer.vocabulary_size(), 10)
 
+    def test_errors_missing_special_tokens(self):
+        bytes_io = io.BytesIO()
+        sentencepiece.SentencePieceTrainer.train(
+            sentence_iterator=iter(["abc"]),
+            model_writer=bytes_io,
+            vocab_size=5,
+            pad_id=-1,
+            eos_id=-1,
+            bos_id=-1,
+        )
+        with self.assertRaises(ValueError):
+            FNetTokenizer(proto=bytes_io.getvalue())
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/gpt2/gpt2_tokenizer_test.py b/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
index 7476bbc2..bbc11cfc 100644
--- a/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
+++ b/keras_nlp/models/gpt2/gpt2_tokenizer_test.py
@@ -66,6 +66,10 @@ class GPT2TokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 10)
 
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            GPT2Tokenizer(vocabulary=["a", "b", "c"], merges=[])
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/opt/opt_tokenizer_test.py b/keras_nlp/models/opt/opt_tokenizer_test.py
index 2e0eb461..b63493a3 100644
--- a/keras_nlp/models/opt/opt_tokenizer_test.py
+++ b/keras_nlp/models/opt/opt_tokenizer_test.py
@@ -66,6 +66,10 @@ class OPTTokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 12)
 
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            OPTTokenizer(vocabulary=["a", "b", "c"], merges=[])
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/models/roberta/roberta_tokenizer_test.py b/keras_nlp/models/roberta/roberta_tokenizer_test.py
index 4f96d455..5da219b9 100644
--- a/keras_nlp/models/roberta/roberta_tokenizer_test.py
+++ b/keras_nlp/models/roberta/roberta_tokenizer_test.py
@@ -66,6 +66,10 @@ class RobertaTokenizerTest(tf.test.TestCase, parameterized.TestCase):
     def test_vocabulary_size(self):
         self.assertEqual(self.tokenizer.vocabulary_size(), 13)
 
+    def test_errors_missing_special_tokens(self):
+        with self.assertRaises(ValueError):
+            RobertaTokenizer(vocabulary=["a", "b", "c"], merges=[])
+
     @parameterized.named_parameters(
         ("tf_format", "tf", "model"),
         ("keras_format", "keras_v3", "model.keras"),
diff --git a/keras_nlp/tokenizers/word_piece_tokenizer_test.py b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
index 86cb3817..fa1b8bd9 100644
--- a/keras_nlp/tokenizers/word_piece_tokenizer_test.py
+++ b/keras_nlp/tokenizers/word_piece_tokenizer_test.py
@@ -226,19 +226,19 @@ class WordPieceTokenizerTest(tf.test.TestCase, parameterized.TestCase):
 
     def test_no_oov_token_in_vocabulary(self):
         vocab_data = ["qu", "@@ick", "br", "@@OWN", "fox"]
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(ValueError):
             WordPieceTokenizer(
                 vocabulary=vocab_data,
             )
 
         vocab_data = ["@UNK@", "qu", "@@ick", "br", "@@OWN", "fox"]
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(ValueError):
             WordPieceTokenizer(
                 vocabulary=vocab_data,
             )
 
         vocab_data = ["UNK", "qu", "@@ick", "br", "@@OWN", "fox"]
-        with self.assertRaises(RuntimeError):
+        with self.assertRaises(ValueError):
             WordPieceTokenizer(
                 vocabulary=vocab_data,
             )