try if multi gpu works

2020-08-19 22:24:26 +02:00
80 changed files with 528 additions and 1836 deletions
@@ -65,7 +65,7 @@ Awesome! Please provide the following information:
 If you are willing to contribute the model yourself, let us know so we can best
 guide you.

-We have added a **detailed guide and templates** to guide you in the process of adding a new model. You can find them
+We have added a **detailed guide and templates** to guide you in the process of adding a new model. You can find them 
 in the [`templates`](https://github.com/huggingface/transformers/tree/master/templates) folder.

 ### Do you want a new feature (that is not a model)?
@@ -87,8 +87,8 @@ A world-class feature request addresses the following points:
 If your issue is well written we're already 80% of the way there by the time you
 post it.

-We have added **templates** to guide you in the process of adding a new example script for training or testing the
-models in the library. You can find them in the [`templates`](https://github.com/huggingface/transformers/tree/master/templates)
+We have added **templates** to guide you in the process of adding a new example script for training or testing the 
+models in the library. You can find them in the [`templates`](https://github.com/huggingface/transformers/tree/master/templates) 
 folder.

 ## Start contributing! (Pull Requests)
@@ -149,14 +149,6 @@ Follow these steps to start contributing:
   $ make test
   ```

-   Note, that this command uses `-n auto` pytest flag, therefore, it will start as many parallel `pytest` processes as the number of your computer's CPU-cores, and if you have lots of those and a few GPUs and not a great amount of RAM, it's likely to overload your computer. Therefore, to run the test suite, you may want to consider using this command instead:
-
-   ```bash
-   $ python -m pytest -n 3 --dist=loadfile -s -v ./tests/
-   ```
-
-   Adjust the value of `-n` to fit the load your hardware can support.
-
   `transformers` relies on `black` and `isort` to format its source code
   consistently. After you make changes, format them with:

@@ -216,21 +208,21 @@ Follow these steps to start contributing:
   are useful to avoid duplicated work, and to differentiate it from PRs ready
   to be merged;
 4. Make sure existing tests pass;
-5. Add high-coverage tests. No quality testing = no merge.
-   - If you are adding a new model, make sure that you use
+5. Add high-coverage tests. No quality testing = no merge. 
+   - If you are adding a new model, make sure that you use 
     `ModelTester.all_model_classes = (MyModel, MyModelWithLMHead,...)`, which triggers the common tests.
-   - If you are adding new `@slow` tests, make sure they pass using
-     `RUN_SLOW=1 python -m pytest tests/test_my_new_model.py`.
-   - If you are adding a new tokenizer, write tests, and make sure
+   - If you are adding new `@slow` tests, make sure they pass using 
+     `RUN_SLOW=1 python -m pytest tests/test_my_new_model.py`. 
+   - If you are adding a new tokenizer, write tests, and make sure 
     `RUN_SLOW=1 python -m pytest tests/test_tokenization_{your_model_name}.py` passes.
   CircleCI does not run the slow tests, but github actions does every night!
-6. All public methods must have informative docstrings that work nicely with sphinx. See `modeling_ctrl.py` for an
+6. All public methods must have informative docstrings that work nicely with sphinx. See `modeling_ctrl.py` for an 
   example.

 ### Tests

-An extensive test suite is included to test the library behavior and several examples. Library tests can be found in
-the [tests folder](https://github.com/huggingface/transformers/tree/master/tests) and examples tests in the
+An extensive test suite is included to test the library behavior and several examples. Library tests can be found in 
+the [tests folder](https://github.com/huggingface/transformers/tree/master/tests) and examples tests in the 
 [examples folder](https://github.com/huggingface/transformers/tree/master/examples).

 We like `pytest` and `pytest-xdist` because it's faster. From the root of the
@@ -76,8 +76,7 @@ exclude_patterns = [u'_build', 'Thumbs.db', '.DS_Store']
 pygments_style = None

 # Remove the prompt when copying examples
-copybutton_prompt_text = r">>> |\.\.\. "
-copybutton_prompt_is_regexp = True
+copybutton_prompt_text = ">>> "

 # -- Options for HTML output -------------------------------------------------

@@ -1,13 +1,6 @@
 Fine-tuning with custom datasets
 ================================

-.. note::
-
-    The datasets used in this tutorial are available and can be more easily accessed using the
-    `🤗 NLP library <https://github.com/huggingface/nlp>`_. We do not use this library to access the datasets here
-    since this tutorial meant to illustrate how to work with your own data. A brief of introduction can be found
-    at the end of the tutorial in the section ":ref:`nlplib`".
-
 This tutorial will take you through several examples of using 🤗 Transformers models with your own datasets. The
 guide shows one of many valid workflows for using these models and is meant to be illustrative rather than
 definitive. We show examples of reading in several data formats, preprocessing the data for several types of tasks,
@@ -21,16 +14,17 @@ We include several examples, each of which demonstrates a different type of comm
  - :ref:`qa_squad`
  - :ref:`resources`

+.. note::
+
+    Many of the datasets used in this tutorial are available and can be more easily accessed using the
+    `🤗 NLP library <https://github.com/huggingface/nlp>`_. We do not use this library to access the datasets here
+    since this tutorial meant to illustrate how to work with your own data.
+
 .. _seq_imdb:

 Sequence Classification with IMDb Reviews
 -----------------------------------------

-.. note::
-
-    This dataset can be explored in the Hugging Face model hub (`IMDb <https://huggingface.co/datasets/imdb>`_), and can
-    be alternatively downloaded with the 🤗 NLP library with ``load_dataset("imdb")``.
-
 In this example, we'll show how to download, tokenize, and train a model on the IMDb reviews dataset. This task
 takes the text of a review and requires the model to predict whether the sentiment of the review is positive or
 negative. Let's start by downloading the dataset from the
@@ -62,8 +56,8 @@ read this in.
    train_texts, train_labels = read_imdb_split('aclImdb/train')
    test_texts, test_labels = read_imdb_split('aclImdb/test')

-We now have a train and test dataset, but let's also also create a validation set which we can use for for
-evaluation and tuning without training our test set results. Sklearn has a convenient utility for creating such
+We now have a train and test dataset, but let's also also create a validation set which we can use for
+for evaluation and tuning without taining our test set results. Sklearn has a convenient utility for creating such
 splits:

 .. code-block:: python
@@ -246,11 +240,6 @@ We can also train use native PyTorch or TensorFlow:
 Token Classification with W-NUT Emerging Entities
 -------------------------------------------------

-.. note::
-
-    This dataset can be explored in the Hugging Face model hub (`WNUT-17 <https://huggingface.co/datasets/wnut_17>`_), and can
-    be alternatively downloaded with the 🤗 NLP library with ``load_dataset("wnut_17")``.
-
 Next we will look at token classification. Rather than classifying an entire sequence, this task classifies token by
 token. We'll demonstrate how to do this with 
 `Named Entity Recognition <http://nlpprogress.com/english/named_entity_recognition.html>`_, which involves
@@ -445,11 +434,6 @@ sequence classification example above.
 Question Answering with SQuAD 2.0
 ---------------------------------

-.. note::
-
-    This dataset can be explored in the Hugging Face model hub (`SQuAD V2 <https://huggingface.co/datasets/squad_v2>`_), and can
-    be alternatively downloaded with the 🤗 NLP library with ``load_dataset("squad_v2")``.
-
 Question answering comes in many forms. In this example, we'll look at the particular type of extractive QA that
 involves answering a question about a passage by highlighting the segment of the passage that answers the question.
 This involves fine-tuning a model which predicts a start position and an end position in the passage. We will use the
@@ -662,54 +646,3 @@ Additional Resources
    masked language model from scratch.
  - :doc:`Preprocessing <preprocessing>`. Docs page on data preprocessing.
  - :doc:`Training <training>`. Docs page on training and fine-tuning.
-
-.. _nlplib:
-
-Using the 🤗 NLP Datasets & Metrics library
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-This tutorial demonstrates how to read in datasets from various raw text formats and prepare them for training with
-🤗 Transformers so that you can do the same thing with your own custom datasets. However, we recommend users use the
-`🤗 NLP library <https://github.com/huggingface/nlp>`_ for working with the 150+ datasets included in the
-`hub <https://huggingface.co/datasets>`_, including the three datasets used in this tutorial. As a very brief overview,
-we will show how to use the NLP library to download and prepare the IMDb dataset from the first example,
-:ref:`seq_imdb`.
-
-Start by downloading the dataset:
-
-.. code-block:: python
-
-    from nlp import load_dataset
-    train = load_dataset("imdb", split="train")
-
-Each dataset has multiple columns corresponding to different features. Let's see what our columns are.
-
-.. code-block:: python
-
-    >>> print(train.column_names)
-    ['label', 'text']
-
-Great. Now let's tokenize the text. We can do this using the ``map`` method. We'll also rename the ``label`` column
-to ``labels`` to match the model's input arguments.
-
-.. code-block:: python
-
-    train = train.map(lambda batch: tokenizer(batch["text"], truncation=True, padding=True), batched=True)
-    train.rename_column_("label", "labels")
-
-Lastly, we can use the ``set_format`` method to determine which columns and in what data format we want to access
-dataset elements.
-
-.. code-block:: python
-
-    ## PYTORCH CODE
-    >>> train.set_format("torch", columns=["input_ids", "attention_mask", "labels"])
-    >>> {key: val.shape for key, val in train[0].items()})
-    {'labels': torch.Size([]), 'input_ids': torch.Size([512]), 'attention_mask': torch.Size([512])}
-    ## TENSORFLOW CODE
-    >>> train.set_format("tensorflow", columns=["input_ids", "attention_mask", "labels"])
-    >>> {key: val.shape for key, val in train[0].items()})
-    {'labels': TensorShape([]), 'input_ids': TensorShape([512]), 'attention_mask': TensorShape([512])}
-
-We now have a fully-prepared dataset. Check out `the 🤗 NLP docs <https://huggingface.co/nlp/processing.html>`_ for
-a more thorough introduction.
@@ -49,13 +49,6 @@ CamembertModel
    :members:


-CamembertForCausalLM
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.CamembertForCausalLM
-    :members:
-
-
 CamembertForMaskedLM
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -478,31 +478,6 @@ pretraining tasks, a composition of the following transformations are applied:

 The library provides a version of this model for conditional generation and sequence classification.

-Pegasus
----------------------------------------------
-
-.. raw:: html
-
-   <a href="https://huggingface.co/models?filter=pegasus">
-       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-pegasus-blueviolet">
-   </a>
-   <a href="model_doc/pegasus.html">
-       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-pegasus-blueviolet">
-   </a>
-
-`PEGASUS: Pre-training with Extracted Gap-sentences forAbstractive Summarization 
-<https://arxiv.org/pdf/1912.08777.pdf>`_, Jingqing Zhang, Yao Zhao, Mohammad Saleh and Peter J. Liu on Dec 18, 2019.
-
-Sequence-to-sequence model with the same encoder-decoder model architecture as BART. Pegasus is pre-trained jointly on two self-supervised objective functions: Masked Language Modeling (MLM) and a novel summarization specific pre-training objective, called Gap Sentence Generation (GSG).
-
-  * MLM: encoder input tokens are randomely replaced by a mask tokens and have to be predicted by the encoder (like in BERT)
-  * GSG: whole encoder input sentences are replaced by a second mask token and fed to the decoder, but which has a causal mask to hide the future words like a regular auto-regressive transformer decoder.
-
-In contrast to BART, Pegasus' pretraining task is intentionally similar to summarization: important sentences are masked and are generated together as one output sequence from the remaining sentences, similar to an extractive summary.
-
-The library provides a version of this model for conditional generation, which should be used for summarization.
-
-
 MarianMT
 ----------------------------------------------

@@ -552,31 +527,6 @@ input becomes “My <x> very <y> .” and the target input becomes “<x> dog is

 The library provides a version of this model for conditional generation.

-MBart
----------------------------------------------
-
-.. raw:: html
-
-   <a href="https://huggingface.co/models?filter=mbart">
-       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-mbart-blueviolet">
-   </a>
-   <a href="model_doc/mbart.html">
-       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-mbart-blueviolet">
-   </a>
-
-`Multilingual Denoising Pre-training for Neural Machine Translation <https://arxiv.org/abs/2001.08210>`_ by Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov
-Marjan Ghazvininejad, Mike Lewis, Luke Zettlemoyer.
-
-The model architecture and pre-training objective is same as BART, but MBart is trained on 25 languages 
-and is intended for supervised and unsupervised machine translation. MBart is one of the first methods 
-for pre-training a complete sequence-to-sequence model by denoising full texts in multiple languages,
-
-The library provides a version of this model for conditional generation.
-
-The `mbart-large-en-ro checkpoint <https://huggingface.co/facebook/mbart-large-en-ro>`_ can be used for english -> romanian translation.
-
-The `mbart-large-cc25 <https://huggingface.co/facebook/mbart-large-cc25>`_ checkpoint can be finetuned for other translation and summarization tasks, using code in ```examples/seq2seq/``` , but is not very useful without finetuning.
-
 .. _multimodal-models:

 Multimodal models
@@ -284,12 +284,6 @@ The tokenizer also accept pre-tokenized inputs. This is particularly useful when
 predictions in `named entity recognition (NER) <https://en.wikipedia.org/wiki/Named-entity_recognition>`__ or
 `part-of-speech tagging (POS tagging) <https://en.wikipedia.org/wiki/Part-of-speech_tagging>`__.

-.. warning::
-
-    Pre-tokenized does not mean your inputs are already tokenized (you wouldn't need to pass them though the tokenizer
-    if that was the case) but just split into words (which is often the first step in subword tokenization algorithms
-    like BPE).
-
 If you want to use pre-tokenized inputs, just set :obj:`is_pretokenized=True` when passing your inputs to the
 tokenizer. For instance, we have:

@@ -52,17 +52,15 @@ Below are some of the operators which can be enabled to speed up inference throu
 * Skip connection LayerNormalization fusing
 * FastGeLU approximation

-Some of the optimizations performed by ONNX runtime can be hardware specific and thus lead to different performances
-if used on another machine with a different hardware configuration than the one used for exporting the model.
-For this reason, when using ``convert_graph_to_onnx.py`` optimizations are not enabled,
-ensuring the model can be easily exported to various hardware.
-Optimizations can then be enabled when loading the model through ONNX runtime for inference.

+Fortunately, you can let ONNXRuntime find all the possible optimized operators for you. Simply add ``--optimize``
+when exporting your model through ``convert_graph_to_onnx.py``.

-.. note::
-    When quantization is enabled (see below), ``convert_graph_to_onnx.py`` script will enable optimizations on the model
-    because quantization would modify the underlying graph making it impossible for ONNX runtime to do the optimizations
-    afterwards.
+Example:
+
+.. code-block:: bash
+
+    python convert_graph_to_onnx.py --framework <pt, tf> --model bert-base-cased --optimize bert-base-cased.onnx

 .. note::
    For more information about the optimizations enabled by ONNXRuntime, please have a look at the (`ONNXRuntime Github <https://github.com/microsoft/onnxruntime/tree/master/onnxruntime/python/tools/transformers>`_)
@@ -114,6 +112,8 @@ Example of quantized BERT model export:
    above command will contain the original ONNX model storing `float32` weights.
    The second one, with ``-quantized`` suffix, will hold the quantized parameters.

+.. note::
+    The quantization export gives the best performances when used in combination with ``--optimize``.

 TorchScript
 =======================================
@@ -106,7 +106,7 @@ The following command should work on a 16GB GPU:
    --train_batch_size=1 \
    --eval_batch_size=1 \
    --output_dir=xsum_results \
-    --num_train_epochs 6 \
+    --num_train_epochs 1 \
    --model_name_or_path facebook/bart-large
 ```

@@ -122,8 +122,8 @@ Best performing command:
 export ENRO_DIR='wmt_en_ro' # Download instructions above
 # export WANDB_PROJECT="MT" # optional
 export MAX_LEN=128
-export BS=8
-./train_mbart_cc25_enro.sh --output_dir enro_finetune_baseline_dropper --label_smoothing 0 --fp16_opt_level=O1 --logger_name wandb --sortish_sampler
+export BS=4
+./train_mbart_cc25_enro.sh --output_dir enro_finetune_baseline --label_smoothing 0.1 --fp16_opt_level=O1 --logger_name wandb --sortish_sampler
 ```
 This should take < 6h/epoch on a 16GB v100 and achieve test BLEU above 26
 To get results in line with fairseq, you need to do some postprocessing. (see `romanian_postprocessing.md`)
@@ -141,7 +141,7 @@ export BS=4
 As you train, `output_dir` will be filled with files, that look kind of like this (comments are mine).
 Some of them are metrics, some of them are checkpoints, some of them are metadata. Here is a quick tour:

-```
+```bash
 output_dir
 ├── best_tfmr  # this is a huggingface checkpoint generated by save_pretrained. It is the same model as the PL .ckpt file below
 │   ├── config.json
@@ -14,7 +14,6 @@ from torch.utils.data import DataLoader

 from lightning_base import BaseTransformer, add_generic_args, generic_train
 from transformers import MarianTokenizer, MBartTokenizer, T5ForConditionalGeneration
-from transformers.modeling_bart import shift_tokens_right


 try:
@@ -37,7 +36,6 @@ try:
    )

    from .callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
-    from .loss_dropper import LossDropper
 except ImportError:
    from utils import (
        Seq2SeqDataset,
@@ -57,21 +55,19 @@ except ImportError:
        label_smoothed_nll_loss,
    )
    from callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
-    from loss_dropper import LossDropper

 logger = logging.getLogger(__name__)


 class SummarizationModule(BaseTransformer):
    mode = "summarization"
-    loss_names = ["loss", "dropper_mask_mean"]
+    loss_names = ["loss"]
    metric_names = ROUGE_KEYS
    val_metric = "rouge2"

    def __init__(self, hparams, **kwargs):
        super().__init__(hparams, num_labels=None, mode=self.mode, **kwargs)
        use_task_specific_params(self.model, "summarization")
-        self.dropper = LossDropper(dropc=.05)
        save_git_info(self.hparams.output_dir)
        self.metrics_save_path = Path(self.output_dir) / "metrics.json"
        self.hparams_save_path = Path(self.output_dir) / "hparams.pkl"
@@ -141,10 +137,7 @@ class SummarizationModule(BaseTransformer):
        pad_token_id = self.tokenizer.pad_token_id
        source_ids, source_mask, target_ids = batch["input_ids"], batch["attention_mask"], batch["decoder_input_ids"]

-        if "labels" in batch:
-            lm_labels = batch["labels"]
-            decoder_input_ids = shift_tokens_right(lm_labels, pad_token_id)
-        elif isinstance(self.model, T5ForConditionalGeneration):
+        if isinstance(self.model, T5ForConditionalGeneration):
            decoder_input_ids = self.model._shift_right(target_ids)
            lm_labels = target_ids
        else:
@@ -152,32 +145,19 @@ class SummarizationModule(BaseTransformer):
            lm_labels = target_ids[:, 1:].clone()  # why clone?

        outputs = self(source_ids, attention_mask=source_mask, decoder_input_ids=decoder_input_ids, use_cache=False)
-        bs = source_ids.shape[0]

        if self.hparams.label_smoothing == 0:
-
            # Same behavior as modeling_bart.py
-            loss_fct = torch.nn.CrossEntropyLoss(reduction='none', ignore_index=pad_token_id)
+            loss_fct = torch.nn.CrossEntropyLoss(ignore_index=pad_token_id)
            lm_logits = outputs[0]
            assert lm_logits.shape[-1] == self.model.config.vocab_size
