Address comments

`ElectraDiscriminatorPredictions.forward` should not need `attention_mask`.

.circleci/config.yml

@@ -57,7 +57,10 @@ jobs:
         steps:
             - checkout
             - run: sudo pip install .[mecab,testing]
-            - run: python -m pytest -sv ./tests/test_tokenization_bert_japanese.py
+            - run: python -m pytest -s ./tests/test_tokenization_bert_japanese.py | tee output.txt
+            - store_artifacts:
+                path: ~/transformers/output.txt
+                destination: test_output.txt
     run_examples_torch:
         working_directory: ~/transformers
         docker:

.github/workflows/self-push.yml

@@ -51,4 +51,11 @@ jobs:
         USE_CUDA: yes
       run: |
         source .env/bin/activate
-        python -m pytest -n 2 --dist=loadfile -s -v ./tests/
+        python -m pytest -n 2 --dist=loadfile -s ./tests/ | tee output.txt
+    - name: cat output.txt
+      run: cat output.txt
+    - name: Upload output.txt
+      uses: actions/upload-artifact@v1
+      with:
+        name: pytest_output
+        path: output.txt

.github/workflows/self-scheduled.yml

@@ -46,5 +46,11 @@ jobs:
         USE_CUDA: yes
       run: |
         source .env/bin/activate
-        python -m pytest -n 1 --dist=loadfile -s -v ./tests/
-        
+        python -m pytest -n 1 --dist=loadfile -s ./tests/  | tee output.txt
+    - name: cat output.txt
+      run: cat output.txt
+    - name: Upload output.txt
+      uses: actions/upload-artifact@v1
+      with:
+        name: pytest_output
+        path: output.txt

docs/source/_static/js/custom.js

