BigBirdTokenizer Serialization Error on Spark
Author: jmwoloso
Pull request: no
Labels: name: bug

### System Info

Hello! Running `transformers-cli env` throws the following error for me on the databricks cluster that I'm using currently (but it works locally).

```python
/databricks/conda/envs/default/lib/python3.8/site-packages/requests/__init__.py:89: RequestsDependencyWarning: urllib3 (1.26.10) or chardet (3.0.4) doesn't match a supported version!
  warnings.warn("urllib3 ({}) or chardet ({}) doesn't match a supported "
Traceback (most recent call last):
  File "/databricks/conda/envs/prism_ai/bin/transformers-cli", line 5, in <module>
    from transformers.commands.transformers_cli import main
  File "/databricks/conda/envs/prism_ai/lib/python3.8/site-packages/transformers/commands/transformers_cli.py", line 26, in <module>
    from .user import UserCommands
  File "/databricks/conda/envs/prism_ai/lib/python3.8/site-packages/transformers/commands/user.py", line 20, in <module>
    from huggingface_hub.hf_api import HfFolder, create_repo, login, logout, whoami
ImportError: cannot import name 'login' from 'huggingface_hub.hf_api' (/databricks/conda/envs/prism_ai/lib/python3.8/site-packages/huggingface_hub/hf_api.py)
```

I can provide some of the info that I know (we maintain the container that we use):

```
transformers: 4.17.0
python: 3.8.12
pt version: 1.11.0+cu113
tf version: 2.8.0
flax: N/A
jax: N/A
jaxlib: N/A
using gpu: nope
using distributed: yes, via pyspark udf
```

The reprex below this produces the following error:

```bash
---------------------------------------------------------------------------
...

/databricks/spark/python/pyspark/sql/dataframe.py in count(self)
    668         2
    669         """
--> 670         return int(self._jdf.count())
    671 
    672     def collect(self):

/databricks/spark/python/lib/py4j-0.10.9-src.zip/py4j/java_gateway.py in __call__(self, *args)
   1302 
   1303         answer = self.gateway_client.send_command(command)
-> 1304         return_value = get_return_value(
   1305             answer, self.gateway_client, self.target_id, self.name)
   1306 

/databricks/spark/python/pyspark/sql/utils.py in deco(*a, **kw)
    121                 # Hide where the exception came from that shows a non-Pythonic
    122                 # JVM exception message.
--> 123                 raise converted from None
    124             else:
    125                 raise

PythonException: An exception was thrown from a UDF: 'pyspark.serializers.SerializationError: Caused by Traceback (most recent call last):
  File "/databricks/spark/python/pyspark/serializers.py", line 165, in _read_with_length
    return self.loads(obj)
  File "/databricks/spark/python/pyspark/serializers.py", line 469, in loads
    return pickle.loads(obj, encoding=encoding)
  File "/databricks/python/lib/python3.8/site-packages/transformers/models/big_bird/tokenization_big_bird.py", line 164, in __setstate__
    self.sp_model.Load(self.vocab_file)
  File "/databricks/python/lib/python3.8/site-packages/sentencepiece/__init__.py", line 367, in Load
    return self.LoadFromFile(model_file)
  File "/databricks/python/lib/python3.8/site-packages/sentencepiece/__init__.py", line 171, in LoadFromFile
    return _sentencepiece.SentencePieceProcessor_LoadFromFile(self, arg)
OSError: Not found: "/.cache/huggingface/d318d7bb69cafb1d8964fc87515592ac3092a2c8fdb305068f9ba4020df3ee3b.271d467a9adc15fb44348481bc75c48b63cba0fd4934bc5377d63a63de052c45": No such file or directory Error #2'. Full traceback below:
Traceback (most recent call last):
  File "/databricks/spark/python/pyspark/serializers.py", line 165, in _read_with_length
    return self.loads(obj)
  File "/databricks/spark/python/pyspark/serializers.py", line 469, in loads
    return pickle.loads(obj, encoding=encoding)
  File "/databricks/python/lib/python3.8/site-packages/transformers/models/big_bird/tokenization_big_bird.py", line 164, in __setstate__
    self.sp_model.Load(self.vocab_file)
  File "/databricks/python/lib/python3.8/site-packages/sentencepiece/__init__.py", line 367, in Load
    return self.LoadFromFile(model_file)
  File "/databricks/python/lib/python3.8/site-packages/sentencepiece/__init__.py", line 171, in LoadFromFile
    return _sentencepiece.SentencePieceProcessor_LoadFromFile(self, arg)
OSError: Not found: "/.cache/huggingface/d318d7bb69cafb1d8964fc87515592ac3092a2c8fdb305068f9ba4020df3ee3b.271d467a9adc15fb44348481bc75c48b63cba0fd4934bc5377d63a63de052c45": No such file or directory Error #2

During handling of the above exception, another exception occurred:

pyspark.serializers.SerializationError: Caused by Traceback (most recent call last):
  File "/databricks/spark/python/pyspark/serializers.py", line 165, in _read_with_length
    return self.loads(obj)
  File "/databricks/spark/python/pyspark/serializers.py", line 469, in loads
    return pickle.loads(obj, encoding=encoding)
  File "/databricks/python/lib/python3.8/site-packages/transformers/models/big_bird/tokenization_big_bird.py", line 164, in __setstate__
    self.sp_model.Load(self.vocab_file)
  File "/databricks/python/lib/python3.8/site-packages/sentencepiece/__init__.py", line 367, in Load
    return self.LoadFromFile(model_file)
  File "/databricks/python/lib/python3.8/site-packages/sentencepiece/__init__.py", line 171, in LoadFromFile
    return _sentencepiece.SentencePieceProcessor_LoadFromFile(self, arg)
OSError: Not found: "/.cache/huggingface/d318d7bb69cafb1d8964fc87515592ac3092a2c8fdb305068f9ba4020df3ee3b.271d467a9adc15fb44348481bc75c48b63cba0fd4934bc5377d63a63de052c45": No such file or directory Error #2
```