-
-            #loss_fct = torch.nn.NLLLoss(reduction='none', ignore_index=pad_token_id)
-            #logit_shape =
-            #weights = torch.ones(logit_shape
            loss = loss_fct(lm_logits.view(-1, lm_logits.shape[-1]), lm_labels.view(-1))
-            loss = loss.view(-1, bs)
-            loss = loss.mean(dim=0)
-            mask = self.dropper(loss)
-            loss *= mask
-            loss = loss.mean()
-            return (loss, 1-mask.mean())
-            #loss = loss.view(-1, bs)
        else:
            lprobs = torch.nn.functional.log_softmax(outputs[0], dim=-1)
            loss, nll_loss = label_smoothed_nll_loss(
                lprobs, lm_labels, self.hparams.label_smoothing, ignore_index=pad_token_id
            )
-        return (loss,torch.tensor(1.))
+        return (loss,)

    @property
    def pad(self) -> int:
@@ -322,7 +302,6 @@ class SummarizationModule(BaseTransformer):
            "--task", type=str, default="summarization", required=False, help="# examples. -1 means use all."
        )
        parser.add_argument("--label_smoothing", type=float, default=0.0, required=False)
-        parser.add_argument("--loss_dropper", type=float, default=0.0, required=False)
        parser.add_argument("--src_lang", type=str, default="", required=False)
        parser.add_argument("--tgt_lang", type=str, default="", required=False)
        parser.add_argument(
@@ -337,6 +316,7 @@ class SummarizationModule(BaseTransformer):

 class TranslationModule(SummarizationModule):
    mode = "translation"
+    loss_names = ["loss"]
    metric_names = ["bleu"]
    val_metric = "bleu"

@@ -1,53 +0,0 @@
-import numpy as np
-import torch.nn as nn
-
-
-class LossDropper(nn.Module):
-    def __init__(
-            self,
-            dropc=0.4,
-            min_count=10000,
-            recompute=10000,
-            verbose=True
-    ):
-        super().__init__()
-        self.keepc = 1. - dropc
-        self.count = 0
-        self.min_count = min_count
-
-        self.recompute = recompute
-        self.last_computed = 0
-        self.percentile_val = 100000000.
-        self.cur_idx = 0
-
-        self.verbose = verbose
-
-        self.vals = np.zeros(self.recompute, dtype=np.float32)
-
-    def forward(self, loss):
-        if loss is None:
-            return loss
-
-        self.last_computed += loss.numel()
-        self.count += loss.numel()
-        if self.count < len(self.vals):
-            self.vals[self.count - loss.numel():self.count] = loss.detach().cpu().numpy().flatten()
-            self.cur_idx += loss.numel()
-            return (loss < np.inf).type(loss.dtype)
-        else:
-            for idx, item in enumerate(loss):
-                self.vals[self.cur_idx] = item
-                self.cur_idx += 1
-                if self.cur_idx >= len(self.vals):
-                    self.cur_idx = 0
-        if self.count < self.min_count:
-            return (loss < np.inf).type(loss.dtype)
-
-        if self.last_computed > self.recompute:
-            self.percentile_val = np.percentile(self.vals, self.keepc * 100)
-            if self.verbose:
-                print('Using cutoff', self.percentile_val)
-            self.last_computed = 0
-
-        mask = (loss < self.percentile_val).type(loss.dtype)
-        return mask
@@ -6,7 +6,7 @@ export GAS=1

 python finetune.py \
    --learning_rate=3e-5 \
-    --fp16 --fp16_opt_level=O1 \
+    --fp16 \
    --gpus 1 \
    --do_train \
    --do_predict \
@@ -6,7 +6,7 @@ python distillation.py \
  --learning_rate=3e-4 \
  --do_train \
  --do_predict \
-  --fp16 --fp16_opt_level=O1 \
+  --fp16 \
  --val_check_interval 0.1 --n_val 1000 \
  --teacher facebook/bart-large-xsum --data_dir $XSUM_DIR \
  --max_target_length=60 --val_max_target_length=60 --test_max_target_length=100 \
@@ -14,4 +14,5 @@ python finetune.py \
    --task translation \
    --warmup_steps 500 \
    --freeze_embeds \
+    --model_name_or_path=facebook/mbart-large-cc25 \
    "$@"
@@ -1,73 +0,0 @@
---
-language: es
-thumbnail: https://i.imgur.com/uxAvBfh.png
-
-
---
-
-## ELECTRICIDAD: The Spanish Electra [Imgur](https://imgur.com/uxAvBfh)
-
-**Electricidad-base-discriminator** (uncased) is a ```base``` Electra like model (discriminator in this case) trained on a + 20 GB of  the [OSCAR](https://oscar-corpus.com/) Spanish corpus.
-
-As mentioned in the original [paper](https://openreview.net/pdf?id=r1xMH1BtvB):
-**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.
-
-For a detailed description and experimental results, please refer the paper [ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators](https://openreview.net/pdf?id=r1xMH1BtvB).
-
-
-## Model details ⚙
-
-|Name| # Value|
-|-----|--------|
-|Layers|	12   |
-|Hidden |768 	|
-|Params| 110M|
-
-## Evaluation metrics (for discriminator) 🧾
-
-|Metric | # Score |
-|-------|---------|
-|Accuracy| 0.985|
-|Precision| 0.726|
-|AUC | 0.922|
-
-
-
-## Fast example of usage 🚀
-
-```python
-from transformers import ElectraForPreTraining, ElectraTokenizerFast
-import torch
-
-discriminator = ElectraForPreTraining.from_pretrained("/content/electricidad-base-discriminator")
-tokenizer = ElectraTokenizerFast.from_pretrained("/content/electricidad-base-discriminator")
-
-sentence = "El rápido zorro marrón salta sobre el perro perezoso"
-fake_sentence = "El rápido zorro marrón amar sobre el perro perezoso"
-
-fake_tokens = tokenizer.tokenize(fake_sentence)
-fake_inputs = tokenizer.encode(fake_sentence, return_tensors="pt")
-discriminator_outputs = discriminator(fake_inputs)
-predictions = torch.round((torch.sign(discriminator_outputs[0]) + 1) / 2)
-
-[print("%7s" % token, end="") for token in fake_tokens]
-
-[print("%7s" % prediction, end="") for prediction in predictions.tolist()]
-
-# Output:
-'''
-el rapido  zorro  marro    ##n   amar  sobre     el  perro   pere ##zoso    0.0    0.0    0.0    0.0    0.0    0.0    1.0    1.0    0.0    0.0    0.0    0.0    0.0[None, None, None, None, None, None, None, None, None, None, None, None, None
-'''
-```
-
-As you can see there are **1s** in the places where the model detected a fake token. So, it works! 🎉
-
-## Acknowledgments
-
-I thank [🤗/transformers team](https://github.com/huggingface/transformers) for allowing me to train the model (specially to [Julien Chaumond](https://twitter.com/julien_c)).
-
-
-
-> Created by [Manuel Romero/@mrm8488](https://twitter.com/mrm8488)
-
-> Made with <span style="color: #e25555;">&hearts;</span> in Spain
@@ -1,48 +0,0 @@
---
-language: es
-thumbnail: https://i.imgur.com/uxAvBfh.png
-
-
---
-
-## ELECTRICIDAD: The Spanish Electra [Imgur](https://imgur.com/uxAvBfh)
-
-**Electricidad-base-generator** (uncased) is a ```base``` Electra like model (generator in this case) trained on a + 20 GB of  the [OSCAR](https://oscar-corpus.com/) Spanish corpus.
-
-As mentioned in the original [paper](https://openreview.net/pdf?id=r1xMH1BtvB):
-**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.
-
-For a detailed description and experimental results, please refer the paper [ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators](https://openreview.net/pdf?id=r1xMH1BtvB).
-
-
-
-
-
-## Fast example of usage 🚀
-
-```python
-from transformers import pipeline
-
-fill_mask = pipeline(
-    "fill-mask",
-    model="mrm8488/electricidad-base-generator",
-    tokenizer="mrm8488/electricidad-base-generator"
-)
-
-print(
-    fill_mask(f"HuggingFace está creando {fill_mask.tokenizer.mask_token} que la comunidad usa para resolver tareas de NLP.")
-)
-
-# Output: [{'sequence': '[CLS] huggingface esta creando herramientas que la comunidad usa para resolver tareas de nlp. [SEP]', 'score': 0.0896105170249939, 'token': 8760, 'token_str': 'herramientas'}, ...]
-
-```
-
-## Acknowledgments
-
-I thank [🤗/transformers team](https://github.com/huggingface/transformers) for allowing me to train the model (specially to [Julien Chaumond](https://twitter.com/julien_c)).
-
-
-
-> Created by [Manuel Romero/@mrm8488](https://twitter.com/mrm8488)
-
-> Made with <span style="color: #e25555;">&hearts;</span> in Spain
@@ -19,20 +19,10 @@ Paper: [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Tr

 Authors: *Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu* 

+
 ## Abstract

 Transfer learning, where a model is first pre-trained on a data-rich task before being fine-tuned on a downstream task, has emerged as a powerful technique in natural language processing (NLP). The effectiveness of transfer learning has given rise to a diversity of approaches, methodology, and practice. In this paper, we explore the landscape of transfer learning techniques for NLP by introducing a unified framework that converts every language problem into a text-to-text format. Our systematic study compares pre-training objectives, architectures, unlabeled datasets, transfer approaches, and other factors on dozens of language understanding tasks. By combining the insights from our exploration with scale and our new “Colossal Clean Crawled Corpus”, we achieve state-of-the-art results on many benchmarks covering summarization, question answering, text classification, and more. To facilitate future work on transfer learning for NLP, we release our dataset, pre-trained models, and code.

 ![model image](https://camo.githubusercontent.com/623b4dea0b653f2ad3f36c71ebfe749a677ac0a1/68747470733a2f2f6d69726f2e6d656469756d2e636f6d2f6d61782f343030362f312a44304a31674e51663876727255704b657944387750412e706e67)

-## Disclaimer
-
-Due do it's immense size, `t5-11b` requires some special treatment. 
-First, `t5-11b` should be loaded with flag `use_cdn` set to `False` as follows:
-
-```python
-t5 = transformers.T5ForConditionalGeneration.from_pretrained('t5-11b', use_cdn = False)
-```
-
-Secondly, a single GPU will most likely not have enough memory to even load the model into memory as the weights alone amount to over 40 GB.
-Model parallelism has to be used here to overcome this problem as is explained in this [PR](https://github.com/huggingface/transformers/pull/3578).
@@ -193,7 +193,7 @@
    "  \n",
    "  assert provider in get_all_providers(), f\"provider {provider} not found, {get_all_providers()}\"\n",
    "\n",
-    "  # Few properties that might have an impact on performances (provided by MS)\n",
+    "  # Few properties than might have an impact on performances (provided by MS)\n",
    "  options = SessionOptions()\n",
    "  options.intra_op_num_threads = 1\n",
    "\n",
@@ -489,4 +489,4 @@
 },
 "nbformat": 4,
 "nbformat_minor": 1
-}
+}
@@ -82,7 +82,6 @@ from .file_utils import (
    add_start_docstrings,
    cached_path,
    is_apex_available,
-    is_nlp_available,
    is_psutil_available,
    is_py3nvml_available,
    is_tf_available,
@@ -299,6 +298,7 @@ if is_torch_available():
        BartModel,
        BartForConditionalGeneration,
        BartForQuestionAnswering,
+        BART_PRETRAINED_MODEL_ARCHIVE_LIST,
    )
    from .modeling_mbart import MBartForConditionalGeneration
    from .modeling_marian import MarianMTModel
@@ -330,7 +330,6 @@ if is_torch_available():
        CamembertForMultipleChoice,
        CamembertForTokenClassification,
        CamembertForQuestionAnswering,
-        CamembertForCausalLM,
        CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
    )
    from .modeling_encoder_decoder import EncoderDecoderModel
@@ -95,8 +95,6 @@ BART_CONFIG_ARGS_DOC = r"""
            for SequenceClassification
        is_encoder_decoder (:obj:`int`, optional, defaults to True):
            True
-        force_bos_token_to_be_generated (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether or not to force BOS token to be generated at step 1 (after ``decoder_start_token_id``), only true for `bart-large-cnn`.

 """

@@ -139,7 +137,6 @@ class BartConfig(PretrainedConfig):
        normalize_embedding=True,
        static_position_embeddings=False,
        add_bias_logits=False,
-        force_bos_token_to_be_generated=False,
        **common_kwargs
    ):
        r"""
@@ -198,8 +195,6 @@ class BartConfig(PretrainedConfig):
        # pos embedding offset
        self.extra_pos_embeddings = self.pad_token_id + 1

-        self.force_bos_token_to_be_generated = force_bos_token_to_be_generated
-
    @property
    def num_attention_heads(self) -> int:
        return self.encoder_attention_heads
@@ -87,7 +87,7 @@ class EncoderDecoderConfig(PretrainedConfig):

    @classmethod
    def from_encoder_decoder_configs(
-        cls, encoder_config: PretrainedConfig, decoder_config: PretrainedConfig, **kwargs
+        cls, encoder_config: PretrainedConfig, decoder_config: PretrainedConfig
    ) -> PretrainedConfig:
        r"""
        Instantiate a :class:`~transformers.EncoderDecoderConfig` (or a derived class) from a pre-trained encoder model configuration and decoder model configuration.
@@ -99,7 +99,7 @@ class EncoderDecoderConfig(PretrainedConfig):
        decoder_config.is_decoder = True
        decoder_config.add_cross_attention = True

-        return cls(encoder=encoder_config.to_dict(), decoder=decoder_config.to_dict(), **kwargs)
+        return cls(encoder=encoder_config.to_dict(), decoder=decoder_config.to_dict())

    def to_dict(self):
        """
@@ -22,7 +22,6 @@ from .file_utils import add_start_docstrings_to_callable

 logger = logging.getLogger(__name__)

-# These config values do not vary between checkpoints
 DEFAULTS = dict(
    vocab_size=96103,
    max_position_embeddings=512,
@@ -47,47 +46,6 @@ DEFAULTS = dict(
    num_beams=8,
    activation_function="relu",
 )
-# Config values that vary between checkpoints: for testing and conversion
-max_gen_length = {
-    # See appendix C of paper
-    "xsum": 64,
-    "cnn_dailymail": 128,
-    "newsroom": 128,
-    "wikihow": 256,
-    "multi_news": 256,
-    "reddit_tifu": 128,
-    "big_patent": 256,
-    "arxiv": 256,
-    "pubmed": 256,
-    "gigaword": 32,
-    "aeslc": 32,
-    "billsum": 256,
-    "large": 256,  # @sshleifer chose arbitrarily
-}
-max_model_length = {
-    "xsum": 512,
-    "cnn_dailymail": 1024,
-    "newsroom": 512,
-    "wikihow": 512,
-    "multi_news": 1024,
-    "reddit_tifu": 512,
-    "big_patent": 1024,
-    "arxiv": 1024,
-    "pubmed": 1024,
-    "gigaword": 128,
-    "aeslc": 512,
-    "billsum": 1024,
-    "large": 1024,
-}
-expected_alpha = {
-    "multinews": 0.9,
-    "wikihow": 0.6,
-    "reddit_tifu": 0.6,
-    "big_patent": 0.7,
-    "gigaword": 0.6,
-    "aeslc": 0.6,
-    "billsum": 0.6,
-}  # otherwise 0.8


@add_start_docstrings_to_callable(BART_CONFIG_ARGS_DOC)
@@ -98,3 +56,7 @@ class PegasusConfig(BartConfig):
    """
    model_type = "pegasus"
    # The implementation of the config object is in BartConfig
+
+    @property
+    def default_config_parameters(self):
+        return DEFAULTS
@@ -58,8 +58,6 @@ class PretrainedConfig(object):
                Whether the model is used as decoder or not (in which case it's used as an encoder).
            add_cross_attention (:obj:`bool`, `optional`, defaults to :obj:`False`):
                Whether cross-attention layers should be added to the model. Note, this option is only relevant for models that can be used as decoder models within the `:class:~transformers.EncoderDecoderModel` class, which consists of all models in ``AUTO_MODELS_FOR_CAUSAL_LM``.
-            tie_encoder_decoder (:obj:`bool`, `optional`, defaults to :obj:`False`)
-                Whether all encoder weights should be tied to their equivalent decoder weights. This requires the encoder and decoder model to have the exact same parameter names.
            prune_heads (:obj:`Dict[int, List[int]]`, `optional`, defaults to :obj:`{}`):
                Pruned heads of the model. The keys are the selected layer indices and the associated values, the list
                of heads to prune in said layer.
@@ -155,7 +153,6 @@ class PretrainedConfig(object):
        self.is_encoder_decoder = kwargs.pop("is_encoder_decoder", False)
        self.is_decoder = kwargs.pop("is_decoder", False)
        self.add_cross_attention = kwargs.pop("add_cross_attention", False)
-        self.tie_encoder_decoder = kwargs.pop("tie_encoder_decoder", False)

        # Parameters for sequence generation
        self.max_length = kwargs.pop("max_length", 20)
@@ -191,7 +188,6 @@ class PretrainedConfig(object):
        self.pad_token_id = kwargs.pop("pad_token_id", None)
        self.eos_token_id = kwargs.pop("eos_token_id", None)
        self.decoder_start_token_id = kwargs.pop("decoder_start_token_id", None)
-        self.chunk_size_feed_forward = kwargs.pop("chunk_size_feed_forwar", 0)

        # task specific arguments
        self.task_specific_params = kwargs.pop("task_specific_params", None)
@@ -22,7 +22,7 @@ import torch
 from tqdm import tqdm

 from transformers import PegasusConfig, PegasusForConditionalGeneration, PegasusTokenizer
-from transformers.configuration_pegasus import DEFAULTS, expected_alpha, max_gen_length, max_model_length
+from transformers.configuration_pegasus import DEFAULTS


 PATTERNS = [
@@ -52,7 +52,47 @@ def rename_state_dict_key(k):


 # See appendix C of paper for all hyperparams
+max_gen_length = {
+    # See appendix C of paper
+    "xsum": 64,
+    "cnn_dailymail": 128,
+    "newsroom": 128,
+    "wikihow": 256,
+    "multi_news": 256,
+    "reddit_tifu": 128,
+    "big_patent": 256,
+    "arxiv": 256,
+    "pubmed": 256,
+    "gigaword": 32,
+    "aeslc": 32,
+    "billsum": 256,
+    "large": 256,  # @sshleifer chose arbitrarily
+}
+max_model_length = {
+    "xsum": 512,
+    "cnn_dailymail": 1024,
+    "newsroom": 512,
+    "wikihow": 512,
+    "multi_news": 1024,
+    "reddit_tifu": 512,
+    "big_patent": 1024,
+    "arxiv": 1024,
+    "pubmed": 1024,
+    "gigaword": 128,
+    "aeslc": 512,
+    "billsum": 1024,
+    "large": 1024,
+}

+expected_alpha = {
+    "multinews": 0.9,
+    "wikihow": 0.6,
+    "reddit_tifu": 0.6,
+    "big_patent": 0.7,
+    "gigaword": 0.6,
+    "aeslc": 0.6,
+    "billsum": 0.6,
+}  # otherwise 0.8
 # TODO(SS): one constant


@@ -111,11 +151,7 @@ def convert_pegasus_ckpt_to_pytorch(ckpt_path, save_dir):

    # convert model
    tf_weights = get_tf_weights_as_numpy(ckpt_path)
-    cfg_updates = dict(
-        max_length=max_gen_length[dataset],
-        length_penalty=expected_alpha.get(dataset, 0.8),
-        max_position_embeddings=desired_max_model_length,
-    )
+    cfg_updates = dict(max_length=max_gen_length[dataset], length_penalty=expected_alpha.get(dataset, 0.8))
    torch_model = convert_pegasus_to_bart(tf_weights, cfg_updates)
    torch_model.save_pretrained(save_dir)

@@ -62,7 +62,7 @@ def default_data_collator(features: List[InputDataClass]) -> Dict[str, torch.Ten
            if isinstance(v, torch.Tensor):
                batch[k] = torch.stack([f[k] for f in features])
            else:
-                batch[k] = torch.tensor([f[k] for f in features])
+                batch[k] = torch.tensor([f[k] for f in features], dtype=torch.long)

    return batch

@@ -64,14 +64,6 @@ except (ImportError, AssertionError):
    _tf_available = False  # pylint: disable=invalid-name


-try:
-    import nlp  # noqa: F401
-
-    _nlp_available = True
-
-except ImportError:
-    _nlp_available = False
-
 try:
    from torch.hub import _get_torch_home

@@ -152,10 +144,6 @@ def is_torch_tpu_available():
    return _torch_tpu_available


-def is_nlp_available():
-    return _nlp_available
-
-
 def is_psutil_available():
    return _psutil_available

@@ -177,7 +177,7 @@ class TFGenerationMixin:

            tokenizer = AutoTokenizer.from_pretrained('gpt2')   # Initialize tokenizer
            model = TFAutoModelWithLMHead.from_pretrained('gpt2')    # Download model and configuration from S3 and cache.
-            input_context = 'My cute dog'
+            input_context = 'My cute dog'  # "Legal" is one of the control codes for ctrl
            bad_words_ids = [tokenizer.encode(bad_word, add_prefix_space=True) for bad_word in ['idiot', 'stupid', 'shut up']]
            input_ids = tokenizer.encode(input_context, return_tensors='tf')  # encode input context
            outputs = model.generate(input_ids=input_ids, max_length=100, do_sample=True, bad_words_ids=bad_words_ids)  # generate sequences without allowing bad_words to be generated
@@ -284,7 +284,7 @@ class TFGenerationMixin:
            pad_token_id = eos_token_id

        # current position and vocab size
-        cur_len = shape_list(input_ids)[1]  # unused
+        cur_len = shape_list(input_ids)[1]
        vocab_size = self.config.vocab_size

        # set effective batch size and effective batch multiplier according to do_sample
@@ -366,8 +366,10 @@ class TFGenerationMixin:
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                bad_words_ids=bad_words_ids,
+                bos_token_id=bos_token_id,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
+                decoder_start_token_id=decoder_start_token_id,
                batch_size=effective_batch_size,
                num_return_sequences=num_return_sequences,
                length_penalty=length_penalty,
@@ -390,8 +392,10 @@ class TFGenerationMixin:
                repetition_penalty=repetition_penalty,
                no_repeat_ngram_size=no_repeat_ngram_size,
                bad_words_ids=bad_words_ids,
+                bos_token_id=bos_token_id,
                pad_token_id=pad_token_id,
                eos_token_id=eos_token_id,
+                decoder_start_token_id=decoder_start_token_id,
                batch_size=effective_batch_size,
                vocab_size=vocab_size,
                encoder_outputs=encoder_outputs,
@@ -414,8 +418,10 @@ class TFGenerationMixin:
        repetition_penalty,
        no_repeat_ngram_size,
        bad_words_ids,
+        bos_token_id,
        pad_token_id,
        eos_token_id,
+        decoder_start_token_id,
        batch_size,
        vocab_size,
        encoder_outputs,
@@ -576,7 +582,9 @@ class TFGenerationMixin:
        repetition_penalty,
        no_repeat_ngram_size,
        bad_words_ids,
+        bos_token_id,
        pad_token_id,
+        decoder_start_token_id,
        eos_token_id,
        batch_size,
        num_return_sequences,
@@ -608,7 +616,6 @@ class TFGenerationMixin:

        # cache compute states
        past = encoder_outputs
-        # to stay similar to torch : past = (encoder_outputs, None) if encoder_outputs is not None else None

        # done sentences
        done = [False for _ in range(batch_size)]
@@ -392,6 +392,7 @@ class GenerationMixin:

            # get encoder and store encoder outputs
            encoder = self.get_encoder()
+
            encoder_outputs: tuple = encoder(input_ids, attention_mask=attention_mask)

        # Expand input ids if num_beams > 1 or num_return_sequences > 1
@@ -43,7 +43,7 @@ from .modeling_outputs import (
    SequenceClassifierOutput,
    TokenClassifierOutput,
 )
-from .modeling_utils import PreTrainedModel, apply_chunking_to_forward, find_pruneable_heads_and_indices
+from .modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices


 logger = logging.getLogger(__name__)
@@ -69,7 +69,6 @@ def load_tf_weights_in_albert(model, config, tf_checkpoint_path):
    """ Load tf checkpoints in a pytorch model."""
    try:
        import re
-
        import numpy as np
        import tensorflow as tf
    except ImportError:
@@ -287,8 +286,6 @@ class AlbertLayer(nn.Module):
        super().__init__()

        self.config = config
-        self.chunk_size_feed_forward = config.chunk_size_feed_forward
-        self.seq_len_dim = 1
        self.full_layer_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
        self.attention = AlbertAttention(config)
        self.ffn = nn.Linear(config.hidden_size, config.intermediate_size)
@@ -300,20 +297,14 @@ class AlbertLayer(nn.Module):
        self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, output_hidden_states=False
    ):
        attention_output = self.attention(hidden_states, attention_mask, head_mask, output_attentions)
-
-        ffn_output = apply_chunking_to_forward(
-            self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output[0],
-        )
+        ffn_output = self.ffn(attention_output[0])
+        ffn_output = self.activation(ffn_output)
+        ffn_output = self.ffn_output(ffn_output)
+        ffn_output = self.dropout(ffn_output)
        hidden_states = self.full_layer_layer_norm(ffn_output + attention_output[0])

        return (hidden_states,) + attention_output[1:]  # add attentions if we output them

-    def ff_chunk(self, attention_output):
-        ffn_output = self.ffn(attention_output)
-        ffn_output = self.activation(ffn_output)
-        ffn_output = self.ffn_output(ffn_output)
-        return ffn_output
-

 class AlbertLayerGroup(nn.Module):
    def __init__(self, config):
@@ -73,7 +73,6 @@ from .modeling_bert import (
    BertModel,
 )
 from .modeling_camembert import (
-    CamembertForCausalLM,
    CamembertForMaskedLM,
    CamembertForMultipleChoice,
    CamembertForQuestionAnswering,
@@ -254,7 +253,6 @@ MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(

 MODEL_FOR_CAUSAL_LM_MAPPING = OrderedDict(
    [
-        (CamembertConfig, CamembertForCausalLM),
        (RobertaConfig, RobertaForCausalLM),
        (BertConfig, BertLMHeadModel),
        (OpenAIGPTConfig, OpenAIGPTLMHeadModel),
@@ -51,7 +51,15 @@ _CONFIG_FOR_DOC = "BartConfig"
 _TOKENIZER_FOR_DOC = "BartTokenizer"


-# See all BART models at https://huggingface.co/models?filter=bart
+BART_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "facebook/bart-base",
+    "facebook/bart-large",
+    "facebook/bart-large-mnli",
+    "facebook/bart-large-cnn",
+    "facebook/bart-large-xsum",
+    "facebook/mbart-large-en-ro",
+    # See all BART models at https://huggingface.co/models?filter=bart
+]


 BART_START_DOCSTRING = r"""
@@ -1015,7 +1023,6 @@ class BartForConditionalGeneration(PretrainedBartModel):

        if labels is not None:
            use_cache = False
-            decoder_input_ids = shift_tokens_right(labels, self.config.pad_token_id)

        outputs = self.model(
            input_ids,
@@ -1066,15 +1073,23 @@ class BartForConditionalGeneration(PretrainedBartModel):
        }

    def adjust_logits_during_generation(self, logits, cur_len, max_length):
-        if cur_len == 1 and self.config.force_bos_token_to_be_generated:
+        if cur_len == 1:
            self._force_token_ids_generation(logits, self.config.bos_token_id)
-        elif cur_len == max_length - 1 and self.config.eos_token_id is not None:
+        if cur_len == max_length - 1 and self.config.eos_token_id is not None:
            self._force_token_ids_generation(logits, self.config.eos_token_id)
        return logits

-    def _force_token_ids_generation(self, scores, token_id) -> None:
-        """force one of token_ids to be generated by setting prob of all other tokens to 0 (logprob=-float("inf"))"""
-        scores[:, [x for x in range(self.config.vocab_size) if x != token_id]] = -float("inf")
+    def _force_token_ids_generation(self, scores, token_ids) -> None:
+        """force one of token_ids to be generated by setting prob of all other tokens to 0"""
+        if isinstance(token_ids, int):
+            token_ids = [token_ids]
+        all_but_token_ids_mask = torch.tensor(
+            [x for x in range(self.config.vocab_size) if x not in token_ids],
+            dtype=torch.long,
+            device=next(self.parameters()).device,
+        )
+        assert len(scores.shape) == 2, "scores should be of rank 2 with shape: [batch_size, vocab_size]"
+        scores[:, all_but_token_ids_mask] = -float("inf")

    @staticmethod
    def _reorder_cache(past, beam_idx):
@@ -424,7 +424,7 @@ class BertLayer(nn.Module):
            outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights

        layer_output = apply_chunking_to_forward(
-            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+            self.chunk_size_feed_forward, self.seq_len_dim, self.feed_forward_chunk, attention_output
        )
        outputs = (layer_output,) + outputs
        return outputs
@@ -20,7 +20,6 @@ import logging
 from .configuration_camembert import CamembertConfig
 from .file_utils import add_start_docstrings
 from .modeling_roberta import (
-    RobertaForCausalLM,
    RobertaForMaskedLM,
    RobertaForMultipleChoice,
    RobertaForQuestionAnswering,
@@ -134,15 +133,3 @@ class CamembertForQuestionAnswering(RobertaForQuestionAnswering):
    """

    config_class = CamembertConfig
-
-
-@add_start_docstrings(
-    """CamemBERT Model with a `language modeling` head on top for CLM fine-tuning. """, CAMEMBERT_START_DOCSTRING
-)
-class CamembertForCausalLM(RobertaForCausalLM):
-    """
-    This class overrides :class:`~transformers.RobertaForCausalLM`. Please check the
-    superclass for the appropriate documentation alongside usage examples.
-    """
-
-    config_class = CamembertConfig
@@ -44,12 +44,7 @@ from .modeling_outputs import (
    SequenceClassifierOutput,
    TokenClassifierOutput,
 )
-from .modeling_utils import (
-    PreTrainedModel,
-    apply_chunking_to_forward,
-    find_pruneable_heads_and_indices,
-    prune_linear_layer,
-)
+from .modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer


 logger = logging.getLogger(__name__)
@@ -213,8 +208,6 @@ class FFN(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.dropout = nn.Dropout(p=config.dropout)
-        self.chunk_size_feed_forward = config.chunk_size_feed_forward
-        self.seq_len_dim = 1
        self.lin1 = nn.Linear(in_features=config.dim, out_features=config.hidden_dim)
        self.lin2 = nn.Linear(in_features=config.hidden_dim, out_features=config.dim)
        assert config.activation in ["relu", "gelu"], "activation ({}) must be in ['relu', 'gelu']".format(
@@ -223,9 +216,6 @@ class FFN(nn.Module):
        self.activation = gelu if config.activation == "gelu" else nn.ReLU()

    def forward(self, input):
-        return apply_chunking_to_forward(self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, input)
-
-    def ff_chunk(self, input):
        x = self.lin1(input)
        x = self.activation(x)
        x = self.lin2(x)
@@ -71,17 +71,9 @@ class EncoderDecoderModel(PreTrainedModel):
            self.encoder.get_output_embeddings() is None
        ), "The encoder {} should not have a LM Head. Please use a model without LM Head"

-        # tie encoder, decoder weights if config set accordingly
-        self.tie_weights()
-
    def tie_weights(self):
-        # tie encoder & decoder if needed
-        if self.config.tie_encoder_decoder:
-            # tie encoder and decoder base model
-            decoder_base_model_prefix = self.decoder.base_model_prefix
-            self._tie_encoder_decoder_weights(
-                self.encoder, self.decoder._modules[decoder_base_model_prefix], self.decoder.base_model_prefix
-            )
+        # for now no weights tying in encoder-decoder
+        pass

    def get_encoder(self):
        return self.encoder
@@ -130,11 +122,7 @@ class EncoderDecoderModel(PreTrainedModel):
                All remaning positional arguments will be passed to the underlying model's ``__init__`` method

            kwargs: (`optional`) Remaining dictionary of keyword arguments.
-                Can be used to update the configuration object (after it being loaded) and initiate the model. (e.g. ``output_attention=True``).
-                - To update the encoder configuration, use the prefix `encoder_` for each configuration parameter
-                - To update the decoder configuration, use the prefix `decoder_` for each configuration parameter
-                - To update the parent model configuration, do not use a prefix for each configuration parameter
-                Behave differently depending on whether a :obj:`config` is provided or automatically loaded.
+                Can be used to update the configuration object (after it being loaded) and initiate the model. (e.g. ``output_attention=True``). Behave differently depending on whether a `config` is provided or automatically loaded:

        Examples::

@@ -156,12 +144,6 @@ class EncoderDecoderModel(PreTrainedModel):
            argument[len("decoder_") :]: value for argument, value in kwargs.items() if argument.startswith("decoder_")
        }

-        # remove encoder, decoder kwargs from kwargs
-        for key in kwargs_encoder.keys():
-            del kwargs["encoder_" + key]
-        for key in kwargs_decoder.keys():
-            del kwargs["decoder_" + key]
-
        # Load and initialize the encoder and decoder
        # The distinction between encoder and decoder at the model level is made
        # by the value of the flag `is_decoder` that we need to set correctly.
@@ -202,9 +184,7 @@ class EncoderDecoderModel(PreTrainedModel):

            decoder = AutoModelForCausalLM.from_pretrained(decoder_pretrained_model_name_or_path, **kwargs_decoder)

-        # instantiate config with corresponding kwargs
-        config = EncoderDecoderConfig.from_encoder_decoder_configs(encoder.config, decoder.config, **kwargs)
-        return cls(encoder=encoder, decoder=decoder, config=config)
+        return cls(encoder=encoder, decoder=decoder)

    def forward(
        self,
@@ -30,7 +30,7 @@ from .file_utils import (
    add_start_docstrings_to_callable,
    replace_return_docstrings,
 )
-from .modeling_bert import BertIntermediate, BertLayerNorm, BertOutput, BertPooler, BertPreTrainedModel, BertSelfOutput
+from .modeling_bert import BertIntermediate, BertLayerNorm, BertOutput, BertPreTrainedModel, BertSelfOutput
 from .modeling_outputs import (
    BaseModelOutput,
    BaseModelOutputWithPooling,
@@ -41,12 +41,7 @@ from .modeling_outputs import (
    TokenClassifierOutput,
 )
 from .modeling_roberta import RobertaEmbeddings, RobertaLMHead
-from .modeling_utils import (
-    PreTrainedModel,
-    apply_chunking_to_forward,
-    find_pruneable_heads_and_indices,
-    prune_linear_layer,
-)
+from .modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer


 logger = logging.getLogger(__name__)
@@ -690,8 +685,6 @@ class LongformerLayer(nn.Module):
        self.attention = LongformerAttention(config, layer_id)
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)
-        self.chunk_size_feed_forward = config.chunk_size_feed_forward
-        self.seq_len_dim = 1

    def forward(
        self, hidden_states, attention_mask=None, output_attentions=False,
@@ -700,16 +693,10 @@ class LongformerLayer(nn.Module):
        attn_output = self_attn_outputs[0]
        outputs = self_attn_outputs[1:]  # add self attentions if we output attention weights

-        layer_output = apply_chunking_to_forward(
-            self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attn_output
-        )
-        outputs = (layer_output,) + outputs
-        return outputs
-
-    def ff_chunk(self, attn_output):
        intermediate_output = self.intermediate(attn_output)
        layer_output = self.output(intermediate_output, attn_output)
-        return layer_output
+        outputs = (layer_output,) + outputs
+        return outputs


 class LongformerEncoder(nn.Module):
@@ -885,7 +872,6 @@ class LongformerModel(LongformerPreTrainedModel):

        self.embeddings = RobertaEmbeddings(config)
        self.encoder = LongformerEncoder(config)
-        self.pooler = BertPooler(config)

        self.init_weights()

@@ -1051,7 +1037,7 @@ class LongformerModel(LongformerPreTrainedModel):
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
-        pooled_output = self.pooler(sequence_output)
+        pooled_output = sequence_output[:, 0]

        # undo padding
        if padding_len > 0:
@@ -47,7 +47,7 @@ class MarianMTModel(BartForConditionalGeneration):
    """

    def adjust_logits_during_generation(self, logits, cur_len, max_length):
-        logits[:, self.config.pad_token_id] = float("-inf")  # never predict pad token.
+        logits[:, self.config.pad_token_id] = float("-inf")
        if cur_len == max_length - 1 and self.config.eos_token_id is not None:
            self._force_token_ids_generation(logits, self.config.eos_token_id)
        return logits
@@ -23,13 +23,6 @@ from .modeling_bart import BART_START_DOCSTRING, BartForConditionalGeneration
@add_start_docstrings("The Pegasus Model for summarization ", BART_START_DOCSTRING)
 class PegasusForConditionalGeneration(BartForConditionalGeneration):
    config_class = PegasusConfig
-    authorized_missing_keys = [
-        r"final_logits_bias",
-        r"encoder\.version",
-        r"decoder\.version",
-        r"model.encoder.embed_positions",
-        "model.decoder.embed_positions",
-    ]
    r"""
    Pytorch version of google's pegasus model for summarization.
    Model API is identical to BartForConditionalGeneration.
@@ -1369,7 +1369,7 @@ class ChunkReformerFeedForward(nn.Module):

    def forward(self, attention_output):
        return apply_chunking_to_forward(
-            self.forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output,
+            self.chunk_size_feed_forward, self.seq_len_dim, self.forward_chunk, attention_output,
        )

    def forward_chunk(self, hidden_states):
@@ -1730,7 +1730,7 @@ class ReformerOnlyLMHead(nn.Module):
        self.decoder.bias = self.bias

    def forward(self, hidden_states):
-        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
+        return apply_chunking_to_forward(self.chunk_size_lm_head, self.seq_len_dim, self.forward_chunk, hidden_states)

    def forward_chunk(self, hidden_states):
        hidden_states = self.decoder(hidden_states)
@@ -887,12 +887,10 @@ class T5Model(T5PreTrainedModel):

        encoder_config = copy.deepcopy(config)
        encoder_config.use_cache = False
-        encoder_config.is_encoder_decoder = False
        self.encoder = T5Stack(encoder_config, self.shared)

        decoder_config = copy.deepcopy(config)
        decoder_config.is_decoder = True
-        decoder_config.is_encoder_decoder = False
        self.decoder = T5Stack(decoder_config, self.shared)

        self.init_weights()
@@ -1042,12 +1040,10 @@ class T5ForConditionalGeneration(T5PreTrainedModel):

        encoder_config = copy.deepcopy(config)
        encoder_config.use_cache = False
-        encoder_config.is_encoder_decoder = False
        self.encoder = T5Stack(encoder_config, self.shared)

        decoder_config = copy.deepcopy(config)
        decoder_config.is_decoder = True
-        decoder_config.is_encoder_decoder = False
        self.decoder = T5Stack(decoder_config, self.shared)

        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
@@ -416,77 +416,6 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
        if output_embeddings is not None:
            self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings())

-        if self.config.is_encoder_decoder and self.config.tie_encoder_decoder:
-            self._tie_encoder_decoder_weights(self.encoder, self.decoder, self.base_model_prefix)
-
-    @staticmethod
-    def _tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str):
-        uninitialized_encoder_weights: List[str] = []
-        assert decoder.__class__ == encoder.__class__, f"{decoder.__class__} and {encoder.__class__} have to be equal."
-
-        def tie_encoder_to_decoder_recursively(
-            decoder_pointer: nn.Module,
-            encoder_pointer: nn.Module,
-            module_name: str,
-            uninitialized_encoder_weights: List[str],
-            depth=0,
-        ):
-            assert isinstance(decoder_pointer, nn.Module) and isinstance(
-                encoder_pointer, nn.Module
-            ), f"{decoder_pointer} and {encoder_pointer} have to be of type torch.nn.Module"
-            if hasattr(decoder_pointer, "weight"):
-                assert hasattr(encoder_pointer, "weight")
-                encoder_pointer.weight = decoder_pointer.weight
-                if hasattr(decoder_pointer, "bias"):
-                    assert hasattr(encoder_pointer, "bias")
-                    encoder_pointer.bias = decoder_pointer.bias
-                return
-
-            encoder_modules = encoder_pointer._modules
-            decoder_modules = decoder_pointer._modules
-            if len(decoder_modules) > 0:
-                assert (
-                    len(encoder_modules) > 0
-                ), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}"
-
-                all_encoder_weights = set([module_name + "/" + sub_name for sub_name in encoder_modules.keys()])
-                encoder_layer_pos = 0
-                for name, module in decoder_modules.items():
-                    if name.isdigit():
-                        encoder_name = str(int(name) + encoder_layer_pos)
-                        decoder_name = name
-                        if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])):
-                            # this can happen if the name corresponds to the position in a list module list of layers
-                            # in this case the decoder has added a cross-attention that the encoder does not have
-                            # thus skip this step and substract one layer pos from encoder
-                            encoder_layer_pos -= 1
-                            continue
-                    elif name not in encoder_modules:
-                        continue
-                    elif depth > 500:
-                        raise ValueError(
-                            "Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is a circular dependency between two or more `nn.Modules` of your model."
-                        )
-                    else:
-                        decoder_name = encoder_name = name
-                    tie_encoder_to_decoder_recursively(
-                        decoder_modules[decoder_name],
-                        encoder_modules[encoder_name],
-                        module_name + "/" + name,
-                        uninitialized_encoder_weights,
-                        depth=depth + 1,
-                    )
-                    all_encoder_weights.remove(module_name + "/" + encoder_name)
-
-                uninitialized_encoder_weights += list(all_encoder_weights)
-
-        # tie weights recursively
-        tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights)
-        if len(uninitialized_encoder_weights) > 0:
-            logger.warning(
-                f"The following encoder weights were not tied to the decoder {uninitialized_encoder_weights}"
-            )
-
    def _tie_or_clone_weights(self, output_embeddings, input_embeddings):
        """ Tie or clone module weights depending of whether we are using TorchScript or not
        """
@@ -965,8 +894,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
                        model.__class__.__name__, "\n\t".join(error_msgs)
                    )
                )
-        # make sure token embedding weights are still tied if needed
-        model.tie_weights()
+        model.tie_weights()  # make sure token embedding weights are still tied if needed

        # Set model in evaluation mode to deactivate DropOut modules by default
        model.eval()
@@ -1519,7 +1447,7 @@ def prune_layer(


 def apply_chunking_to_forward(
-    forward_fn: Callable[..., torch.Tensor], chunk_size: int, chunk_dim: int, *input_tensors
+    chunk_size: int, chunk_dim: int, forward_fn: Callable[..., torch.Tensor], *input_tensors
 ) -> torch.Tensor:
    """
    This function chunks the :obj:`input_tensors` into smaller input tensor parts of size :obj:`chunk_size` over the
@@ -1529,12 +1457,12 @@ def apply_chunking_to_forward(
    directly applying :obj:`forward_fn` to :obj:`input_tensors`.

    Args:
-        forward_fn (:obj:`Callable[..., torch.Tensor]`):
-            The forward function of the model.
        chunk_size (:obj:`int`):
            The chunk size of a chunked tensor: :obj:`num_chunks = len(input_tensors[0]) / chunk_size`.
        chunk_dim (:obj:`int`):
            The dimension over which the :obj:`input_tensors` should be chunked.
+        forward_fn (:obj:`Callable[..., torch.Tensor]`):
+            The forward function of the model.
        input_tensors (:obj:`Tuple[torch.Tensor]`):
            The input tensors of ``forward_fn`` which will be chunked.
    Returns:
@@ -1550,7 +1478,7 @@ def apply_chunking_to_forward(

        # implement a chunked forward function
        def forward(self, hidden_states):
-            return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
+            return apply_chunking_to_forward(self.chunk_size_lm_head, self.seq_len_dim, self.forward_chunk, hidden_states)
    """

    assert len(input_tensors) > 0, "{} has to be a tuple/list of tensors".format(input_tensors)
@@ -50,7 +50,6 @@ from .modeling_utils import (
    PreTrainedModel,
    SequenceSummary,
    SQuADHead,
-    apply_chunking_to_forward,
    find_pruneable_heads_and_indices,
    prune_linear_layer,
 )
@@ -213,13 +212,8 @@ class TransformerFFN(nn.Module):
        self.lin1 = nn.Linear(in_dim, dim_hidden)
        self.lin2 = nn.Linear(dim_hidden, out_dim)
        self.act = gelu if config.gelu_activation else F.relu
-        self.chunk_size_feed_forward = config.chunk_size_feed_forward
-        self.seq_len_dim = 1

    def forward(self, input):
-        return apply_chunking_to_forward(self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, input)
-
-    def ff_chunk(self, input):
        x = self.lin1(input)
        x = self.act(x)
        x = self.lin2(x)
@@ -35,14 +35,7 @@ from .file_utils import (
    add_start_docstrings_to_callable,
    replace_return_docstrings,
 )
-from .modeling_utils import (
-    PoolerAnswerClass,
-    PoolerEndLogits,
-    PoolerStartLogits,
-    PreTrainedModel,
-    SequenceSummary,
-    apply_chunking_to_forward,
-)
+from .modeling_utils import PoolerAnswerClass, PoolerEndLogits, PoolerStartLogits, PreTrainedModel, SequenceSummary


 logger = logging.getLogger(__name__)
@@ -446,8 +439,7 @@ class XLNetRelativeAttention(nn.Module):
            v_head_h = torch.einsum("ibh,hnd->ibnd", cat, self.v)

            # positional heads
-            # type casting for fp16 support
-            k_head_r = torch.einsum("ibh,hnd->ibnd", r.type(self.r.dtype), self.r)
+            k_head_r = torch.einsum("ibh,hnd->ibnd", r, self.r)

            # core attention ops
            attn_vec = self.rel_attn_core(
@@ -503,8 +495,6 @@ class XLNetLayer(nn.Module):
        self.rel_attn = XLNetRelativeAttention(config)
        self.ff = XLNetFeedForward(config)
        self.dropout = nn.Dropout(config.dropout)
-        self.chunk_size_feed_forward = config.chunk_size_feed_forward
-        self.seq_len_dim = 1

    def forward(
        self,
@@ -534,18 +524,12 @@ class XLNetLayer(nn.Module):
        output_h, output_g = outputs[:2]

        if output_g is not None:
-            output_g = apply_chunking_to_forward(
-                self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, output_g
-            )
-        output_h = apply_chunking_to_forward(self.ff_chunk, self.chunk_size_feed_forward, self.seq_len_dim, output_h)
+            output_g = self.ff(output_g)
+        output_h = self.ff(output_h)

        outputs = (output_h, output_g) + outputs[2:]  # Add again attentions if there are there
        return outputs

-    def ff_chunk(self, output_x):
-        output_x = self.ff(output_x)
-        return output_x
-

 class XLNetPreTrainedModel(PreTrainedModel):
    """ An abstract class to handle weights initialization and
@@ -33,7 +33,6 @@ _all_bart_models = [
    "facebook/bart-large-cnn",
    "facebook/bart-large-xsum",
    "yjernite/bart_eli5",
-    # This is not exhaustive: see https://huggingface.co/models?filter=bart
 ]


@@ -134,7 +133,7 @@ class BartTokenizer(RobertaTokenizer):
        # Process tgt_texts
        if max_target_length is None:
            max_target_length = max_length
-        labels = self(
+        decoder_inputs: BatchEncoding = self(
            tgt_texts,
            add_special_tokens=True,
            return_tensors=return_tensors,
@@ -142,8 +141,10 @@ class BartTokenizer(RobertaTokenizer):
            max_length=max_target_length,
            truncation=truncation,
            **kwargs,
-        )["input_ids"]
-        model_inputs["labels"] = labels
+        )
+        for k, v in decoder_inputs.items():
+            model_inputs[f"decoder_{k}"] = v
+
        return model_inputs


@@ -154,104 +155,3 @@ class BartTokenizerFast(RobertaTokenizerFast):
        "vocab_file": {m: vocab_url for m in _all_bart_models},
        "merges_file": {m: merges_url for m in _all_bart_models},
    }
-
-    def prepare_seq2seq_batch(
-        self,
-        src_texts: List[str],
-        tgt_texts: Optional[List[str]] = None,
-        max_length: Optional[int] = None,
-        max_target_length: Optional[int] = None,
-        padding: str = "longest",
-        return_tensors: str = "None",
-        truncation=True,
-        **kwargs,
-    ) -> BatchEncoding:
-        r"""
-
-        Prepare a batch that can be passed directly to an instance of :class:`~transformers.BartModel`.
-
-        Args:
-            src_texts: (:obj:`List[str]`):
-                List of documents to summarize or source language texts.
-            tgt_texts: (:obj:`List[str]`, `optional`):
-                List of summaries or target language texts.
-            max_length (:obj:`int`, `optional`):
-                Controls the maximum length for encoder inputs (documents to summarize or source language texts).
-                If left unset or set to :obj:`None`, this will use the predefined model maximum length if a maximum
-                length is required by one of the truncation/padding parameters. If the model has no specific maximum
-                input length (like XLNet) truncation/padding to a maximum length will be deactivated.
-            max_target_length (:obj:`int`, `optional`):
-                Controls the maximum length of decoder inputs (target language texts or summaries).
-                If left unset or set to :obj:`None`, this will use the max_length value.
-            padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`False`):
-                Activates and controls padding. Accepts the following values:
-
-                * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a
-                  single sequence if provided).
-                * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
-                  maximum acceptable input length for the model if that argument is not provided.
-                * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
-                  different lengths).
-            return_tensors (:obj:`str` or :class:`~transformers.tokenization_utils_base.TensorType`, `optional`, defaults to "pt"):
-                If set, will return tensors instead of list of python integers. Acceptable values are:
-
-                * :obj:`'tf'`: Return TensorFlow :obj:`tf.constant` objects.
-                * :obj:`'pt'`: Return PyTorch :obj:`torch.Tensor` objects.
-                * :obj:`'np'`: Return Numpy :obj:`np.ndarray` objects.
-            truncation (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.TruncationStrategy`, `optional`, defaults to :obj:`True`):
-                Activates and controls truncation. Accepts the following values:
-
-                * :obj:`True` or :obj:`'longest_first'`: Truncate to a maximum length specified with the argument
-                  :obj:`max_length` or to the maximum acceptable input length for the model if that argument is not
-                  provided. This will truncate token by token, removing a token from the longest sequence in the pair
-                  if a pair of sequences (or a batch of pairs) is provided.
-                * :obj:`'only_first'`: Truncate to a maximum length specified with the argument :obj:`max_length` or to
-                  the maximum acceptable input length for the model if that argument is not provided. This will only
-                  truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
-                * :obj:`'only_second'`: Truncate to a maximum length specified with the argument :obj:`max_length` or
-                  to the maximum acceptable input length for the model if that argument is not provided. This will only
-                  truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
-                * :obj:`False` or :obj:`'do_not_truncate'` (default): No truncation (i.e., can output batch with
-                  sequence lengths greater than the model maximum admissible input size).
-            **kwargs:
-                Additional keyword arguments passed along to :obj:`self.__call__`.
-
-        Returns:
-            :class:`~transformers.BatchEncoding`: A :class:`~transformers.BatchEncoding` with the following fields:
-
-            - **input_ids** -- List of token ids to be fed to the encoder.
-            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model.
-            - **decoder_input_ids** -- List of token ids to be fed to the decoder.
-            - **decoder_attention_mask** -- List of indices specifying which tokens should be attended to by the decoder.
-                This does not include causal mask, which is built by the model.
-
-            The full set of keys ``[input_ids, attention_mask, decoder_input_ids,  decoder_attention_mask]``,
-            will only be returned if tgt_texts is passed. Otherwise, input_ids, attention_mask will be the only keys.
-        """
-        if max_length is None:
-            max_length = self.model_max_length
-        model_inputs: BatchEncoding = self(
-            src_texts,
-            add_special_tokens=True,
-            return_tensors=return_tensors,
-            max_length=max_length,
-            padding=padding,
-            truncation=truncation,
-            **kwargs,
-        )
-        if tgt_texts is None:
-            return model_inputs
-        # Process tgt_texts
-        if max_target_length is None:
-            max_target_length = max_length
-        labels = self(
-            tgt_texts,
-            add_special_tokens=True,
-            return_tensors=return_tensors,
-            padding=padding,
-            max_length=max_target_length,
-            truncation=truncation,
-            **kwargs,
-        )["input_ids"]
-        model_inputs["labels"] = labels
-        return model_inputs
@@ -56,15 +56,6 @@ FAIRSEQ_LANGUAGE_CODES = [
 ]


-def shift_tokens_right(input_ids, pad_token_id):
-    """Shift input ids one token to the right, and wrap the last non pad token (usually <eos>)."""
-    prev_output_tokens = input_ids.clone()
-    index_of_eos = (input_ids.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
-    prev_output_tokens[:, 0] = input_ids.gather(1, index_of_eos).squeeze()
-    prev_output_tokens[:, 1:] = input_ids[:, :-1]
-    return prev_output_tokens
-
-
 class MBartTokenizer(XLMRobertaTokenizer):
    """
    This inherits from XLMRobertaTokenizer. ``prepare_seq2seq_batch`` should be used to encode inputs.
@@ -107,6 +98,32 @@ class MBartTokenizer(XLMRobertaTokenizer):
        self._additional_special_tokens = list(self.lang_code_to_id.keys())
        self.set_src_lang_special_tokens(kwargs.get("src_lang", "en_XX"))

+    def build_inputs_with_special_tokens(
+        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
+    ) -> List[int]:
+        """
+        Build model inputs from a sequence or a pair of sequence for sequence classification tasks
+        by concatenating and adding special tokens. The special tokens depend on calling set_lang.
+        An MBART sequence has the following format, where ``X`` represents the sequence:
+        - ``input_ids`` (for encoder) ``X [eos, src_lang_code]``
+        - ``decoder_input_ids``: (for decoder) ``[tgt_lang_code] X [eos]``
+        BOS is never used.
+        Pairs of sequences are not the expected use case, but they will be handled without a separator.
+
+        Args:
+            token_ids_0 (:obj:`List[int]`):
+                List of IDs to which the special tokens will be added
+            token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`):
+                Optional second list of IDs for sequence pairs.
+
+        Returns:
+            :obj:`List[int]`: list of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
+        """
+        if token_ids_1 is None:
+            return self.prefix_tokens + token_ids_0 + self.suffix_tokens
+        # We don't expect to process pairs, but leave the pair logic for API consistency
+        return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
+
    def get_special_tokens_mask(
        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None, already_has_special_tokens: bool = False
    ) -> List[int]:
@@ -139,32 +156,6 @@ class MBartTokenizer(XLMRobertaTokenizer):
            return prefix_ones + ([0] * len(token_ids_0)) + suffix_ones
        return prefix_ones + ([0] * len(token_ids_0)) + ([0] * len(token_ids_1)) + suffix_ones

-    def build_inputs_with_special_tokens(
-        self, token_ids_0: List[int], token_ids_1: Optional[List[int]] = None
-    ) -> List[int]:
-        """
-        Build model inputs from a sequence or a pair of sequence for sequence classification tasks
-        by concatenating and adding special tokens. The special tokens depend on calling set_lang.
-        An MBART sequence has the following format, where ``X`` represents the sequence:
-        - ``input_ids`` (for encoder) ``X [eos, src_lang_code]``
-        - ``decoder_input_ids``: (for decoder) ``[tgt_lang_code] X [eos]``
-        BOS is never used.
-        Pairs of sequences are not the expected use case, but they will be handled without a separator.
-
-        Args:
-            token_ids_0 (:obj:`List[int]`):
-                List of IDs to which the special tokens will be added
-            token_ids_1 (:obj:`List[int]`, `optional`, defaults to :obj:`None`):
-                Optional second list of IDs for sequence pairs.
-
-        Returns:
-            :obj:`List[int]`: list of `input IDs <../glossary.html#input-ids>`__ with the appropriate special tokens.
-        """
-        if token_ids_1 is None:
-            return self.prefix_tokens + token_ids_0 + self.suffix_tokens
-        # We don't expect to process pairs, but leave the pair logic for API consistency
-        return self.prefix_tokens + token_ids_0 + token_ids_1 + self.suffix_tokens
-
    @add_start_docstrings_to_callable(PREPARE_SEQ2SEQ_BATCH_DOCSTRING)
    def prepare_seq2seq_batch(
        self,
@@ -260,8 +251,7 @@ class MBartTokenizer(XLMRobertaTokenizer):
        if max_target_length is None:
            max_target_length = max_length
        self.set_tgt_lang_special_tokens(tgt_lang)
-
-        labels = self(
+        decoder_inputs: BatchEncoding = self(
            tgt_texts,
            add_special_tokens=True,
            return_tensors=return_tensors,
@@ -269,9 +259,10 @@ class MBartTokenizer(XLMRobertaTokenizer):
            max_length=max_target_length,
            truncation=True,
            **kwargs,
-        )["input_ids"]
-        model_inputs["decoder_input_ids"] = shift_tokens_right(labels, self.pad_token_id)
-        model_inputs["labels"] = labels
+        )
+        for k, v in decoder_inputs.items():
+            model_inputs[f"decoder_{k}"] = v
+
        self.set_src_lang_special_tokens(src_lang)  # sets to src_lang
        return model_inputs

@@ -284,5 +275,5 @@ class MBartTokenizer(XLMRobertaTokenizer):
    def set_tgt_lang_special_tokens(self, lang: str) -> None:
        """Reset the special tokens to the target language setting. Prefix [tgt_lang_code], suffix =[eos]."""
        self.cur_lang_code = self.lang_code_to_id[lang]
-        self.prefix_tokens = []
-        self.suffix_tokens = [self.eos_token_id, self.cur_lang_code]
+        self.prefix_tokens = [self.cur_lang_code]
+        self.suffix_tokens = [self.eos_token_id]
@@ -133,8 +133,7 @@ class PegasusTokenizer(ReformerTokenizer):
            return model_inputs
        if max_target_length is not None:
            tokenizer_kwargs["max_length"] = max_target_length
-        labels: BatchEncoding = self(tgt_texts, **tokenizer_kwargs)["input_ids"]
-        model_inputs["labels"] = labels
-        # for k, v in decoder_inputs.items():
-        #    model_inputs[f"decoder_{k}"] = v
+        decoder_inputs: BatchEncoding = self(tgt_texts, **tokenizer_kwargs)
+        for k, v in decoder_inputs.items():
+            model_inputs[f"decoder_{k}"] = v
        return model_inputs
@@ -303,7 +303,7 @@ class T5Tokenizer(PreTrainedTokenizer):
        if max_length is None:
            max_length = self.max_len
        self.prefix_tokens = []
-        model_inputs = self(
+        model_inputs: BatchEncoding = self(
            src_texts,
            add_special_tokens=True,
            return_tensors=return_tensors,
@@ -319,7 +319,7 @@ class T5Tokenizer(PreTrainedTokenizer):
            max_target_length = max_length
        # set prefix_tokens for target text
        self.prefix_tokens = [self.pad_token_id]
-        model_inputs["labels"] = self(
+        decoder_inputs: BatchEncoding = self(
            tgt_texts,
            add_special_tokens=True,
            return_tensors=return_tensors,
@@ -327,7 +327,9 @@ class T5Tokenizer(PreTrainedTokenizer):
            max_length=max_target_length,
            truncation=truncation,
            **kwargs,
-        )["input_ids"]
+        )
+        for k, v in decoder_inputs.items():
+            model_inputs[f"decoder_{k}"] = v

        self.prefix_tokens = []
        return model_inputs
@@ -1088,8 +1088,7 @@ ENCODE_KWARGS_DOCSTRING = r"""
                returned to provide some overlap between truncated and overflowing sequences. The value of this
                argument defines the number of overlapping tokens.
            is_pretokenized (:obj:`bool`, `optional`, defaults to :obj:`False`):
-                Whether or not the input is already pre-tokenized (e.g., split into words), in which case the tokenizer
-                will skip the pre-tokenization step. This is useful for NER or token classification.
+                Whether or not the input is already tokenized.
            pad_to_multiple_of (:obj:`int`, `optional`):
                If set will pad the sequence to a multiple of the provided value. This is especially useful to enable
                the use of Tensor Cores on NVIDIA hardware with compute capability >= 7.5 (Volta).
@@ -1,4 +1,3 @@
-import inspect
 import logging
 import math
 import os
@@ -20,7 +19,7 @@ from torch.utils.data.sampler import RandomSampler, Sampler, SequentialSampler
 from tqdm.auto import tqdm, trange

 from .data.data_collator import DataCollator, default_data_collator
-from .file_utils import is_nlp_available, is_torch_tpu_available
+from .file_utils import is_torch_tpu_available
 from .integrations import is_comet_available, is_tensorboard_available, is_wandb_available
 from .modeling_utils import PreTrainedModel
 from .optimization import AdamW, get_linear_schedule_with_warmup
@@ -42,8 +41,6 @@ else:
    _use_native_amp = True
    from torch.cuda.amp import autocast

-if is_nlp_available():
-    import nlp

 if is_torch_tpu_available():
    import torch_xla.core.xla_model as xm
@@ -143,19 +140,19 @@ class Trainer:
        model (:class:`~transformers.PreTrainedModel`):
            The model to train, evaluate or use for predictions.
        args (:class:`~transformers.TrainingArguments`):
-            The arguments to tweak for training.
+            The arguments to tweak training.
        data_collator (:obj:`DataCollator`, `optional`, defaults to :func:`~transformers.default_data_collator`):
-            The function to use to form a batch from a list of elements of :obj:`train_dataset` or
+            The function to use to from a batch from a list of elements of :obj:`train_dataset` or
            :obj:`eval_dataset`.
        train_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
-            The dataset to use for training. If it is an :obj:`nlp.Dataset`, columns not accepted by the
-            ``model.forward()`` method are automatically removed.
+            The dataset to use for training.
        eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
-            The dataset to use for evaluation. If it is an :obj:`nlp.Dataset`, columns not accepted by the
-            ``model.forward()`` method are automatically removed.
+            The dataset to use for evaluation.
        compute_metrics (:obj:`Callable[[EvalPrediction], Dict]`, `optional`):
            The function that will be used to compute metrics at evaluation. Must take a
            :class:`~transformers.EvalPrediction` and return a dictionary string to metric values.
+        prediction_loss_only (:obj:`bool`, `optional`, defaults to `False`):
+            When performing evaluation and predictions, only returns the loss.
        tb_writer (:obj:`SummaryWriter`, `optional`):
            Object to write to TensorBoard.
        optimizers (:obj:`Tuple[torch.optim.Optimizer, torch.optim.lr_scheduler.LambdaLR`, `optional`):
@@ -231,32 +228,11 @@ class Trainer:
                ),
                FutureWarning,
            )
-
-        if is_nlp_available():
-            if isinstance(train_dataset, nlp.Dataset):
-                self._remove_unused_columns(self.train_dataset, description="training")
-            if isinstance(eval_dataset, nlp.Dataset):
-                self._remove_unused_columns(self.eval_dataset, description="evaluation")
-
        self.global_step = None
        self.epoch = None
        if self.args.fp16 and _use_native_amp:
            self.scaler = torch.cuda.amp.GradScaler()

-    def _remove_unused_columns(self, dataset: "nlp.Dataset", description: Optional[str] = None):
-        # Inspect model forward signature to keep only the arguments it accepts.
-        signature = inspect.signature(self.model.forward)
-        signature_columns = list(signature.parameters.keys())
-        # Labels may be named label or label_ids, the default data collator handles that.
-        signature_columns += ["label", "label_ids"]
-        columns = [k for k in signature_columns if k in dataset.column_names]
-        ignored_columns = list(set(dataset.column_names) - set(signature_columns))
-        dset_description = "" if description is None else f"in the {description} set "
-        logger.info(
-            f"The following columns {dset_description}don't have a corresponding argument in `{self.model.__class__.__name__}.forward` and have been ignored: {', '.join(ignored_columns)}."
-        )
-        dataset.set_format(columns=columns)
-
    def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
        if isinstance(self.train_dataset, torch.utils.data.IterableDataset):
            return None
@@ -311,13 +287,11 @@ class Trainer:

        Args:
            eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
-                If provided, will override :obj:`self.eval_dataset`. If it is an :obj:`nlp.Dataset`, columns not
-                accepted by the ``model.forward()`` method are automatically removed.
+                If provided, will override :obj:`self.eval_dataset`.
        """
        if eval_dataset is None and self.eval_dataset is None:
            raise ValueError("Trainer: evaluation requires an eval_dataset.")
-        elif eval_dataset is not None and is_nlp_available() and isinstance(eval_dataset, nlp.Dataset):
-            self._remove_unused_columns(eval_dataset, description="evaluation")
+
        eval_dataset = eval_dataset if eval_dataset is not None else self.eval_dataset
        eval_sampler = self._get_eval_sampler(eval_dataset)

@@ -340,11 +314,8 @@ class Trainer:

        Args:
            eval_dataset (:obj:`torch.utils.data.dataset.Dataset`, `optional`):
-                The test dataset to use. If it is an :obj:`nlp.Dataset`, columns not accepted by the
-                ``model.forward()`` method are automatically removed.
+                The test dataset to use.
        """
-        if is_nlp_available() and isinstance(test_dataset, nlp.Dataset):
-            self._remove_unused_columns(test_dataset, description="test")
        test_sampler = self._get_eval_sampler(test_dataset)

        # We use the same batch_size as for eval.
@@ -478,7 +449,6 @@ class Trainer:
        else:
            t_total = int(len(train_dataloader) // self.args.gradient_accumulation_steps * self.args.num_train_epochs)
            num_train_epochs = self.args.num_train_epochs
-            self.args.max_steps = t_total

        self.create_optimizer_and_scheduler(num_training_steps=t_total)

@@ -560,7 +530,7 @@ class Trainer:
        logging_loss = 0.0
        model.zero_grad()
        train_iterator = trange(
-            epochs_trained, int(np.ceil(num_train_epochs)), desc="Epoch", disable=not self.is_local_process_zero()
+            epochs_trained, int(num_train_epochs), desc="Epoch", disable=not self.is_local_process_zero()
        )
        for epoch in train_iterator:
            if isinstance(train_dataloader, DataLoader) and isinstance(train_dataloader.sampler, DistributedSampler):
@@ -656,10 +626,10 @@ class Trainer:
                            torch.save(self.optimizer.state_dict(), os.path.join(output_dir, "optimizer.pt"))
                            torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))

-                if self.args.max_steps > 0 and self.global_step >= self.args.max_steps:
+                if self.args.max_steps > 0 and self.global_step > self.args.max_steps:
                    epoch_iterator.close()
                    break
-            if self.args.max_steps > 0 and self.global_step >= self.args.max_steps:
+            if self.args.max_steps > 0 and self.global_step > self.args.max_steps:
                train_iterator.close()
                break
            if self.args.tpu_metrics_debug or self.args.debug:
@@ -734,7 +704,9 @@ class Trainer:
        else:
            print(output)

-    def _prepare_inputs(self, inputs: Dict[str, Union[torch.Tensor, Any]]) -> Dict[str, Union[torch.Tensor, Any]]:
+    def _prepare_inputs(
+        self, inputs: Dict[str, Union[torch.Tensor, Any]], model: nn.Module
+    ) -> Dict[str, Union[torch.Tensor, Any]]:
        """
        Prepare :obj:`inputs` before feeding them to the model, converting them to tensors if they are not already and
        handling potential state.
@@ -774,7 +746,7 @@ class Trainer:
            return self._training_step(model, inputs, self.optimizer)

        model.train()
-        inputs = self._prepare_inputs(inputs)
+        inputs = self._prepare_inputs(inputs, model)

        if self.args.fp16 and _use_native_amp:
            with autocast():
@@ -932,8 +904,7 @@ class Trainer:

        Args:
            eval_dataset (:obj:`Dataset`, `optional`):
-                Pass a dataset if you wish to override :obj:`self.eval_dataset`. If it is an :obj:`nlp.Dataset`,
-                columns not accepted by the ``model.forward()`` method are automatically removed.
+                Pass a dataset if you wish to override :obj:`self.eval_dataset`.

        Returns:
            A dictionary containing the evaluation loss and the potential metrics computed from the predictions.
@@ -959,8 +930,7 @@ class Trainer:

        Args:
            test_dataset (:obj:`Dataset`):
-                Dataset to run the predictions on. If it is an :obj:`nlp.Dataset`, columns not accepted by the
-                ``model.forward()`` method are automatically removed.
+                Dataset to run the predictions on.

        Returns:
            `NamedTuple`:
@@ -1018,13 +988,10 @@ class Trainer:
        if self.args.past_index >= 0:
            self._past = None

-        samples_count = 0
        for inputs in tqdm(dataloader, desc=description):
            loss, logits, labels = self.prediction_step(model, inputs, prediction_loss_only)
-            batch_size = inputs[list(inputs.keys())[0]].shape[0]
-            samples_count += batch_size
            if loss is not None:
-                eval_losses.append(loss * batch_size)
+                eval_losses.append(loss)
            if logits is not None:
                preds = logits if preds is None else torch.cat((preds, logits), dim=0)
            if labels is not None:
@@ -1058,7 +1025,7 @@ class Trainer:
        else:
            metrics = {}
        if len(eval_losses) > 0:
-            metrics["eval_loss"] = np.sum(eval_losses) / samples_count
+            metrics["eval_loss"] = np.mean(eval_losses)

        # Prefix all keys with eval_
        for key in list(metrics.keys()):
@@ -1104,7 +1071,7 @@ class Trainer:
        """
        has_labels = any(inputs.get(k) is not None for k in ["labels", "lm_labels", "masked_lm_labels"])

-        inputs = self._prepare_inputs(inputs)
+        inputs = self._prepare_inputs(inputs, model)

        with torch.no_grad():
            outputs = model(**inputs)
@@ -38,17 +38,9 @@ class TFTrainer:
        args (:class:`~transformers.TFTrainingArguments`):
            The arguments to tweak training.
        train_dataset (:class:`~tf.data.Dataset`, `optional`):
-            The dataset to use for training. The dataset should yield tuples of ``(features, labels)`` where
-            ``features`` is a dict of input features and ``labels`` is the labels. If ``labels`` is a tensor, the loss is
-            calculated by the model by calling ``model(features, labels=labels)``. If ``labels`` is a dict, such as when
-            using a QuestionAnswering head model with multiple targets, the loss is instead calculated by calling
-            ``model(features, **labels)``.
+            The dataset to use for training.
        eval_dataset (:class:`~tf.data.Dataset`, `optional`):
-            The dataset to use for evaluation. The dataset should yield tuples of ``(features, labels)`` where
-            ``features`` is a dict of input features and ``labels`` is the labels. If ``labels`` is a tensor, the loss is
-            calculated by the model by calling ``model(features, labels=labels)``. If ``labels`` is a dict, such as when
-            using a QuestionAnswering head model with multiple targets, the loss is instead calculated by calling
-            ``model(features, **labels)``.
+            The dataset to use for evaluation.
        compute_metrics (:obj:`Callable[[EvalPrediction], Dict]`, `optional`):
            The function that will be used to compute metrics at evaluation. Must take a
            :class:`~transformers.EvalPrediction` and return a dictionary string to metric values.
@@ -153,11 +145,7 @@ class TFTrainer:

        Args:
            eval_dataset (:class:`~tf.data.Dataset`, `optional`):
-                If provided, will override `self.eval_dataset`. The dataset should yield tuples of ``(features,
-                labels)`` where ``features`` is a dict of input features and ``labels`` is the labels. If ``labels``
-                is a tensor, the loss is calculated by the model by calling ``model(features, labels=labels)``. If
-                ``labels`` is a dict, such as when using a QuestionAnswering head model with multiple targets, the
-                loss is instead calculated by calling ``model(features, **labels)``.
+                If provided, will override `self.eval_dataset`.

        Subclass and override this method if you want to inject some custom behavior.
        """
@@ -185,12 +173,7 @@ class TFTrainer:
        Returns a test :class:`~tf.data.Dataset`.

        Args:
-            test_dataset (:class:`~tf.data.Dataset`):
-                The dataset to use. The dataset should yield tuples of ``(features, labels)`` where ``features`` is
-                a dict of input features and ``labels`` is the labels. If ``labels`` is a tensor, the loss is
-                calculated by the model by calling ``model(features, labels=labels)``. If ``labels`` is a dict, such
-                as when using a QuestionAnswering head model with multiple targets, the loss is instead calculated
-                by calling ``model(features, **labels)``.
+            test_dataset (:class:`~tf.data.Dataset`): The dataset to use.

        Subclass and override this method if you want to inject some custom behavior.
        """
@@ -422,18 +405,14 @@ class TFTrainer:

        Args:
            eval_dataset (:class:`~tf.data.Dataset`, `optional`):
-                Pass a dataset if you wish to override :obj:`self.eval_dataset`. The dataset should yield tuples of
-                ``(features, labels)`` where ``features`` is a dict of input features and ``labels`` is the labels.
-                If ``labels`` is a tensor, the loss is calculated by the model by calling ``model(features,
-                labels=labels)``. If ``labels`` is a dict, such as when using a QuestionAnswering head model with
-                multiple targets, the loss is instead calculated by calling ``model(features, **labels)``.
+                Pass a dataset if you wish to override :obj:`self.eval_dataset`.

        Returns:
            A dictionary containing the evaluation loss and the potential metrics computed from the predictions.
        """
        eval_ds, steps, num_examples = self.get_eval_tfdataset(eval_dataset)

-        output = self.prediction_loop(eval_ds, steps, num_examples, description="Evaluation")
+        output = self._prediction_loop(eval_ds, steps, num_examples, description="Evaluation")
        logs = {**output.metrics}
        logs["epoch"] = self.epoch_logging

@@ -687,11 +666,7 @@ class TFTrainer:

        Args:
            test_dataset (:class:`~tf.data.Dataset`):
-                Dataset to run the predictions on. The dataset should yield tuples of ``(features, labels)`` where
-                ``features`` is a dict of input features and ``labels`` is the labels. If ``labels`` is a tensor,
-                the loss is calculated by the model by calling ``model(features, labels=labels)``. If ``labels`` is
-                a dict, such as when using a QuestionAnswering head model with multiple targets, the loss is instead
-                calculated by calling ``model(features, **labels)``.
+                Dataset to run the predictions on.
        Returns:
            `NamedTuple`:
            predictions (:obj:`np.ndarray`):
@@ -69,8 +69,7 @@ class TrainingArguments:
        max_grad_norm (:obj:`float`, `optional`, defaults to 1.0):
            Maximum gradient norm (for gradient clipping).
        num_train_epochs(:obj:`float`, `optional`, defaults to 3.0):
-            Total number of training epochs to perform (if not an integer, will perform the decimal part percents of
-            the last epoch before stopping training).
+            Total number of training epochs to perform.
        max_steps (:obj:`int`, `optional`, defaults to -1):
            If set to a positive number, the total number of training steps to perform. Overrides
            :obj:`num_train_epochs`.
@@ -1,152 +0,0 @@
-import unittest
-
-from transformers import AutoTokenizer, is_torch_available
-from transformers.testing_utils import require_torch
-
-
-if is_torch_available():
-    import torch
-
-    from transformers import (
-        DataCollatorForLanguageModeling,
-        DataCollatorForPermutationLanguageModeling,
-        GlueDataset,
-        GlueDataTrainingArguments,
-        LineByLineTextDataset,
-        TextDataset,
-        default_data_collator,
-    )
-
-
-PATH_SAMPLE_TEXT = "./tests/fixtures/sample_text.txt"
-
-
-@require_torch
-class DataCollatorIntegrationTest(unittest.TestCase):
-    def test_default_with_dict(self):
-        features = [{"label": i, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
-        batch = default_data_collator(features)
-        self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8)))))
-        self.assertEqual(batch["labels"].dtype, torch.long)
-        self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))
-
-        # With label_ids
-        features = [{"label_ids": [0, 1, 2], "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
-        batch = default_data_collator(features)
-        self.assertTrue(batch["labels"].equal(torch.tensor([[0, 1, 2]] * 8)))
-        self.assertEqual(batch["labels"].dtype, torch.long)
-        self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))
-
-        # Features can already be tensors
-        features = [{"label": i, "inputs": torch.randint(10, [10])} for i in range(8)]
-        batch = default_data_collator(features)
-        self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8)))))
-        self.assertEqual(batch["labels"].dtype, torch.long)
-        self.assertEqual(batch["inputs"].shape, torch.Size([8, 10]))
-
-        # Labels can already be tensors
-        features = [{"label": torch.tensor(i), "inputs": torch.randint(10, [10])} for i in range(8)]
-        batch = default_data_collator(features)
-        self.assertEqual(batch["labels"].dtype, torch.long)
-        self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8)))))
-        self.assertEqual(batch["labels"].dtype, torch.long)
-        self.assertEqual(batch["inputs"].shape, torch.Size([8, 10]))
-
-    def test_default_with_no_labels(self):
-        features = [{"label": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
-        batch = default_data_collator(features)
-        self.assertTrue("labels" not in batch)
-        self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))
-
-        # With label_ids
-        features = [{"label_ids": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
-        batch = default_data_collator(features)
-        self.assertTrue("labels" not in batch)
-        self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))
-
-    def test_default_classification(self):
-        MODEL_ID = "bert-base-cased-finetuned-mrpc"
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        data_args = GlueDataTrainingArguments(
-            task_name="mrpc", data_dir="./tests/fixtures/tests_samples/MRPC", overwrite_cache=True
-        )
-        dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev")
-        data_collator = default_data_collator
-        batch = data_collator(dataset.features)
-        self.assertEqual(batch["labels"].dtype, torch.long)
-
-    def test_default_regression(self):
-        MODEL_ID = "distilroberta-base"
-        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
-        data_args = GlueDataTrainingArguments(
-            task_name="sts-b", data_dir="./tests/fixtures/tests_samples/STS-B", overwrite_cache=True
-        )
-        dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev")
-        data_collator = default_data_collator
-        batch = data_collator(dataset.features)
-        self.assertEqual(batch["labels"].dtype, torch.float)
-
-    def test_lm_tokenizer_without_padding(self):
-        tokenizer = AutoTokenizer.from_pretrained("gpt2")
-        data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
-        # ^ causal lm
-
-        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
-        examples = [dataset[i] for i in range(len(dataset))]
-        with self.assertRaises(ValueError):
-            # Expect error due to padding token missing on gpt2:
-            data_collator(examples)
-
-        dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
-        self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
-        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
-
-    def test_lm_tokenizer_with_padding(self):
-        tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
-        data_collator = DataCollatorForLanguageModeling(tokenizer)
-        # ^ masked lm
-
-        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
-        self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((31, 107)))
-        self.assertEqual(batch["labels"].shape, torch.Size((31, 107)))
-
-        dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
-        self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
-        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
-
-    def test_plm(self):
-        tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased")
-        data_collator = DataCollatorForPermutationLanguageModeling(tokenizer)
-        # ^ permutation lm
-
-        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
-        self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((31, 112)))
-        self.assertEqual(batch["perm_mask"].shape, torch.Size((31, 112, 112)))
-        self.assertEqual(batch["target_mapping"].shape, torch.Size((31, 112, 112)))
-        self.assertEqual(batch["labels"].shape, torch.Size((31, 112)))
-
-        dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True)
-        examples = [dataset[i] for i in range(len(dataset))]
-        batch = data_collator(examples)
-        self.assertIsInstance(batch, dict)
-        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
-        self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 512, 512)))
-        self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 512, 512)))
-        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
-
-        example = [torch.randint(5, [5])]
-        with self.assertRaises(ValueError):
-            # Expect error due to odd sequence length
-            data_collator(example)
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -71,7 +71,7 @@ class AlbertModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -18,7 +18,7 @@ import unittest

 import timeout_decorator  # noqa

-from transformers import is_torch_available
+from transformers import BatchEncoding, is_torch_available
 from transformers.file_utils import cached_property
 from transformers.testing_utils import require_torch, slow, torch_device

@@ -38,7 +38,6 @@ if is_torch_available():
        BartForQuestionAnswering,
        BartConfig,
        BartTokenizer,
-        BartTokenizerFast,
        pipeline,
    )
    from transformers.modeling_bart import (
@@ -422,10 +421,6 @@ class BartModelIntegrationTests(unittest.TestCase):
    def default_tokenizer(self):
        return BartTokenizer.from_pretrained("facebook/bart-large")

-    @cached_property
-    def default_tokenizer_fast(self):
-        return BartTokenizerFast.from_pretrained("facebook/bart-large")
-
    @slow
    def test_inference_no_head(self):
        model = BartModel.from_pretrained("facebook/bart-large").to(torch_device)
@@ -445,7 +440,8 @@ class BartModelIntegrationTests(unittest.TestCase):
        pbase = pipeline(task="fill-mask", model="facebook/bart-base")
        src_text = [" I went to the <mask>."]
        results = [x["token_str"] for x in pbase(src_text)]
-        assert "Ġbathroom" in results
+        expected_results = ["Ġbathroom", "Ġrestroom", "Ġhospital", "Ġkitchen", "Ġcar"]
+        self.assertListEqual(results, expected_results)

    @slow
    def test_bart_large_mask_filling(self):
@@ -486,9 +482,9 @@ class BartModelIntegrationTests(unittest.TestCase):
    def test_xsum_summarization_same_as_fairseq(self):
        model = BartForConditionalGeneration.from_pretrained("facebook/bart-large-xsum").to(torch_device)
        self.assertFalse(model.config.is_valid_mbart())
-        tok = self.default_tokenizer
+        tok = BartTokenizer.from_pretrained("facebook/bart-large")

-        EXPECTED_SUMMARY = "California's largest power company has begun shutting off electricity to thousands of customers in the state."
+        EXPECTED_SUMMARY = "California's largest power company has begun shutting off power to tens of thousands of homes and businesses in the state."
        dct = tok.batch_encode_plus(
            [PGE_ARTICLE], max_length=1024, padding="max_length", truncation=True, return_tensors="pt",
        ).to(torch_device)
@@ -568,6 +564,76 @@ class BartModelIntegrationTests(unittest.TestCase):
        # TODO(SS): run fairseq again with num_beams=2, min_len=20.
        # TODO(SS): add test case that hits max_length

+    def test_prepare_seq2seq_batch(self):
+        tokenizer = self.default_tokenizer
+        src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
+        tgt_text = [
+            "Summary of the text.",
+            "Another summary.",
+        ]
+        expected_src_tokens = [0, 250, 251, 17818, 13, 32933, 21645, 1258, 4, 2]
+        batch = tokenizer.prepare_seq2seq_batch(
+            src_text, tgt_texts=tgt_text, max_length=len(expected_src_tokens), return_tensors="pt"
+        )
+        self.assertIsInstance(batch, BatchEncoding)
+
+        self.assertEqual((2, 10), batch.input_ids.shape)
+        self.assertEqual((2, 10), batch.attention_mask.shape)
+        result = batch.input_ids.tolist()[0]
+        self.assertListEqual(expected_src_tokens, result)
+        # Test that special tokens are reset
+
+    def test_empty_target_text(self):
+        tokenizer = self.default_tokenizer
+        src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
+        batch = tokenizer.prepare_seq2seq_batch(src_text, return_tensors="pt")
+        # check if input_ids are returned and no decoder_input_ids
+        self.assertIn("input_ids", batch)
+        self.assertIn("attention_mask", batch)
+        self.assertNotIn("decoder_input_ids", batch)
+        self.assertNotIn("decoder_attention_mask", batch)
+
+    def test_max_target_length(self):
+        tokenizer = self.default_tokenizer
+        src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
+        tgt_text = [
+            "Summary of the text.",
+            "Another summary.",
+        ]
+        batch = tokenizer.prepare_seq2seq_batch(
+            src_text, tgt_texts=tgt_text, max_target_length=32, padding="max_length", return_tensors="pt"
+        )
+        self.assertEqual(32, batch["decoder_input_ids"].shape[1])
+        self.assertEqual(32, batch["decoder_attention_mask"].shape[1])
+
+        # test None max_target_length
+        batch = tokenizer.prepare_seq2seq_batch(
+            src_text, tgt_texts=tgt_text, max_length=32, padding="max_length", return_tensors="pt"
+        )
+        self.assertEqual(32, batch["decoder_input_ids"].shape[1])
+        self.assertEqual(32, batch["decoder_attention_mask"].shape[1])
+
+    def test_outputs_not_longer_than_maxlen(self):
+        tokenizer = self.default_tokenizer
+
+        batch = tokenizer.prepare_seq2seq_batch(["I am a small frog" * 1024, "I am a small frog"], return_tensors="pt")
+        self.assertIsInstance(batch, BatchEncoding)
+        self.assertEqual(batch.input_ids.shape, (2, 1024))
+
+    def test_special_tokens(self):
+        tokenizer = self.default_tokenizer
+        src_text = ["A long paragraph for summrization."]
+        tgt_text = [
+            "Summary of the text.",
+        ]
+        batch = tokenizer.prepare_seq2seq_batch(src_text, tgt_texts=tgt_text, return_tensors="pt")
+        input_ids = batch["input_ids"]
+        decoder_input_ids = batch["decoder_input_ids"]
+        self.assertTrue((input_ids[:, 0] == tokenizer.bos_token_id).all().item())
+        self.assertTrue((decoder_input_ids[:, 0] == tokenizer.bos_token_id).all().item())
+        self.assertTrue((input_ids[:, -1] == tokenizer.eos_token_id).all().item())
+        self.assertTrue((decoder_input_ids[:, -1] == tokenizer.eos_token_id).all().item())
+

@require_torch
 class TestSinusoidalPositionalEmbeddings(unittest.TestCase):
@@ -20,21 +20,21 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


 if is_torch_available():
    from transformers import (
        BertConfig,
+        BertModel,
+        BertLMHeadModel,
        BertForMaskedLM,
-        BertForMultipleChoice,
        BertForNextSentencePrediction,
        BertForPreTraining,
        BertForQuestionAnswering,
        BertForSequenceClassification,
        BertForTokenClassification,
-        BertLMHeadModel,
-        BertModel,
+        BertForMultipleChoice,
    )
    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 = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -370,6 +370,7 @@ class BertModelTest(ModelTesterMixin, unittest.TestCase):
        if is_torch_available()
        else ()
    )
+    test_chunking = True

    def setUp(self):
        self.model_tester = BertModelTester(self)
@@ -25,15 +25,15 @@ from transformers.testing_utils import require_multigpu, require_torch, slow, to


 if is_torch_available():
-    import numpy as np
    import torch
+    import numpy as np

    from transformers import (
        AdaptiveEmbedding,
        PretrainedConfig,
        PreTrainedModel,
-        BertConfig,
        BertModel,
+        BertConfig,
        BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
        MODEL_FOR_QUESTION_ANSWERING_MAPPING,
@@ -65,6 +65,7 @@ class ModelTesterMixin:
    test_resize_embeddings = True
    test_head_masking = True
    test_missing_keys = True
+    test_chunking = False
    is_encoder_decoder = False

    def _prepare_for_class(self, inputs_dict, model_class, return_labels=False):
@@ -551,6 +552,9 @@ class ModelTesterMixin:

    def test_feed_forward_chunking(self):
        (original_config, inputs_dict,) = self.model_tester.prepare_config_and_inputs_for_common()
+        if not self.test_chunking:
+            return
+
        for model_class in self.all_model_classes:
            torch.manual_seed(0)
            config = copy.deepcopy(original_config)
@@ -707,7 +711,7 @@ class ModelTesterMixin:
                    else:
                        self.assertTrue(
                            torch.allclose(tuple_object, dict_object, atol=1e-5),
-                            msg=f"Tuple and dict output are not equal. Difference: {torch.max(torch.abs(tuple_object - dict_object))}. Tuple has `nan`: {torch.isnan(tuple_object).any()} and `inf`: {torch.isinf(tuple_object)}. Dict has `nan`: {torch.isnan(dict_object).any()} and `inf`: {torch.isinf(dict_object)}.",
+                            msg=f"Tuple and dict output are not equal. Difference: {torch.max(torch.abs(tuple_object - dict_object))}",
                        )

                recursive_check(tuple_output, dict_output)
@@ -934,13 +938,6 @@ def ids_tensor(shape, vocab_size, rng=None, name=None):
    return torch.tensor(data=values, dtype=torch.long, device=torch_device).view(shape).contiguous()


-def random_attention_mask(shape, rng=None, name=None):
-    attn_mask = ids_tensor(shape, vocab_size=2, rng=None, name=None)
-    # make sure that at least one token is attended to for each batch
-    attn_mask[:, -1] = 1
-    return attn_mask
-
-
 def floats_tensor(shape, scale=1.0, rng=None, name=None):
    """Creates a random float32 tensor"""
    if rng is None:
@@ -19,7 +19,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -60,7 +60,7 @@ class CTRLModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -89,7 +89,7 @@ if is_torch_available():

            input_mask = None
            if self.use_input_mask:
-                input_mask = random_attention_mask([self.batch_size, self.seq_length])
+                input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

            sequence_labels = None
            token_labels = None
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -88,7 +88,7 @@ class DPRModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -69,7 +69,7 @@ class ElectraModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -268,88 +268,6 @@ class EncoderDecoderMixin:
        )
        self.assertEqual(generated_output.shape, (input_ids.shape[0],) + (decoder_config.max_length,))

-    def create_and_check_encoder_decoder_shared_weights(
-        self,
-        config,
-        input_ids,
-        attention_mask,
-        encoder_hidden_states,
-        decoder_config,
-        decoder_input_ids,
-        decoder_attention_mask,
-        labels,
-        **kwargs
-    ):
-        torch.manual_seed(0)
-        encoder_model, decoder_model = self.get_encoder_decoder_model(config, decoder_config)
-        model = EncoderDecoderModel(encoder=encoder_model, decoder=decoder_model)
-        model.to(torch_device)
-        model.eval()
-        # load state dict copies weights but does not tie them
-        decoder_state_dict = model.decoder._modules[model.decoder.base_model_prefix].state_dict()
-        model.encoder.load_state_dict(decoder_state_dict, strict=False)
-
-        torch.manual_seed(0)
-        tied_encoder_model, tied_decoder_model = self.get_encoder_decoder_model(config, decoder_config)
-        config = EncoderDecoderConfig.from_encoder_decoder_configs(
-            tied_encoder_model.config, tied_decoder_model.config, tie_encoder_decoder=True
-        )
-        tied_model = EncoderDecoderModel(encoder=tied_encoder_model, decoder=tied_decoder_model, config=config)
-        tied_model.to(torch_device)
-        tied_model.eval()
-
-        model_result = model(
-            input_ids=input_ids,
-            decoder_input_ids=decoder_input_ids,
-            attention_mask=attention_mask,
-            decoder_attention_mask=decoder_attention_mask,
-        )
-
-        tied_model_result = tied_model(
-            input_ids=input_ids,
-            decoder_input_ids=decoder_input_ids,
-            attention_mask=attention_mask,
-            decoder_attention_mask=decoder_attention_mask,
-        )
-
-        # check that models has less parameters
-        self.assertLess(sum(p.numel() for p in tied_model.parameters()), sum(p.numel() for p in model.parameters()))
-        random_slice_idx = ids_tensor((1,), model_result[0].shape[-1]).item()
-
-        # check that outputs are equal
-        self.assertTrue(
-            torch.allclose(
-                model_result[0][0, :, random_slice_idx], tied_model_result[0][0, :, random_slice_idx], atol=1e-4
-            )
-        )
-
-        # check that outputs after saving and loading are equal
-        with tempfile.TemporaryDirectory() as tmpdirname:
-            tied_model.save_pretrained(tmpdirname)
-            tied_model = EncoderDecoderModel.from_pretrained(tmpdirname)
-            tied_model.to(torch_device)
-            tied_model.eval()
-
-            # check that models has less parameters
-            self.assertLess(
-                sum(p.numel() for p in tied_model.parameters()), sum(p.numel() for p in model.parameters())
-            )
-            random_slice_idx = ids_tensor((1,), model_result[0].shape[-1]).item()
-
-            tied_model_result = tied_model(
-                input_ids=input_ids,
-                decoder_input_ids=decoder_input_ids,
-                attention_mask=attention_mask,
-                decoder_attention_mask=decoder_attention_mask,
-            )
-
-            # check that outputs are equal
-            self.assertTrue(
-                torch.allclose(
-                    model_result[0][0, :, random_slice_idx], tied_model_result[0][0, :, random_slice_idx], atol=1e-4
-                )
-            )
-
    def test_encoder_decoder_model(self):
        input_ids_dict = self.prepare_config_and_inputs()
        self.check_encoder_decoder_model(**input_ids_dict)
@@ -378,10 +296,6 @@ class EncoderDecoderMixin:
        input_ids_dict = self.prepare_config_and_inputs()
        self.check_encoder_decoder_model_generate(**input_ids_dict)

-    def test_encoder_decoder_model_shared_weights(self):
-        input_ids_dict = self.prepare_config_and_inputs()
-        self.create_and_check_encoder_decoder_shared_weights(**input_ids_dict)
-
    @slow
    def test_real_model_save_load_from_pretrained(self):
        model_2 = self.get_pretrained_model()
@@ -566,6 +480,3 @@ class GPT2EncoderDecoderModelTest(EncoderDecoderMixin, unittest.TestCase):

    def get_pretrained_model(self):
        return EncoderDecoderModel.from_encoder_decoder_pretrained("bert-base-cased", "gpt2")
-
-    def test_encoder_decoder_model_shared_weights(self):
-        pass
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -72,7 +72,7 @@ class FlaubertModelTester(object):

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
-        input_mask = random_attention_mask([self.batch_size, self.seq_length])
+        input_mask = ids_tensor([self.batch_size, self.seq_length], 2).float()

        input_lengths = None
        if self.use_input_lengths:
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


 if is_torch_available():
@@ -92,7 +92,7 @@ class GPT2ModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -82,7 +82,7 @@ class LongformerModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -205,9 +205,9 @@ class TestMarian_MT_EN(MarianIntegrationTest):
        self._assert_generated_batch_equal_expected()


-class TestMarian_en_zh(MarianIntegrationTest):
-    src = "en"
-    tgt = "zh"
+class TestMarian_eng_zho(MarianIntegrationTest):
+    src = "eng"
+    tgt = "zho"
    src_text = ["My name is Wolfgang and I live in Berlin"]
    expected_text = ["我叫沃尔夫冈 我住在柏林"]

@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


 if is_torch_available():
@@ -94,7 +94,7 @@ class MobileBertModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -1,7 +1,6 @@
 import unittest

-from transformers import AutoConfig, AutoTokenizer, is_torch_available
-from transformers.configuration_pegasus import max_gen_length, max_model_length
+from transformers import AutoConfig, is_torch_available
 from transformers.file_utils import cached_property
 from transformers.testing_utils import require_torch, slow, torch_device

@@ -31,7 +30,9 @@ class PegasusXSUMIntegrationTest(AbstractSeq2SeqIntegrationTest):
    @slow
    def test_pegasus_xsum_summary(self):
        assert self.tokenizer.model_max_length == 512
-        inputs = self.tokenizer(self.src_text, return_tensors="pt", truncation=True, padding=True).to(torch_device)
+        inputs = self.tokenizer(self.src_text, return_tensors="pt", truncation=True, max_length=512, padding=True).to(
+            torch_device
+        )
        assert inputs.input_ids.shape == (2, 421)
        translated_tokens = self.model.generate(**inputs)
        decoded = self.tokenizer.batch_decode(translated_tokens, skip_special_tokens=True)
@@ -49,28 +50,28 @@ class PegasusXSUMIntegrationTest(AbstractSeq2SeqIntegrationTest):

 class PegasusConfigTests(unittest.TestCase):
    def test_all_config_max_lengths(self):
+        expected_max_length = {
+            # See appendix C of paper
+            "xsum": 64,
+            "cnn_dailymail": 128,
+            "newsroom": 128,
+            "wikihow": 256,
+            "multi_news": 256,
+            "reddit_tifu": 128,
+            "big_patent": 256,
+            "arxiv": 256,
+            "pubmed": 256,
+            "gigaword": 32,
+            "aeslc": 32,
+            "billsum": 256,
+        }
        failures = []
        pegasus_prefix = "google/pegasus"
-        for dataset, max_len in max_gen_length.items():
+        for dataset, max_len in expected_max_length.items():
            mname = f"{pegasus_prefix}-{dataset}"
            cfg = AutoConfig.from_pretrained(mname)
-
            if cfg.max_length != max_len:
                failures.append(f"config for {mname} had max_length: {cfg.max_length}, expected {max_len}")
-
-            if cfg.max_position_embeddings < max_model_length[dataset]:
-                # otherwise you get IndexError for e.g. position 513
-                # see https://github.com/huggingface/transformers/issues/6599
-                failures.append(
-                    f"config for {mname} had max_position_embeddings: {cfg.max_position_embeddings}, expected {max_model_length[dataset]}"
-                )
-
-            tokenizer = AutoTokenizer.from_pretrained(mname)
-            if max_model_length[dataset] != tokenizer.model_max_length:
-                failures.append(
-                    f"tokenizer.model_max_length {tokenizer.model_max_length} expected {max_model_length[dataset]}"
-                )
-
        if failures == []:
            return
        # error
@@ -19,7 +19,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


 if is_torch_available():
@@ -133,7 +133,7 @@ class ReformerModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        choice_labels = None
        if self.use_labels:
@@ -555,6 +555,7 @@ class ReformerLocalAttnModelTest(ReformerTesterMixin, ModelTesterMixin, unittest
    test_pruning = False
    test_headmasking = False
    test_torchscript = False
+    test_chunking = True

    def prepare_kwargs(self):
        return {
@@ -615,6 +616,7 @@ class ReformerLSHAttnModelTest(ReformerTesterMixin, ModelTesterMixin, unittest.T
    test_pruning = False
    test_headmasking = False
    test_torchscript = False
+    test_chunking = True

    def prepare_kwargs(self):
        return {
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, floats_tensor, ids_tensor


 if is_torch_available():
@@ -71,7 +71,7 @@ class RobertaModelTester:

        input_mask = None
        if self.use_input_mask:
-            input_mask = random_attention_mask([self.batch_size, self.seq_length])
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)

        token_type_ids = None
        if self.use_token_type_ids:
@@ -14,8 +14,6 @@
 # limitations under the License.


-import copy
-import tempfile
 import unittest

 from transformers import is_torch_available
@@ -132,7 +130,7 @@ class T5ModelTester:
                # all items after square
                self.parent.assertListEqual(decoder_input_ids_slice[1:].tolist(), lm_labels_slice[:-1].tolist())

-    def create_and_check_model(
+    def create_and_check_t5_model(
        self, config, input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels,
    ):
        model = T5Model(config=config)
@@ -158,7 +156,7 @@ class T5ModelTester:
        # There should be a self attn key, a self attn value, a cross attn key and a cross attn value stored in each decoder_past[1] tuple
        self.parent.assertEqual(len(decoder_past[1][0]), 4)

-    def create_and_check_with_lm_head(
+    def create_and_check_t5_with_lm_head(
        self, config, input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels,
    ):
        model = T5ForConditionalGeneration(config=config).to(torch_device).eval()
@@ -172,7 +170,7 @@ class T5ModelTester:
        self.parent.assertEqual(outputs["logits"].size(), (self.batch_size, self.decoder_seq_length, self.vocab_size))
        self.parent.assertEqual(outputs["loss"].size(), ())

-    def create_and_check_decoder_model_past(
+    def create_and_check_t5_decoder_model_past(
        self, config, input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels,
    ):
        model = T5Model(config=config).get_decoder().to(torch_device).eval()
@@ -203,7 +201,7 @@ class T5ModelTester:
        # test that outputs are equal for slice
        self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))

-    def create_and_check_decoder_model_attention_mask_past(
+    def create_and_check_t5_decoder_model_attention_mask_past(
        self, config, input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels,
    ):
        model = T5Model(config=config).get_decoder()
@@ -247,7 +245,7 @@ class T5ModelTester:
        # test that outputs are equal for slice
        self.parent.assertTrue(torch.allclose(output_from_past_slice, output_from_no_past_slice, atol=1e-3))

-    def create_and_check_generate_with_past_key_value_states(
+    def create_t5_and_check_t5_generate_with_past_key_value_states(
        self, config, input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels,
    ):
        model = T5ForConditionalGeneration(config=config).to(torch_device).eval()
@@ -259,83 +257,13 @@ class T5ModelTester:
        output_with_past_cache = model.generate(input_ids[:1], num_beams=2, max_length=5, do_sample=True)
        self.parent.assertTrue(torch.all(output_with_past_cache == output_without_past_cache))

-    def create_and_check_model_fp16_forward(
+    def create_and_check_t5_model_fp16_forward(
        self, config, input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels,
    ):
        model = T5Model(config=config).to(torch_device).half().eval()
        output = model(input_ids, decoder_input_ids=input_ids, attention_mask=attention_mask)["last_hidden_state"]
        self.parent.assertFalse(torch.isnan(output).any().item())

-    def create_and_check_encoder_decoder_shared_weights(
-        self, config, input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels,
-    ):
-        for model_class in [T5Model, T5ForConditionalGeneration]:
-            torch.manual_seed(0)
-            model = model_class(config=config).to(torch_device).eval()
-            # load state dict copies weights but does not tie them
-            model.encoder.load_state_dict(model.decoder.state_dict(), strict=False)
-
-            torch.manual_seed(0)
-            tied_config = copy.deepcopy(config)
-            tied_config.tie_encoder_decoder = True
-            tied_model = model_class(config=tied_config).to(torch_device).eval()
-
-            model_result = model(
-                input_ids=input_ids,
-                decoder_input_ids=decoder_input_ids,
-                attention_mask=attention_mask,
-                decoder_attention_mask=decoder_attention_mask,
-            )
-
-            tied_model_result = tied_model(
-                input_ids=input_ids,
-                decoder_input_ids=decoder_input_ids,
-                attention_mask=attention_mask,
-                decoder_attention_mask=decoder_attention_mask,
-            )
-
-            # check that models has less parameters
-            self.parent.assertLess(
-                sum(p.numel() for p in tied_model.parameters()), sum(p.numel() for p in model.parameters())
-            )
-            random_slice_idx = ids_tensor((1,), model_result[0].shape[-1]).item()
-
-            # check that outputs are equal
-            self.parent.assertTrue(
-                torch.allclose(
-                    model_result[0][0, :, random_slice_idx], tied_model_result[0][0, :, random_slice_idx], atol=1e-4
-                )
-            )
-
-            # check that outputs after saving and loading are equal
-            with tempfile.TemporaryDirectory() as tmpdirname:
-                tied_model.save_pretrained(tmpdirname)
-                tied_model = model_class.from_pretrained(tmpdirname)
-                tied_model.to(torch_device)
-                tied_model.eval()
-
-                # check that models has less parameters
-                self.parent.assertLess(
-                    sum(p.numel() for p in tied_model.parameters()), sum(p.numel() for p in model.parameters())
-                )
-                random_slice_idx = ids_tensor((1,), model_result[0].shape[-1]).item()
-
-                tied_model_result = tied_model(
-                    input_ids=input_ids,
-                    decoder_input_ids=decoder_input_ids,
-                    attention_mask=attention_mask,
-                    decoder_attention_mask=decoder_attention_mask,
-                )
-
-                # check that outputs are equal
-                self.parent.assertTrue(
-                    torch.allclose(
-                        model_result[0][0, :, random_slice_idx],
-                        tied_model_result[0][0, :, random_slice_idx],
-                        atol=1e-4,
-                    )
-                )
-
    def prepare_config_and_inputs_for_common(self):
        config_and_inputs = self.prepare_config_and_inputs()
        (config, input_ids, decoder_input_ids, attention_mask, decoder_attention_mask, lm_labels,) = config_and_inputs
@@ -371,34 +299,30 @@ class T5ModelTest(ModelTesterMixin, unittest.TestCase):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        self.model_tester.check_prepare_lm_labels_via_shift_left(*config_and_inputs)

-    def test_model(self):
+    def test_t5_model(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_model(*config_and_inputs)
+        self.model_tester.create_and_check_t5_model(*config_and_inputs)

    def test_with_lm_head(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_with_lm_head(*config_and_inputs)
+        self.model_tester.create_and_check_t5_with_lm_head(*config_and_inputs)

-    def test_decoder_model_past(self):
+    def test_t5_decoder_model_past(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_decoder_model_past(*config_and_inputs)
+        self.model_tester.create_and_check_t5_decoder_model_past(*config_and_inputs)

-    def test_decoder_model_past_with_attn_mask(self):
+    def test_t5_decoder_model_past_with_attn_mask(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_decoder_model_attention_mask_past(*config_and_inputs)
+        self.model_tester.create_and_check_t5_decoder_model_attention_mask_past(*config_and_inputs)

-    def test_generate_with_past_key_value_states(self):
+    def test_t5_generate_with_past_key_value_states(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_generate_with_past_key_value_states(*config_and_inputs)
-
-    def test_encoder_decoder_shared_weights(self):
-        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_encoder_decoder_shared_weights(*config_and_inputs)
+        self.model_tester.create_t5_and_check_t5_generate_with_past_key_value_states(*config_and_inputs)

    @unittest.skipIf(torch_device == "cpu", "Cant do half precision")
-    def test_model_fp16_forward(self):
+    def test_t5_model_fp16_forward(self):
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
-        self.model_tester.create_and_check_model_fp16_forward(*config_and_inputs)
+        self.model_tester.create_and_check_t5_model_fp16_forward(*config_and_inputs)

    @slow
    def test_model_from_pretrained(self):
@@ -407,6 +331,8 @@ class T5ModelTest(ModelTesterMixin, unittest.TestCase):
            self.assertIsNotNone(model)

    def test_export_to_onnx(self):
+        import tempfile
+
        config_and_inputs = self.model_tester.prepare_config_and_inputs()
        model = T5Model(config_and_inputs[0]).to(torch_device)
        with tempfile.TemporaryDirectory() as tmpdirname:
@@ -306,13 +306,13 @@ class TFModelTesterMixin:

            max_diff = np.amax(np.abs(tf_hidden_states - pt_hidden_states))
            # Debug info (remove when fixed)
-            if max_diff >= 4e-2:
+            if max_diff >= 2e-2:
                print("===")
                print(model_class)
                print(config)
                print(inputs_dict)
                print(pt_inputs_dict)
-            self.assertLessEqual(max_diff, 4e-2)
+            self.assertLessEqual(max_diff, 2e-2)

            # Check we can load pt model in tf and vice-versa with checkpoint => model functions
            with tempfile.TemporaryDirectory() as tmpdirname:
@@ -348,7 +348,7 @@ class TFModelTesterMixin:
            tfo[pt_nans] = 0

            max_diff = np.amax(np.abs(tfo - pto))
-            self.assertLessEqual(max_diff, 4e-2)
+            self.assertLessEqual(max_diff, 2e-2)

    def test_compile_tf_model(self):
        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
@@ -20,7 +20,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -73,7 +73,7 @@ class XLMModelTester:

    def prepare_config_and_inputs(self):
        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
-        input_mask = random_attention_mask([self.batch_size, self.seq_length])
+        input_mask = ids_tensor([self.batch_size, self.seq_length], 2).float()

        input_lengths = None
        if self.use_input_lengths:
@@ -21,7 +21,7 @@ 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, random_attention_mask
+from .test_modeling_common import ModelTesterMixin, ids_tensor


 if is_torch_available():
@@ -100,7 +100,7 @@ class XLNetModelTester:
        input_ids_1 = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
        input_ids_2 = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
        segment_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
-        input_mask = random_attention_mask([self.batch_size, self.seq_length])
+        input_mask = ids_tensor([self.batch_size, self.seq_length], 2).float()

        input_ids_q = ids_tensor([self.batch_size, self.seq_length + 1], self.vocab_size)
        perm_mask = torch.zeros(
@@ -1,90 +0,0 @@
-import unittest
-
-from transformers import BartTokenizer, BartTokenizerFast, BatchEncoding
-from transformers.file_utils import cached_property
-
-
-class TestTokenizationBart(unittest.TestCase):
-    @cached_property
-    def default_tokenizer(self):
-        return BartTokenizer.from_pretrained("facebook/bart-large")
-
-    @cached_property
-    def default_tokenizer_fast(self):
-        return BartTokenizerFast.from_pretrained("facebook/bart-large")
-
-    def test_prepare_seq2seq_batch(self):
-        tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
-        src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
-        tgt_text = [
-            "Summary of the text.",
-            "Another summary.",
-        ]
-        expected_src_tokens = [0, 250, 251, 17818, 13, 32933, 21645, 1258, 4, 2]
-
-        for tokenizer in tokenizers:
-            batch = tokenizer.prepare_seq2seq_batch(
-                src_text, tgt_texts=tgt_text, max_length=len(expected_src_tokens), return_tensors="pt"
-            )
-            self.assertIsInstance(batch, BatchEncoding)
-
-            self.assertEqual((2, 10), batch.input_ids.shape)
-            self.assertEqual((2, 10), batch.attention_mask.shape)
-            result = batch.input_ids.tolist()[0]
-            self.assertListEqual(expected_src_tokens, result)
-            # Test that special tokens are reset
-
-    def test_empty_target_text(self):
-        tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
-        src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
-        for tokenizer in tokenizers:
-            batch = tokenizer.prepare_seq2seq_batch(src_text, return_tensors="pt")
-            # check if input_ids are returned and no labels
-            self.assertIn("input_ids", batch)
-            self.assertIn("attention_mask", batch)
-            self.assertNotIn("labels", batch)
-            self.assertNotIn("decoder_attention_mask", batch)
-
-    def test_max_target_length(self):
-        tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
-        src_text = ["A long paragraph for summrization.", "Another paragraph for summrization."]
-        tgt_text = [
-            "Summary of the text.",
-            "Another summary.",
-        ]
-        for tokenizer in tokenizers:
-            batch = tokenizer.prepare_seq2seq_batch(
-                src_text, tgt_texts=tgt_text, max_target_length=32, padding="max_length", return_tensors="pt"
-            )
-            self.assertEqual(32, batch["labels"].shape[1])
-
-            # test None max_target_length
-            batch = tokenizer.prepare_seq2seq_batch(
-                src_text, tgt_texts=tgt_text, max_length=32, padding="max_length", return_tensors="pt"
-            )
-            self.assertEqual(32, batch["labels"].shape[1])
-
-    def test_outputs_not_longer_than_maxlen(self):
-        tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
-
-        for tokenizer in tokenizers:
-            batch = tokenizer.prepare_seq2seq_batch(
-                ["I am a small frog" * 1024, "I am a small frog"], return_tensors="pt"
-            )
-            self.assertIsInstance(batch, BatchEncoding)
-            self.assertEqual(batch.input_ids.shape, (2, 1024))
-
-    def test_special_tokens(self):
-        tokenizers = [self.default_tokenizer, self.default_tokenizer_fast]
-        src_text = ["A long paragraph for summrization."]
-        tgt_text = [
-            "Summary of the text.",
-        ]
-        for tokenizer in tokenizers:
-            batch = tokenizer.prepare_seq2seq_batch(src_text, tgt_texts=tgt_text, return_tensors="pt")
-            input_ids = batch["input_ids"]
-            labels = batch["labels"]
-            self.assertTrue((input_ids[:, 0] == tokenizer.bos_token_id).all().item())
-            self.assertTrue((labels[:, 0] == tokenizer.bos_token_id).all().item())
-            self.assertTrue((input_ids[:, -1] == tokenizer.eos_token_id).all().item())
-            self.assertTrue((labels[:, -1] == tokenizer.eos_token_id).all().item())
@@ -1544,11 +1544,11 @@ class TokenizerTesterMixin:
            src_texts=src_text, tgt_texts=tgt_text, max_length=3, max_target_length=10, return_tensors="pt"
        )
        self.assertEqual(batch.input_ids.shape[1], 3)
-        self.assertEqual(batch.labels.shape[1], 10)
+        self.assertEqual(batch.decoder_input_ids.shape[1], 10)
        # max_target_length will default to max_length if not specified
        batch = tokenizer.prepare_seq2seq_batch(src_text, tgt_texts=tgt_text, max_length=3)
        self.assertEqual(batch.input_ids.shape[1], 3)
-        self.assertEqual(batch.labels.shape[1], 3)
+        self.assertEqual(batch.decoder_input_ids.shape[1], 3)

        batch_encoder_only = tokenizer.prepare_seq2seq_batch(
            src_texts=src_text, max_length=3, max_target_length=10, return_tensors="pt"
@@ -182,17 +182,3 @@ class MBartEnroIntegrationTest(unittest.TestCase):
        self.tokenizer.save_pretrained(tmpdirname)
        new_tok = MBartTokenizer.from_pretrained(tmpdirname)
        self.assertDictEqual(new_tok.fairseq_tokens_to_ids, original_special_tokens)
-
-    def test_batch_fairseq_parity(self):
-        batch: BatchEncoding = self.tokenizer.prepare_seq2seq_batch(
-            self.src_text, tgt_texts=self.tgt_text, return_tensors="pt"
-        )
-        for k in batch:
-            batch[k] = batch[k].tolist()
-        # batch = {k: v.tolist() for k,v in batch.items()}
-        # fairseq batch: https://gist.github.com/sshleifer/cba08bc2109361a74ac3760a7e30e4f4
-        # batch.decoder_inputs_ids[0][0] ==
-        assert batch.input_ids[1][-2:] == [2, EN_CODE]
-        assert batch.decoder_input_ids[1][0] == RO_CODE
-        assert batch.decoder_input_ids[1][-1] == 2
-        assert batch.labels[1][-2:] == [2, RO_CODE]
@@ -63,6 +63,7 @@ class PegasusTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
        batch = self.pegasus_large_tokenizer.prepare_seq2seq_batch(src_texts, tgt_texts=tgt_texts, max_target_length=5)
        assert batch.input_ids.shape == (2, 1024)
        assert batch.attention_mask.shape == (2, 1024)
-        assert "labels" in batch  # because tgt_texts was specified
-        assert batch.labels.shape == (2, 5)
-        assert len(batch) == 3  # input_ids, attention_mask, labels. Other things make by BartModel
+        assert "decoder_input_ids" in batch  # because tgt_texts was specified
+        assert batch.decoder_input_ids.shape == (2, 5)
+        assert batch.decoder_attention_mask.shape == (2, 5)
+        assert len(batch) == 4  # no extra keys
@@ -1,253 +0,0 @@
-# coding=utf-8
-# Copyright 2018 The Google AI Language Team Authors.
-#
-# Licensed under the Apache License, Version 2.0 (the "License");
-# you may not use this file except in compliance with the License.
-# You may obtain a copy of the License at
-#
-#     http://www.apache.org/licenses/LICENSE-2.0
-#
-# Unless required by applicable law or agreed to in writing, software
-# distributed under the License is distributed on an "AS IS" BASIS,
-# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
-# See the License for the specific language governing permissions and
-# limitations under the License.
-
-
-import os
-import unittest
-
-from transformers.file_utils import cached_property
-from transformers.testing_utils import require_torch, slow
-from transformers.tokenization_reformer import SPIECE_UNDERLINE, ReformerTokenizer
-
-from .test_tokenization_common import TokenizerTesterMixin
-
-
-SAMPLE_VOCAB = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fixtures/test_sentencepiece.model")
-
-
-class ReformerTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
-
-    tokenizer_class = ReformerTokenizer
-
-    def setUp(self):
-        super().setUp()
-
-        tokenizer = ReformerTokenizer(SAMPLE_VOCAB, keep_accents=True)
-        tokenizer.save_pretrained(self.tmpdirname)
-
-    def test_full_tokenizer(self):
-        tokenizer = ReformerTokenizer(SAMPLE_VOCAB, keep_accents=True)
-
-        tokens = tokenizer.tokenize("This is a test")
-        self.assertListEqual(tokens, ["▁This", "▁is", "▁a", "▁t", "est"])
-
-        self.assertListEqual(
-            tokenizer.convert_tokens_to_ids(tokens), [285, 46, 10, 170, 382],
-        )
-
-        tokens = tokenizer.tokenize("I was born in 92000, and this is falsé.")
-        self.assertListEqual(
-            tokens,
-            [
-                SPIECE_UNDERLINE + "I",
-                SPIECE_UNDERLINE + "was",
-                SPIECE_UNDERLINE + "b",
-                "or",
-                "n",
-                SPIECE_UNDERLINE + "in",
-                SPIECE_UNDERLINE + "",
-                "9",
-                "2",
-                "0",
-                "0",
-                "0",
-                ",",
-                SPIECE_UNDERLINE + "and",
-                SPIECE_UNDERLINE + "this",
-                SPIECE_UNDERLINE + "is",
-                SPIECE_UNDERLINE + "f",
-                "al",
-                "s",
-                "é",
-                ".",
-            ],
-        )
-        ids = tokenizer.convert_tokens_to_ids(tokens)
-        self.assertListEqual(
-            ids, [8, 21, 84, 55, 24, 19, 7, 0, 602, 347, 347, 347, 3, 12, 66, 46, 72, 80, 6, 0, 4],
-        )
-
-        back_tokens = tokenizer.convert_ids_to_tokens(ids)
-        self.assertListEqual(
-            back_tokens,
-            [
-                SPIECE_UNDERLINE + "I",
-                SPIECE_UNDERLINE + "was",
-                SPIECE_UNDERLINE + "b",
-                "or",
-                "n",
-                SPIECE_UNDERLINE + "in",
-                SPIECE_UNDERLINE + "",
-                "<unk>",
-                "2",
-                "0",
-                "0",
-                "0",
-                ",",
-                SPIECE_UNDERLINE + "and",
-                SPIECE_UNDERLINE + "this",
-                SPIECE_UNDERLINE + "is",
-                SPIECE_UNDERLINE + "f",
-                "al",
-                "s",
-                "<unk>",
-                ".",
-            ],
-        )
-
-    @cached_property
-    def big_tokenizer(self):
-        return ReformerTokenizer.from_pretrained("google/reformer-crime-and-punishment")
-
-    @slow
-    def test_tokenization_base_easy_symbols(self):
-        symbols = "Hello World!"
-        original_tokenizer_encodings = [126, 32, 262, 152, 38, 72, 287]
-
-        self.assertListEqual(original_tokenizer_encodings, self.big_tokenizer.encode(symbols))
-
-    @slow
-    def test_tokenization_base_hard_symbols(self):
-        symbols = 'This is a very long text with a lot of weird characters, such as: . , ~ ? ( ) " [ ] ! : - . Also we will add words that should not exsist and be tokenized to <unk>, such as saoneuhaoesuth'
-        original_tokenizer_encodings = [
-            108,
-            265,
-            24,
-            111,
-            4,
-            258,
-            156,
-            35,
-            28,
-            275,
-            3,
-            259,
-            297,
-            260,
-            84,
-            4,
-            35,
-            110,
-            44,
-            8,
-            259,
-            91,
-            268,
-            21,
-            11,
-            209,
-            274,
-            109,
-            266,
-            277,
-            117,
-            86,
-            93,
-            315,
-            258,
-            278,
-            258,
-            277,
-            258,
-            0,
-            258,
-            288,
-            258,
-            319,
-            258,
-            0,
-            258,
-            0,
-            258,
-            0,
-            258,
-            0,
-            258,
-            287,
-            258,
-            315,
-            258,
-            289,
-            258,
-            278,
-            99,
-            269,
-            266,
-            262,
-            8,
-            259,
-            241,
-            4,
-            217,
-            230,
-            268,
-            266,
-            55,
-            168,
-            106,
-            75,
-            193,
-            266,
-            223,
-            27,
-            49,
-            26,
-            282,
-            25,
-            264,
-            299,
-            19,
-            26,
-            0,
-            258,
-            277,
-            117,
-            86,
-            93,
-            176,
-            183,
-            270,
-            11,
-            262,
-            42,
-            61,
-            265,
-        ]
-
-        self.assertListEqual(original_tokenizer_encodings, self.big_tokenizer.encode(symbols))
-
-    @slow
-    @require_torch
-    def test_torch_encode_plus_sent_to_model(self):
-        import torch
-        from transformers import ReformerModel, ReformerConfig
-
-        # Build sequence
-        first_ten_tokens = list(self.big_tokenizer.get_vocab().keys())[:10]
-        sequence = " ".join(first_ten_tokens)
-        encoded_sequence = self.big_tokenizer.encode_plus(sequence, return_tensors="pt")
-        batch_encoded_sequence = self.big_tokenizer.batch_encode_plus([sequence, sequence], return_tensors="pt")
-
-        config = ReformerConfig()
-        # The input gets padded during training so adjust the axial position encodings from the pretrained model value of (512, 1024)
-        config.axial_pos_shape = encoded_sequence["input_ids"].shape
-        model = ReformerModel(config)
-
-        # Reformer has config.vocab_size == tokenizer.vocab_size == len(tokenizer) - 1 = 320; len(tokenizer) is 321 (including a pad token with id 320)
-        assert model.get_input_embeddings().weight.shape[0] >= self.big_tokenizer.vocab_size
-
-        with torch.no_grad():
-            model(**encoded_sequence)
-            model(**batch_encoded_sequence)
@@ -1,7 +1,5 @@
 import unittest

-import numpy as np
-
 from transformers import AutoTokenizer, TrainingArguments, is_torch_available
 from transformers.testing_utils import require_torch

@@ -12,38 +10,149 @@ if is_torch_available():

    from transformers import (
        AutoModelForSequenceClassification,
+        DataCollatorForLanguageModeling,
+        DataCollatorForPermutationLanguageModeling,
        GlueDataset,
        GlueDataTrainingArguments,
        LineByLineTextDataset,
+        TextDataset,
        Trainer,
+        default_data_collator,
    )


 PATH_SAMPLE_TEXT = "./tests/fixtures/sample_text.txt"


-class RegressionDataset:
-    def __init__(self, a=2, b=3, length=64, seed=42):
-        np.random.seed(seed)
-        self.length = length
-        self.x = np.random.normal(size=(length,)).astype(np.float32)
-        self.y = a * self.x + b + np.random.normal(scale=0.1, size=(length,))
+@require_torch
+class DataCollatorIntegrationTest(unittest.TestCase):
+    def test_default_with_dict(self):
+        features = [{"label": i, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
+        batch = default_data_collator(features)
+        self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8)))))
+        self.assertEqual(batch["labels"].dtype, torch.long)
+        self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))

-    def __len__(self):
-        return self.length
+        # With label_ids
+        features = [{"label_ids": [0, 1, 2], "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
+        batch = default_data_collator(features)
+        self.assertTrue(batch["labels"].equal(torch.tensor([[0, 1, 2]] * 8)))
+        self.assertEqual(batch["labels"].dtype, torch.long)
+        self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))

-    def __getitem__(self, i):
-        return {"input_x": self.x[i], "label": self.y[i]}
+        # Features can already be tensors
+        features = [{"label": i, "inputs": torch.randint(10, [10])} for i in range(8)]
+        batch = default_data_collator(features)
+        self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8)))))
+        self.assertEqual(batch["labels"].dtype, torch.long)
+        self.assertEqual(batch["inputs"].shape, torch.Size([8, 10]))

+        # Labels can already be tensors
+        features = [{"label": torch.tensor(i), "inputs": torch.randint(10, [10])} for i in range(8)]
+        batch = default_data_collator(features)
+        self.assertEqual(batch["labels"].dtype, torch.long)
+        self.assertTrue(batch["labels"].equal(torch.tensor(list(range(8)))))
+        self.assertEqual(batch["labels"].dtype, torch.long)
+        self.assertEqual(batch["inputs"].shape, torch.Size([8, 10]))

-class AlmostAccuracy:
-    def __init__(self, thresh=0.25):
-        self.thresh = thresh
+    def test_default_with_no_labels(self):
+        features = [{"label": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
+        batch = default_data_collator(features)
+        self.assertTrue("labels" not in batch)
+        self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))

-    def __call__(self, eval_pred):
-        predictions, labels = eval_pred
-        true = np.abs(predictions - labels) <= self.thresh
-        return {"accuracy": true.astype(np.float32).mean().item()}
+        # With label_ids
+        features = [{"label_ids": None, "inputs": [0, 1, 2, 3, 4, 5]} for i in range(8)]
+        batch = default_data_collator(features)
+        self.assertTrue("labels" not in batch)
+        self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))
+
+    def test_default_classification(self):
+        MODEL_ID = "bert-base-cased-finetuned-mrpc"
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        data_args = GlueDataTrainingArguments(
+            task_name="mrpc", data_dir="./tests/fixtures/tests_samples/MRPC", overwrite_cache=True
+        )
+        dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev")
+        data_collator = default_data_collator
+        batch = data_collator(dataset.features)
+        self.assertEqual(batch["labels"].dtype, torch.long)
+
+    def test_default_regression(self):
+        MODEL_ID = "distilroberta-base"
+        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
+        data_args = GlueDataTrainingArguments(
+            task_name="sts-b", data_dir="./tests/fixtures/tests_samples/STS-B", overwrite_cache=True
+        )
+        dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev")
+        data_collator = default_data_collator
+        batch = data_collator(dataset.features)
+        self.assertEqual(batch["labels"].dtype, torch.float)
+
+    def test_lm_tokenizer_without_padding(self):
+        tokenizer = AutoTokenizer.from_pretrained("gpt2")
+        data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
+        # ^ causal lm
+
+        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
+        examples = [dataset[i] for i in range(len(dataset))]
+        with self.assertRaises(ValueError):
+            # Expect error due to padding token missing on gpt2:
+            data_collator(examples)
+
+        dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True)
+        examples = [dataset[i] for i in range(len(dataset))]
+        batch = data_collator(examples)
+        self.assertIsInstance(batch, dict)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
+
+    def test_lm_tokenizer_with_padding(self):
+        tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
+        data_collator = DataCollatorForLanguageModeling(tokenizer)
+        # ^ masked lm
+
+        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
+        examples = [dataset[i] for i in range(len(dataset))]
+        batch = data_collator(examples)
+        self.assertIsInstance(batch, dict)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((31, 107)))
+        self.assertEqual(batch["labels"].shape, torch.Size((31, 107)))
+
+        dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True)
+        examples = [dataset[i] for i in range(len(dataset))]
+        batch = data_collator(examples)
+        self.assertIsInstance(batch, dict)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
+
+    def test_plm(self):
+        tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased")
+        data_collator = DataCollatorForPermutationLanguageModeling(tokenizer)
+        # ^ permutation lm
+
+        dataset = LineByLineTextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512)
+        examples = [dataset[i] for i in range(len(dataset))]
+        batch = data_collator(examples)
+        self.assertIsInstance(batch, dict)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((31, 112)))
+        self.assertEqual(batch["perm_mask"].shape, torch.Size((31, 112, 112)))
+        self.assertEqual(batch["target_mapping"].shape, torch.Size((31, 112, 112)))
+        self.assertEqual(batch["labels"].shape, torch.Size((31, 112)))
+
+        dataset = TextDataset(tokenizer, file_path=PATH_SAMPLE_TEXT, block_size=512, overwrite_cache=True)
+        examples = [dataset[i] for i in range(len(dataset))]
+        batch = data_collator(examples)
+        self.assertIsInstance(batch, dict)
+        self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
+        self.assertEqual(batch["perm_mask"].shape, torch.Size((2, 512, 512)))
+        self.assertEqual(batch["target_mapping"].shape, torch.Size((2, 512, 512)))
+        self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
+
+        example = [torch.randint(5, [5])]
+        with self.assertRaises(ValueError):
+            # Expect error due to odd sequence length
+            data_collator(example)


 if is_torch_available():
@@ -59,137 +168,9 @@ if is_torch_available():
        def __iter__(self):
            return iter(self.parse_file())

-    class RegressionModel(torch.nn.Module):
-        def __init__(self, a=0, b=0):
-            super().__init__()
-            self.a = torch.nn.Parameter(torch.tensor(a).float())
-            self.b = torch.nn.Parameter(torch.tensor(b).float())
-
-        def forward(self, input_x=None, labels=None):
-            y = input_x * self.a + self.b
-            if labels is None:
-                return (y,)
-            loss = torch.nn.functional.mse_loss(y, labels)
-            return (loss, y)
-
-    def get_regression_trainer(a=0, b=0, train_len=64, eval_len=64, **kwargs):
-        train_dataset = RegressionDataset(length=train_len)
-        eval_dataset = RegressionDataset(length=eval_len)
-        model = RegressionModel(a, b)
-        compute_metrics = kwargs.pop("compute_metrics", None)
-        data_collator = kwargs.pop("data_collator", None)
-        optimizers = kwargs.pop("optimizers", (None, None))
-        args = TrainingArguments("./regression", **kwargs)
-        return Trainer(
-            model,
-            args,
-            data_collator=data_collator,
-            train_dataset=train_dataset,
-            eval_dataset=eval_dataset,
-            compute_metrics=compute_metrics,
-            optimizers=optimizers,
-        )
-

@require_torch
 class TrainerIntegrationTest(unittest.TestCase):
-    def setUp(self):
-        # Get the default values (in case they change):
-        args = TrainingArguments(".")
-        self.n_epochs = args.num_train_epochs
-        self.batch_size = args.per_device_train_batch_size
-
-    def test_reproducible_training(self):
-        # Checks that training worked, model trained and seed made a reproducible training.
-        trainer = get_regression_trainer(learning_rate=0.1)
-        trainer.train()
-        self.assertTrue(torch.abs(trainer.model.a - 0.6975) < 1e-4)
-        self.assertTrue(torch.abs(trainer.model.b - 1.2415) < 1e-4)
-
-        # Checks that a different seed gets different (reproducible) results.
-        trainer = get_regression_trainer(learning_rate=0.1, seed=314)
-        trainer.train()
-        self.assertTrue(torch.abs(trainer.model.a - 1.0171) < 1e-4)
-        self.assertTrue(torch.abs(trainer.model.b - 1.2494) < 1e-4)
-
-    def test_number_of_steps_in_training(self):
-        # Regular training has n_epochs * len(train_dl) steps
-        trainer = get_regression_trainer(learning_rate=0.1)
-        train_output = trainer.train()
-        self.assertEqual(train_output.global_step, self.n_epochs * 64 / self.batch_size)
-
-        # Check passing num_train_epochs works (and a float version too):
-        trainer = get_regression_trainer(learning_rate=0.1, num_train_epochs=1.5)
-        train_output = trainer.train()
-        self.assertEqual(train_output.global_step, int(1.5 * 64 / self.batch_size))
-
-        # If we pass a max_steps, num_train_epochs is ignored
-        trainer = get_regression_trainer(learning_rate=0.1, max_steps=10)
-        train_output = trainer.train()
-        self.assertEqual(train_output.global_step, 10)
-
-    def test_train_and_eval_dataloaders(self):
-        trainer = get_regression_trainer(learning_rate=0.1, per_device_train_batch_size=16)
-        self.assertEqual(trainer.get_train_dataloader().batch_size, 16)
-        trainer = get_regression_trainer(learning_rate=0.1, per_device_eval_batch_size=16)
-        self.assertEqual(trainer.get_eval_dataloader().batch_size, 16)
-
-        # Check drop_last works
-        trainer = get_regression_trainer(
-            train_len=66, eval_len=74, learning_rate=0.1, per_device_train_batch_size=16, per_device_eval_batch_size=32
-        )
-        self.assertEqual(len(trainer.get_train_dataloader()), 66 // 16 + 1)
-        self.assertEqual(len(trainer.get_eval_dataloader()), 74 // 32 + 1)
-
-        trainer = get_regression_trainer(
-            train_len=66,
-            eval_len=74,
-            learning_rate=0.1,
-            per_device_train_batch_size=16,
-            per_device_eval_batch_size=32,
-            dataloader_drop_last=True,
-        )
-        self.assertEqual(len(trainer.get_train_dataloader()), 66 // 16)
-        self.assertEqual(len(trainer.get_eval_dataloader()), 74 // 32)
-
-        # Check passing a new dataset fpr evaluation wors
-        new_eval_dataset = RegressionDataset(length=128)
-        self.assertEqual(len(trainer.get_eval_dataloader(new_eval_dataset)), 128 // 32)
-
-    def test_evaluate(self):
-        trainer = get_regression_trainer(a=1.5, b=2.5, compute_metrics=AlmostAccuracy())
-        results = trainer.evaluate()
-
-        x, y = trainer.eval_dataset.x, trainer.eval_dataset.y
-        pred = 1.5 * x + 2.5
-        expected_loss = ((pred - y) ** 2).mean()
-        self.assertAlmostEqual(results["eval_loss"], expected_loss)
-        expected_acc = AlmostAccuracy()((pred, y))["accuracy"]
-        self.assertAlmostEqual(results["eval_accuracy"], expected_acc)
-
-        # With a number of elements not a round multiple of the batch size
-        trainer = get_regression_trainer(a=1.5, b=2.5, eval_len=66, compute_metrics=AlmostAccuracy())
-        results = trainer.evaluate()
-
-        x, y = trainer.eval_dataset.x, trainer.eval_dataset.y
-        pred = 1.5 * x + 2.5
-        expected_loss = ((pred - y) ** 2).mean()
-        self.assertAlmostEqual(results["eval_loss"], expected_loss)
-        expected_acc = AlmostAccuracy()((pred, y))["accuracy"]
-        self.assertAlmostEqual(results["eval_accuracy"], expected_acc)
-
-    def test_predict(self):
-        trainer = get_regression_trainer(a=1.5, b=2.5)
-        preds = trainer.predict(trainer.eval_dataset).predictions
-        x = trainer.eval_dataset.x
-        self.assertTrue(np.allclose(preds, 1.5 * x + 2.5))
-
-        # With a number of elements not a round multiple of the batch size
-        trainer = get_regression_trainer(a=1.5, b=2.5, eval_len=66)
-        preds = trainer.predict(trainer.eval_dataset).predictions
-        x = trainer.eval_dataset.x
-        self.assertTrue(np.allclose(preds, 1.5 * x + 2.5))
-
    def test_trainer_eval_mrpc(self):
        MODEL_ID = "bert-base-cased-finetuned-mrpc"
        tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -60,8 +60,7 @@ if is_torch_available():

 if __name__ == "__main__":
    parser = HfArgumentParser((TrainingArguments,))
-    sys.argv += ["--output_dir", "./examples"]
-    training_args = parser.parse_args_into_dataclasses()[0]
+    training_args = parser.parse_args_into_dataclasses(sys.argv + ["--output_dir", "./examples"])[0]

    logger.warning(
        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s",