Tokenization+Transformers works with PyTorch but not TensorFlow on TPU
Author: oja
Pull request: no
Labels: name: wontfix

name: Core: Tokenization

# 🐛 Bug

## Information

Model I am using (Bert, XLNet ...): RoBERTa

Language I am using the model on (English, Chinese ...): English

The problem arises when using:
* [ ] the official example scripts: (give details below)
* [x] my own modified scripts: (give details below)

The tasks I am working on is:
* [ ] an official GLUE/SQUaD task: (give the name)
* [x] my own task or dataset: (give details below)

## To reproduce

Steps to reproduce the behavior:

1. In order to use Huggingface's tokenizer with TensorFlow's data pipeline (`tf.data.Dataset.map()`), the tokenize function must be wrapped in `tf.py_function()` or `tf.numpy_function()` (issue #3851). 
2. Neither of these functions are not supported on TPU (tensorflow/tensorflow#30818).
3. This makes it impossible to run Huggingface transformers + tokenizer on a TPU using TensorFlow
4. However, the same tokenization on a TPU works under PyTorch (!)

One workaround is to do tokenization before entering the TF data pipeline, but unfortunately my dataset is too large for that.

Example of code that is necessary, but fails on a TPU:
```
def tokenize_encode_map_fn(text):
    encoded = tf.py_function(tokenize_encode, # tokenize_encode is a wrapper around the Huggingface tokenizer and encoder
                             inp=[text, hypothesis],
                             Tout=[tf.int32, tf.int32])
    return {"input_ids": encoded[0], "attention_mask": encoded[1]}

tf_dataset.map(map_fn)
```

<!-- If you have code snippets, error messages, stack traces please provide them here as well.
     Important! Use code tags to correctly format your code. See https://help.github.com/en/github/writing-on-github/creating-and-highlighting-code-blocks#syntax-highlighting
     Do not use screenshots, as they are hard to read and (more importantly) don't allow others to copy-and-paste your code.-->

## Expected behavior
Be able to run tokenizer on a TPU with TensorFlow, like you can in PyTorch
<!-- A clear and concise description of what you would expect to happen. -->

## Environment info
<!-- You can run the command `transformers-cli env` and copy-and-paste its output below.
     Don't forget to fill out the missing fields in that output! -->
     
- `transformers` version: transformers-2.11.0
- Platform: Google Colab TPU
- Python version: Python 3.6.9
- PyTorch version (GPU?):
- Tensorflow version (GPU?): TensorFlow 2.2.0 / TPU
- Using GPU in script?:
- Using distributed or parallel set-up in script?:


Comments:
This makes sense! Do you have some examples of some tokenization that can happen on TPUs with TensorFlow, so that we may consider what we have to do to enable this?

I am not totally sure, but one thing to look at would probably be TF's native `tf.keras.preprocessing.text.Tokenizer`, which (I think) works when used within tf.data.Dataset maps on a TPU. 

https://www.tensorflow.org/api_docs/python/tf/keras/preprocessing/text/Tokenizer

Assigning myself to keep this in mind


> I am not totally sure, but one thing to look at would probably be TF's native `tf.keras.preprocessing.text.Tokenizer`, which (I think) works when used within tf.data.Dataset maps on a TPU.
> 
> https://www.tensorflow.org/api_docs/python/tf/keras/preprocessing/text/Tokenizer

Interesting. Is there a way to translate HF's tokenziers to this?

I am wondering it's close to being possible that we don't need to convert text to Ids on the HF's data processors. 

TensorFlow 2.3.x adds the Sentencepiece tokenizer to `tensorflow_text`. You can use this short script to turn a HuggingFace tokenizer into a `tensorflow_text.SentencepieceTokenizer`: https://gist.github.com/noahtren/6f9f6ecf2f81d0975c4f54afaeb95318

I tested it on TPU and it's been working for me. My experience is that HuggingFace tokenizers are wrappers for https://github.com/google/sentencepiece so it's really simple to make it compatible with TensorFlow graph mode. Not sure yet if this works for all huggingface pretrained tokenizers.

Did you try on the hf rust tokenizers as well? 

@Santosh-Gupta No, I haven't tried that

This issue has been automatically marked as stale because it has not had recent activity. It will be closed if no further activity occurs. Thank you for your contributions.