I believe this is related to #15982. If this is related, it seems we need fast tokenizers for `Marian` and now also `BigBird`. Is anyone working on that currently? If not I could take a stab at `BigBird` (or maybe we just need to add `BigBird` functionality to `PreTrainedTokenizerFast`?)

### Who can help?

BigBird: @ydshieh 
Marian: @patil-suraj 

### 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

Here is a (simplified) reprex of what my function(s) look like and I'm running this on a multi-node databricks cluster on azure.

```python
import functools

from pyspark.sql import functions as F, types as T
from transformers import BigBirdTokenizer

tokenizer = BigBirdTokenizer.from_pretrained("google/bigbird-roberta-base")

return_type = T.ArrayType(T.IntegerType())

def create_tokens(tokenizer, text):
    return tokenizer.encode_plus(text)["input_ids"]

create_tokens_partial = functools.partial(create_tokens, tokenizer=tokenizer)

create_tokens_udf = F.udf(lambda text: create_tokens_partial(text=text, returnType=return_type)

df = df.withColumn("InputIds", create_tokens_udf(F.column("Text")))
df.cache().count()
```

### Expected behavior

Serialization of the `BigBirdTokenizer` that can be shipped off to worker nodes.

Comments:
Hey @jmwoloso Thank you for reporting the issue.  I have a few questions though: I see there is [fast big_bird tokenizer](https://github.com/huggingface/transformers/blob/main/src/transformers/models/big_bird/tokenization_big_bird_fast.py) in `transformers`. Have you tried it?

Hi @ydshieh. No I havne't tried that directly. I tried the `AutoTokenizer` which uses the Fast tokenizers by default (so I assume the behavior will be the same), but the problem with the fast tokenizers is that you cannot feed them pre-tokenized text like you can the vanilla tokenizers, and I need to do that for my use case.

So it seems we don't need a fast tokenizer implementation, but the fast tokenizers actually need feature parity with the vanilla tokenizers.

Unless using the `BigBirdTokenizerFast` directly is the solution, rather than `AutoTokenizer`?

Possibly related [#1045](https://github.com/huggingface/tokenizers/issues/1045)

I was able to come up with a workaround @ydshieh, but it still feels like the fast tokenizers should behave the same as the standard tokenizers (i.e. accept tokenized inputs)

the solution I came up with tries to load the tokenizer from the local node where the udf is being run. if it doesn't exist, we download it, and save it for future calls to make use of. 

expanding upon `create_tokens` in my reprex above:
```python

def create_tokens(text):
    from transformers import BigBirdTokenizer
    model = "google/bigbird-roberta-base"
    try:
        # try to access it locally first
        tokenizer = BigBirdTokenizer.from_pretrained(f"/{model}")
    except Exception:
        tokenizer = BigBirdTokenizer.from_pretrained(model)
        tokenizer.save_pretrained(f"/{model}")
    return tokenizer.encode_plus(text)["input_ids"]

```

@jmwoloso I am  glad you found a workaround!

But for me to understand the problem a bit better, on the remote nodes (hopefully this term makes sense - I am not really familiar with Spark ecosystem), they are not able to download/load the tokenizer from `google/bigbird-roberta-base`?

And if we download  and save it (from the local node where the udf is being run), it could be loaded (through a local directory?) when the code is running on the remote nodes?

----

Regarding the pre-tokenization in fast tokenizers, probably there are some limitations in `tokenizers` (which is written in `rust`) and is designed to be so. I am not very familiar with that libraries. If you would like to, maybe you can open a feature request in [tokenizers](https://github.com/huggingface/tokenizers), with a description of the current limitation you encounter.

Hi @ydshieh thank you for the quick reply. The issue I linked to above in the tokenizers library captures the limitation I believe, so I can follow up on that thread to see about what is possible there.

And yes "remote nodes" is a fine term to use, we typically think about Spark clusters in terms of the driver node (the node where our notebooks run) and the worker nodes (which take a serialized form of the commands we specify on the driver node and use those to do work on their slice of the data).

Spark also has user-defined functions (UDFs) which are python functions that can be serialized and sent to the worker nodes so they have what they need to operate on data. I have tried instantiating the tokenizer within the udf, but a spark cluster has 200 partitions (jobs) by default and you can adjust that to suit your workload. The problem is that each job will call it's version of the udf which results (in this case) in 200 calls to download the tokenizer and we get an error that we've made too many requests to the site.

The alternative to this is to use partial application with the udf. So create the tokenizer a single time on the driver node and then we use partial application to apply the tokenizer to the udf so that the tokenizer gets shipped along with the udf to each of the workers. 

This works just fine for (as far as I can tell) the non-sentencepiece based tokenizers. The spiece tokenizers however, seem to try and reload the model after they've been sent to the workers. The problem is that the cache exists on the driver node, but not on the worker nodes, so the tokenizer tries to load its config again and the cache doesn't exist so the job fails.

My solution above does exactly what you suggest, I try to access the tokenizer from the local file system of the worker node first and if that fails, I download it and save the config on the worker node so that subsequent calls to the udf use the locally saved tokenizer instead of making the api call to download it.

Thank you @jmwoloso for the patience and effort to answer my question! 

Actually, I was not giving suggestion in my last reply, but just to understand better your workaround/solution :-)

> The problem is that each job will call it's version of the udf which results (in this case) in 200 calls to download the tokenizer and we get an error that we've made too many requests to the site.

> I download it and save the config on the worker node so that subsequent calls to the udf use the locally saved tokenizer instead of making the api call to download it.

If you still have a bit time, one last question (as I am a bit confused here): 

all those ~200 jobs can access the locally saved tokenizer, downloaded only once on a particular worker? Or each worker still has to download it (so totally ~200 downloads) separately, but subsequent calls on each node won't download it again? 

From the description, I believe you mean the former. But I am a bit surprised that all workers can access the same file system.

Finally, I am going to close the issue as you already provide a working solution (and the issue is not really a bug in the library). Hope this is OK for you.

 

> all those ~200 jobs can access the locally saved tokenizer, downloaded only once on a particular worker?

Hi @ydshieh Yes, this. One download per worker that any jobs on that worker can then access if my solution above is used.


> Finally, I am going to close the issue as you already provide a working solution (and the issue is not really a bug in the library). Hope this is OK for you.

I'm torn on this because I would expect that all tokenizers behave the same and could be interoperable in a distributed setting without having to make the fix above, but the `spiece` ones behave differently. But having said that, the solution above would work for all tokenizers `spiece` and otherwise. It doesn't feel so much like a bug, just unexpected/inconsistent behavior among tokenizers within the library. Is this something I could add to like a "recipes" section in the documentation as a pattern to be used in distributed settings? I'd be happy to work on that if you think it would be useful, but otherwise, I'm fine with closing this issue. I appreciate your help!



Thanks again, @jmwoloso . Let's keep this issue open for now, and I will discuss with 2 colleagues regarding the documentation.

Sounds great @ydshieh, thank you!

Hi, @jmwoloso 

After a discussion with one of my colleagues @Narsil, we think this scenario is too specific to be documented in the doc. I think, for now, having your workaround on this GitHub issue page is already very good and helpful for others having the same issue :-) Thanks again.

Regarding slow/fast tokenizer, I left [a comment](https://github.com/huggingface/tokenizers/issues/1045#issuecomment-1257802082) :-)
 

Thank you @ydshieh!