Repository: pytorch/text

Issue:
Add `never_split` kwarg to BERTTokenizer to achieve parity with `transformers.BertTokenizer`
## 🚀 Feature

Add functionality to BERTTokenizer so that it does not split/tokenize special tokens like regular words.

**Motivation**

The HuggingFace [BertTokenizer](https://github.com/huggingface/transformers/blob/bbbb453e5869a5fec4925b02f1265c9e6bfc3ebb/src/transformers/models/bert/tokenization_bert.py#L137) is able to recognize its own special tokens and ensure that they are not split during the tokenization process. For example, if given the following input:

```
"\tHeLLo!how  \n Are yoU? [UNK]"
```

The HuggingFace (HF) tokenizer returns the "[UNK]" token un-split, whereas the TorchText (TT) tokenizer returns it split:

```
HF (expected):   ['hello', '!', 'how', 'are', 'you', '?', '[UNK]']
TT (actual):     ['hello', '!', 'how', 'are', 'you', '?', '[', 'un', '##k', ']']
```

**Pitch**

A `never_split` keyword argument at the constructor-level which enables the user to specify which tokens to ignore during tokenization. Such a keyword argument would be flexible and give the user the freedom to adapt to various other tokenization schemes. An example interface:

```
tt_tokenizer = BERTTokenizer(
    vocab_path=vocab_file,
    do_lower_case=True,
    return_tokens=True,
    never_split=[
        "[UNK]",
        "[CLS]",
        "[SEP]",
        "[PAD]",
    ],
)
```

This closely matches the existing HuggingFace BertTokenizer interface (constructor level kwarg [here](https://github.com/huggingface/transformers/blob/bbbb453e5869a5fec4925b02f1265c9e6bfc3ebb/src/transformers/models/bert/tokenization_bert.py#L189), `tokenize` function implementation [here](https://github.com/huggingface/transformers/blob/bbbb453e5869a5fec4925b02f1265c9e6bfc3ebb/src/transformers/models/bert/tokenization_bert.py#L244)).

If the user were able to instantiate the HuggingFace tokenizer with which they were attempting to achieve parity, the user could trivially assign the `never_split` keyword argument to its special tokens:

```
import transformers

hf_tokenizer = transformers.AutoTokenizer.from_pretrained("bert-base-uncased")
tt_tokenizer = BERTTokenizer(
    vocab_path=vocab_file,
    do_lower_case=True,
    return_tokens=True,
    never_split=hf_tokenizer.all_special_tokens,
)
```

**Alternatives**

None that are more straightforward than this.

**Additional context**

A full reproduction of the above disparity in functionality, as well as a comment about the requested interface can be found in this Gist: https://gist.github.com/geoffreyangus/7b4c04f1e485c9c4091728bff1414134



Discussion:

Resolving patch:
diff --git a/torchtext/transforms.py b/torchtext/transforms.py
--- a/torchtext/transforms.py
+++ b/torchtext/transforms.py
@@ -564,21 +564,31 @@ class BERTTokenizer(Module):
     :type strip_accents: Optional[bool]
     :param return_tokens: Indicate whether to return tokens. If false, returns corresponding token IDs as strings (default: False)
     :type return_tokens: bool
+    :param never_split: Collection of tokens which will not be split during tokenization. (default: None)
+    :type never_split: Optional[List[str]]
     """
 
     __jit_unused_properties__ = ["is_jitable"]
 
     def __init__(
-        self, vocab_path: str, do_lower_case: bool = True, strip_accents: Optional[bool] = None, return_tokens=False
+        self,
+        vocab_path: str,
+        do_lower_case: bool = True,
+        strip_accents: Optional[bool] = None,
+        return_tokens=False,
+        never_split: Optional[List[str]] = None,
     ) -> None:
         super().__init__()
+        if never_split is None:
+            never_split = []
         self.bert_model = BERTEncoderPyBind(
-            get_asset_local_path(vocab_path, overwite=True), do_lower_case, strip_accents
+            get_asset_local_path(vocab_path, overwite=True), do_lower_case, strip_accents, never_split
         )
         self._return_tokens = return_tokens
         self._vocab_path = vocab_path
         self._do_lower_case = do_lower_case
         self._strip_accents = strip_accents
+        self._never_split = never_split
 
     @property
     def is_jitable(self):
@@ -654,7 +664,7 @@ def __prepare_scriptable__(self):
         if not self.is_jitable:
             tokenizer_copy = deepcopy(self)
             tokenizer_copy.bert_model = torch.classes.torchtext.BERTEncoder(
-                self._vocab_path, self._do_lower_case, self._strip_accents
+                self._vocab_path, self._do_lower_case, self._strip_accents, self._never_split
             )
             return tokenizer_copy
 


Test patch:
diff --git a/test/torchtext_unittest/test_transforms.py b/test/torchtext_unittest/test_transforms.py
--- a/test/torchtext_unittest/test_transforms.py
+++ b/test/torchtext_unittest/test_transforms.py
@@ -1,5 +1,6 @@
 import os
 from collections import OrderedDict
+from typing import List, Optional
 from unittest.mock import patch
 
 import torch
@@ -586,7 +587,9 @@ def test_clip_tokenizer_save_load_torchscript(self) -> None:
 
 
 class TestBERTTokenizer(TorchtextTestCase):
-    def _load_tokenizer(self, test_scripting: bool, do_lower_case: bool, return_tokens: bool):
+    def _load_tokenizer(
+        self, test_scripting: bool, do_lower_case: bool, return_tokens: bool, never_split: Optional[List[str]] = None
+    ):
         if do_lower_case:
             vocab_file = "bert_base_uncased_vocab.txt"
         else:
@@ -596,46 +599,117 @@ def _load_tokenizer(self, test_scripting: bool, do_lower_case: bool, return_toke
             vocab_path=get_asset_path(vocab_file),
             do_lower_case=do_lower_case,
             return_tokens=return_tokens,
+            never_split=never_split,
         )
         if test_scripting:
             tokenizer = torch.jit.script(tokenizer)
         return tokenizer
 
-    def _bert_tokenizer(self, tokenizer, do_lower_case):
+    def _bert_tokenizer(self, tokenizer, do_lower_case, never_split: Optional[List[str]] = None):
         sample_texts = [
             "Hello World!, how are you?",
             "Hélló  WoŕlḊ¿",
             "Respublica superiorem",
             "Avdija Vršajević în",
+            " \tHeLLo!how  \n Are yoU? [UNK]",
+            "hi world [UNK] [CLS]",
+            "testing, [UNK] words! [SEP]",
         ]
 
-        if do_lower_case:
-            expected_tokens = [
-                ["hello", "world", "!", ",", "how", "are", "you", "?"],
-                ["hello", "world", "¿"],
-                ["res", "##pu", "##bl", "##ica", "superior", "##em"],
-                ["av", "##di", "##ja", "vr", "##sa", "##jevic", "in"],
-            ]
-            expected_token_ids = [
-                ["7592", "2088", "999", "1010", "2129", "2024", "2017", "1029"],
-                ["7592", "2088", "1094"],
-                ["24501", "14289", "16558", "5555", "6020", "6633"],
-                ["20704", "4305", "3900", "27830", "3736", "26782", "1999"],
-            ]
+        if not never_split:
+            if do_lower_case:
+                expected_tokens = [
+                    ["hello", "world", "!", ",", "how", "are", "you", "?"],
+                    ["hello", "world", "¿"],
+                    ["res", "##pu", "##bl", "##ica", "superior", "##em"],
+                    ["av", "##di", "##ja", "vr", "##sa", "##jevic", "in"],
+                    ["hello", "!", "how", "are", "you", "?", "[", "un", "##k", "]"],
+                    ["hi", "world", "[", "un", "##k", "]", "[", "cl", "##s", "]"],
+                    ["testing", ",", "[", "un", "##k", "]", "words", "!", "[", "sep", "]"],
+                ]
+                expected_token_ids = [
+                    ["7592", "2088", "999", "1010", "2129", "2024", "2017", "1029"],
+                    ["7592", "2088", "1094"],
+                    ["24501", "14289", "16558", "5555", "6020", "6633"],
+                    ["20704", "4305", "3900", "27830", "3736", "26782", "1999"],
+                    ["7592", "999", "2129", "2024", "2017", "1029", "1031", "4895", "2243", "1033"],
+                    ["7632", "2088", "1031", "4895", "2243", "1033", "1031", "18856", "2015", "1033"],
+                    ["5604", "1010", "1031", "4895", "2243", "1033", "2616", "999", "1031", "19802", "1033"],
+                ]
 
+            else:
+                expected_tokens = [
+                    ["Hello", "World", "!", ",", "how", "are", "you", "?"],
+                    ["H", "##é", "##ll", "##ó", "[UNK]", "¿"],
+                    ["Re", "##sp", "##ub", "##lica", "superior", "##em"],
+                    ["A", "##v", "##di", "##ja", "V", "##r", "##ša", "##je", "##vić", "î", "##n"],
+                    ["He", "##LL", "##o", "!", "how", "Are", "yo", "##U", "?", "[", "UN", "##K", "]"],
+                    ["hi", "world", "[", "UN", "##K", "]", "[", "C", "##LS", "]"],
+                    ["testing", ",", "[", "UN", "##K", "]", "words", "!", "[", "SE", "##P", "]"],
+                ]
+                expected_token_ids = [
+                    ["8667", "1291", "106", "117", "1293", "1132", "1128", "136"],
+                    ["145", "2744", "2339", "7774", "100", "225"],
+                    ["11336", "20080", "10354", "9538", "7298", "5521"],
+                    ["138", "1964", "3309", "3174", "159", "1197", "23834", "5561", "10225", "260", "1179"],
+                    [
+                        "1124",
+                        "23955",
+                        "1186",
+                        "106",
+                        "1293",
+                        "2372",
+                        "26063",
+                        "2591",
+                        "136",
+                        "164",
+                        "7414",
+                        "2428",
+                        "166",
+                    ],
+                    ["20844", "1362", "164", "7414", "2428", "166", "164", "140", "15928", "166"],
+                    ["5193", "117", "164", "7414", "2428", "166", "1734", "106", "164", "12342", "2101", "166"],
+                ]
         else:
-            expected_tokens = [
-                ["Hello", "World", "!", ",", "how", "are", "you", "?"],
-                ["H", "##é", "##ll", "##ó", "[UNK]", "¿"],
-                ["Re", "##sp", "##ub", "##lica", "superior", "##em"],
-                ["A", "##v", "##di", "##ja", "V", "##r", "##ša", "##je", "##vić", "î", "##n"],
-            ]
-            expected_token_ids = [
-                ["8667", "1291", "106", "117", "1293", "1132", "1128", "136"],
-                ["145", "2744", "2339", "7774", "100", "225"],
-                ["11336", "20080", "10354", "9538", "7298", "5521"],
-                ["138", "1964", "3309", "3174", "159", "1197", "23834", "5561", "10225", "260", "1179"],
-            ]
+            if do_lower_case:
+                expected_tokens = [
+                    ["hello", "world", "!", ",", "how", "are", "you", "?"],
+                    ["hello", "world", "¿"],
+                    ["res", "##pu", "##bl", "##ica", "superior", "##em"],
+                    ["av", "##di", "##ja", "vr", "##sa", "##jevic", "in"],
+                    ["hello", "!", "how", "are", "you", "?", "[UNK]"],
+                    ["hi", "world", "[UNK]", "[CLS]"],
+                    ["testing", ",", "[UNK]", "words", "!", "[", "sep", "]"],
+                ]
+                expected_token_ids = [
+                    ["7592", "2088", "999", "1010", "2129", "2024", "2017", "1029"],
+                    ["7592", "2088", "1094"],
+                    ["24501", "14289", "16558", "5555", "6020", "6633"],
+                    ["20704", "4305", "3900", "27830", "3736", "26782", "1999"],
+                    ["7592", "999", "2129", "2024", "2017", "1029", "100"],
+                    ["7632", "2088", "100", "101"],
+                    ["5604", "1010", "100", "2616", "999", "1031", "19802", "1033"],
+                ]
+
+            else:
+                expected_tokens = [
+                    ["Hello", "World", "!", ",", "how", "are", "you", "?"],
+                    ["H", "##é", "##ll", "##ó", "[UNK]", "¿"],
+                    ["Re", "##sp", "##ub", "##lica", "superior", "##em"],
+                    ["A", "##v", "##di", "##ja", "V", "##r", "##ša", "##je", "##vić", "î", "##n"],
+                    ["He", "##LL", "##o", "!", "how", "Are", "yo", "##U", "?", "[UNK]"],
+                    ["hi", "world", "[UNK]", "[CLS]"],
+                    ["testing", ",", "[UNK]", "words", "!", "[", "SE", "##P", "]"],
+                ]
+                expected_token_ids = [
+                    ["8667", "1291", "106", "117", "1293", "1132", "1128", "136"],
+                    ["145", "2744", "2339", "7774", "100", "225"],
+                    ["11336", "20080", "10354", "9538", "7298", "5521"],
+                    ["138", "1964", "3309", "3174", "159", "1197", "23834", "5561", "10225", "260", "1179"],
+                    ["1124", "23955", "1186", "106", "1293", "2372", "26063", "2591", "136", "100"],
+                    ["20844", "1362", "100", "101"],
+                    ["5193", "117", "100", "1734", "106", "164", "12342", "2101", "166"],
+                ]
 
         # test batch of sentences
         if tokenizer._return_tokens:
@@ -650,14 +724,18 @@ def _bert_tokenizer(self, tokenizer, do_lower_case):
             else:
                 self.assertEqual(tokenizer(txt), expected_token_ids[idx])
 
-    @nested_params([True, False], [True, False], [True, False])
-    def test_bert_tokenizer(self, test_scripting, do_lower_case, return_tokens):
+    @nested_params([True, False], [True, False], [True, False], [[], None, ["[UNK]", "[CLS]"]])
+    def test_bert_tokenizer(self, test_scripting, do_lower_case, return_tokens, never_split):
         """test tokenization on single sentence input as well as batch on sentences"""
         self._bert_tokenizer(
             self._load_tokenizer(
-                test_scripting=test_scripting, do_lower_case=do_lower_case, return_tokens=return_tokens
+                test_scripting=test_scripting,
+                do_lower_case=do_lower_case,
+                return_tokens=return_tokens,
+                never_split=never_split,
             ),
             do_lower_case=do_lower_case,
+            never_split=never_split,
         )
 
     @nested_params([True, False], [True, False], [True, False])