@@ -87,7 +87,7 @@ function addVersionControl() {
     const parts = location.toString().split('/');
     let versionIndex = parts.length - 2;
     // Index page may not have a last part with filename.html so we need to go up
-    if (parts[parts.length - 1] != "" && ! parts[parts.length - 1].match(/\.html$/)) {
+    if (parts[parts.length - 1] != "" && ! parts[parts.length - 1].match(/\.html$|^search.html?/)) {
         versionIndex = parts.length - 1;
     }
     // Main classes and models are nested so we need to go deeper

docs/source/glossary.rst

@@ -11,7 +11,7 @@ General terms
   tokens at a certain timestep.
 - MLM: masked language modeling, a pretraining task where the model sees a corrupted version of the texts, usually done
   by masking some tokens randomly, and has to predict the original text.
-- multimodal: a task taht combines texts with another kind of inputs (for instance images).
+- multimodal: a task that combines texts with another kind of inputs (for instance images).
 - NLG: natural language generation, all tasks related to generating text ( for instance talk with transformers,
   translation)
 - NLP: natural language processing, a generic way to say "deal with texts".

docs/source/main_classes/optimizer_schedules.rst

@@ -1,4 +1,4 @@
-Optimizer
+Optimization
 ----------------------------------------------------
 
 The ``.optimization`` module provides:
@@ -7,24 +7,25 @@ The ``.optimization`` module provides:
 - several schedules in the form of schedule objects that inherit from ``_LRSchedule``:
 - a gradient accumulation class to accumulate the gradients of multiple batches
 
-``AdamW``
-~~~~~~~~~~~~~~~~
+``AdamW`` (PyTorch)
+~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: transformers.AdamW
     :members:
 
-``AdamWeightDecay``
-~~~~~~~~~~~~~~~~~~~
+``AdamWeightDecay`` (TensorFlow)
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: transformers.AdamWeightDecay
 
 .. autofunction:: transformers.create_optimizer
 
 Schedules
-----------------------------------------------------
+~~~~~~~~~~~~~~~~~~~
+
+Learning Rate Schedules (Pytorch)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
-Learning Rate Schedules
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 .. autofunction:: transformers.get_constant_schedule
 
 
@@ -56,16 +57,16 @@ Learning Rate Schedules
     :target: /imgs/warmup_linear_schedule.png
     :alt:
 
-``Warmup``
-~~~~~~~~~~~~~~~~
+``Warmup`` (TensorFlow)
+^^^^^^^^^^^^^^^^^^^^^^^
 
 .. autoclass:: transformers.WarmUp
     :members:
 
 Gradient Strategies
-----------------------------------------------------
+~~~~~~~~~~~~~~~~~~~~
 
-``GradientAccumulator``
-~~~~~~~~~~~~~~~~~~~~~~~
+``GradientAccumulator`` (TensorFlow)
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
 .. autoclass:: transformers.GradientAccumulator

docs/source/model_doc/reformer.rst

@@ -112,3 +112,17 @@ ReformerModelWithLMHead
 
 .. autoclass:: transformers.ReformerModelWithLMHead
     :members:
+
+
+ReformerForMaskedLM
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.ReformerForMaskedLM
+    :members:
+
+
+ReformerForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.ReformerForQuestionAnswering
+    :members:

examples/README.md

@@ -1,4 +1,4 @@
-## Examples
+# Examples
 
 Version 2.9 of 🤗 Transformers introduces a new [`Trainer`](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer.py) class for PyTorch, and its equivalent [`TFTrainer`](https://github.com/huggingface/transformers/blob/master/src/transformers/trainer_tf.py) for TF 2.
 Running the examples requires PyTorch 1.3.1+ or TensorFlow 2.1+.
@@ -13,7 +13,7 @@ Here is the list of all our examples:
 This is still a work-in-progress – in particular documentation is still sparse – so please **contribute improvements/pull requests.**
 
 
-# The Big Table of Tasks
+## The Big Table of Tasks
 
 | Task | Example datasets | Trainer support | TFTrainer support | pytorch-lightning | Colab
 |---|---|:---:|:---:|:---:|:---:|

examples/multiple-choice/run_tf_multiple_choice.py

@@ -108,7 +108,10 @@ def main():
         level=logging.INFO,
     )
     logger.warning(
-        "device: %s, n_gpu: %s, 16-bits training: %s", training_args.device, training_args.n_gpu, training_args.fp16,
+        "device: %s, n_replicas: %s, 16-bits training: %s",
+        training_args.device,
+        training_args.n_replicas,
+        training_args.fp16,
     )
     logger.info("Training/evaluation parameters %s", training_args)
 

examples/question-answering/run_tf_squad.py

@@ -137,9 +137,9 @@ def main():
         level=logging.INFO,
     )
     logger.info(
-        "n_gpu: %s, distributed training: %s, 16-bits training: %s",
-        training_args.n_gpu,
-        bool(training_args.n_gpu > 1),
+        "n_replicas: %s, distributed training: %s, 16-bits training: %s",
+        training_args.n_replicas,
+        bool(training_args.n_replicas > 1),
         training_args.fp16,
     )
     logger.info("Training/evaluation parameters %s", training_args)

examples/seq2seq/README.md

@@ -71,7 +71,7 @@ Load it with `BartForConditionalGeneration.from_pretrained(f'{output_dir}/best_t
 - If you want to run experiments on improving the summarization finetuning process, try the XSUM Shared Task (below). It's faster to train than CNNDM because the summaries are shorter.
 - For CNN/DailyMail, the default `val_max_target_length` and `test_max_target_length` will truncate the ground truth labels, resulting in slightly higher rouge scores. To get accurate rouge scores, you should rerun calculate_rouge on the `{output_dir}/test_generations.txt` file saved by `trainer.test()`
 - `--max_target_length=60 --val_max_target_length=60 --test_max_target_length=100 ` is a reasonable setting for XSUM.
-- `wandb` can be used by specifying `--logger wandb_shared` or `--logger wandb`. It is useful for reproducibility. 
+- `wandb` can be used by specifying `--logger wandb`. It is useful for reproducibility. Specify the environment variable `WANDB_PROJECT='hf_xsum'` to do the XSUM shared task. 
 - This warning can be safely ignored: 
     > "Some weights of BartForConditionalGeneration were not initialized from the model checkpoint at facebook/bart-large-xsum and are newly initialized: ['final_logits_bias']"
 - Both finetuning and eval are 30% faster with `--fp16`. For that you need to [install apex](https://github.com/NVIDIA/apex#quick-start).
@@ -111,14 +111,14 @@ Compare XSUM results with others by using `--logger wandb_shared`. This requires
 
 Here is an example command, but you can do whatever you want. Hopefully this will make debugging and collaboration easier!
 ```bash
-./finetune.sh \
+WANDB_PROJECT='hf_xsum' ./finetune.sh \
     --data_dir $XSUM_DIR \
     --output_dir xsum_frozen_embs \
     --model_name_or_path facebook/bart-large \
-    --logger wandb_shared \
     --train_batch_size 16 --eval_batch_size 16 --freeze_embeds --freeze_encoder \
     --num_train_epochs 6 \
-    --max_target_length=60 --val_max_target_length=60 --test_max_target_length=100
+    --max_target_length=60 --val_max_target_length=60 --test_max_target_length=100 \
+    --logger wandb
 ```
 
 You can see your wandb logs [here](https://app.wandb.ai/sshleifer/hf_xsum?workspace=user-)

examples/seq2seq/finetune.py

@@ -298,8 +298,6 @@ def main(args, model=None) -> SummarizationModule:
             model: SummarizationModule = SummarizationModule(args)
         else:
             model: SummarizationModule = TranslationModule(args)
-
-    dataset = Path(args.data_dir).name
     if (
         args.logger == "default"
         or args.fast_dev_run
@@ -310,12 +308,12 @@ def main(args, model=None) -> SummarizationModule:
     elif args.logger == "wandb":
         from pytorch_lightning.loggers import WandbLogger
 
-        logger = WandbLogger(name=model.output_dir.name, project=dataset)
+        logger = WandbLogger(name=model.output_dir.name)
 
     elif args.logger == "wandb_shared":
         from pytorch_lightning.loggers import WandbLogger
 
-        logger = WandbLogger(name=model.output_dir.name, project=f"hf_{dataset}")
+        logger = WandbLogger(name=model.output_dir.name)
     trainer: pl.Trainer = generic_train(
         model,
         args,

examples/seq2seq/finetune_bart_tiny.sh

@@ -12,7 +12,7 @@ export OUTPUT_DIR=${CURRENT_DIR}/${OUTPUT_DIR_NAME}
 # Make output directory if it doesn't exist
 mkdir -p $OUTPUT_DIR
 
-# Add parent directory to python path to access lightning_base.py and utils.py
+# Add parent directory to python path to access lightning_base.py and testing_utils.py
 export PYTHONPATH="../":"${PYTHONPATH}"
 python finetune.py \
 --data_dir=cnn_tiny/ \

examples/seq2seq/test_seq2seq_examples.py

@@ -12,6 +12,7 @@ import torch
 from torch.utils.data import DataLoader
 
 from transformers import AutoTokenizer
+from transformers.testing_utils import require_multigpu
 
 from .distillation import distill_main, evaluate_checkpoint
 from .finetune import main
@@ -107,7 +108,7 @@ class TestSummarizationDistiller(unittest.TestCase):
         logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks
         return cls
 
-    @unittest.skipUnless(torch.cuda.device_count() > 1, "skipping multiGPU test")
+    @require_multigpu
     def test_multigpu(self):
         updates = dict(no_teacher=True, freeze_encoder=True, gpus=2, sortish_sampler=False,)
         self._test_distiller_cli(updates)

examples/text-classification/run_tf_glue.py

@@ -131,9 +131,9 @@ def main():
         level=logging.INFO,
     )
     logger.info(
-        "n_gpu: %s, distributed training: %s, 16-bits training: %s",
-        training_args.n_gpu,
-        bool(training_args.n_gpu > 1),
+        "n_replicas: %s, distributed training: %s, 16-bits training: %s",
+        training_args.n_replicas,
+        bool(training_args.n_replicas > 1),
         training_args.fp16,
     )
     logger.info("Training/evaluation parameters %s", training_args)

examples/token-classification/run_ner.py

@@ -75,7 +75,8 @@ class DataTrainingArguments:
         metadata={"help": "The input data dir. Should contain the .txt files for a CoNLL-2003-formatted task."}
     )
     labels: Optional[str] = field(
-        metadata={"help": "Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."}
+        default=None,
+        metadata={"help": "Path to a file containing all labels. If not specified, CoNLL-2003 labels are used."},
     )
     max_seq_length: int = field(
         default=128,

examples/token-classification/run_tf_ner.py

@@ -109,9 +109,9 @@ def main():
         level=logging.INFO,
     )
     logger.info(
-        "n_gpu: %s, distributed training: %s, 16-bits training: %s",
-        training_args.n_gpu,
-        bool(training_args.n_gpu > 1),
+        "n_replicas: %s, distributed training: %s, 16-bits training: %s",
+        training_args.n_replicas,
+        bool(training_args.n_replicas > 1),
         training_args.fp16,
     )
     logger.info("Training/evaluation parameters %s", training_args)

model_cards/MoseliMotsoehli/TswanaBert/README.md

@@ -0,0 +1,72 @@
+---
+language: setswana
+---
+
+# TswanaBert
+
+## Model Description.
+TswanaBERT is a transformers model pretrained on a corpus of Setswana data in a self-supervised fashion by masking part of the input words and training to predict the masks.
+
+## Intended uses & limitations
+The model can  be used for either masked language modeling or next word prediction. it can also be fine-tuned for a specifict application. 
+
+#### How to use
+
+```python
+>>> from transformers import pipeline
+>>> from transformers import AutoTokenizer, AutoModelWithLMHead
+
+>>> tokenizer = AutoTokenizer.from_pretrained("MoseliMotsoehli/TswanaBert")
+>>> model = AutoModelWithLMHead.from_pretrained("MoseliMotsoehli/TswanaBert")
+>>> unmasker = pipeline('fill-mask', model=model, tokenizer=tokenizer)
+>>> unmasker("Ntshopotse <mask> e godile.")
+
+[{'score': 0.32749542593955994,
+  'sequence': '<s>Ntshopotse setse e godile.</s>',
+  'token': 538,
+  'token_str': 'Ġsetse'},
+ {'score': 0.060260992497205734,
+  'sequence': '<s>Ntshopotse le e godile.</s>',
+  'token': 270,
+  'token_str': 'Ġle'},
+ {'score': 0.058460816740989685,
+  'sequence': '<s>Ntshopotse bone e godile.</s>',
+  'token': 364,
+  'token_str': 'Ġbone'},
+ {'score': 0.05694682151079178,
+  'sequence': '<s>Ntshopotse ga e godile.</s>',
+  'token': 298,
+  'token_str': 'Ġga'},
+ {'score': 0.0565204992890358,
+  'sequence': '<s>Ntshopotse, e godile.</s>',
+  'token': 16,
+  'token_str': ','}]
+```
+
+#### Limitations and bias
+The model is trained on a fairly small collection of setwana, mostly from news articles and creative writtings, and so is not representative enough of the language as yet.
+
+## Training data
+
+The largest portion of this dataset (10k)  lines of text, comes from the [Leipzig Corpora Collection](https://wortschatz.uni-leipzig.de/en/download)
+
+The I then added 200 more phrases and sentences by scrapping following sites. I continue to expand the dataset
+
+* http://setswana.blogspot.com/
+* https://omniglot.com/writing/tswana.php
+* http://www.dailynews.gov.bw/
+* http://www.mmegi.bw/index.php
+* https://tsena.co.bw
+* http://www.botswana.co.za/Cultural_Issues-travel/botswana-country-guide-en-route.html
+
+## Training procedure
+The model was trained on a google colab Tesla T4 GPU for 200 epochs with a batch size of 64, on 13446 learned tokens.
+Other model training configuration setting can be found [here](https://s3.amazonaws.com/models.huggingface.co/bert/MoseliMotsoehli/TswanaBert/config.json)
+
+### BibTeX entry and citation info
+
+```bibtex
+@inproceedings{author = {Moseli Motsoehli},
+  year={2020}
+}
+```

model_cards/mrm8488/electra-base-finetuned-squadv1/README.md

@@ -0,0 +1,85 @@
+---
+language: english
+---
+
+# Electra base ⚡ + SQuAD v1 ❓
+
+[Electra-base-discriminator](https://huggingface.co/google/electra-base-discriminator) fine-tuned on [SQUAD v1.1 dataset](https://rajpurkar.github.io/SQuAD-explorer/explore/1.1/dev/) for **Q&A** downstream task.
+
+## Details of the downstream task (Q&A) - Model 🧠
+
+**ELECTRA** is a new method for self-supervised language representation learning. It can be used to pre-train transformer networks using relatively little compute. ELECTRA models are trained to distinguish "real" input tokens vs "fake" input tokens generated by another neural network, similar to the discriminator of a [GAN](https://arxiv.org/pdf/1406.2661.pdf). At small scale, ELECTRA achieves strong results even when trained on a single GPU. At large scale, ELECTRA achieves state-of-the-art results on the [SQuAD 2.0](https://rajpurkar.github.io/SQuAD-explorer/) dataset.
+
+
+## Details of the downstream task (Q&A) - Dataset 📚
+
+**S**tanford **Q**uestion **A**nswering **D**ataset (SQuAD) is a reading comprehension dataset, consisting of questions posed by crowdworkers on a set of Wikipedia articles, where the answer to every question is a segment of text, or span, from the corresponding reading passage, or the question might be unanswerable.
+SQuAD v1.1 contains **100,000+** question-answer pairs on **500+** articles.
+
+## Model training 🏋️‍
+
+The model was trained on a Tesla P100 GPU and 25GB of RAM with the following command:
+
+```bash
+python transformers/examples/question-answering/run_squad.py \
+  --model_type electra \
+  --model_name_or_path 'google/electra-base-discriminator' \
+  --do_eval \
+  --do_train \
+  --do_lower_case \
+  --train_file '/content/dataset/train-v1.1.json' \
+  --predict_file '/content/dataset/dev-v1.1.json' \
+  --per_gpu_train_batch_size 16 \
+  --learning_rate 3e-5 \
+  --num_train_epochs 10 \
+  --max_seq_length 384 \
+  --doc_stride 128 \
+  --output_dir '/content/output' \
+  --overwrite_output_dir \
+  --save_steps 1000
+```
+
+## Test set Results 🧾
+
+| Metric | # Value   |
+| ------ | --------- |
+| **EM** | **83.03** |
+| **F1** | **90.77** |
+| **Size**| **+ 400 MB** |
+
+Very good metrics for such a "small" model!
+
+```json
+{
+'exact': 83.03689687795648,
+'f1': 90.77486052446231,
+'total': 10570,
+'HasAns_exact': 83.03689687795648,
+'HasAns_f1': 90.77486052446231,
+'HasAns_total': 10570,
+'best_exact': 83.03689687795648,
+'best_exact_thresh': 0.0,
+'best_f1': 90.77486052446231,
+'best_f1_thresh': 0.0
+}
+```
+
+### Model in action 🚀
+
+Fast usage with **pipelines**:
+
+```python
+from transformers import pipeline
+
+QnA_pipeline = pipeline('question-answering', model='mrm8488/electra-base-finetuned-squadv1')
+
+QnA_pipeline({
+    'context': 'A new strain of flu that has the potential to become a pandemic has been identified in China by scientists.',
+    'question': 'What has been discovered by scientists from China ?'
+})
+# Output:
+{'answer': 'A new strain of flu', 'end': 19, 'score': 0.9995211430099182, 'start': 0}
+```
+
+> Created by [Manuel Romero/@mrm8488](https://twitter.com/mrm8488) | [LinkedIn](https://www.linkedin.com/in/manuel-romero-cs/)
+> Made with <span style="color: #e25555;">&hearts;</span> in Spain

model_cards/roberta-large-mnli-README.md

@@ -0,0 +1,22 @@
+---
+license: mit
+widget:
+- text: "I like you. </s></s> I love you."
+---
+
+
+## roberta-large-mnli
+
+Trained by Facebook, [original source](https://github.com/pytorch/fairseq/tree/master/examples/roberta)
+
+```bibtex
+@article{liu2019roberta,
+    title = {RoBERTa: A Robustly Optimized BERT Pretraining Approach},
+    author = {Yinhan Liu and Myle Ott and Naman Goyal and Jingfei Du and
+              Mandar Joshi and Danqi Chen and Omer Levy and Mike Lewis and
+              Luke Zettlemoyer and Veselin Stoyanov},
+    journal={arXiv preprint arXiv:1907.11692},
+    year = {2019},
+}
+```
+

model_cards/schmidek/electra-small-cased/README.md

@@ -0,0 +1,11 @@
+---
+language: english
+license: apache-2.0
+---
+
+## ELECTRA-small-cased
+
+This is a cased version of `google/electra-small-discriminator`, trained on the
+[OpenWebText corpus](https://skylion007.github.io/OpenWebTextCorpus/).
+
+Uses the same tokenizer and vocab from `bert-base-cased`

notebooks/README.md

@@ -39,3 +39,4 @@ Pull Request so it can be included under the Community notebooks.
 |[Fine-tune BERT for Multi-label Classification](https://github.com/abhimishra91/transformers-tutorials/blob/master/transformers_multi_label_classification.ipynb)|How to fine-tune BERT for multi-label classification using PyTorch|[Abhishek Kumar Mishra](https://github.com/abhimishra91) |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abhimishra91/transformers-tutorials/blob/master/transformers_multi_label_classification.ipynb)|
 |[Fine-tune T5 for Summarization](https://github.com/abhimishra91/transformers-tutorials/blob/master/transformers_summarization_wandb.ipynb)|How to fine-tune T5 for summarization in PyTorch and track experiments with WandB|[Abhishek Kumar Mishra](https://github.com/abhimishra91) |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abhimishra91/transformers-tutorials/blob/master/transformers_summarization_wandb.ipynb)|
 |[Speed up Fine-Tuning in Transformers with Dynamic Padding / Bucketing](https://github.com/ELS-RD/transformers-notebook/blob/master/Divide_Hugging_Face_Transformers_training_time_by_2_or_more.ipynb)|How to speed up fine-tuning by a factor of 2 using dynamic padding / bucketing|[Michael Benesty](https://github.com/pommedeterresautee) |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1CBfRU1zbfu7-ijiOqAAQUA-RJaxfcJoO?usp=sharing)|
+|[Pretrain Reformer for Masked Language Modeling](https://github.com/patrickvonplaten/notebooks/blob/master/Reformer_For_Masked_LM.ipynb)| How to train a Reformer model with bi-directional self-attention layers | [Patrick von Platen](https://github.com/patrickvonplaten) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1tzzh0i8PgDQGV3SMFUGxM7_gGae3K-uW?usp=sharing)|

setup.py

@@ -123,7 +123,7 @@ setup(
         # for OpenAI GPT
         "regex != 2019.12.17",
         # for XLNet
-        "sentencepiece",
+        "sentencepiece != 0.1.92",
         # for XLM
         "sacremoses",
     ],

src/transformers/__init__.py

@@ -155,7 +155,7 @@ from .tokenization_xlm_roberta import XLMRobertaTokenizer
 from .tokenization_xlnet import SPIECE_UNDERLINE, XLNetTokenizer
 
 # Trainer
-from .trainer_utils import EvalPrediction
+from .trainer_utils import EvalPrediction, set_seed
 from .training_args import TrainingArguments
 from .training_args_tf import TFTrainingArguments
 
@@ -366,7 +366,9 @@ if is_torch_available():
         ReformerAttention,
         ReformerLayer,
         ReformerModel,
+        ReformerForMaskedLM,
         ReformerModelWithLMHead,
+        ReformerForQuestionAnswering,
         REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
     )
 
@@ -397,7 +399,7 @@ if is_torch_available():
     )
 
     # Trainer
-    from .trainer import Trainer, set_seed, torch_distributed_zero_first
+    from .trainer import Trainer, torch_distributed_zero_first
     from .data.data_collator import default_data_collator, DataCollator, DataCollatorForLanguageModeling
     from .data.datasets import GlueDataset, TextDataset, LineByLineTextDataset, GlueDataTrainingArguments
 

src/transformers/data/data_collator.py

@@ -43,7 +43,8 @@ def default_data_collator(features: List[InputDataClass]) -> Dict[str, torch.Ten
     # Ensure that tensor is created with the correct type
     # (it should be automatically the case, but let's make sure of it.)
     if "label" in first and first["label"] is not None:
-        dtype = torch.long if type(first["label"]) is int else torch.float
+        label = first["label"].item() if isinstance(first["label"], torch.Tensor) else first["label"]
+        dtype = torch.long if isinstance(label, int) else torch.float
         batch["labels"] = torch.tensor([f["label"] for f in features], dtype=dtype)
     elif "label_ids" in first and first["label_ids"] is not None:
         if isinstance(first["label_ids"], torch.Tensor):

src/transformers/data/processors/glue.py

@@ -17,6 +17,7 @@
 
 import logging
 import os
+from dataclasses import asdict
 from enum import Enum
 from typing import List, Optional, Union
 
@@ -81,26 +82,16 @@ if is_tf_available():
 
         def gen():
             for ex in features:
-                yield (
-                    {
-                        "input_ids": ex.input_ids,
-                        "attention_mask": ex.attention_mask,
-                        "token_type_ids": ex.token_type_ids,
-                    },
-                    ex.label,
-                )
+                d = {k: v for k, v in asdict(ex).items() if v is not None}
+                label = d.pop("label")
+                yield (d, label)
+
+        input_names = ["input_ids"] + tokenizer.model_input_names
 
         return tf.data.Dataset.from_generator(
             gen,
-            ({"input_ids": tf.int32, "attention_mask": tf.int32, "token_type_ids": tf.int32}, tf.int64),
-            (
-                {
-                    "input_ids": tf.TensorShape([None]),
-                    "attention_mask": tf.TensorShape([None]),
-                    "token_type_ids": tf.TensorShape([None]),
-                },
-                tf.TensorShape([]),
-            ),
+            ({k: tf.int32 for k in input_names}, tf.int64),
+            ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])),
         )
 
 

src/transformers/data/processors/squad.py

@@ -389,57 +389,102 @@ def squad_convert_examples_to_features(
 
         def gen():
             for i, ex in enumerate(features):
-                yield (
-                    {
-                        "input_ids": ex.input_ids,
-                        "attention_mask": ex.attention_mask,
-                        "token_type_ids": ex.token_type_ids,
-                        "feature_index": i,
-                        "qas_id": ex.qas_id,
-                    },
-                    {
-                        "start_positions": ex.start_position,
-                        "end_positions": ex.end_position,
-                        "cls_index": ex.cls_index,
-                        "p_mask": ex.p_mask,
-                        "is_impossible": ex.is_impossible,
-                    },
-                )
+                if ex.token_type_ids is None:
+                    yield (
+                        {
+                            "input_ids": ex.input_ids,
+                            "attention_mask": ex.attention_mask,
+                            "feature_index": i,
+                            "qas_id": ex.qas_id,
+                        },
+                        {
+                            "start_positions": ex.start_position,
+                            "end_positions": ex.end_position,
+                            "cls_index": ex.cls_index,
+                            "p_mask": ex.p_mask,
+                            "is_impossible": ex.is_impossible,
+                        },
+                    )
+                else:
+                    yield (
+                        {
+                            "input_ids": ex.input_ids,
+                            "attention_mask": ex.attention_mask,
+                            "token_type_ids": ex.token_type_ids,
+                            "feature_index": i,
+                            "qas_id": ex.qas_id,
+                        },
+                        {
+                            "start_positions": ex.start_position,
+                            "end_positions": ex.end_position,
+                            "cls_index": ex.cls_index,
+                            "p_mask": ex.p_mask,
+                            "is_impossible": ex.is_impossible,
+                        },
+                    )
 
         # Why have we split the batch into a tuple? PyTorch just has a list of tensors.
-        train_types = (
-            {
-                "input_ids": tf.int32,
-                "attention_mask": tf.int32,
-                "token_type_ids": tf.int32,
-                "feature_index": tf.int64,
-                "qas_id": tf.string,
-            },
-            {
-                "start_positions": tf.int64,
-                "end_positions": tf.int64,
-                "cls_index": tf.int64,
-                "p_mask": tf.int32,
-                "is_impossible": tf.int32,
-            },
-        )
+        if "token_type_ids" in tokenizer.model_input_names:
+            train_types = (
+                {
+                    "input_ids": tf.int32,
+                    "attention_mask": tf.int32,
+                    "token_type_ids": tf.int32,
+                    "feature_index": tf.int64,
+                    "qas_id": tf.string,
+                },
+                {
+                    "start_positions": tf.int64,
+                    "end_positions": tf.int64,
+                    "cls_index": tf.int64,
+                    "p_mask": tf.int32,
+                    "is_impossible": tf.int32,
+                },
+            )
 
-        train_shapes = (
-            {
-                "input_ids": tf.TensorShape([None]),
-                "attention_mask": tf.TensorShape([None]),
-                "token_type_ids": tf.TensorShape([None]),
-                "feature_index": tf.TensorShape([]),
-                "qas_id": tf.TensorShape([]),
-            },
-            {
-                "start_positions": tf.TensorShape([]),
-                "end_positions": tf.TensorShape([]),
-                "cls_index": tf.TensorShape([]),
-                "p_mask": tf.TensorShape([None]),
-                "is_impossible": tf.TensorShape([]),
-            },
-        )
+            train_shapes = (
+                {
+                    "input_ids": tf.TensorShape([None]),
+                    "attention_mask": tf.TensorShape([None]),
+                    "token_type_ids": tf.TensorShape([None]),
+                    "feature_index": tf.TensorShape([]),
+                    "qas_id": tf.TensorShape([]),
+                },
+                {
+                    "start_positions": tf.TensorShape([]),
+                    "end_positions": tf.TensorShape([]),
+                    "cls_index": tf.TensorShape([]),
+                    "p_mask": tf.TensorShape([None]),
+                    "is_impossible": tf.TensorShape([]),
+                },
+            )
+        else:
+            train_types = (
+                {"input_ids": tf.int32, "attention_mask": tf.int32, "feature_index": tf.int64, "qas_id": tf.string},
+                {
+                    "start_positions": tf.int64,
+                    "end_positions": tf.int64,
+                    "cls_index": tf.int64,
+                    "p_mask": tf.int32,
+                    "is_impossible": tf.int32,
+                },
+            )
+
+            train_shapes = (
+                {
+                    "input_ids": tf.TensorShape([None]),
+                    "attention_mask": tf.TensorShape([None]),
+                    "feature_index": tf.TensorShape([]),
+                    "qas_id": tf.TensorShape([]),
+                },
+                {
+                    "start_positions": tf.TensorShape([]),
+                    "end_positions": tf.TensorShape([]),
+                    "cls_index": tf.TensorShape([]),
+                    "p_mask": tf.TensorShape([None]),
+                    "is_impossible": tf.TensorShape([]),
+                },
+            )
 
         return tf.data.Dataset.from_generator(gen, train_types, train_shapes)
     else:

src/transformers/modeling_auto.py

@@ -122,7 +122,12 @@ from .modeling_mobilebert import (
     MobileBertModel,
 )
 from .modeling_openai import OpenAIGPTLMHeadModel, OpenAIGPTModel
-from .modeling_reformer import ReformerModel, ReformerModelWithLMHead
+from .modeling_reformer import (
+    ReformerForMaskedLM,
+    ReformerForQuestionAnswering,
+    ReformerModel,
+    ReformerModelWithLMHead,
+)
 from .modeling_retribert import RetriBertModel
 from .modeling_roberta import (
     RobertaForMaskedLM,
@@ -266,6 +271,7 @@ MODEL_FOR_MASKED_LM_MAPPING = OrderedDict(
         (FlaubertConfig, FlaubertWithLMHeadModel),
         (XLMConfig, XLMWithLMHeadModel),
         (ElectraConfig, ElectraForMaskedLM),
+        (ReformerConfig, ReformerForMaskedLM),
     ]
 )
 
@@ -310,6 +316,7 @@ MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
         (MobileBertConfig, MobileBertForQuestionAnswering),
         (XLMConfig, XLMForQuestionAnsweringSimple),
         (ElectraConfig, ElectraForQuestionAnswering),
+        (ReformerConfig, ReformerForQuestionAnswering),
     ]
 )
 

src/transformers/modeling_electra.py

@@ -133,7 +133,7 @@ class ElectraDiscriminatorPredictions(nn.Module):
         self.dense_prediction = nn.Linear(config.hidden_size, 1)
         self.config = config
 
-    def forward(self, discriminator_hidden_states, attention_mask):
+    def forward(self, discriminator_hidden_states):
         hidden_states = self.dense(discriminator_hidden_states)
         hidden_states = get_activation(self.config.hidden_act)(hidden_states)
         logits = self.dense_prediction(hidden_states).squeeze()
@@ -518,7 +518,7 @@ class ElectraForPreTraining(ElectraPreTrainedModel):
         )
         discriminator_sequence_output = discriminator_hidden_states[0]
 
-        logits = self.discriminator_predictions(discriminator_sequence_output, attention_mask)
+        logits = self.discriminator_predictions(discriminator_sequence_output)
 
         output = (logits,)
 

src/transformers/modeling_mobilebert.py

@@ -575,7 +575,7 @@ class MobileBertPooler(nn.Module):
             return first_token_tensor
         else:
             pooled_output = self.dense(first_token_tensor)
-            pooled_output = F.tanh(pooled_output)
+            pooled_output = torch.tanh(pooled_output)
             return pooled_output
 
 

src/transformers/modeling_reformer.py

@@ -1704,6 +1704,7 @@ class ReformerModel(ReformerPreTrainedModel):
 class ReformerModelWithLMHead(ReformerPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
+        assert config.is_decoder, "If you want to use `ReformerLMHeadModel` make sure that `is_decoder=True`."
         self.reformer = ReformerModel(config)
         self.lm_head = ReformerOnlyLMHead(config)
 
@@ -1789,3 +1790,190 @@ class ReformerModelWithLMHead(ReformerPreTrainedModel):
             inputs_dict["num_hashes"] = kwargs["num_hashes"]
 
         return inputs_dict
+
+
+@add_start_docstrings("""Reformer Model with a `language modeling` head on top. """, REFORMER_START_DOCSTRING)
+class ReformerForMaskedLM(ReformerPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        assert (
+            not config.is_decoder
+        ), "If you want to use `ReformerForMaskedLM` make sure `config.is_decoder=False` for bi-directional self-attention."
+        self.reformer = ReformerModel(config)
+        self.lm_head = ReformerOnlyLMHead(config)
+
+        self.init_weights()
+
+    def get_output_embeddings(self):
+        return self.lm_head.decoder
+
+    def tie_weights(self):
+        # word embeddings are not tied in Reformer
+        pass
+
+    @add_start_docstrings_to_callable(REFORMER_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(tokenizer_class=_TOKENIZER_FOR_DOC, checkpoint="google/reformer-crime-and-punishment")
+    def forward(
+        self,
+        input_ids=None,
+        position_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        num_hashes=None,
+        labels=None,
+        output_hidden_states=None,
+        output_attentions=None,
+    ):
+        r"""
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+                Labels for computing the masked language modeling loss.
+                Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
+                Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
+
+    Return:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
+            Classification loss (cross entropy).
+        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
+            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+        all_hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        all_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        """
+
+        reformer_outputs = self.reformer(
+            input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            num_hashes=num_hashes,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+
+        sequence_output = reformer_outputs[0]
+        logits = self.lm_head(sequence_output)
+        outputs = (logits,) + reformer_outputs[1:]
+
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()  # -100 index = padding token
+            masked_lm_loss = loss_fct(logits.view(-1, self.config.vocab_size), labels.view(-1))
+            outputs = (masked_lm_loss,) + outputs
+
+        return outputs  # (mlm_loss), lm_logits, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+    """Reformer Model with a span classification head on top for
+    extractive question-answering tasks like SQuAD / TriviaQA ( a linear layer on
+    top of hidden-states output to compute `span start logits` and `span end logits`. """,
+    REFORMER_START_DOCSTRING,
+)
+class ReformerForQuestionAnswering(ReformerPreTrainedModel):
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+
+        self.reformer = ReformerModel(config)
+        # 2 * config.hidden_size because we use reversible residual layers
+        self.qa_outputs = nn.Linear(2 * config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def tie_weights(self):
+        # word embeddings are not tied in Reformer
+        pass
+
+    @add_start_docstrings_to_callable(REFORMER_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(tokenizer_class=_TOKENIZER_FOR_DOC, checkpoint="google/reformer-crime-and-punishment")
+    def forward(
+        self,
+        input_ids=None,
+        position_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        num_hashes=None,
+        start_positions=None,
+        end_positions=None,
+        output_hidden_states=None,
+        output_attentions=None,
+    ):
+        r"""
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+    Return:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.ReformerConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
+            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
+        start_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
+            Span-start scores (before SoftMax).
+        end_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
+            Span-end scores (before SoftMax).
+        all_hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        all_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        """
+
+        reformer_outputs = self.reformer(
+            input_ids,
+            position_ids=position_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            num_hashes=num_hashes,
+            output_hidden_states=output_hidden_states,
+            output_attentions=output_attentions,
+        )
+
+        sequence_output = reformer_outputs[0]
+
+        logits = self.qa_outputs(sequence_output)
+        start_logits, end_logits = logits.split(1, dim=-1)
+        start_logits = start_logits.squeeze(-1)
+        end_logits = end_logits.squeeze(-1)
+
+        outputs = (start_logits, end_logits,) + reformer_outputs[1:]
+
+        if start_positions is not None and end_positions is not None:
+            # If we are on multi-GPU, split add a dimension
+            if len(start_positions.size()) > 1:
+                start_positions = start_positions.squeeze(-1)
+            if len(end_positions.size()) > 1:
+                end_positions = end_positions.squeeze(-1)
+            # sometimes the start/end positions are outside our model inputs, we ignore these terms
+            ignored_index = start_logits.size(1)
+            start_positions.clamp_(0, ignored_index)
+            end_positions.clamp_(0, ignored_index)
+
+            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+            start_loss = loss_fct(start_logits, start_positions)
+            end_loss = loss_fct(end_logits, end_positions)
+            total_loss = (start_loss + end_loss) / 2
+            outputs = (total_loss,) + outputs
+
+        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)

src/transformers/modeling_transfo_xl.py

@@ -1016,11 +1016,14 @@ class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
             return self.crit.out_layers[-1]
 
     def prepare_inputs_for_generation(self, input_ids, past, **model_kwargs):
-        inputs = {"input_ids": input_ids}
+        inputs = {}
 
         # if past is defined in model kwargs then use it for faster decoding
         if past:
             inputs["mems"] = past
+            inputs["input_ids"] = input_ids[:, -1].unsqueeze(-1)
+        else:
+            inputs["input_ids"] = input_ids
 
         return inputs
 

src/transformers/optimization.py

@@ -16,6 +16,7 @@
 
 import logging
 import math
+from typing import Callable, Iterable, Tuple
 
 import torch
 from torch.optim import Optimizer
@@ -25,18 +26,40 @@ from torch.optim.lr_scheduler import LambdaLR
 logger = logging.getLogger(__name__)
 
 
-def get_constant_schedule(optimizer, last_epoch=-1):
-    """ Create a schedule with a constant learning rate.
+def get_constant_schedule(optimizer: Optimizer, last_epoch: int = -1):
+    """
+    Create a schedule with a constant learning rate, using the learning rate set in optimizer.
+
+    Args:
+        optimizer (:class:`~torch.optim.Optimizer`):
+            The optimizer for which to schedule the learning rate.
+        last_epoch (:obj:`int`, `optional`, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
     """
     return LambdaLR(optimizer, lambda _: 1, last_epoch=last_epoch)
 
 
-def get_constant_schedule_with_warmup(optimizer, num_warmup_steps, last_epoch=-1):
-    """ Create a schedule with a constant learning rate preceded by a warmup
-    period during which the learning rate increases linearly between 0 and 1.
+def get_constant_schedule_with_warmup(optimizer: Optimizer, num_warmup_steps: int, last_epoch: int = -1):
+    """
+    Create a schedule with a constant learning rate preceded by a warmup period during which the learning rate
+    increases linearly between 0 and the initial lr set in the optimizer.
+
+    Args:
+        optimizer (:class:`~torch.optim.Optimizer`):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (:obj:`int`):
+            The number of steps for the warmup phase.
+        last_epoch (:obj:`int`, `optional`, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
     """
 
-    def lr_lambda(current_step):
+    def lr_lambda(current_step: int):
         if current_step < num_warmup_steps:
             return float(current_step) / float(max(1.0, num_warmup_steps))
         return 1.0
@@ -45,11 +68,25 @@ def get_constant_schedule_with_warmup(optimizer, num_warmup_steps, last_epoch=-1
 
 
 def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, last_epoch=-1):
-    """ Create a schedule with a learning rate that decreases linearly after
-    linearly increasing during a warmup period.
+    """
+    Create a schedule with a learning rate that decreases linearly from the initial lr set in the optimizer to 0,
+    after a warmup period during which it increases linearly from 0 to the initial lr set in the optimizer.
+
+    Args:
+        optimizer (:class:`~torch.optim.Optimizer`):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (:obj:`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (:obj:`int`):
+            The totale number of training steps.
+        last_epoch (:obj:`int`, `optional`, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
     """
 
-    def lr_lambda(current_step):
+    def lr_lambda(current_step: int):
         if current_step < num_warmup_steps:
             return float(current_step) / float(max(1, num_warmup_steps))
         return max(
@@ -59,10 +96,29 @@ def get_linear_schedule_with_warmup(optimizer, num_warmup_steps, num_training_st
     return LambdaLR(optimizer, lr_lambda, last_epoch)
 
 
-def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_steps, num_cycles=0.5, last_epoch=-1):
-    """ Create a schedule with a learning rate that decreases following the
-    values of the cosine function between 0 and `pi * cycles` after a warmup
-    period during which it increases linearly between 0 and 1.
+def get_cosine_schedule_with_warmup(
+    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: float = 0.5, last_epoch: int = -1
+):
+    """
+    Create a schedule with a learning rate that decreases following the values of the cosine function between the
+    initial lr set in the optimizer to 0, after a warmup period during which it increases linearly between 0 and the
+    initial lr set in the optimizer.
+
+    Args:
+        optimizer (:class:`~torch.optim.Optimizer`):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (:obj:`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (:obj:`int`):
+            The total number of training steps.
+        num_cycles (:obj:`float`, `optional`, defaults to 0.5):
+            The number of waves in the cosine schedule (the defaults is to just decrease from the max value to 0
+            following a half-cosine).
+        last_epoch (:obj:`int`, `optional`, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
     """
 
     def lr_lambda(current_step):
@@ -75,11 +131,27 @@ def get_cosine_schedule_with_warmup(optimizer, num_warmup_steps, num_training_st
 
 
 def get_cosine_with_hard_restarts_schedule_with_warmup(
-    optimizer, num_warmup_steps, num_training_steps, num_cycles=1.0, last_epoch=-1
+    optimizer: Optimizer, num_warmup_steps: int, num_training_steps: int, num_cycles: int = 1, last_epoch: int = -1
 ):
-    """ Create a schedule with a learning rate that decreases following the
-    values of the cosine function with several hard restarts, after a warmup
-    period during which it increases linearly between 0 and 1.
+    """
+    Create a schedule with a learning rate that decreases following the values of the cosine function between the
+    initial lr set in the optimizer to 0, with several hard restarts, after a warmup period during which it increases
+    linearly between 0 and the initial lr set in the optimizer.
+
+    Args:
+        optimizer (:class:`~torch.optim.Optimizer`):
+            The optimizer for which to schedule the learning rate.
+        num_warmup_steps (:obj:`int`):
+            The number of steps for the warmup phase.
+        num_training_steps (:obj:`int`):
+            The total number of training steps.
+        num_cycles (:obj:`int`, `optional`, defaults to 1):
+            The number of hard restarts to use.
+        last_epoch (:obj:`int`, `optional`, defaults to -1):
+            The index of the last epoch when resuming training.
+
+    Return:
+        :obj:`torch.optim.lr_scheduler.LambdaLR` with the appropriate schedule.
     """
 
     def lr_lambda(current_step):
@@ -94,17 +166,34 @@ def get_cosine_with_hard_restarts_schedule_with_warmup(
 
 
 class AdamW(Optimizer):
-    """ Implements Adam algorithm with weight decay fix.
+    """
+    Implements Adam algorithm with weight decay fix as introduced in
+    `Decoupled Weight Decay Regularization <https://arxiv.org/abs/1711.05101>`__.
 
     Parameters:
-        lr (float): learning rate. Default 1e-3.
-        betas (tuple of 2 floats): Adams beta parameters (b1, b2). Default: (0.9, 0.999)
-        eps (float): Adams epsilon. Default: 1e-6
-        weight_decay (float): Weight decay. Default: 0.0
-        correct_bias (bool): can be set to False to avoid correcting bias in Adam (e.g. like in Bert TF repository). Default True.
+        params (:obj:`Iterable[torch.nn.parameter.Parameter]`):
+            Iterable of parameters to optimize or dictionaries defining parameter groups.
+        lr (:obj:`float`, `optional`, defaults to 1e-3):
+            The learning rate to use.
+        betas (:obj:`Tuple[float,float]`, `optional`, defaults to (0.9, 0.999)):
+            Adam's betas parameters (b1, b2).
+        eps (:obj:`float`, `optional`, defaults to 1e-6):
+            Adam's epsilon for numerical stability.
+        weight_decay (:obj:`float`, `optional`, defaults to 0):
+            Decoupled weight decay to apply.
+        correct_bias (:obj:`bool`, `optional`, defaults to `True`):
+            Whether ot not to correct bias in Adam (for instance, in Bert TF repository they use :obj:`False`).
     """
 
-    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.0, correct_bias=True):
+    def __init__(
+        self,
+        params: Iterable[torch.nn.parameter.Parameter],
+        lr: float = 1e-3,
+        betas: Tuple[float, float] = (0.9, 0.999),
+        eps: float = 1e-6,
+        weight_decay: float = 0.0,
+        correct_bias: bool = True,
+    ):
         if lr < 0.0:
             raise ValueError("Invalid learning rate: {} - should be >= 0.0".format(lr))
         if not 0.0 <= betas[0] < 1.0:
@@ -116,12 +205,12 @@ class AdamW(Optimizer):
         defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay, correct_bias=correct_bias)
         super().__init__(params, defaults)
 
-    def step(self, closure=None):
-        """Performs a single optimization step.
+    def step(self, closure: Callable = None):
+        """
+        Performs a single optimization step.
 
         Arguments:
-            closure (callable, optional): A closure that reevaluates the model
-                and returns the loss.
+            closure (:obj:`Callable`, `optional`): A closure that reevaluates the model and returns the loss.
         """
         loss = None
         if closure is not None:

src/transformers/optimization_tf.py

@@ -16,15 +16,36 @@
 
 
 import re
+from typing import Callable, List, Optional, Union
 
 import tensorflow as tf
 
 
 class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):
-    """Applies a warmup schedule on a given learning rate decay schedule."""
+    """
+    Applies a warmup schedule on a given learning rate decay schedule.
+
+    Args:
+        initial_learning_rate (:obj:`float`):
+            The initial learning rate for the schedule after the warmup (so this will be the learning rate at the end
+            of the warmup).
+        decay_schedule_fn (:obj:`Callable`):
+            The schedule function to apply after the warmup for the rest of training.
+        warmup_steps (:obj:`int`):
+            The number of steps for the warmup part of training.
+        power (:obj:`float`, `optional`, defaults to 1):
+            The power to use for the polynomial warmup (defaults is a linear warmup).
+        name (:obj:`str`, `optional`):
+            Optional name prefix for the returned tensors during the schedule.
+    """
 
     def __init__(
-        self, initial_learning_rate, decay_schedule_fn, warmup_steps, power=1.0, name=None,
+        self,
+        initial_learning_rate: float,
+        decay_schedule_fn: Callable,
+        warmup_steps: int,
+        power: float = 1.0,
+        name: str = None,
     ):
         super().__init__()
         self.initial_learning_rate = initial_learning_rate
@@ -59,15 +80,34 @@ class WarmUp(tf.keras.optimizers.schedules.LearningRateSchedule):
 
 
 def create_optimizer(
-    init_lr,
-    num_train_steps,
-    num_warmup_steps,
-    min_lr_ratio=0.0,
-    adam_epsilon=1e-8,
-    weight_decay_rate=0.0,
-    include_in_weight_decay=None,
+    init_lr: float,
+    num_train_steps: int,
+    num_warmup_steps: int,
+    min_lr_ratio: float = 0.0,
+    adam_epsilon: float = 1e-8,
+    weight_decay_rate: float = 0.0,
+    include_in_weight_decay: Optional[List[str]] = None,
 ):
-    """Creates an optimizer with learning rate schedule."""
+    """
+    Creates an optimizer with a learning rate schedule using a warmup phase followed by a linear decay.
+
+    Args:
+        init_lr (:obj:`float`):
+            The desired learning rate at the end of the warmup phase.
+        num_train_step (:obj:`int`):
+            The total number of training steps.
+        num_warmup_steps (:obj:`int`):
+            The number of warmup steps.
+        min_lr_ratio (:obj:`float`, `optional`, defaults to 0):
+            The final learning rate at the end of the linear decay will be :obj:`init_lr * min_lr_ratio`.
+        adam_epsilon (:obj:`float`, `optional`, defaults to 1e-8):
+            The epsilon to use in Adam.
+        weight_decay_rate (:obj:`float`, `optional`, defaults to 0):
+            The weight decay to use.
+        include_in_weight_decay (:obj:`List[str]`, `optional`):
+            List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is
+            applied to all parameters except bias and layer norm parameters.
+    """
     # Implements linear decay of the learning rate.
     lr_schedule = tf.keras.optimizers.schedules.PolynomialDecay(
         initial_learning_rate=init_lr,
@@ -96,26 +136,55 @@ def create_optimizer(
 
 
 class AdamWeightDecay(tf.keras.optimizers.Adam):
-    """Adam enables L2 weight decay and clip_by_global_norm on gradients.
-    Just adding the square of the weights to the loss function is *not* the
-    correct way of using L2 regularization/weight decay with Adam, since that will
-    interact with the m and v parameters in strange ways.
-    Instead we want ot decay the weights in a manner that doesn't interact with
-    the m/v parameters. This is equivalent to adding the square of the weights to
-    the loss with plain (non-momentum) SGD.
+    """
+    Adam enables L2 weight decay and clip_by_global_norm on gradients. Just adding the square of the weights to the
+    loss function is *not* the correct way of using L2 regularization/weight decay with Adam, since that will interact
+    with the m and v parameters in strange ways as shown in
+    `Decoupled Weight Decay Regularization <https://arxiv.org/abs/1711.05101>`__.
+
+    Instead we want ot decay the weights in a manner that doesn't interact with the m/v parameters. This is equivalent
+    to adding the square of the weights to the loss with plain (non-momentum) SGD.
+
+    Args:
+        learning_rate (:obj:`Union[float, tf.keras.optimizers.schedules.LearningRateSchedule]`, `optional`, defaults to 1e-3):
+            The learning rate to use or a schedule.
+        beta_1 (:obj:`float`, `optional`, defaults to 0.9):
+            The beta1 parameter in Adam, which is the exponential decay rate for the 1st momentum estimates.
+        beta_2 (:obj:`float`, `optional`, defaults to 0.999):
+            The beta2 parameter in Adam, which is the exponential decay rate for the 2nd momentum estimates.
+        epsilon (:obj:`float`, `optional`, defaults to 1e-7):
+            The epsilon paramenter in Adam, which is a small constant for numerical stability.
+        amsgrad (:obj:`bool`, `optional`, default to `False`):
+            Wheter to apply AMSGrad varient of this algorithm or not, see
+            `On the Convergence of Adam and Beyond <https://arxiv.org/abs/1904.09237>`__.
+        weight_decay_rate (:obj:`float`, `optional`, defaults to 0):
+            The weight decay to apply.
+        include_in_weight_decay (:obj:`List[str]`, `optional`):
+            List of the parameter names (or re patterns) to apply weight decay to. If none is passed, weight decay is
+            applied to all parameters by default (unless they are in :obj:`exclude_from_weight_decay`).
+        exclude_from_weight_decay (:obj:`List[str]`, `optional`):
+            List of the parameter names (or re patterns) to exclude from applying weight decay to. If a
+            :obj:`include_in_weight_decay` is passed, the names in it will supersede this list.
+        name (:obj:`str`, `optional`, defaults to 'AdamWeightDecay'):
+            Optional name for the operations created when applying gradients.
+        kwargs:
+            Keyward arguments. Allowed to be {``clipnorm``, ``clipvalue``, ``lr``, ``decay``}. ``clipnorm`` is clip
+            gradients by norm; ``clipvalue`` is clip gradients by value, ``decay`` is included for backward
+            compatibility to allow time inverse decay of learning rate. ``lr`` is included for backward compatibility,
+            recommended to use ``learning_rate`` instead.
     """
 
     def __init__(
         self,
-        learning_rate=0.001,
-        beta_1=0.9,
-        beta_2=0.999,
-        epsilon=1e-7,
-        amsgrad=False,
-        weight_decay_rate=0.0,
-        include_in_weight_decay=None,
-        exclude_from_weight_decay=None,
-        name="AdamWeightDecay",
+        learning_rate: Union[float, tf.keras.optimizers.schedules.LearningRateSchedule] = 0.001,
+        beta_1: float = 0.9,
+        beta_2: float = 0.999,
+        epsilon: float = 1e-7,
+        amsgrad: bool = False,
+        weight_decay_rate: float = 0.0,
+        include_in_weight_decay: Optional[List[str]] = None,
+        exclude_from_weight_decay: Optional[List[str]] = None,
+        name: str = "AdamWeightDecay",
         **kwargs
     ):
         super().__init__(learning_rate, beta_1, beta_2, epsilon, amsgrad, name, **kwargs)

src/transformers/pipelines.py

@@ -87,6 +87,18 @@ def get_framework(model=None):
     return framework
 
 
+class PipelineException(Exception):
+    """
+    Raised by pipelines when handling __call__
+    """
+
+    def __init__(self, task: str, model: str, reason: str):
+        super().__init__(reason)
+
+        self.task = task
+        self.model = model
+
+
 class ArgumentHandler(ABC):
     """
     Base interface for handling varargs for each Pipeline
@@ -603,6 +615,28 @@ class TextGenerationPipeline(Pipeline):
         "TFCTRLLMHeadModel",
     ]
 
+    # overriding _parse_and_tokenize to allow for unusual language-modeling tokenizer arguments
+
+    def _parse_and_tokenize(self, *args, padding=True, add_special_tokens=True, **kwargs):
+        """
+        Parse arguments and tokenize
+        """
+        # Parse arguments
+        if self.model.__class__.__name__ in ["TransfoXLLMHeadModel"]:
+            tokenizer_kwargs = {"add_space_before_punct_symbol": True}
+        else:
+            tokenizer_kwargs = {}
+        inputs = self._args_parser(*args, **kwargs)
+        inputs = self.tokenizer(
+            inputs,
+            add_special_tokens=add_special_tokens,
+            return_tensors=self.framework,
+            padding=padding,
+            **tokenizer_kwargs,
+        )
+
+        return inputs
+
     def __call__(
         self, *args, return_tensors=False, return_text=True, clean_up_tokenization_spaces=False, **generate_kwargs
     ):
@@ -808,6 +842,21 @@ class FillMaskPipeline(Pipeline):
 
         self.topk = topk
 
+    def ensure_exactly_one_mask_token(self, masked_index: np.ndarray):
+        numel = np.prod(masked_index.shape)
+        if numel > 1:
+            raise PipelineException(
+                "fill-mask",
+                self.model.base_model_prefix,
+                f"More than one mask_token ({self.tokenizer.mask_token}) is not supported",
+            )
+        elif numel < 1:
+            raise PipelineException(
+                "fill-mask",
+                self.model.base_model_prefix,
+                f"No mask_token ({self.tokenizer.mask_token}) found on the input",
+            )
+
     def __call__(self, *args, **kwargs):
         inputs = self._parse_and_tokenize(*args, **kwargs)
         outputs = self._forward(inputs, return_tensors=True)
@@ -820,14 +869,22 @@ class FillMaskPipeline(Pipeline):
             result = []
 
             if self.framework == "tf":
-                masked_index = tf.where(input_ids == self.tokenizer.mask_token_id).numpy().item()
-                logits = outputs[i, masked_index, :]
+                masked_index = tf.where(input_ids == self.tokenizer.mask_token_id).numpy()
+
+                # Fill mask pipeline supports only one ${mask_token} per sample
+                self.ensure_exactly_one_mask_token(masked_index)
+
+                logits = outputs[i, masked_index.item(), :]
                 probs = tf.nn.softmax(logits)
                 topk = tf.math.top_k(probs, k=self.topk)
                 values, predictions = topk.values.numpy(), topk.indices.numpy()
             else:
-                masked_index = (input_ids == self.tokenizer.mask_token_id).nonzero().item()
-                logits = outputs[i, masked_index, :]
+                masked_index = (input_ids == self.tokenizer.mask_token_id).nonzero()
+
+                # Fill mask pipeline supports only one ${mask_token} per sample
+                self.ensure_exactly_one_mask_token(masked_index.numpy())
+
+                logits = outputs[i, masked_index.item(), :]
                 probs = logits.softmax(dim=0)
                 values, predictions = probs.topk(self.topk)
 
@@ -987,6 +1044,10 @@ class TokenClassificationPipeline(Pipeline):
 
                 entities += [entity]
 
+            # Ensure if an entity is the latest one in the sequence it gets appended to the output
+            if len(entity_group_disagg) > 0:
+                entity_groups.append(self.group_entities(entity_group_disagg))
+
             # Append
             if self.grouped_entities:
                 answers += [entity_groups]

src/transformers/tokenization_marian.py

@@ -129,6 +129,8 @@ class MarianTokenizer(PreTrainedTokenizer):
         max_length: Optional[int] = None,
         pad_to_max_length: bool = True,
         return_tensors: str = "pt",
+        truncation_strategy="only_first",
+        padding="longest",
     ) -> BatchEncoding:
         """Prepare model inputs for translation. For best performance, translate one sentence at a time.
         Arguments:
@@ -147,24 +149,21 @@ class MarianTokenizer(PreTrainedTokenizer):
             raise ValueError(f"found empty string in src_texts: {src_texts}")
         self.current_spm = self.spm_source
         src_texts = [self.normalize(t) for t in src_texts]  # this does not appear to do much
-        model_inputs: BatchEncoding = self.batch_encode_plus(
-            src_texts,
+        tokenizer_kwargs = dict(
             add_special_tokens=True,
             return_tensors=return_tensors,
             max_length=max_length,
             pad_to_max_length=pad_to_max_length,
+            truncation_strategy=truncation_strategy,
+            padding=padding,
         )
+        model_inputs: BatchEncoding = self(src_texts, **tokenizer_kwargs)
+
         if tgt_texts is None:
             return model_inputs
 
         self.current_spm = self.spm_target
-        decoder_inputs: BatchEncoding = self.batch_encode_plus(
-            tgt_texts,
-            add_special_tokens=True,
-            return_tensors=return_tensors,
-            max_length=max_length,
-            pad_to_max_length=pad_to_max_length,
-        )
+        decoder_inputs: BatchEncoding = self(tgt_texts, **tokenizer_kwargs)
         for k, v in decoder_inputs.items():
             model_inputs[f"decoder_{k}"] = v
         self.current_spm = self.spm_source

src/transformers/trainer.py

@@ -1,7 +1,6 @@
 import logging
 import math
 import os
-import random
 import re
 import shutil
 import warnings
@@ -23,7 +22,14 @@ from .data.data_collator import DataCollator, default_data_collator
 from .file_utils import is_apex_available, is_torch_tpu_available
 from .modeling_utils import PreTrainedModel
 from .optimization import AdamW, get_linear_schedule_with_warmup
-from .trainer_utils import PREFIX_CHECKPOINT_DIR, EvalPrediction, PredictionOutput, TrainOutput, is_wandb_available
+from .trainer_utils import (
+    PREFIX_CHECKPOINT_DIR,
+    EvalPrediction,
+    PredictionOutput,
+    TrainOutput,
+    is_wandb_available,
+    set_seed,
+)
 from .training_args import TrainingArguments
 
 
@@ -60,20 +66,6 @@ if is_wandb_available():
 logger = logging.getLogger(__name__)
 
 
-def set_seed(seed: int):
-    """
-    Helper function for reproducible behavior to set the seed in ``random``, ``numpy`` and ``torch``.
-
-    Args:
-        seed (:obj:`int`): The seed to set.
-    """
-    random.seed(seed)
-    np.random.seed(seed)
-    torch.manual_seed(seed)
-    torch.cuda.manual_seed_all(seed)
-    # ^^ safe to call this function even if cuda is not available
-
-
 @contextmanager
 def torch_distributed_zero_first(local_rank: int):
     """
@@ -493,6 +485,10 @@ class Trainer:
             else:
                 epoch_iterator = tqdm(train_dataloader, desc="Iteration", disable=not self.is_local_master())
 
+            # Reset the past mems state at the beginning of each epoch if necessary.
+            if self.args.past_index >= 0:
+                self._past = None
+
             for step, inputs in enumerate(epoch_iterator):
 
                 # Skip past any already trained steps if resuming training
@@ -537,8 +533,8 @@ class Trainer:
 
                         self._log(logs)
 
-                        if self.args.evaluate_during_training:
-                            self.evaluate()
+                    if self.args.evaluate_during_training and self.global_step % self.args.eval_steps == 0:
+                        self.evaluate()
 
                     if self.args.save_steps > 0 and self.global_step % self.args.save_steps == 0:
                         # In all cases (even distributed/parallel), self.model is always a reference
@@ -569,12 +565,15 @@ class Trainer:
             if self.args.max_steps > 0 and self.global_step > self.args.max_steps:
                 train_iterator.close()
                 break
-            if self.args.tpu_metrics_debug:
+            if self.args.tpu_metrics_debug or self.args.debug:
                 # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
                 xm.master_print(met.metrics_report())
 
         if self.tb_writer:
             self.tb_writer.close()
+        if self.args.past_index and hasattr(self, "_past"):
+            # Clean the state at the end of training
+            delattr(self, "_past")
 
         logger.info("\n\nTraining completed. Do not forget to share your model on huggingface.co/models =)\n\n")
         return TrainOutput(self.global_step, tr_loss / self.global_step)
@@ -617,9 +616,15 @@ class Trainer:
             if isinstance(v, torch.Tensor):
                 inputs[k] = v.to(self.args.device)
 
+        if self.args.past_index >= 0 and self._past is not None:
+            inputs["mems"] = self._past
+
         outputs = model(**inputs)
         loss = outputs[0]  # model outputs are always tuple in transformers (see doc)
 
+        if self.args.past_index >= 0:
+            self._past = outputs[self.args.past_index]
+
         if self.args.n_gpu > 1:
             loss = loss.mean()  # mean() to average on multi-gpu parallel training
         if self.args.gradient_accumulation_steps > 1:
@@ -741,7 +746,7 @@ class Trainer:
 
         self._log(output.metrics)
 
-        if self.args.tpu_metrics_debug:
+        if self.args.tpu_metrics_debug or self.args.debug:
             # tpu-comment: Logging debug metrics for PyTorch/XLA (compile, execute times, ops, etc.)
             xm.master_print(met.metrics_report())
 
@@ -802,12 +807,17 @@ class Trainer:
         if is_torch_tpu_available():
             dataloader = pl.ParallelLoader(dataloader, [self.args.device]).per_device_loader(self.args.device)
 
+        if self.args.past_index >= 0:
+            past = None
+
         for inputs in tqdm(dataloader, desc=description):
             has_labels = any(inputs.get(k) is not None for k in ["labels", "lm_labels", "masked_lm_labels"])
 
             for k, v in inputs.items():
                 if isinstance(v, torch.Tensor):
                     inputs[k] = v.to(self.args.device)
+            if self.args.past_index >= 0:
+                inputs["mems"] = past
 
             with torch.no_grad():
                 outputs = model(**inputs)
@@ -816,6 +826,8 @@ class Trainer:
                     eval_losses += [step_eval_loss.mean().item()]
                 else:
                     logits = outputs[0]
+                if self.args.past_index >= 0:
+                    past = outputs[self.args.past_index if has_labels else self.args.past_index - 1]
 
             if not prediction_loss_only:
                 if preds is None:

src/transformers/trainer_tf.py

@@ -3,7 +3,6 @@
 import logging
 import math
 import os
-import random
 from typing import Callable, Dict, Optional, Tuple
 
 import numpy as np
@@ -11,7 +10,7 @@ import tensorflow as tf
 
 from .modeling_tf_utils import TFPreTrainedModel
 from .optimization_tf import GradientAccumulator, create_optimizer
-from .trainer_utils import PREFIX_CHECKPOINT_DIR, EvalPrediction, PredictionOutput, is_wandb_available
+from .trainer_utils import PREFIX_CHECKPOINT_DIR, EvalPrediction, PredictionOutput, is_wandb_available, set_seed
 from .training_args_tf import TFTrainingArguments
 
 
@@ -22,12 +21,6 @@ if is_wandb_available():
 logger = logging.getLogger(__name__)
 
 
-def set_seed(seed: int):
-    random.seed(seed)
-    np.random.seed(seed)
-    tf.random.set_seed(seed)
-
-
 class TFTrainer:
     """
     TFTrainer is a simple but feature-complete training and eval loop for TensorFlow,
@@ -240,6 +233,10 @@ class TFTrainer:
 
         step: int = 1
 
+        # Reset the past mems state at the beginning of the evaluation if necessary.
+        if self.args.past_index >= 0:
+            self._past = None
+
         for features, labels in dataset:
             step = tf.convert_to_tensor(step, dtype=tf.int64)
             loss, logits = self._evaluate_steps(features, labels)
@@ -252,7 +249,7 @@ class TFTrainer:
                 if isinstance(labels, tuple):
                     labels = labels[0]
 
-                if self.args.n_gpu > 1:
+                if self.args.n_replicas > 1:
                     for val in logits.values:
                         if preds is None:
                             preds = val.numpy()
@@ -288,6 +285,10 @@ class TFTrainer:
             if not key.startswith("eval_"):
                 metrics[f"eval_{key}"] = metrics.pop(key)
 
+        if self.args.past_index and hasattr(self, "_past"):
+            # Clean the state at the end of training
+            delattr(self, "_past")
+
         return PredictionOutput(predictions=preds, label_ids=label_ids, metrics=metrics)
 
     def _log(self, logs: Dict[str, float]) -> None:
@@ -405,6 +406,9 @@ class TFTrainer:
         logger.info("  Total optimization steps = %d", t_total)
 
         for epoch_iter in range(epochs_trained, int(epochs + 1)):
+            # Reset the past mems state at the beginning of each epoch if necessary.
+            if self.args.past_index >= 0:
+                self._past = None
             for step, training_loss in enumerate(self._training_steps(train_ds, optimizer)):
                 self.global_step = iterations.numpy()
                 self.epoch_logging = epoch_iter - 1 + (step + 1) / steps_per_epoch
@@ -444,6 +448,10 @@ class TFTrainer:
                 if self.args.max_steps > 0 and self.global_step % self.args.max_steps == 0:
                     break
 
+        if self.args.past_index and hasattr(self, "_past"):
+            # Clean the state at the end of training
+            delattr(self, "_past")
+
     def _training_steps(self, ds, optimizer):
         """
         Returns a generator over training steps (i.e. parameters update).
@@ -518,11 +526,16 @@ class TFTrainer:
           labels: the batched labels.
           training: run the model in training mode or not
         """
+        if self.args.past_index >= 0 and getattr(self, "_past", None) is not None:
+            features["mems"] = self._past
         if isinstance(labels, (dict)):
-            loss, logits = self.model(features, training=training, **labels)[:2]
+            outputs = self.model(features, training=training, **labels)[:2]
         else:
-            loss, logits = self.model(features, labels=labels, training=training)[:2]
-        loss += sum(self.model.losses) * (1.0 / self.args.n_gpu)
+            outputs = self.model(features, labels=labels, training=training)[:2]
+        loss, logits = outputs[:2]
+        if self.args.past_index >= 0:
+            self._past = outputs[self.args.past_index]
+        loss += sum(self.model.losses) * (1.0 / self.args.n_replicas)
 
         return loss, logits
 

src/transformers/trainer_utils.py

@@ -1,8 +1,11 @@
 import os
+import random
 from typing import Dict, NamedTuple, Optional
 
 import numpy as np
 
+from .file_utils import is_tf_available, is_torch_available
+
 
 try:
     import wandb
@@ -21,6 +24,28 @@ def is_wandb_available():
     return _has_wandb
 
 
+def set_seed(seed: int):
+    """
+    Helper function for reproducible behavior to set the seed in ``random``, ``numpy``, ``torch`` and/or ``tf``
+    (if installed).
+
+    Args:
+        seed (:obj:`int`): The seed to set.
+    """
+    random.seed(seed)
+    np.random.seed(seed)
+    if is_torch_available():
+        import torch
+
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        # ^^ safe to call this function even if cuda is not available
+    if is_tf_available():
+        import tensorflow as tf
+
+        tf.random.set_seed(seed)
+
+
 class EvalPrediction(NamedTuple):
     """
     Evaluation output (always contains labels), to be used to compute metrics.

src/transformers/training_args.py

@@ -97,11 +97,18 @@ class TrainingArguments:
             During distributed training, the rank of the process.
         tpu_num_cores (:obj:`int`, `optional`):
             When training on TPU, the mumber of TPU cores (automatically passed by launcher script).
-        tpu_metrics_debug (:obj:`bool`, `optional`, defaults to :obj:`False`):
+        debug (:obj:`bool`, `optional`, defaults to :obj:`False`):
             When training on TPU, whether to print debug metrics or not.
         dataloader_drop_last (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size)
             or not.
+        eval_steps (:obj:`int`, `optional`, defaults to 1000):
+            Number of update steps between two evaluations.
+        past_index (:obj:`int`, `optional`, defaults to -1):
+            Some models like :doc:`TransformerXL <../model_doc/transformerxl>` or :doc`XLNet <../model_doc/xlnet>` can
+            make use of the past hidden states for their predictions. If this argument is set to a positive int, the
+            ``Trainer`` will use the corresponding output (usually index 2) as the past state and feed it to the model
+            at the next training step under the keyword argument ``mems``.
     """
 
     output_dir: str = field(
@@ -197,11 +204,21 @@ class TrainingArguments:
     tpu_num_cores: Optional[int] = field(
         default=None, metadata={"help": "TPU: Number of TPU cores (automatically passed by launcher script)"}
     )
-    tpu_metrics_debug: bool = field(default=False, metadata={"help": "TPU: Whether to print debug metrics"})
+    tpu_metrics_debug: bool = field(
+        default=False,
+        metadata={"help": "Deprecated, the use of `--debug` is preferred. TPU: Whether to print debug metrics"},
+    )
+    debug: bool = field(default=False, metadata={"help": "Whether to print debug metrics on TPU"})
 
     dataloader_drop_last: bool = field(
         default=False, metadata={"help": "Drop the last incomplete batch if it is not divisible by the batch size."}
     )
+    eval_steps: int = field(default=1000, metadata={"help": "Run an evaluation every X steps."})
+
+    past_index: int = field(
+        default=-1,
+        metadata={"help": "If >=0, uses the corresponding part of the output as the past state for next step."},
+    )
 
     @property
     def train_batch_size(self) -> int:

src/transformers/training_args_tf.py

@@ -1,4 +1,5 @@
 import logging
+import warnings
 from dataclasses import dataclass, field
 from typing import Tuple
 
@@ -80,26 +81,25 @@ class TFTrainingArguments(TrainingArguments):
             During distributed training, the rank of the process.
         tpu_num_cores (:obj:`int`, `optional`):
             When training on TPU, the mumber of TPU cores (automatically passed by launcher script).
-        tpu_metrics_debug (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            When training on TPU, whether to print debug metrics or not.
+        debug (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Wheter to activate the trace to record computation graphs and profiling information or not.
         dataloader_drop_last (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size)
             or not.
-        tpu_name (:obj:`str`, `optional`):
-            The name of the TPU the process is running on.
         eval_steps (:obj:`int`, `optional`, defaults to 1000):
             Number of update steps before two evaluations.
-        debug (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Wheter to activate the trace to record computation graphs and profiling information or not.
+        past_index (:obj:`int`, `optional`, defaults to -1):
+            Some models like :doc:`TransformerXL <../model_doc/transformerxl>` or :doc`XLNet <../model_doc/xlnet>` can
+            make use of the past hidden states for their predictions. If this argument is set to a positive int, the
+            ``Trainer`` will use the corresponding output (usually index 2) as the past state and feed it to the model
+            at the next training step under the keyword argument ``mems``.
+        tpu_name (:obj:`str`, `optional`):
+            The name of the TPU the process is running on.
     """
 
     tpu_name: str = field(
         default=None, metadata={"help": "Name of TPU"},
     )
-    eval_steps: int = field(default=1000, metadata={"help": "Run an evaluation every X steps."})
-    debug: bool = field(
-        default=False, metadata={"help": "Activate the trace to record computation graphs and profiling information"}
-    )
 
     @cached_property
     @tf_required
@@ -145,8 +145,46 @@ class TFTrainingArguments(TrainingArguments):
 
     @property
     @tf_required
-    def n_gpu(self) -> int:
+    def n_replicas(self) -> int:
         """
-        The number of replicas (GPUs or TPU cores) used in this training.
+        The number of replicas (CPUs, GPUs or TPU cores) used in this training.
         """
         return self._setup_strategy.num_replicas_in_sync
+
+    @property
+    def train_batch_size(self) -> int:
+        """
+        The actual batch size for training (may differ from :obj:`per_gpu_train_batch_size` in distributed training).
+        """
+        if self.per_gpu_train_batch_size:
+            logger.warning(
+                "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future "
+                "version. Using `--per_device_train_batch_size` is preferred."
+            )
+        per_device_batch_size = self.per_gpu_train_batch_size or self.per_device_train_batch_size
+        return per_device_batch_size * max(1, self.n_replicas)
+
+    @property
+    def eval_batch_size(self) -> int:
+        """
+        The actual batch size for evaluation (may differ from :obj:`per_gpu_eval_batch_size` in distributed training).
+        """
+        if self.per_gpu_eval_batch_size:
+            logger.warning(
+                "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future "
+                "version. Using `--per_device_eval_batch_size` is preferred."
+            )
+        per_device_batch_size = self.per_gpu_eval_batch_size or self.per_device_eval_batch_size
+        return per_device_batch_size * max(1, self.n_replicas)
+
+    @property
+    @tf_required
+    def n_gpu(self) -> int:
+        """
+        The number of replicas (CPUs, GPUs or TPU cores) used in this training.
+        """
+        warnings.warn(
+            "The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead.",
+            FutureWarning,
+        )
+        return self._setup_strategy.num_replicas_in_sync

tests/test_activations.py

@@ -1,8 +1,7 @@
 import unittest
 
 from transformers import is_torch_available
-
-from .utils import require_torch
+from transformers.testing_utils import require_torch
 
 
 if is_torch_available():

tests/test_benchmark.py

@@ -4,8 +4,7 @@ import unittest
 from pathlib import Path
 
 from transformers import AutoConfig, is_torch_available
-
-from .utils import require_torch, torch_device
+from transformers.testing_utils import require_torch, torch_device
 
 
 if is_torch_available():

tests/test_benchmark_tf.py

@@ -4,8 +4,7 @@ import unittest
 from pathlib import Path
 
 from transformers import AutoConfig, is_tf_available
-
-from .utils import require_tf
+from transformers.testing_utils import require_tf
 
 
 if is_tf_available():

tests/test_configuration_auto.py

@@ -19,8 +19,7 @@ import unittest
 from transformers.configuration_auto import CONFIG_MAPPING, AutoConfig
 from transformers.configuration_bert import BertConfig
 from transformers.configuration_roberta import RobertaConfig
-
-from .utils import DUMMY_UNKWOWN_IDENTIFIER
+from transformers.testing_utils import DUMMY_UNKWOWN_IDENTIFIER
 
 
 SAMPLE_ROBERTA_CONFIG = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fixtures/dummy-config.json")

tests/test_doc_samples.py

@@ -21,8 +21,7 @@ from pathlib import Path
 from typing import List, Union
 
 import transformers
-
-from .utils import require_tf, require_torch, slow
+from transformers.testing_utils import require_tf, require_torch, slow
 
 
 logger = logging.getLogger()

tests/test_modeling_albert.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_auto.py

@@ -17,8 +17,7 @@
 import unittest
 
 from transformers import is_torch_available
-
-from .utils import DUMMY_UNKWOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, require_torch, slow
+from transformers.testing_utils import DUMMY_UNKWOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, require_torch, slow
 
 
 if is_torch_available():

tests/test_modeling_bart.py

@@ -20,10 +20,10 @@ import timeout_decorator  # noqa
 
 from transformers import is_torch_available
 from transformers.file_utils import cached_property
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_bert.py

@@ -17,24 +17,24 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
-from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
-from .utils import require_torch, slow, torch_device
+from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask
 
 
 if is_torch_available():
     from transformers import (
         BertConfig,
-        BertModel,
-        BertLMHeadModel,
         BertForMaskedLM,
+        BertForMultipleChoice,
         BertForNextSentencePrediction,
         BertForPreTraining,
         BertForQuestionAnswering,
         BertForSequenceClassification,
         BertForTokenClassification,
-        BertForMultipleChoice,
+        BertLMHeadModel,
+        BertModel,
     )
     from transformers.modeling_bert import BERT_PRETRAINED_MODEL_ARCHIVE_LIST
 
@@ -93,7 +93,7 @@ class BertModelTester:
 
         input_mask = None
         if self.use_input_mask:
-            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)
+            input_mask = random_attention_mask([self.batch_size, self.seq_length])
 
         token_type_ids = None
         if self.use_token_type_ids:
@@ -120,6 +120,7 @@ class BertModelTester:
             type_vocab_size=self.type_vocab_size,
             is_decoder=False,
             initializer_range=self.initializer_range,
+            return_dict=True,
         )
 
         return config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
@@ -151,29 +152,19 @@ class BertModelTester:
             encoder_attention_mask,
         )
 
-    def check_loss_output(self, result):
-        self.parent.assertListEqual(list(result["loss"].size()), [])
-
-    def create_and_check_bert_model(
+    def create_and_check_model(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
         model = BertModel(config=config)
         model.to(torch_device)
         model.eval()
-        sequence_output, pooled_output = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
-        sequence_output, pooled_output = model(input_ids, token_type_ids=token_type_ids)
-        sequence_output, pooled_output = model(input_ids)
+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
+        result = model(input_ids, token_type_ids=token_type_ids)
+        result = model(input_ids)
+        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
+        self.parent.assertEqual(result.pooler_output.shape, (self.batch_size, self.hidden_size))
 
-        result = {
-            "sequence_output": sequence_output,
-            "pooled_output": pooled_output,
-        }
-        self.parent.assertListEqual(
-            list(result["sequence_output"].size()), [self.batch_size, self.seq_length, self.hidden_size]
-        )
-        self.parent.assertListEqual(list(result["pooled_output"].size()), [self.batch_size, self.hidden_size])
-
-    def create_and_check_bert_model_as_decoder(
+    def create_and_check_model_as_decoder(
         self,
         config,
         input_ids,
@@ -185,34 +176,28 @@ class BertModelTester:
         encoder_hidden_states,
         encoder_attention_mask,
     ):
+        config.add_cross_attention = True
         model = BertModel(config)
         model.to(torch_device)
         model.eval()
-        sequence_output, pooled_output = model(
+        result = model(
             input_ids,
             attention_mask=input_mask,
             token_type_ids=token_type_ids,
             encoder_hidden_states=encoder_hidden_states,
             encoder_attention_mask=encoder_attention_mask,
         )
-        sequence_output, pooled_output = model(
+        result = model(
             input_ids,
             attention_mask=input_mask,
             token_type_ids=token_type_ids,
             encoder_hidden_states=encoder_hidden_states,
         )
-        sequence_output, pooled_output = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids)
+        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
+        self.parent.assertEqual(result.pooler_output.shape, (self.batch_size, self.hidden_size))
 
-        result = {
-            "sequence_output": sequence_output,
-            "pooled_output": pooled_output,
-        }
-        self.parent.assertListEqual(
-            list(result["sequence_output"].size()), [self.batch_size, self.seq_length, self.hidden_size]
-        )
-        self.parent.assertListEqual(list(result["pooled_output"].size()), [self.batch_size, self.hidden_size])
-
-    def create_and_check_bert_for_causal_lm(
+    def create_and_check_for_causal_lm(
         self,
         config,
         input_ids,
@@ -227,37 +212,19 @@ class BertModelTester:
         model = BertLMHeadModel(config=config)
         model.to(torch_device)
         model.eval()
-        loss, prediction_scores = model(
-            input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels
-        )
-        result = {
-            "loss": loss,
-            "prediction_scores": prediction_scores,
-        }
-        self.parent.assertListEqual(
-            list(result["prediction_scores"].size()), [self.batch_size, self.seq_length, self.vocab_size]
-        )
-        self.check_loss_output(result)
+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
 
-    def create_and_check_bert_for_masked_lm(
+    def create_and_check_for_masked_lm(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
         model = BertForMaskedLM(config=config)
         model.to(torch_device)
         model.eval()
-        loss, prediction_scores = model(
-            input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels
-        )
-        result = {
-            "loss": loss,
-            "prediction_scores": prediction_scores,
-        }
-        self.parent.assertListEqual(
-            list(result["prediction_scores"].size()), [self.batch_size, self.seq_length, self.vocab_size]
-        )
-        self.check_loss_output(result)
+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
 
-    def create_and_check_bert_model_for_causal_lm_as_decoder(
+    def create_and_check_model_for_causal_lm_as_decoder(
         self,
         config,
         input_ids,
@@ -269,10 +236,11 @@ class BertModelTester:
         encoder_hidden_states,
         encoder_attention_mask,
     ):
+        config.add_cross_attention = True
         model = BertLMHeadModel(config=config)
         model.to(torch_device)
         model.eval()
-        loss, prediction_scores = model(
+        result = model(
             input_ids,
             attention_mask=input_mask,
             token_type_ids=token_type_ids,
@@ -280,117 +248,82 @@ class BertModelTester:
             encoder_hidden_states=encoder_hidden_states,
             encoder_attention_mask=encoder_attention_mask,
         )
-        loss, prediction_scores = model(
+        result = model(
             input_ids,
             attention_mask=input_mask,
             token_type_ids=token_type_ids,
             labels=token_labels,
             encoder_hidden_states=encoder_hidden_states,
         )
-        result = {
-            "loss": loss,
-            "prediction_scores": prediction_scores,
-        }
-        self.parent.assertListEqual(
-            list(result["prediction_scores"].size()), [self.batch_size, self.seq_length, self.vocab_size]
-        )
-        self.check_loss_output(result)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
 
-    def create_and_check_bert_for_next_sequence_prediction(
+    def create_and_check_for_next_sequence_prediction(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
         model = BertForNextSentencePrediction(config=config)
         model.to(torch_device)
         model.eval()
-        loss, seq_relationship_score = model(
-            input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, next_sentence_label=sequence_labels,
+        result = model(
+            input_ids,
+            attention_mask=input_mask,
+            token_type_ids=token_type_ids,
+            next_sentence_label=sequence_labels,
         )
-        result = {
-            "loss": loss,
-            "seq_relationship_score": seq_relationship_score,
-        }
-        self.parent.assertListEqual(list(result["seq_relationship_score"].size()), [self.batch_size, 2])
-        self.check_loss_output(result)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, 2))
 
-    def create_and_check_bert_for_pretraining(
+    def create_and_check_for_pretraining(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
         model = BertForPreTraining(config=config)
         model.to(torch_device)
         model.eval()
-        loss, prediction_scores, seq_relationship_score = model(
+        result = model(
             input_ids,
             attention_mask=input_mask,
             token_type_ids=token_type_ids,
             labels=token_labels,
             next_sentence_label=sequence_labels,
         )
-        result = {
-            "loss": loss,
-            "prediction_scores": prediction_scores,
-            "seq_relationship_score": seq_relationship_score,
-        }
-        self.parent.assertListEqual(
-            list(result["prediction_scores"].size()), [self.batch_size, self.seq_length, self.vocab_size]
-        )
-        self.parent.assertListEqual(list(result["seq_relationship_score"].size()), [self.batch_size, 2])
-        self.check_loss_output(result)
+        self.parent.assertEqual(result.prediction_logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
+        self.parent.assertEqual(result.seq_relationship_logits.shape, (self.batch_size, 2))
 
-    def create_and_check_bert_for_question_answering(
+    def create_and_check_for_question_answering(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
         model = BertForQuestionAnswering(config=config)
         model.to(torch_device)
         model.eval()
-        loss, start_logits, end_logits = model(
+        result = model(
             input_ids,
             attention_mask=input_mask,
             token_type_ids=token_type_ids,
             start_positions=sequence_labels,
             end_positions=sequence_labels,
         )
-        result = {
-            "loss": loss,
-            "start_logits": start_logits,
-            "end_logits": end_logits,
-        }
-        self.parent.assertListEqual(list(result["start_logits"].size()), [self.batch_size, self.seq_length])
-        self.parent.assertListEqual(list(result["end_logits"].size()), [self.batch_size, self.seq_length])
-        self.check_loss_output(result)
+        self.parent.assertEqual(result.start_logits.shape, (self.batch_size, self.seq_length))
+        self.parent.assertEqual(result.end_logits.shape, (self.batch_size, self.seq_length))
 
-    def create_and_check_bert_for_sequence_classification(
+    def create_and_check_for_sequence_classification(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
         config.num_labels = self.num_labels
         model = BertForSequenceClassification(config)
         model.to(torch_device)
         model.eval()
-        loss, logits = model(
-            input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=sequence_labels
-        )
-        result = {
-            "loss": loss,
-            "logits": logits,
-        }
-        self.parent.assertListEqual(list(result["logits"].size()), [self.batch_size, self.num_labels])
-        self.check_loss_output(result)
+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=sequence_labels)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_labels))
 
-    def create_and_check_bert_for_token_classification(
+    def create_and_check_for_token_classification(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
         config.num_labels = self.num_labels
         model = BertForTokenClassification(config=config)
         model.to(torch_device)
         model.eval()
-        loss, logits = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)
-        result = {
-            "loss": loss,
-            "logits": logits,
-        }
-        self.parent.assertListEqual(list(result["logits"].size()), [self.batch_size, self.seq_length, self.num_labels])
-        self.check_loss_output(result)
+        result = model(input_ids, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))
 
-    def create_and_check_bert_for_multiple_choice(
+    def create_and_check_for_multiple_choice(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
         config.num_choices = self.num_choices
@@ -400,18 +333,13 @@ class BertModelTester:
         multiple_choice_inputs_ids = input_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
         multiple_choice_token_type_ids = token_type_ids.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
         multiple_choice_input_mask = input_mask.unsqueeze(1).expand(-1, self.num_choices, -1).contiguous()
-        loss, logits = model(
+        result = model(
             multiple_choice_inputs_ids,
             attention_mask=multiple_choice_input_mask,
             token_type_ids=multiple_choice_token_type_ids,
             labels=choice_labels,
         )
-        result = {
-            "loss": loss,
-            "logits": logits,
-        }
-        self.parent.assertListEqual(list(result["logits"].size()), [self.batch_size, self.num_choices])
-        self.check_loss_output(result)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.num_choices))
 
     def prepare_config_and_inputs_for_common(self):
         config_and_inputs = self.prepare_config_and_inputs()
@@ -453,15 +381,15 @@ class BertModelTest(ModelTesterMixin, unittest.TestCase):
     def test_config(self):
         self.config_tester.run_common_tests()
 
-    def test_bert_model(self):
+    def test_model(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_bert_model(*config_and_inputs)
+        self.model_tester.create_and_check_model(*config_and_inputs)
 
-    def test_bert_model_as_decoder(self):
+    def test_model_as_decoder(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs_for_decoder()
-        self.model_tester.create_and_check_bert_model_as_decoder(*config_and_inputs)
+        self.model_tester.create_and_check_model_as_decoder(*config_and_inputs)
 
-    def test_bert_model_as_decoder_with_default_input_mask(self):
+    def test_model_as_decoder_with_default_input_mask(self):
         # This regression test was failing with PyTorch < 1.3
         (
             config,
@@ -477,7 +405,7 @@ class BertModelTest(ModelTesterMixin, unittest.TestCase):
 
         input_mask = None
 
-        self.model_tester.create_and_check_bert_model_as_decoder(
+        self.model_tester.create_and_check_model_as_decoder(
             config,
             input_ids,
             token_type_ids,
@@ -491,39 +419,115 @@ class BertModelTest(ModelTesterMixin, unittest.TestCase):
 
     def test_for_causal_lm(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs_for_decoder()
-        self.model_tester.create_and_check_bert_for_causal_lm(*config_and_inputs)
+        self.model_tester.create_and_check_for_causal_lm(*config_and_inputs)
 
     def test_for_masked_lm(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_bert_for_masked_lm(*config_and_inputs)
+        self.model_tester.create_and_check_for_masked_lm(*config_and_inputs)
 
     def test_for_causal_lm_decoder(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs_for_decoder()
-        self.model_tester.create_and_check_bert_model_for_causal_lm_as_decoder(*config_and_inputs)
+        self.model_tester.create_and_check_model_for_causal_lm_as_decoder(*config_and_inputs)
 
     def test_for_multiple_choice(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_bert_for_multiple_choice(*config_and_inputs)
+        self.model_tester.create_and_check_for_multiple_choice(*config_and_inputs)
 
     def test_for_next_sequence_prediction(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_bert_for_next_sequence_prediction(*config_and_inputs)
+        self.model_tester.create_and_check_for_next_sequence_prediction(*config_and_inputs)
 
     def test_for_pretraining(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_bert_for_pretraining(*config_and_inputs)
+        self.model_tester.create_and_check_for_pretraining(*config_and_inputs)
 
     def test_for_question_answering(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_bert_for_question_answering(*config_and_inputs)
+        self.model_tester.create_and_check_for_question_answering(*config_and_inputs)
 
     def test_for_sequence_classification(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_bert_for_sequence_classification(*config_and_inputs)
+        self.model_tester.create_and_check_for_sequence_classification(*config_and_inputs)
 
     def test_for_token_classification(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_bert_for_token_classification(*config_and_inputs)
+        self.model_tester.create_and_check_for_token_classification(*config_and_inputs)
+
+    # Copied from test_modeling_common.test_torchscript, but using jit.script, not jit.trace
+    def test_full_torchscript(self):
+        import copy
+        import tempfile
+        import os
+
+        if is_torch_available():
+            import torch
+            from transformers.modeling_bert import (
+                    BertScriptableModel,
+                    BertScriptableForMultipleChoice,
+                    BertScriptableForNextSentencePrediction,
+                    BertScriptableForPreTraining,
+                    BertScriptableForQuestionAnswering,
+                    BertScriptableForSequenceClassification,
+                    BertScriptableForTokenClassification,
+            )
+
+        config, unused_inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        config.return_dict = False
+        configs_no_init = copy.deepcopy(config)
+        for key in configs_no_init.__dict__.keys():
+            if "_range" in key or "_std" in key or "initializer_factor" in key:
+                setattr(configs_no_init, key, 1e-10)
+
+        scriptable_model_classes = (
+            BertScriptableModel,
+            BertScriptableForMultipleChoice,
+            BertScriptableForNextSentencePrediction,
+            BertScriptableForPreTraining,
+            BertScriptableForQuestionAnswering,
+            BertScriptableForSequenceClassification,
+            BertScriptableForTokenClassification,
+        )
+        for model_class in scriptable_model_classes:
+            model = model_class(config=configs_no_init)
+            model.to(torch_device)
+            model.eval()
+
+            try:
+                scripted = torch.jit.script(model)
+            except RuntimeError:
+                self.fail("Couldn't script module.")
+
+            with tempfile.TemporaryDirectory() as tmp_dir_name:
+                pt_file_name = os.path.join(tmp_dir_name, "scripted_model.pt")
+
+                try:
+                    torch.jit.save(scripted, pt_file_name)
+                except Exception:
+                    self.fail("Couldn't save scripted module.")
+
+                try:
+                    loaded_model = torch.jit.load(pt_file_name)
+                except Exception:
+                    self.fail("Couldn't load scripted module.")
+
+            model.to(torch_device)
+            model.eval()
+
+            loaded_model.to(torch_device)
+            loaded_model.eval()
+
+            model_state_dict = model.state_dict()
+            loaded_model_state_dict = loaded_model.state_dict()
+
+            self.assertEqual(set(model_state_dict.keys()), set(loaded_model_state_dict.keys()))
+
+            models_equal = True
+            for layer_name, p1 in model_state_dict.items():
+                p2 = loaded_model_state_dict[layer_name]
+                if p1.data.ne(p2.data).sum() > 0:
+                    models_equal = False
+
+            self.assertTrue(models_equal)
 
     @slow
     def test_model_from_pretrained(self):

tests/test_modeling_camembert.py

@@ -16,8 +16,7 @@
 import unittest
 
 from transformers import is_torch_available
-
-from .utils import require_torch, slow, torch_device
+from transformers.testing_utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_common.py

@@ -21,8 +21,7 @@ import unittest
 from typing import List
 
 from transformers import is_torch_available
-
-from .utils import require_multigpu, require_torch, slow, torch_device
+from transformers.testing_utils import require_multigpu, require_torch, slow, torch_device
 
 
 if is_torch_available():
@@ -812,7 +811,7 @@ class ModelTesterMixin:
             # Wrap model in nn.DataParallel
             model = torch.nn.DataParallel(model)
             with torch.no_grad():
-                _ = model(**inputs_dict)
+                _ = model(**self._prepare_for_class(inputs_dict, model_class))
 
 
 global_rng = random.Random()

tests/test_modeling_ctrl.py

@@ -16,10 +16,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_distilbert.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, torch_device
 
 
 if is_torch_available():

tests/test_modeling_electra.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_encoder_decoder.py

@@ -18,12 +18,12 @@ import tempfile
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 # TODO(PVP): this line reruns all the tests in BertModelTest; not sure whether this can be prevented
 # for now only run module with pytest tests/test_modeling_encoder_decoder.py::EncoderDecoderModelTest
 from .test_modeling_bert import BertModelTester
 from .test_modeling_common import ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_flaubert.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_gpt2.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_longformer.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():
@@ -115,6 +115,18 @@ class LongformerModelTester:
     def check_loss_output(self, result):
         self.parent.assertListEqual(list(result["loss"].size()), [])
 
+    def create_and_check_attention_mask_determinism(
+        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
+    ):
+        model = LongformerModel(config=config)
+        model.to(torch_device)
+        model.eval()
+
+        attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=torch_device)
+        output_with_mask = model(input_ids, attention_mask=attention_mask)[0]
+        output_without_mask = model(input_ids)[0]
+        self.parent.assertTrue(torch.allclose(output_with_mask[0, 0, :5], output_without_mask[0, 0, :5], atol=1e-4))
+
     def create_and_check_longformer_model(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
@@ -134,6 +146,36 @@ class LongformerModelTester:
         )
         self.parent.assertListEqual(list(result["pooled_output"].size()), [self.batch_size, self.hidden_size])
 
+    def create_and_check_longformer_model_with_global_attention_mask(
+        self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
+    ):
+        model = LongformerModel(config=config)
+        model.to(torch_device)
+        model.eval()
+        global_attention_mask = input_mask.clone()
+        global_attention_mask[:, input_mask.shape[-1] // 2] = 0
+        global_attention_mask = global_attention_mask.to(torch_device)
+
+        sequence_output, pooled_output = model(
+            input_ids,
+            attention_mask=input_mask,
+            global_attention_mask=global_attention_mask,
+            token_type_ids=token_type_ids,
+        )
+        sequence_output, pooled_output = model(
+            input_ids, token_type_ids=token_type_ids, global_attention_mask=global_attention_mask
+        )
+        sequence_output, pooled_output = model(input_ids, global_attention_mask=global_attention_mask)
+
+        result = {
+            "sequence_output": sequence_output,
+            "pooled_output": pooled_output,
+        }
+        self.parent.assertListEqual(
+            list(result["sequence_output"].size()), [self.batch_size, self.seq_length, self.hidden_size]
+        )
+        self.parent.assertListEqual(list(result["pooled_output"].size()), [self.batch_size, self.hidden_size])
+
     def create_and_check_longformer_for_masked_lm(
         self, config, input_ids, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
     ):
@@ -243,7 +285,13 @@ class LongformerModelTester:
             token_labels,
             choice_labels,
         ) = config_and_inputs
-        inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": input_mask}
+        global_attention_mask = torch.zeros_like(input_ids)
+        inputs_dict = {
+            "input_ids": input_ids,
+            "token_type_ids": token_type_ids,
+            "attention_mask": input_mask,
+            "global_attention_mask": global_attention_mask,
+        }
         return config, inputs_dict
 
     def prepare_config_and_inputs_for_question_answering(self):
@@ -277,11 +325,10 @@ class LongformerModelTest(ModelTesterMixin, unittest.TestCase):
         (
             LongformerModel,
             LongformerForMaskedLM,
-            # TODO: make tests pass for those models
-            # LongformerForSequenceClassification,
-            # LongformerForQuestionAnswering,
-            # LongformerForTokenClassification,
-            # LongformerForMultipleChoice,
+            LongformerForSequenceClassification,
+            LongformerForQuestionAnswering,
+            LongformerForTokenClassification,
+            LongformerForMultipleChoice,
         )
         if is_torch_available()
         else ()
@@ -298,6 +345,14 @@ class LongformerModelTest(ModelTesterMixin, unittest.TestCase):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_longformer_model(*config_and_inputs)
 
+    def test_longformer_model_attention_mask_determinism(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_attention_mask_determinism(*config_and_inputs)
+
+    def test_longformer_model_global_attention_mask(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_longformer_model_with_global_attention_mask(*config_and_inputs)
+
     def test_longformer_for_masked_lm(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_longformer_for_masked_lm(*config_and_inputs)
@@ -325,15 +380,31 @@ class LongformerModelIntegrationTest(unittest.TestCase):
         model = LongformerModel.from_pretrained("allenai/longformer-base-4096")
         model.to(torch_device)
 
+        # 'Hello world!'
+        input_ids = torch.tensor([[0, 20920, 232, 328, 1437, 2]], dtype=torch.long, device=torch_device)
+        attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=torch_device)
+        output = model(input_ids, attention_mask=attention_mask)[0]
+        output_without_mask = model(input_ids)[0]
+
+        expected_output_slice = torch.tensor([0.0549, 0.1087, -0.1119, -0.0368, 0.0250], device=torch_device)
+        self.assertTrue(torch.allclose(output[0, 0, -5:], expected_output_slice, atol=1e-4))
+        self.assertTrue(torch.allclose(output_without_mask[0, 0, -5:], expected_output_slice, atol=1e-4))
+
+    @slow
+    def test_inference_no_head_long(self):
+        model = LongformerModel.from_pretrained("allenai/longformer-base-4096")
+        model.to(torch_device)
+
         # 'Hello world! ' repeated 1000 times
         input_ids = torch.tensor(
             [[0] + [20920, 232, 328, 1437] * 1000 + [2]], dtype=torch.long, device=torch_device
         )  # long input
 
         attention_mask = torch.ones(input_ids.shape, dtype=torch.long, device=input_ids.device)
-        attention_mask[:, [1, 4, 21]] = 2  # Set global attention on a few random positions
+        global_attention_mask = torch.zeros(input_ids.shape, dtype=torch.long, device=input_ids.device)
+        global_attention_mask[:, [1, 4, 21]] = 1  # Set global attention on a few random positions
 
-        output = model(input_ids, attention_mask=attention_mask)[0]
+        output = model(input_ids, attention_mask=attention_mask, global_attention_mask=global_attention_mask)[0]
 
         expected_output_sum = torch.tensor(74585.8594, device=torch_device)
         expected_output_mean = torch.tensor(0.0243, device=torch_device)
@@ -341,7 +412,7 @@ class LongformerModelIntegrationTest(unittest.TestCase):
         self.assertTrue(torch.allclose(output.mean(), expected_output_mean, atol=1e-4))
 
     @slow
-    def test_inference_masked_lm(self):
+    def test_inference_masked_lm_long(self):
         model = LongformerForMaskedLM.from_pretrained("allenai/longformer-base-4096")
         model.to(torch_device)
 
@@ -352,9 +423,9 @@ class LongformerModelIntegrationTest(unittest.TestCase):
 
         loss, prediction_scores = model(input_ids, labels=input_ids)
 
-        expected_loss = torch.tensor(0.0620, device=torch_device)
-        expected_prediction_scores_sum = torch.tensor(-6.1599e08, device=torch_device)
-        expected_prediction_scores_mean = torch.tensor(-3.0622, device=torch_device)
+        expected_loss = torch.tensor(0.0074, device=torch_device)
+        expected_prediction_scores_sum = torch.tensor(-6.1048e08, device=torch_device)
+        expected_prediction_scores_mean = torch.tensor(-3.0348, device=torch_device)
         input_ids = input_ids.to(torch_device)
 
         self.assertTrue(torch.allclose(loss, expected_loss, atol=1e-4))

tests/test_modeling_marian.py

@@ -19,8 +19,7 @@ import unittest
 from transformers import is_torch_available
 from transformers.file_utils import cached_property
 from transformers.hf_api import HfApi
-
-from .utils import require_torch, slow, torch_device
+from transformers.testing_utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_mobilebert.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_openai.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_reformer.py

@@ -16,19 +16,21 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_multigpu, require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor
-from .utils import require_multigpu, require_torch, slow, torch_device
 
 
 if is_torch_available():
     from transformers import (
         ReformerConfig,
+        ReformerForMaskedLM,
         ReformerModel,
         ReformerModelWithLMHead,
         ReformerTokenizer,
         ReformerLayer,
+        ReformerForQuestionAnswering,
         REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
     )
     import torch
@@ -43,6 +45,7 @@ class ReformerModelTester:
         is_training=None,
         is_decoder=None,
         use_input_mask=None,
+        use_labels=None,
         vocab_size=None,
         attention_head_size=None,
         hidden_size=None,
@@ -81,6 +84,7 @@ class ReformerModelTester:
         self.is_training = is_training
         self.is_decoder = is_decoder
         self.use_input_mask = use_input_mask
+        self.use_labels = use_labels
         self.vocab_size = vocab_size
         self.attention_head_size = attention_head_size
         self.hidden_size = hidden_size
@@ -128,6 +132,10 @@ class ReformerModelTester:
         if self.use_input_mask:
             input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)
 
+        choice_labels = None
+        if self.use_labels:
+            choice_labels = ids_tensor([self.batch_size], 2)
+
         config = ReformerConfig(
             vocab_size=self.vocab_size,
             hidden_size=self.hidden_size,
@@ -160,14 +168,13 @@ class ReformerModelTester:
             config,
             input_ids,
             input_mask,
+            choice_labels,
         )
 
     def check_loss_output(self, result):
         self.parent.assertListEqual(list(result["loss"].size()), [])
 
-    def create_and_check_reformer_model(
-        self, config, input_ids, input_mask,
-    ):
+    def create_and_check_reformer_model(self, config, input_ids, input_mask, choice_labels):
         model = ReformerModel(config=config)
         model.to(torch_device)
         model.eval()
@@ -182,18 +189,14 @@ class ReformerModelTester:
             list(result["sequence_output"].size()), [self.batch_size, self.seq_length, 2 * self.hidden_size],
         )
 
-    def create_and_check_reformer_model_with_lm_backward(
-        self, config, input_ids, input_mask,
-    ):
+    def create_and_check_reformer_model_with_lm_backward(self, config, input_ids, input_mask, choice_labels):
         model = ReformerModelWithLMHead(config=config)
         model.to(torch_device)
         model.eval()
         loss = model(input_ids, attention_mask=input_mask, labels=input_ids)[0]
         loss.backward()
 
-    def create_and_check_reformer_with_lm(
-        self, config, input_ids, input_mask,
-    ):
+    def create_and_check_reformer_with_lm(self, config, input_ids, input_mask, choice_labels):
         model = ReformerModelWithLMHead(config=config)
         model.to(torch_device)
         model.eval()
@@ -207,7 +210,24 @@ class ReformerModelTester:
         )
         self.check_loss_output(result)
 
-    def create_and_check_reformer_model_with_attn_mask(self, config, input_ids, input_mask, is_decoder):
+    def create_and_check_reformer_with_mlm(self, config, input_ids, input_mask, choice_labels):
+        config.is_decoder = False
+        model = ReformerForMaskedLM(config=config)
+        model.to(torch_device)
+        model.eval()
+        loss, prediction_scores = model(input_ids, attention_mask=input_mask, labels=input_ids)
+        result = {
+            "loss": loss,
+            "prediction_scores": prediction_scores,
+        }
+        self.parent.assertListEqual(
+            list(result["prediction_scores"].size()), [self.batch_size, self.seq_length, self.vocab_size],
+        )
+        self.check_loss_output(result)
+
+    def create_and_check_reformer_model_with_attn_mask(
+        self, config, input_ids, input_mask, choice_labels, is_decoder=False
+    ):
         # no special position embeddings
         config.axial_pos_embds = False
         config.is_decoder = is_decoder
@@ -248,7 +268,9 @@ class ReformerModelTester:
 
         self.parent.assertTrue(torch.allclose(output_padded, output_padded_rolled, atol=1e-3))
 
-    def create_and_check_reformer_layer_dropout_seed(self, config, input_ids, input_mask, is_decoder):
+    def create_and_check_reformer_layer_dropout_seed(
+        self, config, input_ids, input_mask, choice_labels, is_decoder=False
+    ):
         config.is_decoder = is_decoder
         layer = ReformerLayer(config).to(torch_device)
         layer.train()
@@ -281,7 +303,7 @@ class ReformerModelTester:
             torch.allclose(next_hidden_states, hidden_states + feed_forward_hidden_states, atol=1e-3,)
         )
 
-    def create_and_check_reformer_feed_forward_chunking(self, config, input_ids, input_mask):
+    def create_and_check_reformer_feed_forward_chunking(self, config, input_ids, input_mask, choice_labels):
         torch.manual_seed(0)
         model = ReformerModel(config=config)
         model.to(torch_device)
@@ -299,7 +321,7 @@ class ReformerModelTester:
         hidden_states_with_chunk = model(input_ids, attention_mask=input_mask)[0]
         self.parent.assertTrue(torch.allclose(hidden_states_no_chunk, hidden_states_with_chunk, atol=1e-3))
 
-    def create_and_check_reformer_feed_backward_chunking(self, config, input_ids, input_mask):
+    def create_and_check_reformer_feed_backward_chunking(self, config, input_ids, input_mask, choice_labels):
         if not self.is_training:
             return
 
@@ -341,7 +363,7 @@ class ReformerModelTester:
             torch.allclose(grad_slice_position_factor_2_chunk, grad_slice_position_factor_2_no_chunk, atol=1e-3)
         )
 
-    def create_and_check_reformer_random_seed(self, config, input_ids, input_mask):
+    def create_and_check_reformer_random_seed(self, config, input_ids, input_mask, choice_labels):
         layer = ReformerLayer(config).to(torch_device)
         layer.train()
 
@@ -372,7 +394,7 @@ class ReformerModelTester:
             seeds.append(layer.feed_forward_seed)
         self.parent.assertGreater(len(set(seeds)), 70)
 
-    def create_and_check_reformer_model_fp16_forward(self, config, input_ids, input_mask):
+    def create_and_check_reformer_model_fp16_forward(self, config, input_ids, input_mask, choice_labels):
         model = ReformerModel(config=config)
         model.to(torch_device)
         model.half()
@@ -380,7 +402,7 @@ class ReformerModelTester:
         output = model(input_ids, attention_mask=input_mask)[0]
         self.parent.assertFalse(torch.isnan(output).any().item())
 
-    def create_and_check_reformer_model_fp16_generate(self, config, input_ids, input_mask):
+    def create_and_check_reformer_model_fp16_generate(self, config, input_ids, input_mask, choice_labels):
         model = ReformerModelWithLMHead(config=config)
         model.to(torch_device)
         model.half()
@@ -388,7 +410,7 @@ class ReformerModelTester:
         output = model.generate(input_ids, attention_mask=input_mask, do_sample=False)
         self.parent.assertFalse(torch.isnan(output).any().item())
 
-    def create_and_check_reformer_no_chunking(self, config, input_ids, input_mask):
+    def create_and_check_reformer_no_chunking(self, config, input_ids, input_mask, choice_labels):
         # force chunk length to be bigger than input_ids
         config.lsh_attn_chunk_length = 2 * input_ids.shape[-1]
         config.local_attn_chunk_length = 2 * input_ids.shape[-1]
@@ -398,9 +420,25 @@ class ReformerModelTester:
         output_logits = model(input_ids, attention_mask=input_mask)[0]
         self.parent.assertTrue(output_logits.shape[1] == input_ids.shape[-1])
 
+    def create_and_check_longformer_for_question_answering(self, config, input_ids, input_mask, choice_labels):
+        model = ReformerForQuestionAnswering(config=config)
+        model.to(torch_device)
+        model.eval()
+        loss, start_logits, end_logits = model(
+            input_ids, attention_mask=input_mask, start_positions=choice_labels, end_positions=choice_labels,
+        )
+        result = {
+            "loss": loss,
+            "start_logits": start_logits,
+            "end_logits": end_logits,
+        }
+        self.parent.assertListEqual(list(result["start_logits"].size()), [self.batch_size, self.seq_length])
+        self.parent.assertListEqual(list(result["end_logits"].size()), [self.batch_size, self.seq_length])
+        self.check_loss_output(result)
+
     def prepare_config_and_inputs_for_common(self):
         config_and_inputs = self.prepare_config_and_inputs()
-        (config, input_ids, input_mask,) = config_and_inputs
+        (config, input_ids, input_mask, choice_labels) = config_and_inputs
         inputs_dict = {"input_ids": input_ids, "attention_mask": input_mask}
         return config, inputs_dict
 
@@ -423,17 +461,21 @@ class ReformerTesterMixin:
 
     def test_reformer_model_attn_masking(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_reformer_model_with_attn_mask(*config_and_inputs, True)
-        self.model_tester.create_and_check_reformer_model_with_attn_mask(*config_and_inputs, False)
+        self.model_tester.create_and_check_reformer_model_with_attn_mask(*config_and_inputs, is_decoder=True)
+        self.model_tester.create_and_check_reformer_model_with_attn_mask(*config_and_inputs, is_decoder=False)
 
     def test_reformer_with_lm(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
         self.model_tester.create_and_check_reformer_with_lm(*config_and_inputs)
 
+    def test_reformer_with_mlm(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_reformer_with_mlm(*config_and_inputs)
+
     def test_reformer_layer_training_dropout(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_reformer_layer_dropout_seed(*config_and_inputs, True)
-        self.model_tester.create_and_check_reformer_layer_dropout_seed(*config_and_inputs, False)
+        self.model_tester.create_and_check_reformer_layer_dropout_seed(*config_and_inputs, is_decoder=True)
+        self.model_tester.create_and_check_reformer_layer_dropout_seed(*config_and_inputs, is_decoder=False)
 
     def test_reformer_chunking_forward_equality(self):
         config_and_inputs = self.model_tester.prepare_config_and_inputs()
@@ -470,7 +512,9 @@ class ReformerTesterMixin:
 
 @require_torch
 class ReformerLocalAttnModelTest(ReformerTesterMixin, ModelTesterMixin, unittest.TestCase):
-    all_model_classes = (ReformerModel, ReformerModelWithLMHead) if is_torch_available() else ()
+    all_model_classes = (
+        (ReformerModel, ReformerModelWithLMHead, ReformerForQuestionAnswering) if is_torch_available() else ()
+    )
     all_generative_model_classes = (ReformerModelWithLMHead,) if is_torch_available() else ()
     test_pruning = False
     test_headmasking = False
@@ -481,8 +525,9 @@ class ReformerLocalAttnModelTest(ReformerTesterMixin, ModelTesterMixin, unittest
             "batch_size": 13,
             "seq_length": 32,
             "is_training": True,
-            "is_decoder": False,
+            "is_decoder": True,
             "use_input_mask": True,
+            "use_labels": True,
             "vocab_size": 32,
             "attention_head_size": 16,
             "hidden_size": 32,
@@ -524,7 +569,9 @@ class ReformerLocalAttnModelTest(ReformerTesterMixin, ModelTesterMixin, unittest
 
 @require_torch
 class ReformerLSHAttnModelTest(ReformerTesterMixin, ModelTesterMixin, unittest.TestCase):
-    all_model_classes = (ReformerModel, ReformerModelWithLMHead) if is_torch_available() else ()
+    all_model_classes = (
+        (ReformerModel, ReformerModelWithLMHead, ReformerForQuestionAnswering) if is_torch_available() else ()
+    )
     all_generative_model_classes = (ReformerModelWithLMHead,) if is_torch_available() else ()
     test_pruning = False
     test_headmasking = False
@@ -535,8 +582,9 @@ class ReformerLSHAttnModelTest(ReformerTesterMixin, ModelTesterMixin, unittest.T
             "batch_size": 13,
             "seq_length": 13,
             "use_input_mask": True,
+            "use_labels": True,
             "is_training": False,
-            "is_decoder": False,
+            "is_decoder": True,
             "vocab_size": 32,
             "attention_head_size": 16,
             "hidden_size": 64,
@@ -886,7 +934,7 @@ class ReformerIntegrationTests(unittest.TestCase):
         config["num_buckets"] = [2, 4]
         config["is_decoder"] = False
         torch.manual_seed(0)
-        model = ReformerModelWithLMHead(ReformerConfig(**config)).to(torch_device)
+        model = ReformerForMaskedLM(ReformerConfig(**config)).to(torch_device)
         model.eval()
         input_ids, attn_mask = self._get_input_ids_and_mask()
         hidden_states = model(input_ids=input_ids, attention_mask=attn_mask)[0]

tests/test_modeling_roberta.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_t5.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_tf_albert.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import AlbertConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_auto.py

@@ -17,8 +17,7 @@
 import unittest
 
 from transformers import is_tf_available
-
-from .utils import DUMMY_UNKWOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, require_tf, slow
+from transformers.testing_utils import DUMMY_UNKWOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_bert.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import BertConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_camembert.py

@@ -16,8 +16,7 @@
 import unittest
 
 from transformers import is_tf_available
-
-from .utils import require_tf, slow
+from transformers.testing_utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_common.py

@@ -23,8 +23,7 @@ import unittest
 from importlib import import_module
 
 from transformers import is_tf_available, is_torch_available
-
-from .utils import _tf_gpu_memory_limit, require_tf
+from transformers.testing_utils import _tf_gpu_memory_limit, require_tf
 
 
 if is_tf_available():

tests/test_modeling_tf_ctrl.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import CTRLConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_distilbert.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import DistilBertConfig, is_tf_available
+from transformers.testing_utils import require_tf
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf
 
 
 if is_tf_available():

tests/test_modeling_tf_electra.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import ElectraConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_flaubert.py

@@ -16,8 +16,7 @@
 import unittest
 
 from transformers import is_tf_available
-
-from .utils import require_tf, slow
+from transformers.testing_utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_gpt2.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import GPT2Config, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_mobilebert.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import MobileBertConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_openai_gpt.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import OpenAIGPTConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_roberta.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import RobertaConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_t5.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import T5Config, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_transfo_xl.py

@@ -18,10 +18,10 @@ import random
 import unittest
 
 from transformers import TransfoXLConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_xlm.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_xlm_roberta.py

@@ -16,8 +16,7 @@
 import unittest
 
 from transformers import is_tf_available
-
-from .utils import require_tf, slow
+from transformers.testing_utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_tf_xlnet.py

@@ -18,10 +18,10 @@ import random
 import unittest
 
 from transformers import XLNetConfig, is_tf_available
+from transformers.testing_utils import require_tf, slow
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_tf_common import TFModelTesterMixin, ids_tensor
-from .utils import require_tf, slow
 
 
 if is_tf_available():

tests/test_modeling_transfo_xl.py

@@ -17,10 +17,10 @@ import random
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_multigpu, require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_multigpu, require_torch, slow, torch_device
 
 
 if is_torch_available():
@@ -162,7 +162,6 @@ class TransfoXLModelTester:
 
 @require_torch
 class TransfoXLModelTest(ModelTesterMixin, unittest.TestCase):
-
     all_model_classes = (TransfoXLModel, TransfoXLLMHeadModel) if is_torch_available() else ()
     all_generative_model_classes = (TransfoXLLMHeadModel,) if is_torch_available() else ()
     test_pruning = False
@@ -448,7 +447,6 @@ class TransfoXLModelLanguageGenerationTest(unittest.TestCase):
         #  father initially slaps him for making such an accusation , Rasputin watches as the
         #  man is chased outside and beaten . Twenty years later , Rasputin sees a vision of
         #  the Virgin Mary , prompting him to become a priest . Rasputin quickly becomes famous ,
-
         #  with people , even a bishop , begging for his blessing . <eod> </s> <eos>
 
         expected_output_ids = [
@@ -595,54 +593,77 @@ class TransfoXLModelLanguageGenerationTest(unittest.TestCase):
             0,
             33,
             1,
-            1857,
+            142,
+            1298,
+            188,
             2,
-            1,
-            1009,
+            29546,
+            113,
+            8,
+            3654,
             4,
+            1,
             1109,
-            11739,
-            4762,
-            358,
-            5,
-            25,
-            245,
-            28,
-            1110,
+            7136,
+            833,
             3,
             13,
-            1041,
+            1645,
             4,
+            29546,
+            11,
+            104,
+            7,
+            1,
+            1109,
+            532,
+            7129,
+            2,
+            10,
+            83507,
+            2,
+            1162,
+            1123,
+            2,
+            6,
+            7245,
+            10,
+            2,
+            5,
+            11,
+            104,
+            7,
+            1,
+            1109,
+            532,
+            7129,
+            2,
+            10,
             24,
-            603,
-            490,
-            2,
-            71477,
-            20098,
-            104447,
-            2,
-            20961,
-            1,
-            2604,
+            24,
+            10,
+            22,
+            10,
+            13,
+            770,
+            5863,
             4,
-            1,
-            329,
-            3,
-            0,
+            7245,
+            10,
         ]
-        #  In 1991, the remains of Russian Tsar Nicholas II and his family (
-        #  except for Alexei and Maria ) are discovered. The voice of young son,
-        #  Tsarevich Alexei Nikolaevich, narrates the remainder of the story.
-        #  1883 Western Siberia, a young Grigori Rasputin is asked by his father
-        #  and a group of men to perform magic. Rasputin has a vision and
-        #  denounces one of the men as a horse thief. Although his father initially
-        #  slaps him for making such an accusation, Rasputin watches as the man
-        #  is chased outside and beaten. Twenty years later, Rasputin sees a vision
-        #  of the Virgin Mary, prompting him to become a priest.
-        #  Rasputin quickly becomes famous, with people, even a bishop, begging for
-        #  his blessing. <unk> <unk> <eos> In the 1990s, the remains of Russian Tsar
-        # Nicholas II and his family were discovered. The voice of <unk> young son,
-        # Tsarevich Alexei Nikolaevich, narrates the remainder of the story.<eos>
+        #  In 1991, the remains of Russian Tsar Nicholas II and his family ( except for
+        #  Alexei and Maria ) are discovered. The voice of young son, Tsarevich Alexei
+        #  Nikolaevich, narrates the remainder of the story. 1883 Western Siberia, a young
+        #  Grigori Rasputin is asked by his father and a group of men to perform magic.
+        #  Rasputin has a vision and denounces one of the men as a horse thief. Although
+        #  his father initially slaps him for making such an accusation, Rasputin watches
+        #  as the man is chased outside and beaten. Twenty years later, Rasputin sees a
+        #  vision of the Virgin Mary, prompting him to become a priest. Rasputin quickly
+        #  becomes famous, with people, even a bishop, begging for his blessing. In the
+        #  early 20th century, Rasputin became a symbol of the Russian Orthodox Church.
+        #  The image of Rasputin was used in the Russian national anthem, " Nearer, My God,
+        #  to Heaven ", and was used in the Russian national anthem, " " ( " The Great Spirit
+        #  of Heaven "
 
         output_ids = model.generate(input_ids, max_length=200, do_sample=False)
         self.assertListEqual(output_ids[0].tolist(), expected_output_ids)

tests/test_modeling_xlm.py

@@ -17,10 +17,10 @@
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_modeling_xlm_roberta.py

@@ -17,8 +17,7 @@
 import unittest
 
 from transformers import is_torch_available
-
-from .utils import slow
+from transformers.testing_utils import slow
 
 
 if is_torch_available():

tests/test_modeling_xlnet.py

@@ -18,10 +18,10 @@ import random
 import unittest
 
 from transformers import is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
 
 from .test_configuration_common import ConfigTester
 from .test_modeling_common import ModelTesterMixin, ids_tensor
-from .utils import require_torch, slow, torch_device
 
 
 if is_torch_available():

tests/test_onnx.py

@@ -3,9 +3,9 @@ from os.path import dirname, exists
 from shutil import rmtree
 from tempfile import NamedTemporaryFile, TemporaryDirectory
 
-from tests.utils import require_tf, require_torch, slow
 from transformers import BertConfig, BertTokenizerFast, FeatureExtractionPipeline
 from transformers.convert_graph_to_onnx import convert, ensure_valid_input, infer_shapes
+from transformers.testing_utils import require_tf, require_torch, slow
 
 
 class FuncContiguousArgs:

tests/test_optimization.py

@@ -19,8 +19,7 @@ import tempfile
 import unittest
 
 from transformers import is_torch_available
-
-from .utils import require_torch
+from transformers.testing_utils import require_torch
 
 
 if is_torch_available():

tests/test_optimization_tf.py

@@ -1,8 +1,7 @@
 import unittest
 
 from transformers import is_tf_available
-
-from .utils import require_tf
+from transformers.testing_utils import require_tf
 
 
 if is_tf_available():

tests/test_pipelines.py

@@ -3,8 +3,7 @@ from typing import Iterable, List, Optional
 
 from transformers import pipeline
 from transformers.pipelines import SUPPORTED_TASKS, DefaultArgumentHandler, Pipeline
-
-from .utils import require_tf, require_torch, slow, torch_device
+from transformers.testing_utils import require_tf, require_torch, slow, torch_device
 
 
 DEFAULT_DEVICE_NUM = -1 if torch_device == "cpu" else 0
@@ -31,12 +30,12 @@ TF_TRANSLATION_FINETUNED_MODELS = [("patrickvonplaten/t5-tiny-random", "translat
 
 expected_fill_mask_result = [
     [
-        {"sequence": "<s> My name is:</s>", "score": 0.009954338893294334, "token": 35},
-        {"sequence": "<s> My name is John</s>", "score": 0.0080940006300807, "token": 610},
+        {"sequence": "<s>My name is John</s>", "score": 0.00782308354973793, "token": 610, "token_str": "ĠJohn"},
+        {"sequence": "<s>My name is Chris</s>", "score": 0.007475061342120171, "token": 1573, "token_str": "ĠChris"},
     ],
     [
-        {"sequence": "<s> The largest city in France is Paris</s>", "score": 0.3185044229030609, "token": 2201},
-        {"sequence": "<s> The largest city in France is Lyon</s>", "score": 0.21112334728240967, "token": 12790},
+        {"sequence": "<s>The largest city in France is Paris</s>", "score": 0.3185044229030609, "token": 2201},
+        {"sequence": "<s>The largest city in France is Lyon</s>", "score": 0.21112334728240967, "token": 12790},
     ],
 ]
 SUMMARIZATION_KWARGS = dict(num_beams=2, min_length=2, max_length=5)
@@ -218,9 +217,15 @@ class MonoColumnInputTestCase(unittest.TestCase):
             "My name is <mask>",
             "The largest city in France is <mask>",
         ]
+        invalid_inputs = [
+            "This is <mask> <mask>"  # More than 1 mask_token in the input is not supported
+            "This is"  # No mask_token is not supported
+        ]
         for model_name in FILL_MASK_FINETUNED_MODELS:
             nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="pt", topk=2,)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys, expected_check_keys=["sequence"])
+            self._test_mono_column_pipeline(
+                nlp, valid_inputs, mandatory_keys, invalid_inputs, expected_check_keys=["sequence"]
+            )
 
     @require_tf
     def test_tf_fill_mask(self):
@@ -229,9 +234,15 @@ class MonoColumnInputTestCase(unittest.TestCase):
             "My name is <mask>",
             "The largest city in France is <mask>",
         ]
+        invalid_inputs = [
+            "This is <mask> <mask>"  # More than 1 mask_token in the input is not supported
+            "This is"  # No mask_token is not supported
+        ]
         for model_name in FILL_MASK_FINETUNED_MODELS:
             nlp = pipeline(task="fill-mask", model=model_name, tokenizer=model_name, framework="tf", topk=2,)
-            self._test_mono_column_pipeline(nlp, valid_inputs, mandatory_keys, expected_check_keys=["sequence"])
+            self._test_mono_column_pipeline(
+                nlp, valid_inputs, mandatory_keys, invalid_inputs, expected_check_keys=["sequence"]
+            )
 
     @require_torch
     @slow

tests/test_tokenization_auto.py

@@ -27,10 +27,9 @@ from transformers import (
     RobertaTokenizer,
     RobertaTokenizerFast,
 )
+from transformers.testing_utils import DUMMY_UNKWOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER  # noqa: F401
 from transformers.tokenization_auto import TOKENIZER_MAPPING
 
-from .utils import DUMMY_UNKWOWN_IDENTIFIER, SMALL_MODEL_IDENTIFIER, slow  # noqa: F401
-
 
 class AutoTokenizerTest(unittest.TestCase):
     # @slow

tests/test_tokenization_bert.py

@@ -17,6 +17,7 @@
 import os
 import unittest
 
+from transformers.testing_utils import slow
 from transformers.tokenization_bert import (
     VOCAB_FILES_NAMES,
     BasicTokenizer,
@@ -29,7 +30,6 @@ from transformers.tokenization_bert import (
 )
 
 from .test_tokenization_common import TokenizerTesterMixin
-from .utils import slow
 
 
 class BertTokenizationTest(TokenizerTesterMixin, unittest.TestCase):

tests/test_tokenization_bert_japanese.py

@@ -17,6 +17,7 @@
 import os
 import unittest
 
+from transformers.testing_utils import custom_tokenizers
 from transformers.tokenization_bert import WordpieceTokenizer
 from transformers.tokenization_bert_japanese import (
     VOCAB_FILES_NAMES,
@@ -26,7 +27,6 @@ from transformers.tokenization_bert_japanese import (
 )
 
 from .test_tokenization_common import TokenizerTesterMixin
-from .utils import custom_tokenizers
 
 
 @custom_tokenizers

tests/test_tokenization_common.py

@@ -22,8 +22,8 @@ import tempfile
 from collections import OrderedDict
 from typing import TYPE_CHECKING, Dict, List, Tuple, Union
 
-from tests.utils import require_tf, require_torch, slow
 from transformers import PreTrainedTokenizer, PreTrainedTokenizerBase, PreTrainedTokenizerFast
+from transformers.testing_utils import require_tf, require_torch, slow
 
 
 if TYPE_CHECKING:

tests/test_tokenization_distilbert.py

@@ -14,10 +14,10 @@
 # limitations under the License.
 
 
+from transformers.testing_utils import slow
 from transformers.tokenization_distilbert import DistilBertTokenizer, DistilBertTokenizerFast
 
 from .test_tokenization_bert import BertTokenizationTest
-from .utils import slow
 
 
 class DistilBertTokenizationTest(BertTokenizationTest):