Add FNet
Author: gchhablani
Pull request: yes

# What does this PR do?

This PR adds the [FNet](https://arxiv.org/abs/2105.03824) model in PyTorch. I was working on it in another PR #12335 which got closed due to inactivity ;-;. This PR closes issue #12411.

## Checklist
- [x] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests),
      Pull Request section?
- [x] Did you make sure to update the documentation with your changes? Here are the
      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
- [x] Did you write any new necessary tests?

Requesting @LysandreJik to review.

~**Note**:This model uses a SentencePiece tokenizer. They have provided the sentence-piece .model file which can be loaded. While creating FNetTokenizer should I inherit from some other existing tokenizer? Alternatively, I can copy the tokenizer from `ALBERT` (which is what I am doing right now). Wdyt?~

**Note**: I am trying to make this model as similar to Bert is possible. The original implementation has slightly different layers. For example, `FNetIntermediate` and `FNetOutput` equivalents are combined into a single layer in original FNet code, but I keep them separate. Hope this is okay?

EDIT 1:
------
I have made necessary changes for the model. And since the model compares against Bert, it makes sense to have (almost) all tasks - MultipleChoice, QuestionAnswering, etc. I am still working on:
- [x] Tokenizer (regular and fast)
- [x] Documentation
- [x] Checkpoint Conversion 
- [x] Tests

EDIT 2:
------
~We also need to skip `attention_mask` totally from the tokenizer. The user, ideally, should not have an option to get the `attention_mask` using `FNetTokenizer`. I am using `model_input_names` for this.~

EDIT 3:
------
~One more concern is that, since I am implementing in PyTorch, do we expect the user to run this on TPU? The reason is that the original implementation changes the way they calculate FFT on TPU, based on the sequence length (they found some optimal rules for faster processing). Currently, I have only used `torch.fft.fftn` directly (they use `jnp.fft.fftn` in the CPU/GPU case). Please let me know what you think.~

EDIT 4:
------
One more thing to consider is that the original code allows `type_vocab_size` of 4, which is used only for GLUE tasks. During pre-training they only use `0` and `1`. 
But, the checkpoints also have the shape of embedding weights as `(4, 768)` . Does that mean that the tokenizer might need to support something like:
```python
tokenizer = FNetTokenizer.from_pretrained('fnet-base')
inputs = tokenizer(text1, text2, text3, text4)
```
?

