Add support for Scriptable BERT tokenizer

[parmeet]

This PR adds support for scriptable BERT tokenizer.

Initial Implementation: Our implementation is derived from the work of https://github.com/LieluoboAi/radish. We have made following major amendments in their implementation:

• Replaced usage of utfcpp with utfproc itself for converting string to and from unicode. This reduces additional dependency on utfcpp.
• Replaced usage of std::unordered_map with torchtext's Vocab implementation to perform look-up.
• Fixed wordpiece (max_seg_) algorithm to match HuggingFace (HF) implementation
• Fixed stripping issue by sending stripped text directly from python (\u2048 at end of string cannot be removed trivially in C++)
• Replaced their implementation of splitting strings based on whitespace with torchtext's split_
• Perform to_lower directly on unicode strings. Also corrected the logic of combining flags to perform lowering and stripping accents to match HF implementation
• Fixed _is_control implementation to match HF implementation
• Remove comparison with kChinesePunts to match HF's implementation of _is_punctuation
• Changed UString type from uint16_t to uint32_t. On rare occasions when a unicode code point cannot fit in a uint16_t type it causes errors.

Testing

Verified that the results matches with HF BERT Tokenizer on EnWik9 dataset (13147026 rows).

usage

bert_base_uncased_vocab_file = "https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt"
import torchtext.transforms as T
from torchtext.utils import get_asset_local_path
# Instantiate tokenizer with lower case, and return tokens=True (we also support return token IDs instead)
bert_tokenizer = T.BERTTokenizer(get_asset_local_path(bert_base_uncased_vocab_file),
                                    do_lower_case=True, strip_accents=None, return_tokens=True)

# non-batch API
tokens = bert_tokenizer("Hello world")
# out: ['hello', 'world']

# batch API 
tokens = bert_tokenizer(["Hello world","How are you!"])
# out: [['hello', 'world'], ['how', 'are', 'you', '!']]

Follow-up:

• Perform batch processing directly in C++ instead of iterating over input sentences in python

[Nayef211]

Just left a couple of noob questions and suggestions. Overall LGTM. Thanks for adding the BERT tokenizer to torchtext @parmeet! This looks like it was a very complex class to implement 🚀

[Nayef211]

Is there a way we could pass in the batch input directly to the C++ kernel and do the for-loop in the kernel itself? As we've seen in previous benchmarking efforts, a lot of time is spent on passing data back and forth between Python and C++ and we may be able to get significant perf gains just by passing the entire list in one go.

[parmeet]

Yes, that's a great idea. Let's do it in follow-up PR.

[Nayef211]

Are we using std::basic_string here because the text being passed in from Python contains unicode which isn't compatible with std::string?

[parmeet]

The string passed from python is UTF-8 encoded bytes. UString is the container to store the unicode code points when converting string to unicode and vice-versa.

[Nayef211]

Noob question: why do we make kUnkToken a static property rather than a constant?

[parmeet]

ohh, just following original implementation :).

[parmeet]

Thanks @Nayef211 for the thorough review and feedback. I have address the comments :).

[parmeet]

Torchtext CI is failing for windows on python 3.7 Seems like following is the culprit:

error: can't copy 'build\lib.win-amd64-3.7\pyd': doesn't exist or not a regular file

Any suggestions @mthrok, @atalman what's going on in here?

cc: @Nayef211

[Nayef211]

@parmeet I'm also seeing failures for test_bert_tokenizer on several platforms. Let's try to fix these before landing

[parmeet]

@parmeet I'm also seeing failures for test_bert_tokenizer on several platforms. Let's try to fix these before landing

yaa, looking at them. Not sure what went wrong. Now locally the tests are passing. Will look into CI results.

[philschmid]

Hello 🙋🏻‍♂️

I was trying to script/trace the new BERTTokenizer but without any success. What does scriptable for you mean in this context?

Here is my example on how tried to trace it

bert_base_uncased_vocab_file = "https://huggingface.co/bert-base-uncased/resolve/main/vocab.txt"
import torchtext.transforms as T
from torchtext.utils import get_asset_local_path
# Instantiate tokenizer with lower case, and return tokens=True (we also support return token IDs instead)
bert_tokenizer = T.BERTTokenizer(get_asset_local_path(bert_base_uncased_vocab_file),
                                    do_lower_case=True, strip_accents=None, return_tokens=True)

traced_tokenizer = torch.jit.trace(bert_tokenizer, "test")

here is the error

RuntimeError: 
Module 'BERTTokenizer' has no attribute 'bert_model' (This attribute exists on the Python module, but we failed to convert Python type: 'torchtext._torchtext.BERTEncoder' to a TorchScript type. Only tensors and (possibly nested) tuples of tensors, lists, or dictsare supported as inputs or outputs of traced functions, but instead got value of type BERTEncoder.. Its type was inferred; try adding a type annotation for the attribute.):
  File "/home/ubuntu/miniconda3/envs/optimum/lib/python3.8/site-packages/torchtext/transforms.py", line 603
    def _batch_encode(self, text: List[str]) -> List[List[str]]:
        """Batch version of _encode i.e operate on list of str"""
        token_ids: List[List[int]] = self.bert_model.batch_encode([t.strip() for t in text])
                                     ~~~~~~~~~~~~~~~ <--- HERE
        tokens_ids_str: List[List[str]] = [[str(t) for t in token_id] for token_id in token_ids]
        return tokens_ids_str

[parmeet]

@philschmid could you try with torch.jit.script for scripting?