Adding tokens to `RobertaTokenizer` is fast, but loading the extended tokenizer from disk takes tens of minutes
Author: Witiko
Pull request: no
Labels: name: bug

### System Info

```shell
- `transformers` version: 4.18.0
- Platform: Linux-5.10.0-0.bpo.9-amd64-x86_64-with-debian-10.12
- Python version: 3.7.3
- Huggingface_hub version: 0.5.1
- PyTorch version (GPU?): 1.11.0+cu102 (False)
- Tensorflow version (GPU?): not installed (NA)
- Flax version (CPU?/GPU?/TPU?): not installed (NA)
- Jax version: not installed
- JaxLib version: not installed
- Using GPU in script?: false
- Using distributed or parallel set-up in script?: false
```


### Who can help?

@SaulLu @Narsil

### Information

- [ ] The official example scripts
- [X] My own modified scripts

### Tasks

- [ ] An officially supported task in the `examples` folder (such as GLUE/SQuAD, ...)
- [X] My own task or dataset (give details below)

### Reproduction

I train a BPE tokenizer on a domain-specific dataset and save it as [`tokenizer-latex.json`](https://github.com/huggingface/transformers/files/8557562/tokenizer-latex.json.txt).

``` python
>>> from tokenizers import Tokenizer, normalizers, pre_tokenizers
>>> from tokenizers.models import BPE
>>> from tokenizers.trainers import BpeTrainer
>>> 
>>> latex_model = BPE(unk_token='[UNK]')
>>> latex_tokenizer = Tokenizer(latex_model)
>>> latex_tokenizer.pre_tokenizer = pre_tokenizers.WhitespaceSplit()
>>> latex_tokenizer.normalizer = normalizers.Sequence([normalizers.Strip()])
>>> latex_tokenizer_trainer = BpeTrainer(special_tokens=['[UNK]'])
>>> latex_tokenizer.train(['dataset-latex.txt'], latex_tokenizer_trainer)
>>> latex_tokenizer.save('tokenizer-latex.json')
```

Then, I extend [the pre-trained `roberta-base` tokenizer][1] with 28,141 new tokens from the vocabulary of my BPE tokenizer and I save the result to the directory `./extended-roberta-base/`. This finishes in a matter of seconds:

``` python
>>> from tokenizers import Tokenizer
>>> from transformers import RobertaTokenizer
>>> 
>>> latex_tokenizer = Tokenizer.from_file('tokenizer-latex.json')
>>> 
>>> text_latex_tokenizer = RobertaTokenizer.from_pretrained('roberta-base', add_prefix_space=True)
>>> text_latex_tokenizer.add_tokens(list(latex_tokenizer.get_vocab()))
28141
>>> text_latex_tokenizer.save_pretrained('./extended-roberta-base/')
('./extended-roberta-base/tokenizer_config.json', './extended-roberta-base/special_tokens_map.json',
 './extended-roberta-base/vocab.json', './extended-roberta-base/merges.txt',
 './extended-roberta-base/added_tokens.json', './extended-roberta-base/tokenizer.json')
```

However, when I load the extended `roberta-base` tokenizer from the directory `./extended-roberta-base/`, the library constructs a trie (see https://github.com/huggingface/transformers/pull/13220) over the course of ca 20 minutes:

``` python
>>> from transformers import RobertaTokenizer
>>>
>>> text_latex_tokenizer = RobertaTokenizer.from_pretrained('./extended-roberta-base/')
^C
Traceback (most recent call last):
  File "<stdin>", line 2, in <module>
    text_latex_tokenizer = RobertaTokenizer.from_pretrained('./extended-roberta-base/')
  File "***/python3.7/site-packages/transformers/tokenization_utils_base.py", line 1787, in from_pretrained
    **kwargs,
  File "***/python3.7/site-packages/transformers/tokenization_utils_base.py", line 1971, in _from_pretrained
    tokenizer.add_tokens(token, special_tokens=bool(token in special_tokens))
  File "***/python3.7/site-packages/transformers/tokenization_utils_base.py", line 945, in add_tokens
    return self._add_tokens(new_tokens, special_tokens=special_tokens)
  File "***/python3.7/site-packages/transformers/tokenization_utils.py", line 444, in _add_tokens
    self._create_trie(self.unique_no_split_tokens)
  File "***/python3.7/site-packages/transformers/tokenization_utils.py", line 454, in _create_trie
    trie.add(token)
  File "***/python3.7/site-packages/transformers/tokenization_utils.py", line 87, in add
    ref = ref[char]
KeyboardInterrupt
```

The time disparity leads me to believe that when `RobertaTokenizer.add_tokens()` is called, a trie is either not created or is created extremely fast, whereas when `RobertaTokenizer.from_pretrained()` is called, a trie is created (slowly). Using `RobertaTokenizerFast` instead of `RobertaTokenizer` produces similar results at a similar timescale.

 [1]: https://huggingface.co/roberta-base

### Expected behavior

Both `add_tokens()` and `from_pretrained()` should take comparable amount of time. Either building the trie is important and cannot be sped up, in which case `add_tokens()` should also take roughly 20 minutes, or building the trie is unimportant or can be sped up, in which case `from_pretrained()` should finish in a matter of seconds.


Comments:
Hi, pretty sure this is because `add_tokens` and therefore the `trie` creation is done N times for all the N tokens, which is indeed excruciatingly slow (and completely uncessary).

I think we can create the `trie` only once, wdyt @SaulLu 

Hi @Witiko,

Thanks for sharing this issue!

I share your analysis @Narsil ! When the `from_pretrained` method is called, the tokens are added 1 by 1 in this loop.
https://github.com/huggingface/transformers/blob/fa322474060beb3673cf5a3e39ccd3c8ad57ecd3/src/transformers/tokenization_utils_base.py#L1948-L1974

My memory may be faulty, but I had the impression that I had already read in an issue / PR that there could be a difficulty to circumvent to achieve this type of change - I can't find it back unfortunately. For the moment, I think that it is necessary to see in this loop that the order of addition has its importance and that we can alternate between addition of normal and special tokens.

We did it in  tokenizers` since the `Trie` insertion order of added tokens should not be important (this is also currently the case in slow tokenizers)

https://github.com/huggingface/tokenizers/blob/main/tokenizers/src/tokenizer/serialization.rs#L172

There might be other things to deal with in the python code, but the `Trie` itself doesn't care about insertion order, so we can create it only once.

Yes I absolutely agree! It just seemed important to mention it because the code that generates the multiple `Trie` builds currently is code that is shared between the fast and python tokenizers. :smile: 

Thank you for investigating. Should I try and open a PR, or are you planning to tackle this, @Narsil?

Hi @Witiko ,

I don't have a lot of bandwidth atm to handle this. If you can try and open a PR that would be awesome.
Feel free to ping me if you want early feedback (doesn't matter if PR is not ready).

Cheers,
Nicolas

Hello @Narsil,

neither do I, but I can take a stab at it sometime the next month. It seems to me that a simple fix might be to add a boolean parameter `_postpone_optimization` to `add_tokens()`, so that we can prevent the trie from being constructed in `from_pretrained()`. However, this does not solve the problem for users who would manually call `add_tokens()` with many small batches of tokens in their code. A more robust fix would be to construct the trie lazily at the point where is is needed.

> However, this does not solve the problem for users who would manually call add_tokens() with many small batches of tokens in their code.

`add_tokens` already accepts lists, so sending the maximum possible amount of tokens in one go is the way to go, laziness is not the solution to this problem here I think.

> `add_tokens` already accepts lists, so sending the maximum possible amount of tokens in one go is the way to go

The current code of `from_pretrained()` calls `add_tokens()` repeatedly with single tokens, so that it can persist the information about whether the token is special or not. Perhaps the way to go would be to first build a list of special and non-special tokens and then call `add_tokens()` once for special and once for non-special tokens?

> laziness is not the solution to this problem here I think.

I agree that laziness makes it more difficult to predict performance and reason about the code, especially in multiprocessing settings. Having `add_tokens()` that behaves optimally when you add tokens in bulk seems more straightforward.

@SaulLu @Narsil I opened PR #17119 that fixes this issue. I would appreciate your review at your earliest convenience.

Hey @Witiko I have seen your PR, but I am not sure I can do a proper review regarding the implications of this code, I did a small quality improvements suggestions in terms of pure code.

Pinging @SaulLu for visibility (If you don't have time I can look more into this btw, I just figured you would review faster than me)

@Narsil Thanks to your suggestions, the diff in #17119 against the current code is now tiny. Furthermore, the code causes a net speedup in loading tokenizers, so hopefully we can have the PR merged soon. 🤞

Thanks a lot for working on this fix.

On my side, I'm trying to look at your PR tomorrow. As this is a change that will impact all tokenizers, this is a contribution that requires a very attentive review on our part, that's why it can be a bit long. 

@SaulLu Thank you. Your caution is appreciated. we wouldn't want to break tokenizers. 😓