EDIT 5:
------
~The colab link to outputs on checkpoint conversion: [Flax to PyTorch](https://colab.research.google.com/drive/1CxxDwaH4Tei9cUBHRaMYWPHCpS2El2He?usp=sharing).
The model outputs, embedding layer, encoder layer 0 outputs match up to `1e-2`, except masked lm output for masked token, which matches to `1e-1`. Any idea on how I can improve this?~

~One reason I can think of this reduction is precision is the difference in precision in `torch.fft.fftn` and `jnp.fft.fftn` which is atmost `1e-4`. From a difference of  atmost `1e-6`, after applying the corresponding transforms, the difference becomes atmost `1e-3` in the real part. Over layers, this might accumulate. Just a guess, however.~

~This was fixed by using `gelu_new` instead of `gelu`.~

EDIT 6:
------
They use a projection layer in the embeddings, and hence the embedding size and hidden size for the model are provided separately in the config. In their experiments, they keep it same, but the flexibility is still there. Do we want to keep both the sizes separate?

EDIT 7:
------
~The FastTokenizer requires a `tokenizer.json` file which I have created using `convert_slow_tokenizer`. I used `AlbertConverter` for this model. I don't know (in-detail) how SentencePiece and FastTokenizers work. Please let me know if I'm missing anything.~

EDIT 8:
------
Just realized that the original model is denoted as `f_net`. I am using `fnet` everywhere, is this acceptable?

EDIT 9:
------
~I am not sure about special tokens in the tokenizer. The original model gives some special tokens as empty string ``. Using the current tokenizer code to load these gives `
` and <\s> for those tokens (bos, eos), and <**unk**> for unknown and <**pad**> for pad tokens. Not sure which is the right way to go. Any suggestions?~

Comments:
Not sure why this test fails:
```python
=========================== short test summary info ============================
FAILED tests/extended/test_trainer_ext.py::TestTrainerExt::test_run_seq2seq_no_dist
==== 1 failed, 7169 passed, 3466 skipped, 708 warnings in 832.52s (0:13:52) ====

```

The failure is due to the new release of sacrebleu. If you rebase on master to get the commit that pins it to < 2.0.0, the failure will go away (but it's not necessary for this PR to be merged as we know it has nothing to do with it).

Thanks for reviews @sgugger @patil-suraj
I'll address them quickly. 

One more concern apart from the ones I have mentioned above:

~I have removed the slow integration testing from tokenization tests as it expects `attention_mask`. I'll take a look and update the test accordingly.~

EDIT:
------
This test has been updated.


Hey @gchhablani :-) 

We've just added you to the Google org, so that you can move the model weights there. If you find some time, it would also be very nice to add some model cards (I can definitely help you with that). 

Regarding the failing doc test, you can just rebase to current master and it'll be fixed

I found two issues with the fourier transform.

### Issue 1
The actual implementation uses `jax.vmap` on the `self.fourier_transform`. I made a mistake earlier in the implementation and do it for all dimensions - `hidden_size`, `sequence_length`, and `batch_size`, but it is just `sequence_length` and `batch_size`.

This leads to a mismatch issue when using `batch_size` more than one. I have fixed this issue by passing in the correct dimensions to `torch.fft.fftn` and using `functools.partial`.

Please check the Flax/Torch output match for `batch_size=2` [here](https://colab.research.google.com/drive/13cqOgP4DNrYbBdjwD0NwSxORCUTRxiZ-?usp=sharing).

Unfortunately, there is no `vmap` in torch as of now in the stable version, but only in the nightly version [here](https://pytorch.org/tutorials/prototype/vmap_recipe.html).

### Issue 2
Following @sgugger's suggestion to add the optimizations for TPUs, I tried adding the `einsum` version of fourier transform where they use DFT matrix multiplication and the axis-wise FFT. I have had to make changes and few additions to support them in PyTorch. Currently, the outputs from those don't match (but they should, at least to some extent). So I am fixing that as well.


EDIT:
------
I understand the issue with the `einsum` implementation. The original code uses the maximum sequence length possible as their sequence length during training - 512. Hence, during the initialization, they specify this maximum sequence length, and then use this variable to initialize the `DFT` matrix for sequence length. While that may have made sense for them, I'm not sure if it makes sense here?

I think we can take in another parameter `sequence_length` during config initialization. This will be used to specify the sequence length (because `max_position_embeddings` is used to initialize the `self.position_embeddings`, so that shouldn't be changed). Along with this, a check that throws an error if the `sequence_length` does not match sequence length passed to the model. 


With the latest changes, an error occurs:

```python
ImportError while importing test module '/home/circleci/transformers/tests/test_modeling_fnet.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
/usr/local/lib/python3.7/importlib/__init__.py:127: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
tests/test_modeling_fnet.py:23: in <module>
    from transformers.models.fnet.modeling_fnet import FNetBasicFourierTransform
src/transformers/models/fnet/modeling_fnet.py:28: in <module>
    from scipy import linalg
E   ModuleNotFoundError: No module named 'scipy'
```
I am trying to use `scipy.linalg.dft` to get `DFT matrix`. Any chance this can be a dependency?

EDIT
------
I have added a variable called `_scipy_available` which is used when initializing the fourier transform, and if it is not available, I add a warning. The users can install SciPy if they want?

I don't see a problem with using `scipy` as an optional dependency for this specific model

Let me know if need help making the tests pass with the dependency - I can fix this in your PR if you want :-)

Hi @patrickvonplaten

I have pushed the code where I used the global variable `_scipy_available`, does that seem okay? The tests are working fine locally.

Also, in model tests I'm verifying whether `fourier_transform` implementations match or not in the test: `create_and_check_fourier_transform` for which I need to access `modeling_fnet`.

I get this error on CircleCI:

```python
_________________ ERROR collecting tests/test_modeling_fnet.py _________________
ImportError while importing test module '/home/circleci/transformers/tests/test_modeling_fnet.py'.
Hint: make sure your test modules/packages have valid Python names.
Traceback:
/usr/local/lib/python3.7/importlib/__init__.py:127: in import_module
    return _bootstrap._gcd_import(name[level:], package, level)
tests/test_modeling_fnet.py:23: in <module>
    from transformers.models.fnet.modeling_fnet import FNetBasicFourierTransform, _scipy_available
src/transformers/models/fnet/modeling_fnet.py:22: in <module>
    import torch
E   ModuleNotFoundError: No module named 'torch'
```

Any idea how do I fix this?

EDIT
------
Test is fixed. I followed `fsmt` tests. Had to add the imports under `is_torch_available()`. 

I have updated the checkpoints and added basic model cards. The model performance isn't great on MLM, not sure why. The accuracy scores are low, though.

Checkpoints
- [fnet-base](https://huggingface.co/google/fnet-base)
- [fnet-large](https://huggingface.co/google/fnet-large)

Also, just to check :-) The reported eval metrics on GLUE - did you run them once with `run_glue.py` or is it a copy-paste of the paper? 

@patrickvonplaten No, I just copy pasted from the paper 🙈. Should I try fine-tuning it?

Maybe, that itself can be the demo?

I am checking the checkpoint conversion. Ideally, there should be less than `1e-3`/`1e-4` differences in the outputs. I'm not sure how to exactly fix this, but the arg-sorted order of the predictions is different for the PyTorch and the Flax model. :/ 

For different fourier transforms, I matched them against `np.fft.fftn` and `jnp.fft.fftn`, both give at best `1e-4` match, which means the problem is not the fourier transform.

I'll do a layer-wise debugging and update here.

Nonetheless, the original masked LM weights lead to similar predictions, so fine-tuning example will be helpful.


EDIT
------
The issue was that the original implementation uses gelu from BERT, which is equivalent to `gelu_new`, I suppose. Changing the activation to `gelu_new` leads to a `1e-4` match on all logits and sequence output ^_^

I am still working on verifying model outputs.

The original MLM model performs decently for the following: "the man worked as a [MASK]." The masked token top-10 predictions are:
```
man
person
use
guide
work
example
reason
source
one
right
```
I had to modify the tokens as expected by the model. The tokenizer is having issues. The original one gives this output for the text above:
```python
[13, 283, 2479, 106, 8, 16657, 6, 16678]
[ '▁the', '▁man', '▁worked', '▁as', '▁a', '▁', '[MASK]', '.' ]
```
The tokenizer I wrote is returning this:
```python
[13, 283, 2479, 106, 8, 6, 845]
['▁the', '▁man', '▁worked', '▁as', '▁a', '[MASK]', '▁.']
```
Notice how the space - `▁` is missing and that the period is actually `.` but becomes `▁.` in our tokenizer.

Any ideas on why this might be happening?

When I change `[MASK]` to `mask`, both lead to same output:

```python
[13, 283, 2479, 106, 8, 10469, 16678]
['▁the', '▁man', '▁worked', '▁as', '▁a', '▁mask', '.']
```
In their [input_pipeline](https://github.com/google-research/google-research/blob/master/f_net/input_pipeline.py#L258), they add the mask, cls and sep ids manually. Hence, they never use `[MASK]` in the text input. So, maybe, it's okay if we get `▁a`, `[MASK]`?
But in either case, we shouldn't get `▁.`? How do I handle this?

The problem happens in `tokenize`, where we split based on the `[MASK]` token. But if we don't do that, then `[MASK]` is broken into several tokens. `tokenize('.')` results in `▁.` 

I tried fixing the issue using a fix which basically skips first character after a mask token, only if it is not a `no_split_token`.

I'm not sure if this is 100% correct.

Also, there is an error with `FNetTokenizerFast`, the `[MASK]` token is not working as expected:

```python
[4, 13, 283, 2479, 106, 8, 1932, 2594, 16681, 6154, 5]
['[CLS]', '▁the', '▁man', '▁worked', '▁as', '▁a', '▁[', 'mas', 'k', '].', '[SEP]']
```

@patrickvonplaten @LysandreJik @sgugger
Can you please check the tokenizer once when you get a chance?

Once that is working, I can proceed with the fine-tuning without any issues.

Checking now!

Here @gchhablani, I looked into it and the tokenizer actually looks correct to me. See this colab: https://colab.research.google.com/drive/1QC4yvSHk0DSOObD6U2fbUE-9-6W3D3_F?usp=sharing 
Note that in the original code tokens are just "manually" replaced by the "[MASK]" token. So in the colab above, if the token for "guide" (3106) is replaced by the mask token id in the original code then the current tokenizer would be correct

I'm wondering whether the model is actually the same. Checking this now...

@patrickvonplaten  The tokenizer is working as expected because of the fixed I pushed in the previous commit. It handles the mask token, but I am not 100% sure if it is correct or if there is a better way to deal with this.

BTW, to fix the pipelines torch tests I think you just have to rebase to current master (or merge master into your branch :-) ) 

@patrickvonplaten

There is an issue with `FNetTokenizerFast`:

```python
>>> from src.transformers.models.fnet.tokenization_fnet_fast import FNetTokenizerFast
2021-08-30 20:30:09.070691: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
2021-08-30 20:30:09.070758: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
>>> tokenizer = FNetTokenizerFast.from_pretrained('google/fnet-base')
>>> text = "the man worked as a [MASK]."
>>> tokenizer.encode(text)
[4, 13, 283, 2479, 106, 8, 1932, 2594, 16681, 6154, 5]
>>> tokenizer.tokenize(text)
['▁the', '▁man', '▁worked', '▁as', '▁a', '▁[', 'mas', 'k', '].']
```

The `[MASK]` should not get tokenized. Any idea why this might be happening?

@gchhablani it seems that FNet was trained with a SPM vocab, so the corect masking token should be `<mask>` :)

@stefan-it 
I haven't worked with sentencepiece before so I'm not sure. But, in the [original code](https://github.com/google-research/google-research/blob/8077479d91cca79b16417055511b7744c155c344/f_net/input_pipeline.py#L256-L258), they specify `[CLS], [SEP], [MASK]` explicitly. However, they do not use the `[MASK]` string token anywhere, but only the id.

What do you think about this?

If changing to `<mask>` will fix things, then we can go with it. I will try it out.

Hi @gchhablani , oh I'm sorry I haven't yet read the official implementation. But it seems that they're really using `[MASK]` as the masking token (as previously done in [ALBERT](https://github.com/google-research/albert#sentencepiece)).

> @patrickvonplaten
> 
> There is an issue with `FNetTokenizerFast`:
> 
> ```python
> >>> from src.transformers.models.fnet.tokenization_fnet_fast import FNetTokenizerFast
> 2021-08-30 20:30:09.070691: W tensorflow/stream_executor/platform/default/dso_loader.cc:64] Could not load dynamic library 'libcudart.so.11.0'; dlerror: libcudart.so.11.0: cannot open shared object file: No such file or directory
> 2021-08-30 20:30:09.070758: I tensorflow/stream_executor/cuda/cudart_stub.cc:29] Ignore above cudart dlerror if you do not have a GPU set up on your machine.
> >>> tokenizer = FNetTokenizerFast.from_pretrained('google/fnet-base')
> >>> text = "the man worked as a [MASK]."
> >>> tokenizer.encode(text)
> [4, 13, 283, 2479, 106, 8, 1932, 2594, 16681, 6154, 5]
> >>> tokenizer.tokenize(text)
> ['▁the', '▁man', '▁worked', '▁as', '▁a', '▁[', 'mas', 'k', '].']
> ```
> 
> The `[MASK]` should not get tokenized. Any idea why this might be happening?

I can fix this once the changes to how `token_type_ids` are generated are applied :-)