Better warning

warn with FutureWarning when using output_attentions in the configuration
uninstalled wandb raises AttributeError
2020-06-09 19:08:27 -04:00 · 2020-06-09 19:06:49 -04:00 · 2020-06-09 18:50:56 -04:00 · 2020-06-09 23:39:06 +02:00 · 2020-06-09 17:38:28 -04:00 · 2020-06-09 23:12:43 +02:00
147 changed files with 6414 additions and 2101 deletions
@@ -8,6 +8,10 @@ __pycache__/
 # C extensions
 *.so

+# tests and logs
+tests/fixtures
+logs/
+
 # Distribution / packaging
 .Python
 build/
@@ -116,6 +120,7 @@ dmypy.json
 .pyre/

 # vscode
+.vs
 .vscode

 # Pycharm
@@ -0,0 +1,6 @@
+coverage:
+  status:
+    project:
+      default:
+        informational: true
+    patch: off
@@ -7,6 +7,14 @@ you can install them with the following command, at the root of the code reposit
 pip install -e ".[docs]"
 ```

+---
+**NOTE**
+
+You only need to generate the documentation to inspect it locally (if you're planning changes and want to 
+check how they look like before committing for instance). You don't have to commit the built documentation.
+
+---
+
 ## Packages installed

 Here's an overview of all the packages installed. If you ran the previous command installing all packages from
@@ -68,6 +76,18 @@ It should build the static app that will be available under `/docs/_build/html`
 Accepted files are reStructuredText (.rst) and Markdown (.md). Create a file with its extension and put it
 in the source directory. You can then link it to the toc-tree by putting the filename without the extension.

+## Preview the documentation in a pull request
+
+Once you have made your pull request, you can check what the documentation will look like after it's merged by
+following these steps:
+
+- Look at the checks at the bottom of the conversation page of your PR (you may need to click on "show all checks" to
+  expand them).
+- Click on "details" next to the `ci/circleci: build_doc` check.
+- In the new window, click on the "Artifacts" tab.
+- Locate the file "docs/_build/html/index.html" (or any specific page you want to check) and click on it to get a 
+  preview.
+
 ## Writing Documentation - Specification

 The `huggingface/transformers` documentation follows the
@@ -60,6 +60,7 @@ The library currently contains PyTorch and Tensorflow implementations, pre-train
    installation
    quickstart
    glossary
+    summary
    pretrained_models
    usage
    model_sharing
@@ -8,7 +8,7 @@ Recognition, Masked Language Modeling, Sentiment Analysis, Feature Extraction an
 There are two categories of pipeline abstractions to be aware about:

 - The :class:`~transformers.pipeline` which is the most powerful object encapsulating all other pipelines
- The other task-specific pipelines, such as :class:`~transformers.NerPipeline`
+- The other task-specific pipelines, such as :class:`~transformers.TokenClassificationPipeline`
  or :class:`~transformers.QuestionAnsweringPipeline`

 The pipeline abstraction
@@ -30,15 +30,15 @@ Parent class: Pipeline
 .. autoclass:: transformers.Pipeline
    :members: predict, transform, save_pretrained

-NerPipeline
-==========================================
-
-.. autoclass:: transformers.NerPipeline
-
 TokenClassificationPipeline
 ==========================================

-This class is an alias of the :class:`~transformers.NerPipeline` defined above. Please refer to that pipeline for
+.. autoclass:: transformers.TokenClassificationPipeline
+
+NerPipeline
+==========================================
+
+This class is an alias of the :class:`~transformers.TokenClassificationPipeline` defined above. Please refer to that pipeline for
 documentation and usage examples.

 FillMaskPipeline
@@ -67,6 +67,12 @@ AlbertForSequenceClassification
 .. autoclass:: transformers.AlbertForSequenceClassification
    :members:

+AlbertForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.AlbertForTokenClassification
+    :members:
+

 AlbertForQuestionAnswering
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -103,6 +109,13 @@ TFAlbertForMultipleChoice
    :members:


+TFAlbertForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFAlbertForTokenClassification
+    :members:
+
+
 TFAlbertForQuestionAnswering
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -4,8 +4,9 @@ Bart
 file a `Github Issue <https://github.com/huggingface/transformers/issues/new?assignees=&labels=&template=bug-report.md&title>`__ and assign
@sshleifer

-Paper
-~~~~~
+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
 The Bart model was `proposed <https://arxiv.org/abs/1910.13461>`_ by Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov and Luke Zettlemoyer on 29 Oct, 2019.
 According to the abstract,

@@ -16,14 +17,26 @@ According to the abstract,
 The Authors' code can be found `here <https://github.com/pytorch/fairseq/tree/master/examples/bart>`_


-Implementation Notes
-~~~~~~~~~~~~~~~~~~~~
+Implementation Notes:
+
 - Bart doesn't use :obj:`token_type_ids` for sequence classification. Use BartTokenizer.encode to get the proper splitting.
 - The forward pass of ``BartModel`` will create decoder inputs (using the helper function ``transformers.modeling_bart._prepare_bart_decoder_inputs``)  if they are not passed. This is different than some other modeling APIs.
 - Model predictions are intended to be identical to the original implementation. This only works, however, if the string you pass to ``fairseq.encode`` starts with a space.
 - ``BartForConditionalGeneration.generate`` should be used for conditional generation tasks like summarization, see the example in that docstrings
 - Models that load the ``"facebook/bart-large-cnn"`` weights will not have a ``mask_token_id``, or be able to perform mask filling tasks.

+BartConfig
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.BartConfig
+    :members:
+
+
+BartTokenizer
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.BartTokenizer
+    :members:


 BartModel
@@ -35,6 +48,13 @@ BartModel
 .. autofunction:: transformers.modeling_bart._prepare_bart_decoder_inputs


+BartForSequenceClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.BartForSequenceClassification
+    :members: forward
+
+
 BartForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -42,15 +62,3 @@ BartForConditionalGeneration
    :members: generate, forward


-BartForSequenceClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.BartForSequenceClassification
-    :members: forward
-
-BartConfig
-~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.BartConfig
-    :members:
-
@@ -1,6 +1,9 @@
 CamemBERT
 ----------------------------------------------------

+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
 The CamemBERT model was proposed in `CamemBERT: a Tasty French Language Model <https://arxiv.org/abs/1911.03894>`__
 by Louis Martin, Benjamin Muller, Pedro Javier Ortiz Suárez, Yoann Dupont, Laurent Romary, Éric Villemonte de la
 Clergerie, Djamé Seddah, and Benoît Sagot. It is based on Facebook's RoBERTa model released in 2019. It is a model
@@ -74,6 +77,13 @@ CamembertForTokenClassification
    :members:


+CamembertForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.CamembertForQuestionAnswering
+    :members:
+
+
 TFCamembertModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -100,3 +110,10 @@ TFCamembertForTokenClassification

 .. autoclass:: transformers.TFCamembertForTokenClassification
    :members:
+
+
+TFCamembertForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFCamembertForQuestionAnswering
+    :members:
@@ -1,6 +1,9 @@
 CTRL
 ----------------------------------------------------

+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
 CTRL model was proposed in `CTRL: A Conditional Transformer Language Model for Controllable Generation <https://arxiv.org/abs/1909.05858>`_
 by Nitish Shirish Keskar*, Bryan McCann*, Lav R. Varshney, Caiming Xiong and Richard Socher.
 It's a causal (unidirectional) transformer pre-trained using language modeling on a very large
@@ -1,6 +1,9 @@
 DistilBERT
 ----------------------------------------------------

+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
 The DistilBERT model was proposed in the blog post
 `Smaller, faster, cheaper, lighter: Introducing DistilBERT, a distilled version of BERT <https://medium.com/huggingface/distilbert-8cf3380435b5>`__,
 and the paper `DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter <https://arxiv.org/abs/1910.01108>`__.
@@ -72,6 +75,13 @@ DistilBertForSequenceClassification
    :members:


+DistilBertForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.DistilBertForTokenClassification
+    :members:
+
+
 DistilBertForQuestionAnswering
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -99,6 +109,22 @@ TFDistilBertForSequenceClassification
    :members:


+
+TFDistilBertForMultipleChoice
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFDistilBertForMultipleChoice
+    :members:
+
+
+
+TFDistilBertForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFDistilBertForTokenClassification
+    :members:
+
+
 TFDistilBertForQuestionAnswering
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -1,6 +1,9 @@
 ELECTRA
 ----------------------------------------------------

+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
 The ELECTRA model was proposed in the paper.
 `ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators <https://openreview.net/pdf?id=r1xMH1BtvB>`__.
 ELECTRA is a new pre-training approach which trains two transformer models: the generator and the discriminator. The
@@ -89,6 +92,13 @@ ElectraForMaskedLM
    :members:


+ElectraForSequenceClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.ElectraForSequenceClassification
+    :members:
+
+
 ElectraForTokenClassification
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -122,3 +132,10 @@ TFElectraForTokenClassification

 .. autoclass:: transformers.TFElectraForTokenClassification
    :members:
+
+
+TFElectraForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFElectraForQuestionAnswering
+    :members:
@@ -1,6 +1,9 @@
 FlauBERT
 ----------------------------------------------------

+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
 The FlauBERT model was proposed in the paper
 `FlauBERT: Unsupervised Language Model Pre-training for French <https://arxiv.org/abs/1912.05372>`__ by Hang Le et al.
 It's a transformer pre-trained using a masked language modeling (MLM) objective (BERT-like).
@@ -72,3 +75,43 @@ FlaubertForQuestionAnswering
    :members:


+TFFlaubertModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFFlaubertModel
+    :members:
+
+
+TFFlaubertWithLMHeadModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFFlaubertWithLMHeadModel
+    :members:
+
+
+TFFlaubertForSequenceClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFFlaubertForSequenceClassification
+    :members:
+
+
+TFFlaubertForMultipleChoice
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFFlaubertForMultipleChoice
+    :members:
+
+
+TFFlaubertForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFFlaubertForTokenClassification
+    :members:
+
+
+TFFlaubertForQuestionAnsweringSimple
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFFlaubertForQuestionAnsweringSimple
+    :members:
@@ -4,7 +4,7 @@ Longformer
 file a `Github Issue <https://github.com/huggingface/transformers/issues/new?assignees=&labels=&template=bug-report.md&title>`_

 Overview
-~~~~~
+~~~~~~~~~
 The Longformer model was presented in `Longformer: The Long-Document Transformer <https://arxiv.org/pdf/2004.05150.pdf>`_ by Iz Beltagy, Matthew E. Peters, Arman Cohan.
 Here the abstract: 

@@ -13,7 +13,7 @@ Here the abstract:
 The Authors' code can be found `here <https://github.com/allenai/longformer>`_ .

 Longformer Self Attention
-~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~
 Longformer self attention employs self attention on both a "local" context and a "global" context.
 Most tokens only attend "locally" to each other meaning that each token attends to its :math:`\frac{1}{2} w` previous tokens and :math:`\frac{1}{2} w` succeding tokens with :math:`w` being the window length as defined in `config.attention_window`. Note that `config.attention_window` can be of type ``list`` to define a different :math:`w` for each layer. 
 A selecetd few tokens attend "globally" to all other tokens, as it is conventionally done for all tokens in *e.g.* `BertSelfAttention`.
@@ -55,6 +55,13 @@ LongformerTokenizer
    :members: 


+LongformerTokenizerFast
+~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LongformerTokenizerFast
+    :members: 
+
+
 LongformerModel
 ~~~~~~~~~~~~~~~~~~~~

@@ -69,10 +76,10 @@ LongformerForMaskedLM
    :members:


-LongformerForQuestionAnswering
+LongformerForSequenceClassification
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

-.. autoclass:: transformers.LongformerForQuestionAnswering
+.. autoclass:: transformers.LongformerForSequenceClassification
    :members:


@@ -89,3 +96,9 @@ LongformerForTokenClassification
 .. autoclass:: transformers.LongformerForTokenClassification
    :members:

+
+LongformerForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LongformerForQuestionAnswering
+    :members:
@@ -6,11 +6,11 @@ file a `Github Issue <https://github.com/huggingface/transformers/issues/new?ass

 Implementation Notes
 ~~~~~~~~~~~~~~~~~~~~
- each model is about 298 MB on disk, there are 1,000+ models.
+- Each model is about 298 MB on disk, there are 1,000+ models.
 - The list of supported language pairs can be found `here <https://huggingface.co/Helsinki-NLP>`__.
 - The 1,000+ models were originally trained by `Jörg Tiedemann <https://researchportal.helsinki.fi/en/persons/j%C3%B6rg-tiedemann>`__ using the `Marian <https://marian-nmt.github.io/>`_ C++ library, which supports fast training and translation.
 - All models are transformer encoder-decoders with 6 layers in each component. Each model's performance is documented in a model card.
- the 80 opus models that require BPE preprocessing are not supported.
+- The 80 opus models that require BPE preprocessing are not supported.
 - The modeling code is the same as ``BartForConditionalGeneration`` with a few minor modifications:
    - static (sinusoid) positional embeddings (``MarianConfig.static_position_embeddings=True``)
    - a new final_logits_bias (``MarianConfig.add_bias_logits=True``)
@@ -86,6 +86,19 @@ Code to see available pretrained models:
    suffix = [x.split('/')[1] for x in model_ids]
    multi_models = [f'{org}/{s}' for s in suffix if s != s.lower()]

+MarianConfig
+~~~~~~~~~~~~~~~~~~~
+.. autoclass:: transformers.MarianConfig
+    :members:
+
+
+MarianTokenizer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.MarianTokenizer
+    :members: prepare_translation_batch
+
+
 MarianMTModel
 ~~~~~~~~~~~~~

@@ -96,10 +109,3 @@ This class inherits all functionality from ``BartForConditionalGeneration``, see

 .. autoclass:: transformers.MarianMTModel
    :members:
-
-
-MarianTokenizer
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.MarianTokenizer
-    :members: prepare_translation_batch
@@ -1,6 +1,9 @@
 RoBERTa
 ----------------------------------------------------

+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
 The RoBERTa model was proposed in `RoBERTa: A Robustly Optimized BERT Pretraining Approach <https://arxiv.org/abs/1907.11692>`_
 by Yinhan Liu, Myle Ott, Naman Goyal, Jingfei Du, Mandar Joshi, Danqi Chen, Omer Levy, Mike Lewis, Luke Zettlemoyer,
 Veselin Stoyanov. It is based on Google's BERT model released in 2018.
@@ -87,6 +90,14 @@ RobertaForTokenClassification
 .. autoclass:: transformers.RobertaForTokenClassification
    :members:

+
+RobertaForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.RobertaForQuestionAnswering
+    :members:
+
+
 TFRobertaModel
 ~~~~~~~~~~~~~~~~~~~~

@@ -108,8 +119,22 @@ TFRobertaForSequenceClassification
    :members:


+TFRobertaForMultipleChoice
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFRobertaForMultipleChoice
+    :members:
+
+
 TFRobertaForTokenClassification
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFRobertaForTokenClassification
    :members:
+
+
+TFRobertaForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFRobertaForQuestionAnswering
+    :members:
@@ -4,7 +4,8 @@ T5
 file a `Github Issue <https://github.com/huggingface/transformers/issues/new?assignees=&labels=&template=bug-report.md&title>`_

 Overview
-~~~~~
+~~~~~~~~~~~~~~~~~~~~~
+
 The T5 model was presented in `Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer <https://arxiv.org/pdf/1910.10683.pdf>`_ by Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu in 
 Here the abstract: 

@@ -14,10 +15,20 @@ Our systematic study compares pre-training objectives, architectures, unlabeled
 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.*

-The Authors' code can be found `here <https://github.com/google-research/text-to-text-transfer-transformer>`_ .
+Tips:
+
+- T5 is an encoder-decoder model pre-trained on a multi-task mixture of unsupervised 
+  and supervised tasks and for which each task is converted into a text-to-text format.
+  T5 works well on a variety of tasks out-of-the-box by prepending a different prefix to the input corresponding to each task, e.g.: for translation: *translate English to German: ..., summarize: ...*.
+  For more information about which prefix to use, it is easiest to look into Appendix D of the `paper <https://arxiv.org/pdf/1910.10683.pdf>`_ .
+- For sequence to sequence generation, it is recommended to use ``T5ForConditionalGeneration.generate()``. The method takes care of feeding the encoded input via cross-attention layers to the decoder and auto-regressively generates the decoder output.
+- T5 uses relative scalar embeddings. Encoder input padding can be done on the left and on the right.
+
+The original code can be found `here <https://github.com/google-research/text-to-text-transfer-transformer>`_.

 Training
-~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~
+
 T5 is an encoder-decoder model and converts all NLP problems into a text-to-text format. It is trained using teacher forcing.
 This means that for training we always need an input sequence and a target sequence. 
 The input sequence is fed to the model using ``input_ids``. The target sequence is shifted to the right, *i.e.* prepended by a start-sequence token and fed to the decoder using the `decoder_input_ids`. In teacher-forcing style, the target sequence is then appended by the EOS token and corresponds to the ``lm_labels``. The PAD token is hereby used as the start-sequence token.
@@ -50,17 +61,6 @@ T5 can be trained / fine-tuned both in a supervised and unsupervised fashion.
  # the forward function automatically creates the correct decoder_input_ids
  model(input_ids=input_ids, lm_labels=lm_labels)

-Tips
-~~~~~~~~~~~~~~~~~~~~
- T5 is an encoder-decoder model pre-trained on a multi-task mixture of unsupervised 
-  and supervised tasks and for which each task is converted into a text-to-text format.
-  T5 works well on a variety of tasks out-of-the-box by prepending a different prefix to the input corresponding to each task, e.g.: for translation: *translate English to German: ..., summarize: ...*.
-  For more information about which prefix to use, it is easiest to look into Appendix D of the `paper <https://arxiv.org/pdf/1910.10683.pdf>`_ .
- For sequence to sequence generation, it is recommended to use ``T5ForConditionalGeneration.generate()``. The method takes care of feeding the encoded input via cross-attention layers to the decoder and auto-regressively generates the decoder output.
- T5 uses relative scalar embeddings. Encoder input padding can be done on the left and on the right.
-
-The original code can be found `here <https://github.com/google-research/text-to-text-transfer-transformer>`_.
-

 T5Config
 ~~~~~~~~~~~~~~~~~~~~~
@@ -99,7 +99,7 @@ TFT5Model


 TFT5ForConditionalGeneration
-~~~~~~~~~~~~~~~~~~~~~~~~~~
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFT5ForConditionalGeneration
    :members:
@@ -102,6 +102,21 @@ TFXLMForSequenceClassification
    :members:


+TFXLMForMultipleChoice
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFXLMForMultipleChoice
+    :members:
+
+
+TFXLMForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFXLMForTokenClassification
+    :members:
+
+
+
 TFXLMForQuestionAnsweringSimple
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -1,6 +1,9 @@
 XLM-RoBERTa
 ------------------------------------------

+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
 The XLM-RoBERTa model was proposed in `Unsupervised Cross-lingual Representation Learning at Scale <https://arxiv.org/abs/1911.02116>`__
 by Alexis Conneau, Kartikay Khandelwal, Naman Goyal, Vishrav Chaudhary, Guillaume Wenzek, Francisco Guzmán,
 Edouard Grave, Myle Ott, Luke Zettlemoyer and Veselin Stoyanov. It is based on Facebook's RoBERTa model released in 2019.
@@ -81,6 +84,13 @@ XLMRobertaForTokenClassification
    :members:


+XLMRobertaForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.XLMRobertaForQuestionAnswering
+    :members:
+
+
 TFXLMRobertaModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -102,8 +112,22 @@ TFXLMRobertaForSequenceClassification
    :members:


+TFXLMRobertaForMultipleChoice
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFXLMRobertaForMultipleChoice
+    :members:
+
+
 TFXLMRobertaForTokenClassification
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFXLMRobertaForTokenClassification
    :members:
+
+
+TFXLMRobertaForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFXLMRobertaForQuestionAnswering
+    :members:
@@ -71,13 +71,6 @@ XLNetForSequenceClassification
    :members:


-XLNetForTokenClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.XLNetForTokenClassification
-    :members:
-
-
 XLNetForMultipleChoice
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -85,6 +78,13 @@ XLNetForMultipleChoice
    :members:


+XLNetForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.XLNetForTokenClassification
+    :members:
+
+
 XLNetForQuestionAnsweringSimple
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -120,6 +120,20 @@ TFXLNetForSequenceClassification
    :members:


+TFLNetForMultipleChoice
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFXLNetForMultipleChoice
+    :members:
+
+
+TFXLNetForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFXLNetForTokenClassification
+    :members:
+
+
 TFXLNetForQuestionAnsweringSimple
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -0,0 +1,618 @@
+Summary of the models
+================================================
+
+This is a summary of the models available in the transformers library. It assumes you’re familiar with the original 
+`transformer model <https://arxiv.org/abs/1706.03762>`_. For a gentle introduction check the `annotated transformer 
+<http://nlp.seas.harvard.edu/2018/04/03/attention.html>`_. Here we focus on the high-level differences between the
+models. You can check them more in detail in their respective documentation. Also checkout the 
+:doc:`pretrained model page </pretrained_models>` to see the checkpoints available for each type of model and all `the 
+community models <https://huggingface.co/models>`_.
+
+Each one of the models in the library falls into one of the following categories:
+
+  * :ref:`autoregressive-models`
+  * :ref:`autoencoding-models`
+  * :ref:`seq-to-seq-models`
+  * :ref:`multimodal-models`
+
+Autoregressive models are pretrained on the classic language modeling task: guess the next token having read all the 
+previous ones. They correspond to the decoder of the original transformer model, and a mask is used on top of the full 
+sentence so that the attention heads can only see what was before in the next, and not what’s after. Although those 
+models can be fine-tuned and achieve great results on many tasks, the most natural application is text generation. 
+A typical example of such models is GPT.
+
+Autoencoding models are pretrained by corrupting the input tokens in some way and trying to reconstruct the original 
+sentence. They correspond to the encoder of the original transformer model in the sense that they get access to the 
+full inputs without any mask. Those models usually build a bidirectional representation of the whole sentence. They can 
+be fine-tuned and achieve great results on many tasks such as text generation, but their most natural application is 
+sentence classification or token classification. A typical example of such models is BERT.
+
+Note that the only difference between autoregressive models and autoencoding models is in the way the model is 
+pretrained. Therefore, the same architecture can be used for both autoregressive and autoencoding models. When a given
+model has been used for both pretraining, we have put it in the category corresponding to the article it was first
+introduced.
+
+Sequence-to-sequence models use both the encoder and the decoder of the original transformer, either for translation 
+tasks or by transforming other tasks to sequence-to-sequence problems. They can be fine-tuned to many tasks but their 
+most natural applications are translation, summarization and question answering. The original transformer model is an 
+example of such a model (only for translation), T5 is an example that can be fine-tuned on other tasks.
+
+Multimodal models mix text inputs with other kinds (like image) and are more specific to a given task.
+
+.. _autoregressive-models:
+
+Autoregressive models
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+As mentioned before, these models rely on the decoder part of the original transformer and use an attention mask so 
+that at each position, the model can only look at the tokens before in the attention heads.
+
+Original GPT
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=openai-gpt">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-openai--gpt-blueviolet">
+   </a>
+   <a href="/model_doc/gpt">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-openai--gpt-blueviolet">
+   </a>
+
+`Improving Language Understanding by Generative Pre-Training <https://cdn.openai.com/research-covers/language-unsupervised/language_understanding_paper.pdf>`_, 
+Alec Radford et al.
+
+The first autoregressive model based on the transformer architecture, pretrained on the Book Corpus dataset.
+
+The library provides versions of the model for language modeling and multitask language modeling/multiple choice 
+classification.
+
+GPT-2
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=gpt2">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-gpt2-blueviolet">
+   </a>
+   <a href="/model_doc/gpt2">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-gpt2-blueviolet">
+   </a>
+
+`Language Models are Unsupervised Multitask Learners <https://d4mucfpksywv.cloudfront.net/better-language-models/language_models_are_unsupervised_multitask_learners.pdf>`_, 
+Alec Radford et al.
+
+A bigger and better version of GPT, pretrained on WebText (web pages from outgoing links in Reddit with 3 karmas or 
+more).
+
+The library provides versions of the model for language modeling and multitask language modeling/multiple choice 
+classification.
+
+CTRL
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=ctrl">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-ctrl-blueviolet">
+   </a>
+   <a href="/model_doc/ctrl">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-ctrl-blueviolet">
+   </a>
+
+`CTRL: A Conditional Transformer Language Model for Controllable Generation <https://arxiv.org/abs/1909.05858>`_, 
+Nitish Shirish Keskar et al.
+
+Same as the GPT model but adds the idea of control codes. Text is generated from a prompt (can be empty) and one (or 
+several) of those control codes which are then used to influence the text generation: generate with the style of 
+wikipedia article, a book or a movie review.
+
+The library provides a version of the model for language modeling only.
+
+Transformer-XL
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=transfo-xl">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-transfo--xl-blueviolet">
+   </a>
+   <a href="/model_doc/transformerxl">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-transfo--xl-blueviolet">
+   </a>
+
+`Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context <https://arxiv.org/abs/1901.02860>`_, 
+Zihang Dai et al.
+
+Same as a regular GPT model, but introduces a recurrence mechanism for two consecutive segments (similar to a regular 
+RNNs with two consecutive inputs). In this context, a segment is a number of consecutive tokens (for instance 512) that 
+may span across multiple documents, and segments are fed in order to the model.
+
+Basically, the hidden states of the previous segment are concatenated to the current input to compute the attention 
+scores. This allows the model to pay attention to information that was in the previous segment as well as the current 
+one. By stacking multiple attention layers, the receptive field can be increased to multiple previous segments.
+
+This changes the positional embeddings to positional relative embeddings (as the regular positional embeddings would 
+give the same results in the current input and the current hidden state at a given position) and needs to make some 
+adjustments in the way attention scores are computed.
+
+The library provides a version of the model for language modeling only.
+
+.. _reformer:
+
+Reformer
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=reformer">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-reformer-blueviolet">
+   </a>
+   <a href="/model_doc/reformer">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-reformer-blueviolet">
+   </a>
+
+`Reformer: The Efficient Transformer <https://arxiv.org/abs/2001.04451>`_,
+Nikita Kitaev et al .
+
+An autoregressive transformer model with lots of tricks to reduce memory footprint and compute time. Those tricks 
+include:
+
+  * Use :ref:`Axial position encoding <axial-pos-encoding>` (see below for more details). It’s a mechanism to avoid 
+    having a huge positional encoding matrix (when the sequence length is very big) by factorizing it in smaller 
+    matrices.
+  * Replace traditional attention by :ref:`LSH (local-sensitive hashing) attention <lsh-attention>` (see below for more 
+    details). It's a technique to avoid compute the full product query-key in the attention layers.
+  * Avoid storing the intermediate results of each layer by using reversible transformer layers to obtain them during 
+    the backward pass (subtracting the residuals from the input of the next layer gives them back) or recomputing them 
+    for results inside a given layer (less efficient than storing them but saves memory).
+  * Compute the feedforward operations by chunks and not on the whole batch.
+
+With those tricks, the model can be fed much larger sentences than traditional transformer autoregressive models.
+
+**Note:** This model could be very well be used in an autoencoding setting, there is no checkpoint for such a
+pretraining yet, though.
+
+The library provides a version of the model for language modeling only.
+
+XLNet
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=xlnet">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-xlnet-blueviolet">
+   </a>
+   <a href="/model_doc/xlnet">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-xlnet-blueviolet">
+   </a>
+
+`XLNet: Generalized Autoregressive Pretraining for Language Understanding <https://arxiv.org/abs/1906.08237>`_,
+Zhilin Yang et al.
+
+XLNet is not a traditional autoregressive model but uses a training strategy that builds on that. It permutes the 
+tokens in the sentence, then allows the model to use the last n tokens to predict the token n+1. Since this is all done 
+with a mask, the sentence is actually fed in the model in the right order, but instead of masking the first n tokens 
+for n+1, XLNet uses a mask that hides the previous tokens in some given permutation of 1,...,sequence length.
+
+XLNet also uses the same recurrence mechanism as TransformerXL to build long-term dependencies. 
+
+The library provides a version of the model for language modeling, token classification, sentence classification, 
+multiple choice classification and question answering.
+
+.. _autoencoding-models:
+
+Autoencoding models
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+As mentioned before, these models rely on the encoder part of the original transformer and use no mask so the model can
+look at all the tokens in the attention heads. For pretraining, inputs are a corrupted version of the sentence, usually 
+obtained by masking tokens, and targets are the original sentences.
+
+BERT
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=bert">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-bert-blueviolet">
+   </a>
+   <a href="/model_doc/bert">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-bert-blueviolet">
+   </a>
+
+`BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding <https://arxiv.org/abs/1810.04805>`_,
+Jacob Devlin et al.
+
+Corrupts the inputs by using random masking, more precisely, during pretraining, a given percentage of tokens (usually 
+15%) are masked by
+ 
+  * a special mask token with probability 0.8
+  * a random token different from the one masked with probability 0.1
+  * the same token with probability 0.1
+
+The model must predict the original sentence, but has a second objective: inputs are two sentences A and B (with a 
+separation token in between). With probability 50%, the sentences are consecutive in the corpus, in the remaining 50% 
+they are not related. The model has to predict if the sentences are consecutive or not.
+
+The library provides a version of the model for language modeling (traditional or masked), next sentence prediction, 
+token classification, sentence classification, multiple choice classification and question answering.
+
+ALBERT
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=albert">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-albert-blueviolet">
+   </a>
+   <a href="/model_doc/albert">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-albert-blueviolet">
+   </a>
+
+`ALBERT: A Lite BERT for Self-supervised Learning of Language Representations <https://arxiv.org/abs/1909.11942>`_,
+Zhenzhong Lan et al.
+
+Same as BERT but with a few tweaks:
+
+  * Embedding size E is different from hidden size H justified because the embeddings are context independent (one 
+    embedding vector represents one token) whereas hidden states are context dependent (one hidden state represents a 
+    sequence of tokens) so it's more logical to have H >> E. Als, the embedding matrix is large since it's V x E (V 
+    being the vocab size). If E < H, it has less parameters.
+  * Layers are split in groups that share parameters (to save memory).
+  * Next sentence prediction is replaced by a sentence ordering prediction: in the inputs, we have two sentences A et B 
+    (that are consecutive) and we either feed A followed by B or B followed by A. The model must predict if they have 
+    been swapped or not.
+
+The library provides a version of the model for masked language modeling, token classification, sentence 
+classification, multiple choice classification and question answering.
+
+RoBERTa
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=roberta">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-roberta-blueviolet">
+   </a>
+   <a href="/model_doc/roberta">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-roberta-blueviolet">
+   </a>
+
+`RoBERTa: A Robustly Optimized BERT Pretraining Approach <https://arxiv.org/abs/1907.11692>`_,
+Yinhan Liu et al.
+
+Same as BERT with better pretraining tricks:
+
+  * dynamic masking: tokens are masked differently at each epoch whereas BERT does it once and for all
+  * no NSP (next sentence prediction) loss and instead of putting just two sentences together, put a chunk of 
+    contiguous texts together to reach 512 tokens (so sentences in in an order than may span other several documents)
+  * train with larger batches
+  * use BPE with bytes as a subunit and not characters (because of unicode characters)
+
+The library provides a version of the model for masked language modeling, token classification, sentence 
+classification, multiple choice classification and question answering.
+
+DistilBERT
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=distilbert">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-distilbert-blueviolet">
+   </a>
+   <a href="/model_doc/distilbert">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-distilbert-blueviolet">
+   </a>
+
+`DistilBERT, a distilled version of BERT: smaller, faster, cheaper and lighter <https://arxiv.org/abs/1910.01108>`_,
+Victor Sanh et al.
+
+Same as BERT but smaller. Trained by distillation of the pretrained BERT model, meaning it's been trained to predict 
+the same probabilities as the larger model. The actual objective is a combination of:
+
+  * finding the same probabilities as the teacher model
+  * predicting the masked tokens correctly (but no next-sentence objective)
+  * a cosine similarity between the hidden states of the student and the teacher model
+
+The library provides a version of the model for masked language modeling, token classification, sentence classification 
+and question answering.
+
+XLM
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=xlm">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-xlm-blueviolet">
+   </a>
+   <a href="/model_doc/xlm">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-xlm-blueviolet">
+   </a>
+
+`Cross-lingual Language Model Pretraining <https://arxiv.org/abs/1901.07291>`_, Guillaume Lample and Alexis Conneau
+
+A transformer model trained on several languages. There are three different type of training for this model and the 
+library provides checkpoints for all of them:
+
+  * Causal language modeling (CLM) which is the traditional autoregressive training (so this model could be in the 
+    previous section as well). One of the languages is selected for each training sample, and the model input is a 
+    sentence of 256 tokens that may span on several documents in one one those languages.
+  * Masked language modeling (MLM) which is like RoBERTa. One of the languages is selected for each training sample, 
+    and the model input is a sentence of 256 tokens that may span on several documents in one one those languages, with
+    dynamic masking of the tokens.
+  * A combination of MLM and translation language modeling (TLM). This consists of concatenating a sentence in two 
+    different languages, with random masking. To predict one of the masked token, the model can use both the 
+    surrounding context in language 1 as well as the context given by language 2.
+
+Checkpoints refer to which method was used for pretraining by having `clm`, `mlm` or `mlm-tlm` in their names. On top
+of positional embeddings, the model has language embeddings. When training using MLM/CLM, this gives the model an
+indication of the language used, and when training using MLM+TLM, an indication of which part of the input is in which
+language.
+
+The library provides a version of the model for language modeling, token classification, sentence classification and 
+question answering.
+
+XLM-RoBERTa
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=xlm-roberta">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-xlm--roberta-blueviolet">
+   </a>
+   <a href="/model_doc/xlmroberta">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-xlm--roberta-blueviolet">
+   </a>
+
+`Unsupervised Cross-lingual Representation Learning at Scale <https://arxiv.org/abs/1911.02116>`_, Alexis Conneau et 
+al.
+
+Uses RoBERTa tricks on the XLM approach, but does not use the translation language modeling objective, only using 
+masked language modeling on sentences coming from one language. However, the model is trained on many more languages 
+(100) and doesn't use the language embeddings, so it's capable of detecting the input language by itself.
+
+The library provides a version of the model for masked language modeling, token classification, sentence 
+classification, multiple choice classification and question answering.
+
+FlauBERT
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=flaubert">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-flaubert-blueviolet">
+   </a>
+   <a href="/model_doc/flaubert">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-flaubert-blueviolet">
+   </a>
+
+`FlauBERT: Unsupervised Language Model Pre-training for French <https://arxiv.org/abs/1912.05372>`_, Hang Le et al.
+
+Like RoBERTa, without the sentence ordering prediction (so just trained on the MLM objective).
+
+The library provides a version of the model for language modeling and sentence classification.
+
+ELECTRA
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=electra">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-electra-blueviolet">
+   </a>
+   <a href="/model_doc/electra">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-electra-blueviolet">
+   </a>
+
+`ELECTRA: Pre-training Text Encoders as Discriminators Rather Than Generators <https://arxiv.org/abs/2003.10555>`_, 
+Kevin Clark et al.
+
+ELECTRA is a transformer model pretrained with the use of another (small) masked language model. The inputs are 
+corrupted by that language model, which takes an input text that is randomly masked and outputs a text in which ELECTRA 
+has to predict which token is an original and which one has been replaced. Like for GAN training, the small language 
+model is trained for a few steps (but with the original texts as objective, not to fool the ELECTRA model like in a 
+traditional GAN setting) then the ELECTRA model is trained for a few steps.
+
+The library provides a version of the model for masked language modeling, token classification and sentence 
+classification.
+
+.. _longformer:
+
+Longformer
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=longformer">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-longformer-blueviolet">
+   </a>
+   <a href="/model_doc/longformer">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-longformer-blueviolet">
+   </a>
+
+`Longformer: The Long-Document Transformer <https://arxiv.org/abs/2004.05150>`_, Iz Beltagy et al.
+
+A transformer model replacing the attention matrices by sparse matrices to go faster. Often, the local context (e.g., 
+what are the two tokens left and right?) is enough to take action for a given token. Some preselected input tokens are 
+still given global attention, but the attention matrix has way less parameters, resulting in a speed-up. See the 
+:ref:`local attention section <local-attention>` for more information.
+
+It is pretrained the same way a RoBERTa otherwise.
+
+**Note:** This model could be very well be used in an autoregressive setting, there is no checkpoint for such a
+pretraining yet, though.
+
+The library provides a version of the model for masked language modeling, token classification, sentence 
+classification, multiple choice classification and question answering.
+
+.. _seq-to-seq-models:
+
+Sequence-to-sequence models
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+As mentioned before, these models keep both the encoder and the decoder of the original transformer.
+
+BART
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=bart">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-bart-blueviolet">
+   </a>
+   <a href="/model_doc/bart">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-bart-blueviolet">
+   </a>
+
+`BART: Denoising Sequence-to-Sequence Pre-training for Natural Language Generation, Translation, and Comprehension 
+<https://arxiv.org/abs/1910.13461>`_, Mike Lewis et al.
+
+Sequence-to-sequence model with an encoder and a decoder. Encoder is fed a corrupted version of the tokens, decoder is 
+fed the tokens (but has a mask to hide the future words like a regular transformers decoder). For the encoder, on the 
+pretraining tasks, a composition of the following transformations are applied:
+
+  * mask random tokens (like in BERT)
+  * delete random tokens
+  * mask a span of k tokens with a single mask token (a span of 0 tokens is an insertion of a mask token)
+  * permute sentences
+  * rotate the document to make it start by a specific token
+
+The library provides a version of this model for conditional generation and sequence classification.
+
+MarianMT
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=marian">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-marian-blueviolet">
+   </a>
+   <a href="/model_doc/marian">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-marian-blueviolet">
+   </a>
+
+`Marian: Fast Neural Machine Translation in C++ <https://arxiv.org/abs/1804.00344>`_, Marcin Junczys-Dowmunt et al.
+
+A framework for translation models, using the same models as BART
+
+The library provides a version of this model for conditional generation.
+
+T5
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=t5">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-t5-blueviolet">
+   </a>
+   <a href="/model_doc/t5">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-t5-blueviolet">
+   </a>
+
+`Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer <https://arxiv.org/abs/1910.10683>`_, 
+Colin Raffel et al.
+
+Uses the traditional transformer model (except a slight change with the positional embeddings, which are learned at 
+each layer). To be able to operate on all NLP tasks, it transforms them in text-to-text problems by using certain 
+prefixes: “Summarize: …”, “question: …”, “translate English to German: …” and so forth.
+
+The pretraining includes both supervised and self-supervised training. Supervised training is conducted on downstream 
+tasks provided by the GLUE and SuperGLUE benchmarks (changing them to text-to-text tasks as explained above).
+
+Self-supervised training consists of corrupted pretrained, which means randomly removing 15% of the tokens and 
+replacing them by individual sentinel tokens (if several consecutive tokens are marked for removal, they are replaced 
+by one single sentinel token). The input of the encoder is the corrupted sentence, the input of the decoder the 
+original sentence and the target is then the dropped out tokens delimited by their sentinel tokens.
+
+For instance, if we have the sentence “My dog is very cute .”, and we decide to remove the token dog, is and cute, the 
+input becomes “My <x> very <y> .” and the target is “<x> dog is <y> . <z>”
+
+The library provides a version of this model for conditional generation.
+
+.. _multimodal-models:
+
+Multimodal models
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+There is one multimodal model in the library which has not been pretrained in the self-supervised fashion like the 
+others.
+
+MMBT
+----------------------------------------------
+
+`Supervised Multimodal Bitransformers for Classifying Images and Text <https://arxiv.org/abs/1909.02950>`_, Douwe Kiela 
+et al.
+
+A transformers model used in multimodal settings, combining a text and an image to make predictions. The transformer 
+model takes as inputs the embeddings of the tokenized text and a the final activations of a pretrained resnet on the 
+images (after the pooling layer) that goes through a linear layer (to go from number of features at the end of the 
+resnet to the hidden state dimension of the transformer).
+
+The different inputs are concatenated, and on top of the positional embeddings, a segment embedding is added to let the 
+model know which part of the input vector corresponds to the text or the image.
+
+The pretrained model only works for classification.
+
+..
+    More information in this :doc:`model documentation </model_doc/mmbt>`.
+    TODO: write this page
+
+More technical aspects
+^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
+
+Full vs sparse attention
+----------------------------------------------
+
+Most transformer models use full attention in the sense that the attention matrix is square. It can be a big 
+computational bottleneck when you have long texts. Longformer and reformer are models that try to be more efficient and 
+use a sparse version of the attention matrix to speed up training.
+
+.. _lsh-attention:
+
+**LSH attention**
+
+:ref:`Reformer <reformer>` uses LSH attention. In the softmax(QK^t), only the biggest elements (in the softmax 
+dimension) of the matrix QK^t are going to give useful contributions. So for each query q in Q, we can only consider 
+the keys k in K that are close to q. A hash function is used to determine if q and k are close. The attention mask is 
+modified to mask the current token (except at the first position) because it will give a query and key equal (so very 
+similar to each other). Since the hash can be a bit random, several hash functions are used in practice (determined by 
+a n_rounds parameter) then are averaged together.
+
+.. _local-attention:
+
+**Local attention**
+
+:ref:`Longformer <longformer>` uses local attention: often, the local context (e.g., what are the two tokens left and 
+right?) is enough to take action for a given token. Also, by stacking attention layers that have a small window, the 
+last layer will have a receptive field of more than just the tokens on the window, allowing them to build a 
+representation of the whole sentence.
+
+Some preselected input tokens are also given global attention: for those few tokens, the attention matrix can access 
+all tokens and this process is symmetric: all other tokens have access to those specific tokens (on top of the ones in 
+their local window). This is shown in Figure 2d of the paper, see below for a sample attention mask:
+
+.. image:: imgs/local_attention_mask.png
+   :scale: 50 %
+   :align: center
+
+Using those attention matrices with less parameters then allows the model to have inputs having a bigger sequence 
+length.
+
+Other tricks
+----------------------------------------------
+
+.. _axial-pos-encoding:
+
+**Axial positional encodings**
+
+:ref:`Reformer <reformer>` uses axial positional encodings: in traditional transformer models, the positional encoding 
+E is a matrix of size :math:`l` by :math:`d`, :math:`l` being the sequence length and :math:`d` the dimension of the 
+hidden state. If you have very long texts, this matrix can be huge and take way too much space on the GPU.
+
+To alleviate that, axial positional encodings consists in factorizing that big matrix E in two smaller matrices E1 and 
+E2, with dimensions :math:`l_{1} \times d_{1}` and :math:`l_{2} \times d_{2}`, such that :math:`l_{1} \times l_{2} = l`
+and :math:`d_{1} + d_{2} = d` (with the product for the lengths, this ends up being way smaller). The embedding for 
+time step :math:`j` in E is obtained by concatenating the embeddings for timestep :math:`j \% l1` in E1 and 
+:math:`j // l1` in E2.
+
@@ -22,7 +22,7 @@ This is still a work-in-progress – in particular documentation is still sparse
 | [**`token-classification`**](https://github.com/huggingface/transformers/tree/master/examples/token-classification) | CoNLL NER       | ✅ | ✅ | ✅ | -
 | [**`multiple-choice`**](https://github.com/huggingface/transformers/tree/master/examples/multiple-choice)           | SWAG, RACE, ARC | ✅ | ✅ | -  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ViktorAlm/notebooks/blob/master/MPC_GPU_Demo_for_TF_and_PT.ipynb)
 | [**`question-answering`**](https://github.com/huggingface/transformers/tree/master/examples/question-answering)     | SQuAD           | -  | ✅ | -  | -
-| [**`text-generation`**](https://github.com/huggingface/transformers/tree/master/examples/text-generation)     | -           | -  | - | -  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)
+| [**`text-generation`**](https://github.com/huggingface/transformers/tree/master/examples/text-generation)           | -               | n/a | n/a | n/a | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)
 | [**`distillation`**](https://github.com/huggingface/transformers/tree/master/examples/distillation)       | All               | -  | -  | -  | -
 | [**`summarization`**](https://github.com/huggingface/transformers/tree/master/examples/summarization)     | CNN/Daily Mail    | -  | -  | -  | -
 | [**`translation`**](https://github.com/huggingface/transformers/tree/master/examples/translation)         | WMT               | -  | -  | -  | -
@@ -3,9 +3,9 @@ from collections import defaultdict
 from dataclasses import dataclass, field
 from typing import Optional

+import matplotlib.pyplot as plt
 import numpy as np

-import matplotlib.pyplot as plt
 from transformers import HfArgumentParser


@@ -165,17 +165,15 @@ Larger batch size may improve the performance while costing more memory.
 python run_tf_squad.py \
    --model_name_or_path bert-base-uncased \
    --output_dir model \
-    --max-seq-length 384 \
+    --max_seq_length 384 \
    --num_train_epochs 2 \
    --per_gpu_train_batch_size 8 \
    --per_gpu_eval_batch_size 16 \
    --do_train \
-    --logging_dir logs \
-    --mode question-answering \
+    --logging_dir logs \    
    --logging_steps 10 \
    --learning_rate 3e-5 \
-    --doc_stride 128 \
-    --optimizer_name adamw
+    --doc_stride 128    
 ```

-For the moment the evaluation is not available in the Tensorflow Trainer only the training.
+For the moment the evaluation is not available in the Tensorflow Trainer only the training.
@@ -1,4 +1,4 @@
-### Get Preprocessed CNN Data
+### Get CNN Data
 To be able to reproduce the authors' results on the CNN/Daily Mail dataset you first need to download both CNN and Daily Mail datasets [from Kyunghyun Cho's website](https://cs.nyu.edu/~kcho/DMQA/) (the links next to "Stories") in the same folder. Then uncompress the archives by running:

 ```bash
@@ -6,24 +6,19 @@ wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/cnn_dm.tgz
 tar -xzvf cnn_dm.tgz
 ```

-this should make a directory called cnn_dm/ with files like `test.source`. 
+this should make a directory called cnn_dm/ with files like `test.source`.
 To use your own data, copy that files format. Each article to be summarized is on its own line.

 ### Evaluation
+
 To create summaries for each article in dataset, run:
 ```bash
-python evaluate_cnn.py <path_to_test.source> cnn_test_summaries.txt
+python evaluate_cnn.py <path_to_test.source> test_generations.txt <model-name>  --score_path rouge_scores.txt
 ```
-the default batch size, 8, fits in 16GB GPU memory, but may need to be adjusted to fit your system.
-
+The default batch size, 8, fits in 16GB GPU memory, but may need to be adjusted to fit your system.

 ### Training
-Run/modify `run_train.sh`
-    
-### Where is the code?
-The core model is in `src/transformers/modeling_bart.py`. This directory only contains examples.
-
-## (WIP) Rouge Scores
+Run/modify `finetune_bart.sh` or `finetune_t5.sh`

 ### Stanford CoreNLP Setup
 ```
@@ -1,71 +0,0 @@
-import argparse
-from pathlib import Path
-
-import torch
-from tqdm import tqdm
-
-from transformers import BartForConditionalGeneration, BartTokenizer
-
-
-DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
-
-
-def chunks(lst, n):
-    """Yield successive n-sized chunks from lst."""
-    for i in range(0, len(lst), n):
-        yield lst[i : i + n]
-
-
-def generate_summaries(
-    examples: list, out_file: str, model_name: str, batch_size: int = 8, device: str = DEFAULT_DEVICE
-):
-    fout = Path(out_file).open("w")
-    model = BartForConditionalGeneration.from_pretrained(model_name).to(device)
-    tokenizer = BartTokenizer.from_pretrained("facebook/bart-large")
-
-    max_length = 140
-    min_length = 55
-
-    for batch in tqdm(list(chunks(examples, batch_size))):
-        dct = tokenizer.batch_encode_plus(batch, max_length=1024, return_tensors="pt", pad_to_max_length=True)
-        summaries = model.generate(
-            input_ids=dct["input_ids"].to(device),
-            attention_mask=dct["attention_mask"].to(device),
-            num_beams=4,
-            length_penalty=2.0,
-            max_length=max_length + 2,  # +2 from original because we start at step=1 and stop before max_length
-            min_length=min_length + 1,  # +1 from original because we start at step=1
-            no_repeat_ngram_size=3,
-            early_stopping=True,
-            decoder_start_token_id=model.config.eos_token_id,
-        )
-        dec = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summaries]
-        for hypothesis in dec:
-            fout.write(hypothesis + "\n")
-            fout.flush()
-
-
-def run_generate():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "source_path", type=str, help="like cnn_dm/test.source",
-    )
-    parser.add_argument(
-        "output_path", type=str, help="where to save summaries",
-    )
-    parser.add_argument(
-        "model_name", type=str, default="facebook/bart-large-cnn", help="like bart-large-cnn",
-    )
-    parser.add_argument(
-        "--device", type=str, required=False, default=DEFAULT_DEVICE, help="cuda, cuda:1, cpu etc.",
-    )
-    parser.add_argument(
-        "--bs", type=int, default=8, required=False, help="batch size: how many to summarize at a time",
-    )
-    args = parser.parse_args()
-    examples = [" " + x.rstrip() for x in open(args.source_path).readlines()]
-    generate_summaries(examples, args.output_path, args.model_name, batch_size=args.bs, device=args.device)
-
-
-if __name__ == "__main__":
-    run_generate()
@@ -0,0 +1,100 @@
+import argparse
+from pathlib import Path
+
+import torch
+from rouge_score import rouge_scorer, scoring
+from tqdm import tqdm
+
+from transformers import AutoModelWithLMHead, AutoTokenizer
+
+
+DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
+
+
+def chunks(lst, n):
+    """Yield successive n-sized chunks from lst."""
+    for i in range(0, len(lst), n):
+        yield lst[i : i + n]
+
+
+def generate_summaries(
+    examples: list, out_file: str, model_name: str, batch_size: int = 8, device: str = DEFAULT_DEVICE
+):
+    fout = Path(out_file).open("w", encoding="utf-8")
+    model = AutoModelWithLMHead.from_pretrained(model_name).to(device)
+
+    tokenizer = AutoTokenizer.from_pretrained(model_name)
+
+    # update config with summarization specific params
+    task_specific_params = model.config.task_specific_params
+    if task_specific_params is not None:
+        model.config.update(task_specific_params.get("summarization", {}))
+
+    for batch in tqdm(list(chunks(examples, batch_size))):
+        if "t5" in model_name:
+            batch = [model.config.prefix + text for text in batch]
+        dct = tokenizer.batch_encode_plus(batch, max_length=1024, return_tensors="pt", pad_to_max_length=True).to(
+            device
+        )
+        summaries = model.generate(**dct)
+
+        dec = tokenizer.batch_decode(summaries, skip_special_tokens=True, clean_up_tokenization_spaces=False)
+        for hypothesis in dec:
+            fout.write(hypothesis + "\n")
+            fout.flush()
+
+
+def calculate_rouge(output_lns, reference_lns, score_path):
+    score_file = Path(score_path).open("w")
+    scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
+    aggregator = scoring.BootstrapAggregator()
+
+    for reference_ln, output_ln in zip(reference_lns, output_lns):
+        scores = scorer.score(reference_ln, output_ln)
+        aggregator.add_scores(scores)
+
+    result = aggregator.aggregate()
+    score_file.write(
+        "ROUGE_1: \n{} \n\n ROUGE_2: \n{} \n\n ROUGE_L: \n{} \n\n".format(
+            result["rouge1"], result["rouge2"], result["rougeL"]
+        )
+    )
+
+
+def run_generate():
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "input_path", type=str, help="like cnn_dm/test.source or cnn_dm/test_articles_input.txt",
+    )
+    parser.add_argument(
+        "output_path", type=str, help="where to save summaries",
+    )
+    parser.add_argument(
+        "model_name",
+        type=str,
+        default="facebook/bart-large-cnn",
+        help="like bart-large-cnn,'t5-small', 't5-base', 't5-large', 't5-3b', 't5-11b",
+    )
+    parser.add_argument("--reference_path", type=str, required=False, help="like cnn_dm/test_reference_summaries.txt")
+    parser.add_argument(
+        "--score_path", type=str, required=False, help="where to save the rouge score",
+    )
+    parser.add_argument(
+        "--device", type=str, required=False, default=DEFAULT_DEVICE, help="cuda, cuda:1, cpu etc.",
+    )
+    parser.add_argument(
+        "--bs", type=int, default=8, required=False, help="batch size: how many to summarize at a time",
+    )
+    args = parser.parse_args()
+    examples = [" " + x.rstrip() if "t5" in args.model_name else x.rstrip() for x in open(args.input_path).readlines()]
+
+    generate_summaries(examples, args.output_path, args.model_name, batch_size=args.bs, device=args.device)
+    if args.score_path is not None:
+        output_lns = [x.rstrip() for x in open(args.output_path).readlines()]
+        reference_lns = [x.rstrip() for x in open(args.reference_path).readlines()]
+
+        calculate_rouge(output_lns, reference_lns, args.score_path)
+
+
+if __name__ == "__main__":
+    run_generate()
@@ -6,7 +6,7 @@ export OUTPUT_DIR=${CURRENT_DIR}/${OUTPUT_DIR_NAME}
 mkdir -p $OUTPUT_DIR

 # Add parent directory to python path to access lightning_base.py
-export PYTHONPATH="../../":"${PYTHONPATH}"
+export PYTHONPATH="../":"${PYTHONPATH}"

 python finetune.py \
 --data_dir=./cnn-dailymail/cnn_dm \
@@ -13,7 +13,7 @@ export OUTPUT_DIR=${CURRENT_DIR}/${OUTPUT_DIR_NAME}
 mkdir -p $OUTPUT_DIR

 # Add parent directory to python path to access lightning_base.py and utils.py
-export PYTHONPATH="../../":"${PYTHONPATH}"
+export PYTHONPATH="../":"${PYTHONPATH}"
 python finetune.py \
 --data_dir=cnn_tiny/ \
 --model_type=bart \
@@ -0,0 +1,19 @@
+export OUTPUT_DIR_NAME=t5
+export CURRENT_DIR=${PWD}
+export OUTPUT_DIR=${CURRENT_DIR}/${OUTPUT_DIR_NAME}
+
+# Make output directory if it doesn't exist
+mkdir -p $OUTPUT_DIR
+
+# Add parent directory to python path to access lightning_base.py
+export PYTHONPATH="../":"${PYTHONPATH}"
+
+python finetune.py \
+--data_dir=./cnn-dailymail/cnn_dm \
+--model_name_or_path=t5-large \
+--model_type=t5
+--learning_rate=3e-5 \
+--train_batch_size=4 \
+--eval_batch_size=4 \
+--output_dir=$OUTPUT_DIR \
+--do_train  $@
@@ -1,29 +0,0 @@
-***This script evaluates the the multitask pre-trained checkpoint for ``t5-base`` (see paper [here](https://arxiv.org/pdf/1910.10683.pdf)) on the CNN/Daily Mail test dataset. Please note that the results in the paper were attained using a model fine-tuned on summarization, so that results will be worse here by approx. 0.5 ROUGE points***
-
-### Get the CNN Data
-First, you need to download the CNN data. It's about ~400 MB and can be downloaded by 
-running 
-
-```bash
-python download_cnn_daily_mail.py cnn_articles_input_data.txt cnn_articles_reference_summaries.txt
-```
-
-You should confirm that each file has 11490 lines:
-
-```bash
-wc -l cnn_articles_input_data.txt # should print 11490
-wc -l cnn_articles_reference_summaries.txt # should print 11490
-```
-
-### Generating Summaries
-
-To create summaries for each article in dataset, run:
-```bash
-python evaluate_cnn.py cnn_articles_input_data.txt cnn_generated_articles_summaries.txt cnn_articles_reference_summaries.txt rouge_score.txt
-```
-The default batch size, 8, fits in 16GB GPU memory, but may need to be adjusted to fit your system.
-The rouge scores "rouge1, rouge2, rougeL" are automatically created and saved in ``rouge_score.txt``.
-
-
-### Finetuning
-Pass model_type=t5 and model `examples/summarization/bart/finetune.py`
@@ -1,31 +0,0 @@
-import argparse
-from pathlib import Path
-
-import tensorflow_datasets as tfds
-
-
-def main(input_path, reference_path, data_dir):
-    cnn_ds = tfds.load("cnn_dailymail", split="test", shuffle_files=False, data_dir=data_dir)
-    cnn_ds_iter = tfds.as_numpy(cnn_ds)
-
-    test_articles_file = Path(input_path).open("w")
-    test_summaries_file = Path(reference_path).open("w")
-
-    for example in cnn_ds_iter:
-        test_articles_file.write(example["article"].decode("utf-8") + "\n")
-        test_articles_file.flush()
-        test_summaries_file.write(example["highlights"].decode("utf-8").replace("\n", " ") + "\n")
-        test_summaries_file.flush()
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser.add_argument("input_path", type=str, help="where to save the articles input data")
-    parser.add_argument(
-        "reference_path", type=str, help="where to save the reference summaries",
-    )
-    parser.add_argument(
-        "--data_dir", type=str, default="~/tensorflow_datasets", help="where to save the tensorflow datasets.",
-    )
-    args = parser.parse_args()
-    main(args.input_path, args.reference_path, args.data_dir)
@@ -1,101 +0,0 @@
-import argparse
-from pathlib import Path
-
-import torch
-from rouge_score import rouge_scorer, scoring
-from tqdm import tqdm
-
-from transformers import T5ForConditionalGeneration, T5Tokenizer
-
-
-def chunks(lst, n):
-    """Yield successive n-sized chunks from lst."""
-    for i in range(0, len(lst), n):
-        yield lst[i : i + n]
-
-
-def generate_summaries(lns, output_file_path, model_size, batch_size, device):
-    output_file = Path(output_file_path).open("w")
-
-    model = T5ForConditionalGeneration.from_pretrained(model_size)
-    model.to(device)
-
-    tokenizer = T5Tokenizer.from_pretrained(model_size)
-
-    # update config with summarization specific params
-    task_specific_params = model.config.task_specific_params
-    if task_specific_params is not None:
-        model.config.update(task_specific_params.get("summarization", {}))
-
-    for batch in tqdm(list(chunks(lns, batch_size))):
-        batch = [model.config.prefix + text for text in batch]
-
-        dct = tokenizer.batch_encode_plus(batch, max_length=512, return_tensors="pt", pad_to_max_length=True)
-        input_ids = dct["input_ids"].to(device)
-        attention_mask = dct["attention_mask"].to(device)
-
-        summaries = model.generate(input_ids=input_ids, attention_mask=attention_mask)
-        dec = [tokenizer.decode(g, skip_special_tokens=True, clean_up_tokenization_spaces=False) for g in summaries]
-
-        for hypothesis in dec:
-            output_file.write(hypothesis + "\n")
-            output_file.flush()
-
-
-def calculate_rouge(output_lns, reference_lns, score_path):
-    score_file = Path(score_path).open("w")
-    scorer = rouge_scorer.RougeScorer(["rouge1", "rouge2", "rougeL"], use_stemmer=True)
-    aggregator = scoring.BootstrapAggregator()
-
-    for reference_ln, output_ln in zip(reference_lns, output_lns):
-        scores = scorer.score(reference_ln, output_ln)
-        aggregator.add_scores(scores)
-
-    result = aggregator.aggregate()
-    score_file.write(
-        "ROUGE_1: \n{} \n\n ROUGE_2: \n{} \n\n ROUGE_L: \n{} \n\n".format(
-            result["rouge1"], result["rouge2"], result["rougeL"]
-        )
-    )
-
-
-def run_generate():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "model_size",
-        type=str,
-        help="T5 model size, either 't5-small', 't5-base', 't5-large', 't5-3b', 't5-11b'. Defaults to 't5-base'.",
-        default="t5-base",
-    )
-    parser.add_argument(
-        "input_path", type=str, help="like cnn_dm/test_articles_input.txt",
-    )
-    parser.add_argument(
-        "output_path", type=str, help="where to save summaries",
-    )
-    parser.add_argument("reference_path", type=str, help="like cnn_dm/test_reference_summaries.txt")
-    parser.add_argument(
-        "score_path", type=str, help="where to save the rouge score",
-    )
-    parser.add_argument(
-        "--batch_size", type=int, default=8, required=False, help="batch size: how many to summarize at a time",
-    )
-    parser.add_argument(
-        "--no_cuda", default=False, type=bool, help="Whether to force the execution on CPU.",
-    )
-
-    args = parser.parse_args()
-    args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
-
-    source_lns = [x.rstrip() for x in open(args.input_path).readlines()]
-
-    generate_summaries(source_lns, args.output_path, args.model_size, args.batch_size, args.device)
-
-    output_lns = [x.rstrip() for x in open(args.output_path).readlines()]
-    reference_lns = [x.rstrip() for x in open(args.reference_path).readlines()]
-
-    calculate_rouge(output_lns, reference_lns, args.score_path)
-
-
-if __name__ == "__main__":
-    run_generate()
@@ -1,44 +0,0 @@
-import logging
-import sys
-import tempfile
-import unittest
-from pathlib import Path
-from unittest.mock import patch
-
-from .evaluate_cnn import run_generate
-
-
-output_file_name = "output_t5_sum.txt"
-score_file_name = "score_t5_sum.txt"
-
-articles = ["New York (CNN)When Liana Barrientos was 23 years old, she got married in Westchester County."]
-
-logging.basicConfig(level=logging.DEBUG)
-
-logger = logging.getLogger()
-
-
-class TestT5Examples(unittest.TestCase):
-    def test_t5_cli(self):
-        stream_handler = logging.StreamHandler(sys.stdout)
-        logger.addHandler(stream_handler)
-        tmp = Path(tempfile.gettempdir()) / "utest_generations_t5_sum.hypo"
-        with tmp.open("w") as f:
-            f.write("\n".join(articles))
-
-        output_file_name = Path(tempfile.gettempdir()) / "utest_output_t5_sum.hypo"
-        score_file_name = Path(tempfile.gettempdir()) / "utest_score_t5_sum.hypo"
-
-        testargs = [
-            "evaluate_cnn.py",
-            "patrickvonplaten/t5-tiny-random",
-            str(tmp),
-            str(output_file_name),
-            str(tmp),
-            str(score_file_name),
-        ]
-
-        with patch.object(sys, "argv", testargs):
-            run_generate()
-            self.assertTrue(Path(output_file_name).exists())
-            self.assertTrue(Path(score_file_name).exists())
@@ -146,3 +146,34 @@ class TestBartExamples(unittest.TestCase):
            # show that targets were truncated
            self.assertEqual(batch["target_ids"].shape[1], trunc_target)  # Truncated
            self.assertGreater(max_len_target, trunc_target)  # Truncated
+
+
+class TestT5Examples(unittest.TestCase):
+    def test_t5_cli(self):
+        output_file_name = "output_t5_sum.txt"
+        score_file_name = "score_t5_sum.txt"
+        articles = ["New York (CNN)When Liana Barrientos was 23 years old, she got married in Westchester County."]
+        stream_handler = logging.StreamHandler(sys.stdout)
+        logger.addHandler(stream_handler)
+        tmp = Path(tempfile.gettempdir()) / "utest_generations_t5_sum.hypo"
+        with tmp.open("w", encoding="utf-8") as f:
+            f.write("\n".join(articles))
+
+        output_file_name = Path(tempfile.gettempdir()) / "utest_output_t5_sum.hypo"
+        score_file_name = Path(tempfile.gettempdir()) / "utest_score_t5_sum.hypo"
+
+        testargs = [
+            "evaluate_cnn.py",
+            str(tmp),
+            str(output_file_name),
+            "patrickvonplaten/t5-tiny-random",
+            "--reference_path",
+            str(tmp),
+            "--score_path",
+            str(score_file_name),
+        ]
+
+        with patch.object(sys, "argv", testargs):
+            run_generate()
+            self.assertTrue(Path(output_file_name).exists())
+            self.assertTrue(Path(score_file_name).exists())
@@ -61,7 +61,6 @@ export GLUE_DIR=/path/to/glue
 export TASK_NAME=MRPC

 python run_glue.py \
-  --model_type bert \
  --model_name_or_path bert-base-cased \
  --task_name $TASK_NAME \
  --do_train \
@@ -21,7 +21,7 @@ import logging
 import os
 import sys
 from dataclasses import dataclass, field
-from typing import Dict, Optional
+from typing import Callable, Dict, Optional

 import numpy as np

@@ -134,16 +134,29 @@ def main():
    )

    # Get datasets
-    train_dataset = GlueDataset(data_args, tokenizer=tokenizer) if training_args.do_train else None
-    eval_dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="dev") if training_args.do_eval else None
-    test_dataset = GlueDataset(data_args, tokenizer=tokenizer, mode="test") if training_args.do_predict else None
+    train_dataset = (
+        GlueDataset(data_args, tokenizer=tokenizer, cache_dir=model_args.cache_dir) if training_args.do_train else None
+    )
+    eval_dataset = (
+        GlueDataset(data_args, tokenizer=tokenizer, mode="dev", cache_dir=model_args.cache_dir)
+        if training_args.do_eval
+        else None
+    )
+    test_dataset = (
+        GlueDataset(data_args, tokenizer=tokenizer, mode="test", cache_dir=model_args.cache_dir)
+        if training_args.do_predict
+        else None
+    )

-    def compute_metrics(p: EvalPrediction) -> Dict:
-        if output_mode == "classification":
-            preds = np.argmax(p.predictions, axis=1)
-        elif output_mode == "regression":
-            preds = np.squeeze(p.predictions)
-        return glue_compute_metrics(data_args.task_name, preds, p.label_ids)
+    def build_compute_metrics_fn(task_name: str) -> Callable[[EvalPrediction], Dict]:
+        def compute_metrics_fn(p: EvalPrediction):
+            if output_mode == "classification":
+                preds = np.argmax(p.predictions, axis=1)
+            elif output_mode == "regression":
+                preds = np.squeeze(p.predictions)
+            return glue_compute_metrics(task_name, preds, p.label_ids)
+
+        return compute_metrics_fn

    # Initialize our Trainer
    trainer = Trainer(
@@ -151,7 +164,7 @@ def main():
        args=training_args,
        train_dataset=train_dataset,
        eval_dataset=eval_dataset,
-        compute_metrics=compute_metrics,
+        compute_metrics=build_compute_metrics_fn(data_args.task_name),
    )

    # Training
@@ -174,9 +187,12 @@ def main():
        eval_datasets = [eval_dataset]
        if data_args.task_name == "mnli":
            mnli_mm_data_args = dataclasses.replace(data_args, task_name="mnli-mm")
-            eval_datasets.append(GlueDataset(mnli_mm_data_args, tokenizer=tokenizer, mode="dev"))
+            eval_datasets.append(
+                GlueDataset(mnli_mm_data_args, tokenizer=tokenizer, mode="dev", cache_dir=model_args.cache_dir)
+            )

        for eval_dataset in eval_datasets:
+            trainer.compute_metrics = build_compute_metrics_fn(eval_dataset.args.task_name)
            eval_result = trainer.evaluate(eval_dataset=eval_dataset)

            output_eval_file = os.path.join(
@@ -196,7 +212,9 @@ def main():
        test_datasets = [test_dataset]
        if data_args.task_name == "mnli":
            mnli_mm_data_args = dataclasses.replace(data_args, task_name="mnli-mm")
-            test_datasets.append(GlueDataset(mnli_mm_data_args, tokenizer=tokenizer, mode="test"))
+            test_datasets.append(
+                GlueDataset(mnli_mm_data_args, tokenizer=tokenizer, mode="test", cache_dir=model_args.cache_dir)
+            )

        for test_dataset in test_datasets:
            predictions = trainer.predict(test_dataset=test_dataset).predictions
@@ -17,7 +17,7 @@ Please check out the repo under uber-research for more information: https://gith
 git clone https://github.com/huggingface/transformers && cd transformers
 pip install .
 pip install nltk torchtext # additional requirements.
-cd examples/pplm
+cd examples/text-generation/pplm
 ```

 ## PPLM-BoW 
@@ -31,7 +31,6 @@ from typing import List, Optional, Tuple, Union
 import numpy as np
 import torch
 import torch.nn.functional as F
-from torch.autograd import Variable
 from tqdm import trange

 from pplm_classification_head import ClassificationHead
@@ -76,14 +75,6 @@ DISCRIMINATOR_MODELS_PARAMS = {
 }


-def to_var(x, requires_grad=False, volatile=False, device="cuda"):
-    if torch.cuda.is_available() and device == "cuda":
-        x = x.cuda()
-    elif device != "cuda":
-        x = x.to(device)
-    return Variable(x, requires_grad=requires_grad, volatile=volatile)
-
-
 def top_k_filter(logits, k, probs=False):
    """
    Masks everything but the k top entries as -infinity (1e10).
@@ -156,9 +147,7 @@ def perturb_past(
    new_accumulated_hidden = None
    for i in range(num_iterations):
        print("Iteration ", i + 1)
-        curr_perturbation = [
-            to_var(torch.from_numpy(p_), requires_grad=True, device=device) for p_ in grad_accumulator
-        ]
+        curr_perturbation = [torch.from_numpy(p_).requires_grad_(True).to(device=device) for p_ in grad_accumulator]

        # Compute hidden using perturbed past
        perturbed_past = list(map(add, past, curr_perturbation))
@@ -247,7 +236,7 @@ def perturb_past(
        past = new_past

    # apply the accumulated perturbations to the past
-    grad_accumulator = [to_var(torch.from_numpy(p_), requires_grad=True, device=device) for p_ in grad_accumulator]
+    grad_accumulator = [torch.from_numpy(p_).requires_grad_(True).to(device=device) for p_ in grad_accumulator]
    pert_past = list(map(add, past, grad_accumulator))

    return pert_past, new_accumulated_hidden, grad_norms, loss_per_iter
@@ -266,7 +255,7 @@ def get_classifier(
    elif "path" in params:
        resolved_archive_file = params["path"]
    else:
-        raise ValueError("Either url or path have to be specified " "in the discriminator model parameters")
+        raise ValueError("Either url or path have to be specified in the discriminator model parameters")
    classifier.load_state_dict(torch.load(resolved_archive_file, map_location=device))
    classifier.eval()

@@ -569,9 +558,9 @@ def generate_text_pplm(

 def set_generic_model_params(discrim_weights, discrim_meta):
    if discrim_weights is None:
-        raise ValueError("When using a generic discriminator, " "discrim_weights need to be specified")
+        raise ValueError("When using a generic discriminator, discrim_weights need to be specified")
    if discrim_meta is None:
-        raise ValueError("When using a generic discriminator, " "discrim_meta need to be specified")
+        raise ValueError("When using a generic discriminator, discrim_meta need to be specified")

    with open(discrim_meta, "r") as discrim_meta_file:
        meta = json.load(discrim_meta_file)
@@ -619,7 +608,7 @@ def run_pplm_example(

    if discrim is not None:
        pretrained_model = DISCRIMINATOR_MODELS_PARAMS[discrim]["pretrained_model"]
-        print("discrim = {}, pretrained_model set " "to discriminator's = {}".format(discrim, pretrained_model))
+        print("discrim = {}, pretrained_model set to discriminator's = {}".format(discrim, pretrained_model))

    # load pretrained model
    model = GPT2LMHeadModel.from_pretrained(pretrained_model, output_hidden_states=True)
@@ -706,7 +695,7 @@ def run_pplm_example(
                for word_id in pert_gen_tok_text.tolist()[0]:
                    if word_id in bow_word_ids:
                        pert_gen_text += "{}{}{}".format(
-                            colorama.Fore.RED, tokenizer.decode([word_id]), colorama.Style.RESET_ALL
+                            colorama.Fore.RED, tokenizer.decode([word_id]), colorama.Style.RESET_ALL,
                        )
                    else:
                        pert_gen_text += tokenizer.decode([word_id])
@@ -744,9 +733,11 @@ if __name__ == "__main__":
        "-B",
        type=str,
        default=None,
-        help="Bags of words used for PPLM-BoW. "
-        "Either a BOW id (see list in code) or a filepath. "
-        "Multiple BoWs separated by ;",
+        help=(
+            "Bags of words used for PPLM-BoW. "
+            "Either a BOW id (see list in code) or a filepath. "
+            "Multiple BoWs separated by ;"
+        ),
    )
    parser.add_argument(
        "--discrim",
@@ -756,9 +747,11 @@ if __name__ == "__main__":
        choices=("clickbait", "sentiment", "toxicity", "generic"),
        help="Discriminator to use",
    )
-    parser.add_argument("--discrim_weights", type=str, default=None, help="Weights for the generic discriminator")
    parser.add_argument(
-        "--discrim_meta", type=str, default=None, help="Meta information for the generic discriminator"
+        "--discrim_weights", type=str, default=None, help="Weights for the generic discriminator",
+    )
+    parser.add_argument(
+        "--discrim_meta", type=str, default=None, help="Meta information for the generic discriminator",
    )
    parser.add_argument(
        "--class_label", type=int, default=-1, help="Class label used for the discriminator",
@@ -774,7 +767,7 @@ if __name__ == "__main__":
        "--window_length",
        type=int,
        default=0,
-        help="Length of past which is being optimized; " "0 corresponds to infinite window length",
+        help="Length of past which is being optimized; 0 corresponds to infinite window length",
    )
    parser.add_argument(
        "--horizon_length", type=int, default=1, help="Length of future to optimize over",
@@ -2,10 +2,17 @@

 Based on the scripts [`run_ner.py`](https://github.com/huggingface/transformers/blob/master/examples/token-classification/run_ner.py) for Pytorch and
 [`run_tf_ner.py`](https://github.com/huggingface/transformers/blob/master/examples/token-classification/run_tf_ner.py) for Tensorflow 2.
-This example fine-tune Bert Multilingual on GermEval 2014 (German NER).
+
+The following examples are covered in this section:
+
+* NER on the GermEval 2014 (German NER) dataset
+* Emerging and Rare Entities task: WNUT’17 (English NER) dataset
+
 Details and results for the fine-tuning provided by @stefan-it.

-### Data (Download and pre-processing steps)
+### GermEval 2014 (German NER) dataset
+
+#### Data (Download and pre-processing steps)

 Data can be obtained from the [GermEval 2014](https://sites.google.com/site/germeval2014ner/data) shared task page.

@@ -20,11 +27,10 @@ curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-test.tsv?attr
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > test.txt.tmp
 ```

-The GermEval 2014 dataset contains some strange "control character" tokens like `'\x96', '\u200e', '\x95', '\xad' or '\x80'`. One problem with these tokens is, that `BertTokenizer` returns an empty token for them, resulting in misaligned `InputExample`s. I wrote a script that a) filters these tokens and b) splits longer sentences into smaller ones (once the max. subtoken length is reached).
+The GermEval 2014 dataset contains some strange "control character" tokens like `'\x96', '\u200e', '\x95', '\xad' or '\x80'`.
+One problem with these tokens is, that `BertTokenizer` returns an empty token for them, resulting in misaligned `InputExample`s.
+The `preprocess.py` script located in the `scripts` folder a) filters these tokens and b) splits longer sentences into smaller ones (once the max. subtoken length is reached).

-```bash
-wget "https://raw.githubusercontent.com/stefan-it/fine-tuned-berts-seq/master/scripts/preprocess.py"
-```
 Let's define some variables that we need for further pre-processing steps and training the model:

 ```bash
@@ -35,9 +41,9 @@ export BERT_MODEL=bert-base-multilingual-cased
 Run the pre-processing script on training, dev and test datasets:

 ```bash
-python3 preprocess.py train.txt.tmp $BERT_MODEL $MAX_LENGTH > train.txt
-python3 preprocess.py dev.txt.tmp $BERT_MODEL $MAX_LENGTH > dev.txt
-python3 preprocess.py test.txt.tmp $BERT_MODEL $MAX_LENGTH > test.txt
+python3 scripts/preprocess.py train.txt.tmp $BERT_MODEL $MAX_LENGTH > train.txt
+python3 scripts/preprocess.py dev.txt.tmp $BERT_MODEL $MAX_LENGTH > dev.txt
+python3 scripts/preprocess.py test.txt.tmp $BERT_MODEL $MAX_LENGTH > test.txt
 ```

 The GermEval 2014 dataset has much more labels than CoNLL-2002/2003 datasets, so an own set of labels must be used:
@@ -46,7 +52,7 @@ The GermEval 2014 dataset has much more labels than CoNLL-2002/2003 datasets, so
 cat train.txt dev.txt test.txt | cut -d " " -f 2 | grep -v "^$"| sort | uniq > labels.txt
 ```

-### Prepare the run
+#### Prepare the run

 Additional environment variables must be set:

@@ -58,7 +64,7 @@ export SAVE_STEPS=750
 export SEED=1
 ```

-### Run the Pytorch version
+#### Run the Pytorch version

 To start training, just run:

@@ -79,7 +85,7 @@ python3 run_ner.py --data_dir ./ \

 If your GPU supports half-precision training, just add the `--fp16` flag. After training, the model will be both evaluated on development and test datasets.

-### JSON-based configuration file
+#### JSON-based configuration file

 Instead of passing all parameters via commandline arguments, the `run_ner.py` script also supports reading parameters from a json-based configuration file:

@@ -124,16 +130,6 @@ On the test dataset the following results could be achieved:
 10/04/2019 00:42:42 - INFO - __main__ -     recall = 0.8624150210424085
 ```

-#### Comparing BERT (large, cased), RoBERTa (large, cased) and DistilBERT (base, uncased)
-
-Here is a small comparison between BERT (large, cased), RoBERTa (large, cased) and DistilBERT (base, uncased) with the same hyperparameters as specified in the [example documentation](https://huggingface.co/transformers/examples.html#named-entity-recognition) (one run):
-
-| Model | F-Score Dev | F-Score Test
-| --------------------------------- | ------- | --------
-| `bert-large-cased`            | 95.59 | 91.70
-| `roberta-large`                  | 95.96 | 91.87
-| `distilbert-base-uncased` | 94.34 | 90.32
-
 #### Run PyTorch version using PyTorch-Lightning

 Run `bash run_pl.sh` from the `ner` directory. This would also install `pytorch-lightning` and the `examples/requirements.txt`. It is a shell pipeline which would automatically download, pre-process the data and run the models in `germeval-model` directory. Logs are saved in `lightning_logs` directory.
@@ -141,7 +137,7 @@ Run `bash run_pl.sh` from the `ner` directory. This would also install `pytorch-
 Pass `--n_gpu` flag to change the number of GPUs. Default uses 1. At the end, the expected results are: `TEST RESULTS {'val_loss': tensor(0.0707), 'precision': 0.852427800698191, 'recall': 0.869537067011978, 'f1': 0.8608974358974358}`


-### Run the Tensorflow 2 version
+#### Run the Tensorflow 2 version

 To start training, just run:

@@ -205,3 +201,102 @@ On the test dataset the following results could be achieved:
 micro avg     0.8722    0.8774    0.8748     13869
 macro avg     0.8712    0.8774    0.8740     13869
 ```
+
+### Emerging and Rare Entities task: WNUT’17 (English NER) dataset
+
+Description of the WNUT’17 task from the [shared task website](http://noisy-text.github.io/2017/index.html):
+
+> The WNUT’17 shared task focuses on identifying unusual, previously-unseen entities in the context of emerging discussions.
+> Named entities form the basis of many modern approaches to other tasks (like event clustering and summarization), but recall on
+> them is a real problem in noisy text - even among annotators. This drop tends to be due to novel entities and surface forms.
+
+Six labels are available in the dataset. An overview can be found on this [page](http://noisy-text.github.io/2017/files/).
+
+#### Data (Download and pre-processing steps)
+
+The dataset can be downloaded from the [official GitHub](https://github.com/leondz/emerging_entities_17) repository.
+
+The following commands show how to prepare the dataset for fine-tuning:
+
+```bash
+mkdir -p data_wnut_17
+
+curl -L 'https://github.com/leondz/emerging_entities_17/raw/master/wnut17train.conll'  | tr '\t' ' ' > data_wnut_17/train.txt.tmp
+curl -L 'https://github.com/leondz/emerging_entities_17/raw/master/emerging.dev.conll' | tr '\t' ' ' > data_wnut_17/dev.txt.tmp
+curl -L 'https://raw.githubusercontent.com/leondz/emerging_entities_17/master/emerging.test.annotated' | tr '\t' ' ' > data_wnut_17/test.txt.tmp
+```
+
+Let's define some variables that we need for further pre-processing steps:
+
+```bash
+export MAX_LENGTH=128
+export BERT_MODEL=bert-large-cased
+```
+
+Here we use the English BERT large model for fine-tuning.
+The `preprocess.py` scripts splits longer sentences into smaller ones (once the max. subtoken length is reached):
+
+```bash
+python3 scripts/preprocess.py data_wnut_17/train.txt.tmp $BERT_MODEL $MAX_LENGTH > data_wnut_17/train.txt
+python3 scripts/preprocess.py data_wnut_17/dev.txt.tmp $BERT_MODEL $MAX_LENGTH > data_wnut_17/dev.txt
+python3 scripts/preprocess.py data_wnut_17/test.txt.tmp $BERT_MODEL $MAX_LENGTH > data_wnut_17/test.txt
+```
+
+In the last pre-processing step, the `labels.txt` file needs to be generated. This file contains all available labels:
+
+```bash
+cat data_wnut_17/train.txt data_wnut_17/dev.txt data_wnut_17/test.txt | cut -d " " -f 2 | grep -v "^$"| sort | uniq > data_wnut_17/labels.txt
+```
+
+#### Run the Pytorch version
+
+Fine-tuning with the PyTorch version can be started using the `run_ner.py` script. In this example we use a JSON-based configuration file.
+
+This configuration file looks like:
+
+```json
+{
+    "data_dir": "./data_wnut_17",
+    "labels": "./data_wnut_17/labels.txt",
+    "model_name_or_path": "bert-large-cased",
+    "output_dir": "wnut-17-model-1",
+    "max_seq_length": 128,
+    "num_train_epochs": 3,
+    "per_device_train_batch_size": 32,
+    "save_steps": 425,
+    "seed": 1,
+    "do_train": true,
+    "do_eval": true,
+    "do_predict": true,
+    "fp16": false
+}
+```
+
+If your GPU supports half-precision training, please set `fp16` to `true`.
+
+Save this JSON-based configuration under `wnut_17.json`. The fine-tuning can be started with `python3 run_ner.py wnut_17.json`.
+
+#### Evaluation
+
+Evaluation on development dataset outputs the following:
+
+```bash
+05/29/2020 23:33:44 - INFO - __main__ -   ***** Eval results *****
+05/29/2020 23:33:44 - INFO - __main__ -     eval_loss = 0.26505235286212275
+05/29/2020 23:33:44 - INFO - __main__ -     eval_precision = 0.7008264462809918
+05/29/2020 23:33:44 - INFO - __main__ -     eval_recall = 0.507177033492823
+05/29/2020 23:33:44 - INFO - __main__ -     eval_f1 = 0.5884802220680084
+05/29/2020 23:33:44 - INFO - __main__ -     epoch = 3.0
+```
+
+On the test dataset the following results could be achieved:
+
+```bash
+05/29/2020 23:33:44 - INFO - transformers.trainer -   ***** Running Prediction *****
+05/29/2020 23:34:02 - INFO - __main__ -     eval_loss = 0.30948806500973547
+05/29/2020 23:34:02 - INFO - __main__ -     eval_precision = 0.5840108401084011
+05/29/2020 23:34:02 - INFO - __main__ -     eval_recall = 0.3994439295644115
+05/29/2020 23:34:02 - INFO - __main__ -     eval_f1 = 0.47440836543753434
+```
+
+WNUT’17 is a very difficult task. Current state-of-the-art results on this dataset can be found [here](http://nlpprogress.com/english/named_entity_recognition.html).
@@ -4,12 +4,12 @@ curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-dev.tsv?attre
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > dev.txt.tmp
 curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-test.tsv?attredirects=0&d=1' \
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > test.txt.tmp
-wget "https://raw.githubusercontent.com/stefan-it/fine-tuned-berts-seq/master/scripts/preprocess.py"
+
 export MAX_LENGTH=128
 export BERT_MODEL=bert-base-multilingual-cased
-python3 preprocess.py train.txt.tmp $BERT_MODEL $MAX_LENGTH > train.txt
-python3 preprocess.py dev.txt.tmp $BERT_MODEL $MAX_LENGTH > dev.txt
-python3 preprocess.py test.txt.tmp $BERT_MODEL $MAX_LENGTH > test.txt
+python3 scripts/preprocess.py train.txt.tmp $BERT_MODEL $MAX_LENGTH > train.txt
+python3 scripts/preprocess.py dev.txt.tmp $BERT_MODEL $MAX_LENGTH > dev.txt
+python3 scripts/preprocess.py test.txt.tmp $BERT_MODEL $MAX_LENGTH > test.txt
 cat train.txt dev.txt test.txt | cut -d " " -f 2 | grep -v "^$"| sort | uniq > labels.txt
 export OUTPUT_DIR=germeval-model
 export BATCH_SIZE=32
@@ -11,12 +11,12 @@ curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-dev.tsv?attre
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > dev.txt.tmp
 curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-test.tsv?attredirects=0&d=1' \
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > test.txt.tmp
-wget "https://raw.githubusercontent.com/stefan-it/fine-tuned-berts-seq/master/scripts/preprocess.py"
+
 export MAX_LENGTH=128
 export BERT_MODEL=bert-base-multilingual-cased
-python3 preprocess.py train.txt.tmp $BERT_MODEL $MAX_LENGTH > train.txt
-python3 preprocess.py dev.txt.tmp $BERT_MODEL $MAX_LENGTH > dev.txt
-python3 preprocess.py test.txt.tmp $BERT_MODEL $MAX_LENGTH > test.txt
+python3 scripts/preprocess.py train.txt.tmp $BERT_MODEL $MAX_LENGTH > train.txt
+python3 scripts/preprocess.py dev.txt.tmp $BERT_MODEL $MAX_LENGTH > dev.txt
+python3 scripts/preprocess.py test.txt.tmp $BERT_MODEL $MAX_LENGTH > test.txt
 cat train.txt dev.txt test.txt | cut -d " " -f 2 | grep -v "^$"| sort | uniq > labels.txt
 export BATCH_SIZE=32
 export NUM_EPOCHS=3
@@ -0,0 +1,41 @@
+import sys
+
+from transformers import AutoTokenizer
+
+
+dataset = sys.argv[1]
+model_name_or_path = sys.argv[2]
+max_len = int(sys.argv[3])
+
+subword_len_counter = 0
+
+tokenizer = AutoTokenizer.from_pretrained(model_name_or_path)
+max_len -= tokenizer.num_special_tokens_to_add()
+
+with open(dataset, "rt") as f_p:
+    for line in f_p:
+        line = line.rstrip()
+
+        if not line:
+            print(line)
+            subword_len_counter = 0
+            continue
+
+        token = line.split()[0]
+
+        current_subwords_len = len(tokenizer.tokenize(token))
+
+        # Token contains strange control characters like \x96 or \x95
+        # Just filter out the complete line
+        if current_subwords_len == 0:
+            continue
+
+        if (subword_len_counter + current_subwords_len) > max_len:
+            print("")
+            print(line)
+            subword_len_counter = current_subwords_len
+            continue
+
+        subword_len_counter += current_subwords_len
+
+        print(line)
@@ -0,0 +1,26 @@
+# BERT-Small CORD-19 fine-tuned on SQuAD 2.0
+
+[bert-small-cord19 model](https://huggingface.co/NeuML/bert-small-cord19) fine-tuned on SQuAD 2.0
+
+## Building the model
+
+```bash
+python run_squad.py
+    --model_type bert
+    --model_name_or_path bert-small-cord19
+    --do_train
+    --do_eval
+    --do_lower_case
+    --version_2_with_negative
+    --train_file train-v2.0.json
+    --predict_file dev-v2.0.json
+    --per_gpu_train_batch_size 8
+    --learning_rate 3e-5
+    --num_train_epochs 3.0
+    --max_seq_length 384
+    --doc_stride 128
+    --output_dir bert-small-cord19-squad2
+    --save_steps 0
+    --threads 8
+    --overwrite_cache
+    --overwrite_output_dir
@@ -0,0 +1,25 @@
+# BERT-Small fine-tuned on CORD-19 dataset
+
+[BERT L6_H-512_A-8 model](https://huggingface.co/google/bert_uncased_L-6_H-512_A-8) fine-tuned on the [CORD-19 dataset](https://www.semanticscholar.org/cord19).
+
+## CORD-19 data subset
+The training data for this dataset is stored as a [Kaggle dataset](https://www.kaggle.com/davidmezzetti/cord19-qa?select=cord19.txt). The training
+data is a subset of the full corpus, focusing on high-quality, study-design detected articles.
+
+## Building the model
+
+```bash
+python run_language_modeling.py
+    --model_type bert
+    --model_name_or_path google/bert_uncased_L-6_H-512_A-8
+    --do_train
+    --mlm
+    --line_by_line
+    --block_size 512
+    --train_data_file cord19.txt
+    --per_gpu_train_batch_size 4
+    --learning_rate 3e-5
+    --num_train_epochs 3.0
+    --output_dir bert-small-cord19
+    --save_steps 0
+    --overwrite_output_dir
@@ -0,0 +1,63 @@
+# BERT-Small fine-tuned on CORD-19 QA dataset
+
+[bert-small-cord19-squad model](https://huggingface.co/NeuML/bert-small-cord19-squad2) fine-tuned on the [CORD-19 QA dataset](https://www.kaggle.com/davidmezzetti/cord19-qa?select=cord19-qa.json).
+
+## CORD-19 QA dataset
+The CORD-19 QA dataset is a SQuAD 2.0 formatted list of question, context, answer combinations covering the [CORD-19 dataset](https://www.semanticscholar.org/cord19).
+
+## Building the model
+
+```bash
+python run_squad.py \
+    --model_type bert \
+    --model_name_or_path bert-small-cord19-squad \
+    --do_train \
+    --do_lower_case \
+    --version_2_with_negative \
+    --train_file cord19-qa.json \
+    --per_gpu_train_batch_size 8 \
+    --learning_rate 5e-5 \
+    --num_train_epochs 10.0 \
+    --max_seq_length 384 \
+    --doc_stride 128 \
+    --output_dir bert-small-cord19qa \
+    --save_steps 0 \
+    --threads 8 \
+    --overwrite_cache \
+    --overwrite_output_dir
+```
+
+## Testing the model
+
+Example usage below:
+
+```python
+from transformers import pipeline
+
+qa = pipeline(
+    "question-answering",
+    model="NeuML/bert-small-cord19qa",
+    tokenizer="NeuML/bert-small-cord19qa"
+)
+
+qa({
+    "question": "What is the median incubation period?",
+    "context": "The incubation period is around 5 days (range: 4-7 days) with a maximum of 12-13 day"
+})
+
+qa({
+    "question": "What is the incubation period range?",
+    "context": "The incubation period is around 5 days (range: 4-7 days) with a maximum of 12-13 day"
+})
+
+qa({
+    "question": "What type of surfaces does it persist?",
+    "context": "The virus can survive on surfaces for up to 72 hours such as plastic and stainless steel ."
+})
+```
+
+```json
+{"score": 0.5970273583242793, "start": 32, "end": 38, "answer": "5 days"}
+{"score": 0.999555868193891, "start": 39, "end": 56, "answer": "(range: 4-7 days)"}
+{"score": 0.9992726505196998, "start": 61, "end": 88, "answer": "plastic and stainless steel"}
+```
@@ -0,0 +1,10 @@
+The Bart model was proposed by Mike Lewis, Yinhan Liu, Naman Goyal, Marjan Ghazvininejad, Abdelrahman Mohamed, Omer Levy, Ves Stoyanov and Luke Zettlemoyer on 29 Oct, 2019. According to the abstract,
+
+Bart uses a standard seq2seq/machine translation architecture with a bidirectional encoder (like BERT) and a left-to-right decoder (like GPT).
+
+The pretraining task involves randomly shuffling the order of the original sentences and a novel in-filling scheme, where spans of text are replaced with a single mask token.
+
+BART is particularly effective when fine tuned for text generation but also works well for comprehension tasks. It matches the performance of RoBERTa with comparable training resources on GLUE and SQuAD, achieves new state-of-the-art results on a range of abstractive dialogue, question answering, and summarization tasks, with gains of up to 6 ROUGE.
+
+The Authors’ code can be found here:
+https://github.com/pytorch/fairseq/tree/master/examples/bart
@@ -0,0 +1,73 @@
+---
+language: english
+thumbnail:
+---
+
+# T5-base fine-tuned for Sentiment Span Extraction
+
+All credits to [Lorenzo Ampil](https://twitter.com/AND__SO)
+
+[Google's T5](https://ai.googleblog.com/2020/02/exploring-transfer-learning-with-t5.html) base fine-tuned on [Tweet Sentiment Extraction Dataset](https://www.kaggle.com/c/tweet-sentiment-extraction) for **Span Sentiment Extraction** downstream task.
+
+## Details of T5
+
+The **T5** model was presented in [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/pdf/1910.10683.pdf) by *Colin Raffel, Noam Shazeer, Adam Roberts, Katherine Lee, Sharan Narang, Michael Matena, Yanqi Zhou, Wei Li, Peter J. Liu* in Here the 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.
+
+
+## Details of the downstream task (Span Sentiment Extraction) - Dataset 📚
+
+[Tweet Sentiment Extraction Dataset](https://www.kaggle.com/c/tweet-sentiment-extraction)
+
+"My ridiculous dog is amazing." [sentiment: positive]
+
+With all of the tweets circulating every second it is hard to tell whether the sentiment behind a specific tweet will impact a company, or a person's, brand for being viral (positive), or devastate profit because it strikes a negative tone. Capturing sentiment in language is important in these times where decisions and reactions are created and updated in seconds. But, which words actually lead to the sentiment description? In this competition you will need to pick out the part of the tweet (word or phrase) that reflects the sentiment.
+
+Help build your skills in this important area with this broad dataset of tweets. Work on your technique to grab a top spot in this competition. What words in tweets support a positive, negative, or neutral sentiment? How can you help make that determination using machine learning tools?
+
+In this competition we've extracted support phrases from Figure Eight's Data for Everyone platform. The dataset is titled Sentiment Analysis: Emotion in Text tweets with existing sentiment labels, used here under creative commons attribution 4.0. international licence. Your objective in this competition is to construct a model that can do the same - look at the labeled sentiment for a given tweet and figure out what word or phrase best supports it.
+
+Disclaimer: The dataset for this competition contains text that may be considered profane, vulgar, or offensive.
+
+| Dataset  | Split | # samples |
+| -------- | ----- | --------- |
+| TSE | train | 23907      |
+| TSE | eval  | 3573     |
+
+
+
+## Model fine-tuning 🏋️‍
+
+The training script is a slightly modified version of [this Colab Notebook](https://github.com/enzoampil/t5-intro/blob/master/t5_qa_training_pytorch_span_extraction.ipynb) created by [Lorenzo Ampil](https://github.com/enzoampil), so all credits to him!
+
+
+## Model in Action 🚀
+
+```python
+from transformers import AutoModelWithLMHead, AutoTokenizer
+
+tokenizer = AutoTokenizer.from_pretrained("mrm8488/t5-base-finetuned-span-sentiment-extraction")
+model = AutoModelWithLMHead.from_pretrained("mrm8488/t5-base-finetuned-span-sentiment-extraction")
+
+def get_sentiment_span(text):
+  input_ids = tokenizer.encode(text, return_tensors="pt", add_special_tokens=True)  # Batch size 1
+  
+  generated_ids = model.generate(input_ids=input_ids, num_beams=1, max_length=80).squeeze()
+  
+  predicted_span = tokenizer.decode(generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True)
+    
+  return predicted_span
+  
+get_sentiment_span("question: negative context: My bike was put on hold...should have known that.... argh total bummer")
+
+# output: 'argh total bummer'
+
+get_sentiment_span("question: positive context: On the monday, so i wont be able to be with you! i love you")
+
+# output: 'i love you'
+```
+
+> Created by [Manuel Romero/@mrm8488](https://twitter.com/mrm8488) | [LinkedIn](https://www.linkedin.com/in/manuel-romero-cs/)
+
+> Made with <span style="color: #e25555;">&hearts;</span> in Spain
@@ -0,0 +1,99 @@
+---
+language: sanskrit
+---
+
+# RoBERTa trained on Sanskrit (SanBERTa)
+
+**Mode size** (after training): **340MB**
+
+### Dataset:
+
+[Wikipedia articles](https://www.kaggle.com/disisbig/sanskrit-wikipedia-articles) (used in [iNLTK](https://github.com/goru001/nlp-for-sanskrit)).
+It contains evaluation set.
+
+[Sanskrit scraps from CLTK](http://cltk.org/)
+
+### Configuration
+
+| Parameter | Value |
+|---|---|
+| `num_attention_heads` | 12 |
+| `num_hidden_layers` | 6 |
+| `hidden_size` | 768 |
+| `vocab_size` | 29407 |
+
+### Training :
+- On TPU
+- For language modelling
+- Iteratively increasing `--block_size` from 128 to 256 over epochs
+
+### Evaluation
+
+|Metric| # Value |
+|---|---|
+|Perplexity (`block_size=256`)|4.04|
+
+## Example of usage:
+
+### For Embeddings
+
+```
+
+tokenizer = AutoTokenizer.from_pretrained("surajp/SanBERTa")
+model = RobertaModel.from_pretrained("surajp/SanBERTa")
+
+op = tokenizer.encode("इयं भाषा न केवलं भारतस्य अपि तु विश्वस्य प्राचीनतमा भाषा इति मन्यते।", return_tensors="pt")
+ps = model(op)
+ps[0].shape
+
+```
+```
+'''
+Output:
+--------
+torch.Size([1, 47, 768])
+
+```
+
+
+### For \<mask\> Prediction
+
+```
+from transformers import pipeline
+
+fill_mask = pipeline(
+    "fill-mask",
+    model="surajp/SanBERTa",
+    tokenizer="surajp/SanBERTa"
+)
+
+## इयं भाषा न केवलं भारतस्य अपि तु विश्वस्य प्राचीनतमा भाषा इति मन्यते।
+fill_mask("इयं भाषा न केवल<mask> भारतस्य अपि तु विश्वस्य प्राचीनतमा भाषा इति मन्यते।")
+
+ps = model(torch.tensor(enc).unsqueeze(1))
+print(ps[0].shape)
+```
+```
+'''
+Output:
+--------
+[{'score': 0.7516744136810303,
+  'sequence': '<s> इयं भाषा न केवलं भारतस्य अपि तु विश्वस्य प्राचीनतमा भाषा इति मन्यते।</s>',
+  'token': 280,
+  'token_str': 'à¤Ĥ'},
+ {'score': 0.06230105459690094,
+  'sequence': '<s> इयं भाषा न केवली भारतस्य अपि तु विश्वस्य प्राचीनतमा भाषा इति मन्यते।</s>',
+  'token': 289,
+  'token_str': 'à¥Ģ'},
+ {'score': 0.055410224944353104,
+  'sequence': '<s> इयं भाषा न केवला भारतस्य अपि तु विश्वस्य प्राचीनतमा भाषा इति मन्यते।</s>',
+  'token': 265,
+  'token_str': 'à¤¾'},
+  ...]
+```
+
+### It works!! 🎉 🎉 🎉
+
+> Created by [Suraj Parmar/@parmarsuraj99](https://twitter.com/parmarsuraj99) | [LinkedIn](https://www.linkedin.com/in/parmarsuraj99/)
+
+> Made with <span style="color: #e25555;">&hearts;</span> in India
@@ -0,0 +1,40 @@
+---
+language: french
+---
+
+# tf-allociné
+
+A french sentiment analysis model, based on [CamemBERT](https://camembert-model.fr/), and finetuned on a large-scale dataset scraped from [Allociné.fr](http://www.allocine.fr/) user reviews.
+
+## Results
+
+| Validation Accuracy | Validation F1-Score | Test Accuracy | Test F1-Score |
+|--------------------:| -------------------:| -------------:|--------------:|
+|               97.39 |               97.36 |         97.44 |         97.34 |
+
+The dataset and the evaluation code are available on [this repo](https://github.com/TheophileBlard/french-sentiment-analysis-with-bert).
+
+## Usage
+
+```python
+from transformers import AutoTokenizer, TFAutoModelForSequenceClassification
+from transformers import pipeline
+
+tokenizer = AutoTokenizer.from_pretrained("tblard/tf-allocine")
+model = TFAutoModelForSequenceClassification.from_pretrained("tblard/tf-allocine")
+
+nlp = pipeline('sentiment-analysis', model=model, tokenizer=tokenizer)
+
+print(nlp("Alad'2 est clairement le meilleur film de l'année 2018.")) # POSITIVE
+print(nlp("Juste whoaaahouuu !")) # POSITIVE
+print(nlp("NUL...A...CHIER ! FIN DE TRANSMISSION.")) # NEGATIVE
+print(nlp("Je m'attendais à mieux de la part de Franck Dubosc !")) # NEGATIVE
+```
+
+## Author
+
+Théophile Blard – :email: theophile.blard@gmail.com
+
+If you use this work (code, model or dataset), please cite as:
+
+> Théophile Blard, French sentiment analysis with BERT, (2020), GitHub repository, <https://github.com/TheophileBlard/french-sentiment-analysis-with-bert>
@@ -34,3 +34,6 @@ Pull Request so it can be included under the Community notebooks.
 | [Fine-tune Longformer for QA](https://github.com/patil-suraj/Notebooks/blob/master/longformer_qa_training.ipynb) | How to fine-tune longformer model for QA task | [Suraj Patil](https://github.com/patil-suraj) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/patil-suraj/Notebooks/blob/master/longformer_qa_training.ipynb) |
 | [Evaluate Model with 🤗nlp](https://github.com/patrickvonplaten/notebooks/blob/master/How_to_evaluate_Longformer_on_TriviaQA_using_NLP.ipynb) | How to evaluate longformer on TriviaQA with `nlp` | [Patrick von Platen](https://github.com/patrickvonplaten) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/drive/1m7eTGlPmLRgoPkkA7rkhQdZ9ydpmsdLE?usp=sharing) |
 | [Fine-tune T5 for Sentiment Span Extraction](https://github.com/enzoampil/t5-intro/blob/master/t5_qa_training_pytorch_span_extraction.ipynb)  | How to fine-tune T5 for sentiment span extraction using a text-to-text format with PyTorch Lightning |  [Lorenzo Ampil](https://github.com/enzoampil) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/enzoampil/t5-intro/blob/master/t5_qa_training_pytorch_span_extraction.ipynb) |
+| [Fine-tune DistilBert for Multiclass Classification](https://github.com/abhimishra91/transformers-tutorials/blob/master/transformers_multiclass_classification.ipynb) | How to fine-tune DistilBert for multiclass classification with PyTorch | [Abhishek Kumar Mishra](https://github.com/abhimishra91) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abhimishra91/transformers-tutorials/blob/master/transformers_multiclass_classification.ipynb)|
+|[Fine-tune BERT for Multi-label Classification](https://github.com/abhimishra91/transformers-tutorials/blob/master/transformers_multi_label_classification.ipynb)|How to fine-tune BERT for multi-label classification using PyTorch|[Abhishek Kumar Mishra](https://github.com/abhimishra91) |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abhimishra91/transformers-tutorials/blob/master/transformers_multi_label_classification.ipynb)|
+|[Fine-tune T5 for Summarization](https://github.com/abhimishra91/transformers-tutorials/blob/master/transformers_summarization_wandb.ipynb)|How to fine-tune T5 for summarization in PyTorch and track experiments with WandB|[Abhishek Kumar Mishra](https://github.com/abhimishra91) |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/abhimishra91/transformers-tutorials/blob/master/transformers_summarization_wandb.ipynb)|
@@ -9,6 +9,7 @@ known_third_party =
    fastprogress
    git
    h5py
+    matplotlib
    MeCab
    nltk
    numpy
@@ -71,24 +71,24 @@ extras["sklearn"] = ["scikit-learn"]
 # keras2onnx and onnxconverter-common version is specific through a commit until 1.7.0 lands on pypi
 extras["tf"] = [
    "tensorflow",
-    "onnxconverter-common",
-    "keras2onnx"
+    "onnxconverter-common @ git+git://github.com/microsoft/onnxconverter-common.git@f64ca15989b6dc95a1f3507ff6e4c395ba12dff5#egg=onnxconverter-common",
+    "keras2onnx @ git+git://github.com/onnx/keras-onnx.git@cbdc75cb950b16db7f0a67be96a278f8d2953b48#egg=keras2onnx"
 ]
 extras["tf-cpu"] = [
    "tensorflow-cpu",
-    "onnxconverter-common",
-    "keras2onnx"
+    "onnxconverter-common @ git+git://github.com/microsoft/onnxconverter-common.git@f64ca15989b6dc95a1f3507ff6e4c395ba12dff5#egg=onnxconverter-common",
+    "keras2onnx @ git+git://github.com/onnx/keras-onnx.git@cbdc75cb950b16db7f0a67be96a278f8d2953b48#egg=keras2onnx"
 ]
 extras["torch"] = ["torch"]

 extras["serving"] = ["pydantic", "uvicorn", "fastapi", "starlette"]
 extras["all"] = extras["serving"] + ["tensorflow", "torch"]

-extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator"]
+extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator", "psutil"]
 extras["docs"] = ["recommonmark", "sphinx", "sphinx-markdown-tables", "sphinx-rtd-theme"]
 extras["quality"] = [
    "black",
-    "isort",
+    "isort @ git+git://github.com/timothycrosley/isort.git@e63ae06ec7d70b06df9e528357650281a3d3ec22#egg=isort",
    "flake8",
 ]
 extras["dev"] = extras["testing"] + extras["quality"] + ["mecab-python3", "scikit-learn", "tensorflow", "torch"]
@@ -78,6 +78,7 @@ from .file_utils import (
    cached_path,
    is_tf_available,
    is_torch_available,
+    is_torch_tpu_available,
 )
 from .hf_argparser import HfArgumentParser

@@ -132,10 +133,18 @@ from .tokenization_reformer import ReformerTokenizer
 from .tokenization_roberta import RobertaTokenizer, RobertaTokenizerFast
 from .tokenization_t5 import T5Tokenizer
 from .tokenization_transfo_xl import TransfoXLCorpus, TransfoXLTokenizer, TransfoXLTokenizerFast
-from .tokenization_utils import PreTrainedTokenizer
+from .tokenization_utils import (
+    BatchEncoding,
+    PreTrainedTokenizer,
+    PreTrainedTokenizerFast,
+    SpecialTokensMixin,
+    TensorType,
+)
 from .tokenization_xlm import XLMTokenizer
 from .tokenization_xlm_roberta import XLMRobertaTokenizer
 from .tokenization_xlnet import SPIECE_UNDERLINE, XLNetTokenizer
+
+# Trainer
 from .trainer_utils import EvalPrediction
 from .training_args import TrainingArguments
 from .training_args_tf import TFTrainingArguments
@@ -230,6 +239,7 @@ if is_torch_available():
        XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
    )
    from .modeling_bart import (
+        PretrainedBartModel,
        BartForSequenceClassification,
        BartModel,
        BartForConditionalGeneration,
@@ -289,6 +299,7 @@ if is_torch_available():
        XLMRobertaForMultipleChoice,
        XLMRobertaForSequenceClassification,
        XLMRobertaForTokenClassification,
+        XLMRobertaForQuestionAnswering,
        XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
    )
    from .modeling_mmbt import ModalEmbeddings, MMBTModel, MMBTForClassification
@@ -352,173 +363,193 @@ if is_torch_available():
 # TensorFlow
 if is_tf_available():
    from .modeling_tf_utils import (
-        TFPreTrainedModel,
-        TFSharedEmbeddings,
-        TFSequenceSummary,
        shape_list,
        tf_top_k_top_p_filtering,
+        TFPreTrainedModel,
+        TFSequenceSummary,
+        TFSharedEmbeddings,
    )
    from .modeling_tf_auto import (
-        TFAutoModel,
-        TFAutoModelForPreTraining,
-        TFAutoModelForMultipleChoice,
-        TFAutoModelForSequenceClassification,
-        TFAutoModelForQuestionAnswering,
-        TFAutoModelWithLMHead,
-        TFAutoModelForTokenClassification,
        TF_MODEL_MAPPING,
+        TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING,
        TF_MODEL_FOR_PRETRAINING_MAPPING,
-        TF_MODEL_WITH_LM_HEAD_MAPPING,
-        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING,
+        TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
        TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
-    )
-
-    from .modeling_tf_bert import (
-        TFBertPreTrainedModel,
-        TFBertMainLayer,
-        TFBertEmbeddings,
-        TFBertModel,
-        TFBertForPreTraining,
-        TFBertForMaskedLM,
-        TFBertForNextSentencePrediction,
-        TFBertForSequenceClassification,
-        TFBertForMultipleChoice,
-        TFBertForTokenClassification,
-        TFBertForQuestionAnswering,
-        TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_gpt2 import (
-        TFGPT2PreTrainedModel,
-        TFGPT2MainLayer,
-        TFGPT2Model,
-        TFGPT2LMHeadModel,
-        TFGPT2DoubleHeadsModel,
-        TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_openai import (
-        TFOpenAIGPTPreTrainedModel,
-        TFOpenAIGPTMainLayer,
-        TFOpenAIGPTModel,
-        TFOpenAIGPTLMHeadModel,
-        TFOpenAIGPTDoubleHeadsModel,
-        TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_transfo_xl import (
-        TFTransfoXLPreTrainedModel,
-        TFTransfoXLMainLayer,
-        TFTransfoXLModel,
-        TFTransfoXLLMHeadModel,
-        TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
-        TFAdaptiveEmbedding,
-    )
-
-    from .modeling_tf_xlnet import (
-        TFXLNetPreTrainedModel,
-        TFXLNetMainLayer,
-        TFXLNetModel,
-        TFXLNetLMHeadModel,
-        TFXLNetForSequenceClassification,
-        TFXLNetForTokenClassification,
-        TFXLNetForQuestionAnsweringSimple,
-        TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_xlm import (
-        TFXLMPreTrainedModel,
-        TFXLMMainLayer,
-        TFXLMModel,
-        TFXLMWithLMHeadModel,
-        TFXLMForSequenceClassification,
-        TFXLMForQuestionAnsweringSimple,
-        TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_xlm_roberta import (
-        TFXLMRobertaForMaskedLM,
-        TFXLMRobertaModel,
-        TFXLMRobertaForSequenceClassification,
-        TFXLMRobertaForTokenClassification,
-        TF_XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_roberta import (
-        TFRobertaPreTrainedModel,
-        TFRobertaMainLayer,
-        TFRobertaModel,
-        TFRobertaForMaskedLM,
-        TFRobertaForSequenceClassification,
-        TFRobertaForTokenClassification,
-        TFRobertaForQuestionAnswering,
-        TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_camembert import (
-        TFCamembertModel,
-        TFCamembertForMaskedLM,
-        TFCamembertForSequenceClassification,
-        TFCamembertForTokenClassification,
-        TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_flaubert import (
-        TFFlaubertModel,
-        TFFlaubertWithLMHeadModel,
-        TFFlaubertForSequenceClassification,
-        TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_distilbert import (
-        TFDistilBertPreTrainedModel,
-        TFDistilBertMainLayer,
-        TFDistilBertModel,
-        TFDistilBertForMaskedLM,
-        TFDistilBertForSequenceClassification,
-        TFDistilBertForTokenClassification,
-        TFDistilBertForQuestionAnswering,
-        TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
-    )
-
-    from .modeling_tf_ctrl import (
-        TFCTRLPreTrainedModel,
-        TFCTRLModel,
-        TFCTRLLMHeadModel,
-        TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TF_MODEL_WITH_LM_HEAD_MAPPING,
+        TFAutoModel,
+        TFAutoModelForMultipleChoice,
+        TFAutoModelForPreTraining,
+        TFAutoModelForQuestionAnswering,
+        TFAutoModelForSequenceClassification,
+        TFAutoModelForTokenClassification,
+        TFAutoModelWithLMHead,
    )

    from .modeling_tf_albert import (
-        TFAlbertPreTrainedModel,
-        TFAlbertMainLayer,
-        TFAlbertModel,
-        TFAlbertForPreTraining,
+        TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFAlbertForMaskedLM,
        TFAlbertForMultipleChoice,
-        TFAlbertForSequenceClassification,
+        TFAlbertForPreTraining,
        TFAlbertForQuestionAnswering,
-        TF_ALBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFAlbertForSequenceClassification,
+        TFAlbertForTokenClassification,
+        TFAlbertMainLayer,
+        TFAlbertModel,
+        TFAlbertPreTrainedModel,
    )

-    from .modeling_tf_t5 import (
-        TFT5PreTrainedModel,
-        TFT5Model,
-        TFT5ForConditionalGeneration,
-        TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST,
+    from .modeling_tf_bert import (
+        TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFBertEmbeddings,
+        TFBertForMaskedLM,
+        TFBertForMultipleChoice,
+        TFBertForNextSentencePrediction,
+        TFBertForPreTraining,
+        TFBertForQuestionAnswering,
+        TFBertForSequenceClassification,
+        TFBertForTokenClassification,
+        TFBertMainLayer,
+        TFBertModel,
+        TFBertPreTrainedModel,
+    )
+
+    from .modeling_tf_camembert import (
+        TF_CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFCamembertForMaskedLM,
+        TFCamembertModel,
+        TFCamembertForMultipleChoice,
+        TFCamembertForQuestionAnswering,
+        TFCamembertForSequenceClassification,
+        TFCamembertForTokenClassification,
+    )
+
+    from .modeling_tf_ctrl import (
+        TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFCTRLLMHeadModel,
+        TFCTRLModel,
+        TFCTRLPreTrainedModel,
+    )
+
+    from .modeling_tf_distilbert import (
+        TF_DISTILBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFDistilBertForMaskedLM,
+        TFDistilBertForMultipleChoice,
+        TFDistilBertForQuestionAnswering,
+        TFDistilBertForSequenceClassification,
+        TFDistilBertForTokenClassification,
+        TFDistilBertMainLayer,
+        TFDistilBertModel,
+        TFDistilBertPreTrainedModel,
    )

    from .modeling_tf_electra import (
-        TFElectraPreTrainedModel,
-        TFElectraModel,
-        TFElectraForPreTraining,
-        TFElectraForMaskedLM,
-        TFElectraForTokenClassification,
        TF_ELECTRA_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFElectraForMaskedLM,
+        TFElectraForPreTraining,
+        TFElectraForQuestionAnswering,
+        TFElectraForTokenClassification,
+        TFElectraModel,
+        TFElectraPreTrainedModel,
+    )
+
+    from .modeling_tf_flaubert import (
+        TF_FLAUBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFFlaubertForMultipleChoice,
+        TFFlaubertForQuestionAnsweringSimple,
+        TFFlaubertForSequenceClassification,
+        TFFlaubertForTokenClassification,
+        TFFlaubertWithLMHeadModel,
+        TFFlaubertModel,
+    )
+
+    from .modeling_tf_gpt2 import (
+        TF_GPT2_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFGPT2DoubleHeadsModel,
+        TFGPT2LMHeadModel,
+        TFGPT2MainLayer,
+        TFGPT2Model,
+        TFGPT2PreTrainedModel,
+    )
+
+    from .modeling_tf_openai import (
+        TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFOpenAIGPTDoubleHeadsModel,
+        TFOpenAIGPTLMHeadModel,
+        TFOpenAIGPTMainLayer,
+        TFOpenAIGPTModel,
+        TFOpenAIGPTPreTrainedModel,
+    )
+
+    from .modeling_tf_roberta import (
+        TF_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFRobertaForMaskedLM,
+        TFRobertaForMultipleChoice,
+        TFRobertaForQuestionAnswering,
+        TFRobertaForSequenceClassification,
+        TFRobertaForTokenClassification,
+        TFRobertaMainLayer,
+        TFRobertaModel,
+        TFRobertaPreTrainedModel,
+    )
+
+    from .modeling_tf_t5 import (
+        TF_T5_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFT5ForConditionalGeneration,
+        TFT5Model,
+        TFT5PreTrainedModel,
+    )
+
+    from .modeling_tf_transfo_xl import (
+        TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFAdaptiveEmbedding,
+        TFTransfoXLLMHeadModel,
+        TFTransfoXLMainLayer,
+        TFTransfoXLModel,
+        TFTransfoXLPreTrainedModel,
+    )
+
+    from .modeling_tf_xlm import (
+        TF_XLM_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFXLMForMultipleChoice,
+        TFXLMForQuestionAnsweringSimple,
+        TFXLMForSequenceClassification,
+        TFXLMForTokenClassification,
+        TFXLMWithLMHeadModel,
+        TFXLMMainLayer,
+        TFXLMModel,
+        TFXLMPreTrainedModel,
+    )
+
+    from .modeling_tf_xlm_roberta import (
+        TF_XLM_ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFXLMRobertaForMaskedLM,
+        TFXLMRobertaForMultipleChoice,
+        TFXLMRobertaForQuestionAnswering,
+        TFXLMRobertaForSequenceClassification,
+        TFXLMRobertaForTokenClassification,
+        TFXLMRobertaModel,
+    )
+
+    from .modeling_tf_xlnet import (
+        TF_XLNET_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFXLNetForMultipleChoice,
+        TFXLNetForQuestionAnsweringSimple,
+        TFXLNetForSequenceClassification,
+        TFXLNetForTokenClassification,
+        TFXLNetLMHeadModel,
+        TFXLNetMainLayer,
+        TFXLNetModel,
+        TFXLNetPreTrainedModel,
    )

    # Optimization
-    from .optimization_tf import WarmUp, create_optimizer, AdamWeightDecay, GradientAccumulator
+    from .optimization_tf import (
+        AdamWeightDecay,
+        create_optimizer,
+        GradientAccumulator,
+        WarmUp,
+    )

    # Trainer
    from .trainer_tf import TFTrainer
@@ -18,13 +18,18 @@
 """


-import inspect
 import logging
 import timeit

-from transformers import MODEL_MAPPING, MODEL_WITH_LM_HEAD_MAPPING, PretrainedConfig, is_torch_available
+from transformers import (
+    MODEL_MAPPING,
+    MODEL_WITH_LM_HEAD_MAPPING,
+    PretrainedConfig,
+    is_torch_available,
+    is_torch_tpu_available,
+)

-from .benchmark_utils import Benchmark, Memory, start_memory_tracing, stop_memory_tracing
+from .benchmark_utils import Benchmark, Memory, measure_peak_memory_cpu, start_memory_tracing, stop_memory_tracing


 if is_torch_available():
@@ -48,70 +53,48 @@ class PyTorchBenchmark(Benchmark):
    def train(self, model_name, batch_size, sequence_length, trace_memory=False):
        try:
            config = self.config_dict[model_name]
+
+            if self.args.torchscript:
+                config.torchscript = True
+
            model = MODEL_WITH_LM_HEAD_MAPPING[config.__class__](config)
            model.to(self.args.device)
            model.train()

+            # encoder-decoder has vocab size saved differently
+            vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
            input_ids = torch.randint(
-                model.config.vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device
+                vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device
            )

-            def compute_loss_and_backprob():
-                # TODO: Not all models call labels argument labels => this hack using the function signature should be corrected once all models have a common name for labels
-                function_argument_names = inspect.getfullargspec(model.forward).args
-                if "labels" in function_argument_names:
-                    loss = model(input_ids, labels=input_ids)[0]
-                elif "lm_labels" in function_argument_names:
-                    loss = model(input_ids, lm_labels=input_ids)[0]
-                elif "masked_lm_labels" in function_argument_names:
-                    loss = model(input_ids, masked_lm_labels=input_ids)[0]
-                else:
-                    NotImplementedError(f"{model_name} does not seem to allow training with labels")
-
-                loss.backward()
-                model.zero_grad()
-
-            if trace_memory is True:
-                if self.args.trace_memory_line_by_line or self.args.n_gpu == 0:
-                    trace = start_memory_tracing("transformers")
-                else:
-                    # clear cuda cache
-                    torch.cuda.empty_cache()
-                    torch.cuda.reset_peak_memory_stats()
-
-                # calculate loss and do backpropagation
-                compute_loss_and_backprob()
-
-                if self.args.trace_memory_line_by_line or self.args.n_gpu == 0:
-                    summary = stop_memory_tracing(trace)
-                    memory = summary.total
-                else:
-                    memory = Memory(torch.cuda.max_memory_reserved())
-
-                return memory
+            if self.args.torchscript:
+                raise NotImplementedError("Training for torchscript is currently not implemented")
            else:
-                # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
-                runtimes = timeit.repeat(lambda: compute_loss_and_backprob(), repeat=self.args.repeat, number=10,)
-                return min(runtimes) / 10.0
-        except RuntimeError as e:
-            self.print_fn("Doesn't fit on GPU. {}".format(e))
-            return "N/A"
+                train_model = model

-    def inference(self, model_name, batch_size, sequence_length, trace_memory=False):
-        try:
-            config = self.config_dict[model_name]
-            model = MODEL_MAPPING[config.__class__](config)
-            model.to(self.args.device)
-            model.eval()
+            def compute_loss_and_backprob_encoder():
+                loss = train_model(input_ids, labels=input_ids)[0]
+                loss.backward()
+                train_model.zero_grad()

-            input_ids = torch.randint(
-                config.vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device
+            def compute_loss_and_backprob_encoder_decoder():
+                loss = train_model(input_ids, decoder_input_ids=input_ids, labels=input_ids)[0]
+                loss.backward()
+                train_model.zero_grad()
+
+            _train = (
+                compute_loss_and_backprob_encoder_decoder
+                if config.is_encoder_decoder
+                else compute_loss_and_backprob_encoder
            )
+
            if trace_memory is True:
-                if self.args.trace_memory_line_by_line or self.args.n_gpu == 0:
+                if self.args.trace_memory_line_by_line:
                    trace = start_memory_tracing("transformers")
-                else:
-                    # clear cuda cache
+
+                if self.args.n_gpu > 0:
+                    # gpu
+                    # clear gpu cache
                    torch.cuda.empty_cache()
                    if hasattr(torch.cuda, "max_memory_reserved"):
                        torch.cuda.reset_peak_memory_stats()
@@ -121,12 +104,136 @@ class PyTorchBenchmark(Benchmark):
                        )
                        torch.cuda.reset_max_memory_cached()

-                model(input_ids)
-
-                if self.args.trace_memory_line_by_line or self.args.n_gpu == 0:
-                    summary = stop_memory_tracing(trace)
-                    memory = summary.total
+                        # calculate loss and do backpropagation
+                        _train()
+                elif not self.args.no_tpu and is_torch_tpu_available():
+                    # tpu
+                    raise NotImplementedError(
+                        "Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with `args.no_memory=True`"
+                    )
                else:
+                    # cpu
+                    memory_bytes = measure_peak_memory_cpu(_train)
+                    memory = Memory(memory_bytes) if isinstance(memory_bytes, int) else memory_bytes
+
+                if self.args.trace_memory_line_by_line:
+                    summary = stop_memory_tracing(trace)
+                else:
+                    summary = None
+
+                if self.args.n_gpu > 0:
+                    # gpu
+                    if hasattr(torch.cuda, "max_memory_reserved"):
+                        memory = Memory(torch.cuda.max_memory_reserved())
+                    else:
+                        logger.info(
+                            "Please consider updating PyTorch to version 1.4 to get more accuracy on GPU memory usage"
+                        )
+                        memory = Memory(torch.cuda.max_memory_cached())
+                    memory = Memory(torch.cuda.max_memory_reserved())
+
+                return memory, summary
+            else:
+                if (not self.args.no_tpu and is_torch_tpu_available()) or self.args.torchscript:
+                    # run additional 10 times to stabilize compilation for tpu and torchscript
+                    logger.info("Do inference on TPU or torchscript. Running model 5 times to stabilize compilation")
+                    timeit.repeat(
+                        _train, repeat=1, number=5,
+                    )
+
+                # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
+                runtimes = timeit.repeat(_train, repeat=self.args.repeat, number=10,)
+
+                if not self.args.no_tpu and is_torch_tpu_available() and self.args.tpu_print_metrics:
+                    import torch_xla.debug.metrics as met
+
+                    self.print_fn(met.metrics_report())
+
+                return min(runtimes) / 10.0
+        except RuntimeError as e:
+            self.print_fn("Doesn't fit on GPU. {}".format(e))
+            if trace_memory:
+                return "N/A", None
+            else:
+                return "N/A"
+
+    def inference(self, model_name, batch_size, sequence_length, trace_memory=False):
+        try:
+            config = self.config_dict[model_name]
+            model = None
+
+            if self.args.torchscript:
+                config.torchscript = True
+
+            if self.args.with_lm_head:
+                model = MODEL_WITH_LM_HEAD_MAPPING[config.__class__](config)
+            else:
+                model = MODEL_MAPPING[config.__class__](config)
+
+            model.eval()
+            model.to(self.args.device)
+
+            # encoder-decoder has vocab size saved differently
+            vocab_size = config.vocab_size if hasattr(config, "vocab_size") else config.encoder.vocab_size
+
+            input_ids = torch.randint(
+                vocab_size, (batch_size, sequence_length), dtype=torch.long, device=self.args.device
+            )
+
+            if self.args.torchscript:
+                with torch.no_grad():
+                    if config.is_encoder_decoder:
+                        raise NotImplementedError("Torchscript is currently not supported for EncoderDecoder models")
+                    else:
+                        inference_model = torch.jit.trace(model, input_ids)
+            else:
+                inference_model = model
+
+            def encoder_decoder_forward():
+                with torch.no_grad():
+                    inference_model(input_ids, decoder_input_ids=input_ids)
+
+            def encoder_forward():
+                with torch.no_grad():
+                    inference_model(input_ids)
+
+            _forward = encoder_decoder_forward if config.is_encoder_decoder else encoder_forward
+
+            if trace_memory is True:
+                if self.args.trace_memory_line_by_line:
+                    trace = start_memory_tracing("transformers")
+
+                if self.args.n_gpu > 0:
+                    # gpu
+                    # clear gpu cache
+                    torch.cuda.empty_cache()
+                    if hasattr(torch.cuda, "max_memory_reserved"):
+                        torch.cuda.reset_peak_memory_stats()
+                    else:
+                        logger.info(
+                            "Please consider updating PyTorch to version 1.4 to get more accuracy on GPU memory usage"
+                        )
+                        torch.cuda.reset_max_memory_cached()
+
+                        # run forward
+                        _forward()
+                elif not self.args.no_tpu and is_torch_tpu_available():
+                    # tpu
+                    raise NotImplementedError(
+                        "Memory Benchmarking is currently not implemented for TPU. Please disable memory benchmarking with `args.no_memory=True`"
+                    )
+                else:
+                    # cpu
+                    memory_bytes = measure_peak_memory_cpu(_forward)
+                    memory = Memory(memory_bytes) if isinstance(memory_bytes, int) else memory_bytes
+
+                if self.args.trace_memory_line_by_line:
+                    summary = stop_memory_tracing(trace)
+                else:
+                    summary = None
+
+                if self.args.n_gpu > 0:
+                    # gpu
                    if hasattr(torch.cuda, "max_memory_reserved"):
                        memory = Memory(torch.cuda.max_memory_reserved())
                    else:
@@ -135,12 +242,29 @@ class PyTorchBenchmark(Benchmark):
                        )
                        memory = Memory(torch.cuda.max_memory_cached())

-                return memory
+                return memory, summary
            else:
+
+                if (not self.args.no_tpu and is_torch_tpu_available()) or self.args.torchscript:
+                    # run additional 10 times to stabilize compilation for tpu and torchscript
+                    logger.info("Do inference on TPU or torchscript. Running model 5 times to stabilize compilation")
+                    timeit.repeat(
+                        _forward, repeat=1, number=5,
+                    )
+
                # as written in https://docs.python.org/2/library/timeit.html#timeit.Timer.repeat, min should be taken rather than the average
-                runtimes = timeit.repeat(lambda: model(input_ids), repeat=self.args.repeat, number=10,)
+                runtimes = timeit.repeat(_forward, repeat=self.args.repeat, number=10,)
+
+                if not self.args.no_tpu and is_torch_tpu_available() and self.args.tpu_print_metrics:
+                    import torch_xla.debug.metrics as met
+
+                    self.print_fn(met.metrics_report())
+
                return min(runtimes) / 10.0

        except RuntimeError as e:
            self.print_fn("Doesn't fit on GPU. {}".format(e))
-            return "N/A"
+            if trace_memory:
+                return "N/A", None
+            else:
+                return "N/A"
@@ -18,25 +18,16 @@ import logging
 from dataclasses import dataclass, field
 from typing import Tuple

-from ..file_utils import cached_property, is_torch_available, torch_required
+from ..file_utils import cached_property, is_torch_available, is_torch_tpu_available, torch_required
 from .benchmark_args_utils import BenchmarkArguments


 if is_torch_available():
    import torch

-try:
+if is_torch_tpu_available():
    import torch_xla.core.xla_model as xm

-    _has_tpu = True
-except ImportError:
-    _has_tpu = False
-
-
-@torch_required
-def is_tpu_available():
-    return _has_tpu
-

 logger = logging.getLogger(__name__)

@@ -45,7 +36,9 @@ logger = logging.getLogger(__name__)
 class PyTorchBenchmarkArguments(BenchmarkArguments):
    no_cuda: bool = field(default=False, metadata={"help": "Whether to run on available cuda devices"})
    torchscript: bool = field(default=False, metadata={"help": "Trace the models using torchscript"})
+    no_tpu: bool = field(default=False, metadata={"help": "Whether to run on available tpu devices"})
    fp16: bool = field(default=False, metadata={"help": "Use FP16 to accelerate inference."})
+    tpu_print_metrics: bool = field(default=False, metadata={"help": "Use FP16 to accelerate inference."})

    @cached_property
    @torch_required
@@ -54,7 +47,7 @@ class PyTorchBenchmarkArguments(BenchmarkArguments):
        if self.no_cuda:
            device = torch.device("cpu")
            n_gpu = 0
-        elif is_tpu_available():
+        elif is_torch_tpu_available():
            device = xm.xla_device()
            n_gpu = 0
        else:
@@ -61,6 +61,12 @@ class BenchmarkArguments:
    save_to_csv: bool = field(default=False, metadata={"help": "Save result to a CSV file"})
    log_print: bool = field(default=False, metadata={"help": "Save all print statements in a log file"})
    no_env_print: bool = field(default=False, metadata={"help": "Don't print environment information"})
+    with_lm_head: bool = field(
+        default=False,
+        metadata={
+            "help": "Use model with its language model head (MODEL_WITH_LM_HEAD_MAPPING instead of MODEL_MAPPING)"
+        },
+    )
    inference_time_csv_file: str = field(
        default=f"inference_time_{round(time())}.csv",
        metadata={"help": "CSV filename used if saving time results to csv."},
@@ -14,12 +14,15 @@ import sys
 from abc import ABC, abstractmethod
 from collections import defaultdict, namedtuple
 from datetime import datetime
-from typing import Iterable, List, NamedTuple, Optional, Union
+from multiprocessing import Pipe, Process
+from multiprocessing.connection import Connection
+from signal import SIGKILL
+from typing import Callable, Iterable, List, NamedTuple, Optional, Union

 from transformers import AutoConfig, PretrainedConfig
 from transformers import __version__ as version

-from ..file_utils import is_tf_available, is_torch_available
+from ..file_utils import is_tf_available, is_torch_available, is_torch_tpu_available
 from .benchmark_args_utils import BenchmarkArguments


@@ -36,7 +39,15 @@ logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
 _is_memory_tracing_enabled = False

 BenchmarkOutput = namedtuple(
-    "BenchmarkOutput", ["time_inference_result", "memory_inference_result", "time_train_result", "memory_train_result"]
+    "BenchmarkOutput",
+    [
+        "time_inference_result",
+        "memory_inference_result",
+        "time_train_result",
+        "memory_train_result",
+        "inference_summary",
+        "train_summary",
+    ],
 )


@@ -120,6 +131,127 @@ class MemorySummary(NamedTuple):
 MemoryTrace = List[UsedMemoryState]


+def measure_peak_memory_cpu(function: Callable[[], None], interval=0.5) -> int:
+    """
+        measures peak cpu memory consumption of a given `function`
+        running the function for at least interval seconds
+        and at most 20 * interval seconds.
+        This function is heavily inspired by: `memory_usage`
+        of the package `memory_profiler`: https://github.com/pythonprofilers/memory_profiler/blob/895c4ac7a08020d66ae001e24067da6dcea42451/memory_profiler.py#L239
+
+        Args:
+            - `function`: (`callable`): function() -> ...
+                function without any arguments to measure for which to measure the peak memory
+
+            - `interval`: (`float`)
+                interval in second for which to measure the memory usage
+
+        Returns:
+            - `max_memory`: (`int`)
+                cosumed memory peak in Bytes
+    """
+    try:
+        import psutil
+    except (ImportError):
+        logger.warning(
+            "Psutil not installed, we won't log CPU memory usage. "
+            "Install Psutil (pip install psutil) to use CPU memory tracing."
+        )
+        max_memory = "N/A"
+    else:
+
+        def _get_memory(process_id: int) -> int:
+            """
+                measures current cpu memory usage of a given `process_id`
+
+                Args:
+                    - `process_id`: (`int`)
+                        process_id for which to measure memory
+
+                Returns
+                    - `memory`: (`int`)
+                        cosumed memory in Bytes
+            """
+            process = psutil.Process(process_id)
+            try:
+                meminfo_attr = "memory_info" if hasattr(process, "memory_info") else "get_memory_info"
+                memory = getattr(process, meminfo_attr)()[0]
+            except psutil.AccessDenied:
+                raise ValueError("Error with Psutil.")
+            return memory
+
+        class MemoryMeasureProcess(Process):
+
+            """
+                `MemoryMeasureProcess` inherits from `Process` and overwrites
+                its `run()` method. Used to measure the memory usage of a process
+            """
+
+            def __init__(self, process_id: int, child_connection: Connection, interval: float):
+                super().__init__()
+                self.process_id = process_id
+                self.interval = interval
+                self.connection = child_connection
+                self.num_measurements = 1
+                self.mem_usage = _get_memory(process_id)
+
+            def run(self):
+                self.connection.send(0)
+                stop = False
+                while True:
+                    self.mem_usage = max(self.mem_usage, _get_memory(self.process_id))
+                    self.num_measurements += 1
+
+                    if stop:
+                        break
+
+                    stop = self.connection.poll(self.interval)
+
+                # send results to parent pipe
+                self.connection.send(self.mem_usage)
+                self.connection.send(self.num_measurements)
+
+        while True:
+            # create child, parent connection
+            child_connection, parent_connection = Pipe()
+
+            # instantiate process
+            mem_process = MemoryMeasureProcess(os.getpid(), child_connection, interval)
+            mem_process.start()
+
+            # wait until we get memory
+            parent_connection.recv()
+
+            try:
+                # execute function
+                function()
+
+                # start parent connection
+                parent_connection.send(0)
+
+                # receive memory and num measurements
+                max_memory = parent_connection.recv()
+                num_measurements = parent_connection.recv()
+            except Exception:
+                # kill process in a clean way
+                parent = psutil.Process(os.getpid())
+                for child in parent.children(recursive=True):
+                    os.kill(child.pid, SIGKILL)
+                mem_process.join(0)
+                raise RuntimeError("Process killed. Error in Process")
+
+            # run process at least 20 * interval or until it finishes
+            mem_process.join(20 * interval)
+
+            if (num_measurements > 4) or (interval < 1e-6):
+                break
+
+            # reduce interval
+            interval /= 10
+
+        return max_memory
+
+
 def start_memory_tracing(
    modules_to_trace: Optional[Union[str, Iterable[str]]] = None,
    modules_not_to_trace: Optional[Union[str, Iterable[str]]] = None,
@@ -401,15 +533,10 @@ class Benchmark(ABC):
    def print_fn(self):
        if self._print_fn is None:
            if self.args.log_print:
-                logging.basicConfig(
-                    level=logging.DEBUG,
-                    filename=self.args.log_filename,
-                    filemode="a+",
-                    format="%(asctime)-15s %(levelname)-8s %(message)s",
-                )

                def print_and_log(*args):
-                    logging.info(*args)
+                    with open(self.args.log_filename, "a") as log_file:
+                        log_file.write(str(*args) + "\n")
                    print(*args)

                self._print_fn = print_and_log
@@ -421,6 +548,10 @@ class Benchmark(ABC):
    def is_gpu(self):
        return self.args.n_gpu > 0

+    @property
+    def is_tpu(self):
+        return is_torch_tpu_available() and not self.args.no_tpu
+
    @property
    @abstractmethod
    def framework_version(self):
@@ -454,11 +585,15 @@ class Benchmark(ABC):
            train_result_time[model_name] = copy.deepcopy(model_dict)
            train_result_memory[model_name] = copy.deepcopy(model_dict)

+            inference_summary = train_summary = None
+
            for batch_size in self.args.batch_sizes:
                for sequence_length in self.args.sequence_lengths:
                    if not self.args.no_inference:
                        if not self.args.no_memory:
-                            memory = self.inference(model_name, batch_size, sequence_length, trace_memory=True)
+                            memory, inference_summary = self.inference(
+                                model_name, batch_size, sequence_length, trace_memory=True
+                            )
                            inference_result_memory[model_name]["result"][batch_size][sequence_length] = memory
                        if not self.args.no_speed:
                            time = self.inference(model_name, batch_size, sequence_length, trace_memory=False)
@@ -466,7 +601,9 @@ class Benchmark(ABC):

                    if self.args.training:
                        if not self.args.no_memory:
-                            memory = self.train(model_name, batch_size, sequence_length, trace_memory=True)
+                            memory, train_summary = self.train(
+                                model_name, batch_size, sequence_length, trace_memory=True
+                            )
                            train_result_memory[model_name]["result"][batch_size][sequence_length] = memory
                        if not self.args.no_speed:
                            time = self.inference(model_name, batch_size, sequence_length, trace_memory=False)
@@ -477,23 +614,39 @@ class Benchmark(ABC):
                self.print_fn("======= INFERENCE - SPEED - RESULT =======")
                self.print_results(inference_result_time)
                self.save_to_csv(inference_result_time, self.args.inference_time_csv_file)
+                if self.is_tpu:
+                    self.print_fn(
+                        "TPU was used for inference. Note that the time after compilation stabilized (after ~10 inferences model.forward(..) calls) was measured."
+                    )

            if not self.args.no_memory:
                self.print_fn("======= INFERENCE - MEMORY - RESULT =======")
                self.print_results(inference_result_memory)
                self.save_to_csv(inference_result_memory, self.args.inference_memory_csv_file)

+            if self.args.trace_memory_line_by_line:
+                self.print_fn("======= INFERENCE - MEMORY LINE BY LINE TRACE - SUMMARY =======")
+                self.print_memory_trace_statistics(inference_summary)
+
        if self.args.training:
            if not self.args.no_speed:
                self.print_fn("======= TRAIN - SPEED - RESULT =======")
                self.print_results(train_result_time)
                self.save_to_csv(train_result_time, self.args.train_time_csv_file)
+                if self.is_tpu:
+                    self.print_fn(
+                        "TPU was used for training. Note that the time after compilation stabilized (after ~10 train loss=model.forward(...) + loss.backward() calls) was measured."
+                    )

            if not self.args.no_memory:
                self.print_fn("======= TRAIN - MEMORY - RESULT =======")
                self.print_results(train_result_memory)
                self.save_to_csv(train_result_memory, self.args.train_memory_csv_file)

+            if self.args.trace_memory_line_by_line:
+                self.print_fn("======= TRAIN - MEMORY LINE BY LINE TRACE - SUMMARY =======")
+                self.print_memory_trace_statistics(train_summary)
+
        if not self.args.no_env_print:
            self.print_fn("\n======== ENVIRONMENT - INFORMATION ========")
            self.print_fn(
@@ -506,7 +659,14 @@ class Benchmark(ABC):
                for key, value in self.environment_info.items():
                    writer.writerow([key, value])

-        return BenchmarkOutput(inference_result_time, inference_result_memory, train_result_time, train_result_memory)
+        return BenchmarkOutput(
+            inference_result_time,
+            inference_result_memory,
+            train_result_time,
+            train_result_memory,
+            inference_summary,
+            train_summary,
+        )

    @property
    def environment_info(self):
@@ -514,6 +674,8 @@ class Benchmark(ABC):
            info = {}
            info["transformers_version"] = version
            info["framework"] = self.framework
+            if self.framework == "PyTorch":
+                info["use_torchscript"] = self.args.torchscript
            info["framework_version"] = self.framework_version
            info["python_version"] = platform.python_version()
            info["system"] = platform.system()
@@ -566,6 +728,10 @@ class Benchmark(ABC):
                    info["gpu_performance_state"] = py3nvml.nvmlDeviceGetPerformanceState(handle)
                    py3nvml.nvmlShutdown()

+            info["use_tpu"] = self.is_tpu
+            # TODO(PVP): See if we can add more information about TPU
+            # see: https://github.com/pytorch/xla/issues/2180
+
            self._environment_info = info
        return self._environment_info

@@ -85,6 +85,9 @@ class EncoderDecoderConfig(PretrainedConfig):
        Returns:
            :class:`EncoderDecoderConfig`: An instance of a configuration object
        """
+        logger.info("Set `config.is_decoder=True` for decoder_config")
+        decoder_config.is_decoder = True
+
        return cls(encoder=encoder_config.to_dict(), decoder=decoder_config.to_dict())

    def to_dict(self):
@@ -20,6 +20,7 @@ import copy
 import json
 import logging
 import os
+import warnings
 from typing import Dict, Tuple

 from .file_utils import CONFIG_NAME, cached_path, hf_bucket_url, is_remote_url
@@ -44,8 +45,6 @@ class PretrainedConfig(object):
                Name of the task used to fine-tune the model. This can be used when converting from an original (TensorFlow or PyTorch) checkpoint.
            num_labels (:obj:`int`, `optional`, defaults to `2`):
                Number of classes to use when the model is a classification model (sequences/tokens)
-            output_attentions (:obj:`bool`, `optional`, defaults to :obj:`False`):
-                Should the model returns attentions weights.
            output_hidden_states (:obj:`string`, `optional`, defaults to :obj:`False`):
                Should the model returns all hidden-states.
            torchscript (:obj:`bool`, `optional`, defaults to :obj:`False`):
@@ -55,8 +54,15 @@ class PretrainedConfig(object):

    def __init__(self, **kwargs):
        # Attributes with defaults
-        self.output_attentions = kwargs.pop("output_attentions", False)
        self.output_hidden_states = kwargs.pop("output_hidden_states", False)
+
+        if "output_attentions" in kwargs:
+            warnings.warn(
+                "The `output_attentions` in the configuration is deprecated and will be removed in a later version. "
+                "Please use `output_attention` as an argument to your model's foward method instead.",
+                FutureWarning,
+            )
+        self.output_attentions = kwargs.pop("output_attentions", False)
        self.use_cache = kwargs.pop("use_cache", True)  # Not used by all models
        self.torchscript = kwargs.pop("torchscript", False)  # Only used by PyTorch models
        self.use_bfloat16 = kwargs.pop("use_bfloat16", False)
@@ -289,6 +295,13 @@ class PretrainedConfig(object):
        if hasattr(config, "pruned_heads"):
            config.pruned_heads = dict((int(key), value) for key, value in config.pruned_heads.items())

+        if "output_attentions" in kwargs:
+            warnings.warn(
+                "The `output_attentions` in the configuration is deprecated and will be removed in a later version. "
+                "Please use `output_attention` as an argument to your model's forward method instead.",
+                FutureWarning,
+            )
+
        # Update config with kwargs if needed
        to_remove = []
        for key, value in kwargs.items():
@@ -163,6 +163,10 @@ class XLNetConfig(PretrainedConfig):
        self.n_layer = n_layer
        self.n_head = n_head
        assert d_model % n_head == 0
+        if "d_head" in kwargs:
+            assert (
+                kwargs["d_head"] == d_model // n_head
+            ), f"`d_head` ({kwargs['d_head']}) should be equal to `d_model // n_head` ({d_model // n_head})"
        self.d_head = d_model // n_head
        self.ff_activation = ff_activation
        self.d_inner = d_inner
@@ -91,7 +91,7 @@ class DataCollatorForLanguageModeling(DataCollator):
        batch = self._tensorize_batch(examples)
        if self.mlm:
            inputs, labels = self.mask_tokens(batch)
-            return {"input_ids": inputs, "masked_lm_labels": labels}
+            return {"input_ids": inputs, "labels": labels}
        else:
            return {"input_ids": batch, "labels": batch}

@@ -70,6 +70,7 @@ class GlueDataset(Dataset):
        tokenizer: PreTrainedTokenizer,
        limit_length: Optional[int] = None,
        mode: Union[str, Split] = Split.train,
+        cache_dir: Optional[str] = None,
    ):
        self.args = args
        self.processor = glue_processors[args.task_name]()
@@ -81,7 +82,7 @@ class GlueDataset(Dataset):
                raise KeyError("mode is not a valid split name")
        # Load data features from cache or dataset file
        cached_features_file = os.path.join(
-            args.data_dir,
+            cache_dir if cache_dir is not None else args.data_dir,
            "cached_{}_{}_{}_{}".format(
                mode.value, tokenizer.__class__.__name__, str(args.max_seq_length), args.task_name,
            ),
@@ -63,9 +63,9 @@ if _has_sklearn:
        elif task_name == "qqp":
            return acc_and_f1(preds, labels)
        elif task_name == "mnli":
-            return {"acc": simple_accuracy(preds, labels)}
+            return {"mnli/acc": simple_accuracy(preds, labels)}
        elif task_name == "mnli-mm":
-            return {"acc": simple_accuracy(preds, labels)}
+            return {"mnli-mm/acc": simple_accuracy(preds, labels)}
        elif task_name == "qnli":
            return {"acc": simple_accuracy(preds, labels)}
        elif task_name == "rte":
@@ -394,8 +394,8 @@ def squad_convert_examples_to_features(
                        "qas_id": ex.qas_id,
                    },
                    {
-                        "start_position": ex.start_position,
-                        "end_position": ex.end_position,
+                        "start_positions": ex.start_position,
+                        "end_positions": ex.end_position,
                        "cls_index": ex.cls_index,
                        "p_mask": ex.p_mask,
                        "is_impossible": ex.is_impossible,
@@ -412,8 +412,8 @@ def squad_convert_examples_to_features(
                "qas_id": tf.string,
            },
            {
-                "start_position": tf.int64,
-                "end_position": tf.int64,
+                "start_positions": tf.int64,
+                "end_positions": tf.int64,
                "cls_index": tf.int64,
                "p_mask": tf.int32,
                "is_impossible": tf.int32,
@@ -429,8 +429,8 @@ def squad_convert_examples_to_features(
                "qas_id": tf.TensorShape([]),
            },
            {
-                "start_position": tf.TensorShape([]),
-                "end_position": tf.TensorShape([]),
+                "start_positions": tf.TensorShape([]),
+                "end_positions": tf.TensorShape([]),
                "cls_index": tf.TensorShape([]),
                "p_mask": tf.TensorShape([None]),
                "is_impossible": tf.TensorShape([]),
@@ -68,6 +68,21 @@ except ImportError:
    torch_cache_home = os.path.expanduser(
        os.getenv("TORCH_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch"))
    )
+
+
+try:
+    import torch_xla.core.xla_model as xm
+
+    tpu_device = xm.xla_device()
+
+    if _torch_available:
+        _torch_tpu_available = True  # pylint: disable=
+    else:
+        _torch_tpu_available = False
+except ImportError:
+    _torch_tpu_available = False
+
+
 default_cache_path = os.path.join(torch_cache_home, "transformers")


@@ -98,6 +113,10 @@ def is_tf_available():
    return _tf_available


+def is_torch_tpu_available():
+    return _torch_tpu_available
+
+
 def add_start_docstrings(*docstr):
    def docstring_decorator(fn):
        fn.__doc__ = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")
@@ -64,6 +64,8 @@ class S3Object:
        self.lastModified = lastModified
        self.size = size
        self.rfilename = rfilename
+        for k, v in kwargs.items():
+            setattr(self, k, v)


 class ModelInfo:
@@ -78,7 +80,7 @@ class ModelInfo:
        author: Optional[str] = None,
        downloads: Optional[int] = None,
        tags: List[str] = [],
-        siblings: List[Dict] = [],  # list of files that constitute the model
+        siblings: Optional[List[Dict]] = None,  # list of files that constitute the model
        **kwargs
    ):
        self.modelId = modelId
@@ -86,7 +88,9 @@ class ModelInfo:
        self.author = author
        self.downloads = downloads
        self.tags = tags
-        self.siblings = [S3Object(**x) for x in siblings]
+        self.siblings = [S3Object(**x) for x in siblings] if siblings is not None else None
+        for k, v in kwargs.items():
+            setattr(self, k, v)


 class HfApi:
@@ -81,6 +81,8 @@ class HfArgumentParser(ArgumentParser):
                kwargs["type"] = field.type
                if field.default is not dataclasses.MISSING:
                    kwargs["default"] = field.default
+                elif field.default_factory is not dataclasses.MISSING:
+                    kwargs["default"] = field.default_factory()
                else:
                    kwargs["required"] = True
            self.add_argument(field_name, **kwargs)
@@ -149,9 +149,7 @@ class ModelCard:

        try:
            # Load from URL or cache if already cached
-            resolved_model_card_file = cached_path(
-                model_card_file, cache_dir=cache_dir, force_download=True, proxies=proxies, resume_download=False
-            )
+            resolved_model_card_file = cached_path(model_card_file, cache_dir=cache_dir, proxies=proxies)
            if resolved_model_card_file is None:
                raise EnvironmentError
            if resolved_model_card_file == model_card_file:
@@ -17,6 +17,7 @@
 import logging
 import math
 import os
+import warnings

 import torch
 import torch.nn as nn
@@ -25,7 +26,7 @@ from torch.nn import CrossEntropyLoss, MSELoss
 from .configuration_albert import AlbertConfig
 from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
 from .modeling_bert import ACT2FN, BertEmbeddings, BertSelfAttention, prune_linear_layer
-from .modeling_utils import PreTrainedModel
+from .modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices


 logger = logging.getLogger(__name__)
@@ -186,7 +187,6 @@ class AlbertAttention(BertSelfAttention):
    def __init__(self, config):
        super().__init__(config)

-        self.output_attentions = config.output_attentions
        self.num_attention_heads = config.num_attention_heads
        self.hidden_size = config.hidden_size
        self.attention_head_size = config.hidden_size // config.num_attention_heads
@@ -198,14 +198,9 @@ class AlbertAttention(BertSelfAttention):
    def prune_heads(self, heads):
        if len(heads) == 0:
            return
-        mask = torch.ones(self.num_attention_heads, self.attention_head_size)
-        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
-        for head in heads:
-            # Compute how many pruned heads are before the head and move the index accordingly
-            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
-            mask[head] = 0
-        mask = mask.view(-1).contiguous().eq(1)
-        index = torch.arange(len(mask))[mask].long()
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.num_attention_heads, self.attention_head_size, self.pruned_heads
+        )

        # Prune linear layers
        self.query = prune_linear_layer(self.query, index)
@@ -218,7 +213,7 @@ class AlbertAttention(BertSelfAttention):
        self.all_head_size = self.attention_head_size * self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

-    def forward(self, input_ids, attention_mask=None, head_mask=None):
+    def forward(self, input_ids, attention_mask=None, head_mask=None, output_attentions=False):
        mixed_query_layer = self.query(input_ids)
        mixed_key_layer = self.key(input_ids)
        mixed_value_layer = self.value(input_ids)
@@ -260,7 +255,7 @@ class AlbertAttention(BertSelfAttention):
        projected_context_layer = torch.einsum("bfnd,ndh->bfh", context_layer, w) + b
        projected_context_layer_dropout = self.dropout(projected_context_layer)
        layernormed_context_layer = self.LayerNorm(input_ids + projected_context_layer_dropout)
-        return (layernormed_context_layer, attention_probs) if self.output_attentions else (layernormed_context_layer,)
+        return (layernormed_context_layer, attention_probs) if output_attentions else (layernormed_context_layer,)


 class AlbertLayer(nn.Module):
@@ -274,8 +269,8 @@ class AlbertLayer(nn.Module):
        self.ffn_output = nn.Linear(config.intermediate_size, config.hidden_size)
        self.activation = ACT2FN[config.hidden_act]

-    def forward(self, hidden_states, attention_mask=None, head_mask=None):
-        attention_output = self.attention(hidden_states, attention_mask, head_mask)
+    def forward(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False):
+        attention_output = self.attention(hidden_states, attention_mask, head_mask, output_attentions)
        ffn_output = self.ffn(attention_output[0])
        ffn_output = self.activation(ffn_output)
        ffn_output = self.ffn_output(ffn_output)
@@ -288,19 +283,18 @@ class AlbertLayerGroup(nn.Module):
    def __init__(self, config):
        super().__init__()

-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.albert_layers = nn.ModuleList([AlbertLayer(config) for _ in range(config.inner_group_num)])

-    def forward(self, hidden_states, attention_mask=None, head_mask=None):
+    def forward(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False):
        layer_hidden_states = ()
        layer_attentions = ()

        for layer_index, albert_layer in enumerate(self.albert_layers):
-            layer_output = albert_layer(hidden_states, attention_mask, head_mask[layer_index])
+            layer_output = albert_layer(hidden_states, attention_mask, head_mask[layer_index], output_attentions)
            hidden_states = layer_output[0]

-            if self.output_attentions:
+            if output_attentions:
                layer_attentions = layer_attentions + (layer_output[1],)

            if self.output_hidden_states:
@@ -309,7 +303,7 @@ class AlbertLayerGroup(nn.Module):
        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (layer_hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (layer_attentions,)
        return outputs  # last-layer hidden state, (layer hidden states), (layer attentions)

@@ -319,12 +313,11 @@ class AlbertTransformer(nn.Module):
        super().__init__()

        self.config = config
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.embedding_hidden_mapping_in = nn.Linear(config.embedding_size, config.hidden_size)
        self.albert_layer_groups = nn.ModuleList([AlbertLayerGroup(config) for _ in range(config.num_hidden_groups)])

-    def forward(self, hidden_states, attention_mask=None, head_mask=None):
+    def forward(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False):
        hidden_states = self.embedding_hidden_mapping_in(hidden_states)

        all_attentions = ()
@@ -343,10 +336,11 @@ class AlbertTransformer(nn.Module):
                hidden_states,
                attention_mask,
                head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group],
+                output_attentions,
            )
            hidden_states = layer_group_output[0]

-            if self.output_attentions:
+            if output_attentions:
                all_attentions = all_attentions + layer_group_output[-1]

            if self.output_hidden_states:
@@ -355,7 +349,7 @@ class AlbertTransformer(nn.Module):
        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)

@@ -492,6 +486,7 @@ class AlbertModel(AlbertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -512,7 +507,7 @@ class AlbertModel(AlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -532,6 +527,8 @@ class AlbertModel(AlbertPreTrainedModel):

        """

+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -556,7 +553,9 @@ class AlbertModel(AlbertPreTrainedModel):
        embedding_output = self.embeddings(
            input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
        )
-        encoder_outputs = self.encoder(embedding_output, extended_attention_mask, head_mask=head_mask)
+        encoder_outputs = self.encoder(
+            embedding_output, extended_attention_mask, head_mask=head_mask, output_attentions=output_attentions,
+        )

        sequence_output = encoder_outputs[0]

@@ -599,11 +598,13 @@ class AlbertForPreTraining(AlbertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
-        masked_lm_labels=None,
+        labels=None,
        sentence_order_label=None,
+        output_attentions=None,
+        **kwargs,
    ):
        r"""
-        masked_lm_labels (``torch.LongTensor`` of shape ``(batch_size, sequence_length)``, `optional`, defaults to :obj:`None`):
+        labels (``torch.LongTensor`` of shape ``(batch_size, sequence_length)``, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
@@ -613,10 +614,12 @@ class AlbertForPreTraining(AlbertPreTrainedModel):
            Indices should be in ``[0, 1]``.
            ``0`` indicates original order (sequence A, then sequence B),
            ``1`` indicates switched order (sequence B, then sequence A).
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
-        loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -628,7 +631,7 @@ class AlbertForPreTraining(AlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -651,6 +654,14 @@ class AlbertForPreTraining(AlbertPreTrainedModel):

        """

+        if "masked_lm_labels" in kwargs:
+            warnings.warn(
+                "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("masked_lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
+
        outputs = self.albert(
            input_ids,
            attention_mask=attention_mask,
@@ -658,6 +669,7 @@ class AlbertForPreTraining(AlbertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output, pooled_output = outputs[:2]
@@ -667,9 +679,9 @@ class AlbertForPreTraining(AlbertPreTrainedModel):

        outputs = (prediction_scores, sop_scores,) + outputs[2:]  # add hidden states and attention if they are here

-        if masked_lm_labels is not None and sentence_order_label is not None:
+        if labels is not None and sentence_order_label is not None:
            loss_fct = CrossEntropyLoss()
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            sentence_order_loss = loss_fct(sop_scores.view(-1, 2), sentence_order_label.view(-1))
            total_loss = masked_lm_loss + sentence_order_loss
            outputs = (total_loss,) + outputs
@@ -742,18 +754,22 @@ class AlbertForMaskedLM(AlbertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
-        masked_lm_labels=None,
+        labels=None,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
-        masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with
            labels in ``[0, ..., config.vocab_size]``
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
-        loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -762,7 +778,7 @@ class AlbertForMaskedLM(AlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -777,10 +793,18 @@ class AlbertForMaskedLM(AlbertPreTrainedModel):
        tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
        model = AlbertForMaskedLM.from_pretrained('albert-base-v2')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
-        outputs = model(input_ids, masked_lm_labels=input_ids)
+        outputs = model(input_ids, labels=input_ids)
        loss, prediction_scores = outputs[:2]

        """
+        if "masked_lm_labels" in kwargs:
+            warnings.warn(
+                "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("masked_lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
+
        outputs = self.albert(
            input_ids=input_ids,
            attention_mask=attention_mask,
@@ -788,15 +812,16 @@ class AlbertForMaskedLM(AlbertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        sequence_outputs = outputs[0]

        prediction_scores = self.predictions(sequence_outputs)

        outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention if they are here
-        if masked_lm_labels is not None:
+        if labels is not None:
            loss_fct = CrossEntropyLoss()
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            outputs = (masked_lm_loss,) + outputs

        return outputs
@@ -828,6 +853,7 @@ class AlbertForSequenceClassification(AlbertPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -847,7 +873,7 @@ class AlbertForSequenceClassification(AlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -875,6 +901,7 @@ class AlbertForSequenceClassification(AlbertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        pooled_output = outputs[1]
@@ -923,6 +950,7 @@ class AlbertForTokenClassification(AlbertPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
@@ -940,7 +968,7 @@ class AlbertForTokenClassification(AlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -970,6 +998,7 @@ class AlbertForTokenClassification(AlbertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -1020,6 +1049,7 @@ class AlbertForQuestionAnswering(AlbertPreTrainedModel):
        inputs_embeds=None,
        start_positions=None,
        end_positions=None,
+        output_attentions=None,
    ):
        r"""
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -1044,7 +1074,7 @@ class AlbertForQuestionAnswering(AlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1074,6 +1104,7 @@ class AlbertForQuestionAnswering(AlbertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -121,6 +121,7 @@ from .modeling_xlm import (
 from .modeling_xlm_roberta import (
    XLMRobertaForMaskedLM,
    XLMRobertaForMultipleChoice,
+    XLMRobertaForQuestionAnswering,
    XLMRobertaForSequenceClassification,
    XLMRobertaForTokenClassification,
    XLMRobertaModel,
@@ -230,6 +231,7 @@ MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
        (DistilBertConfig, DistilBertForQuestionAnswering),
        (AlbertConfig, AlbertForQuestionAnswering),
        (LongformerConfig, LongformerForQuestionAnswering),
+        (XLMRobertaConfig, XLMRobertaForQuestionAnswering),
        (RobertaConfig, RobertaForQuestionAnswering),
        (BertConfig, BertForQuestionAnswering),
        (XLNetConfig, XLNetForQuestionAnsweringSimple),
@@ -16,6 +16,7 @@
 import logging
 import math
 import random
+import warnings
 from typing import Dict, List, Optional, Tuple

 import numpy as np
@@ -182,7 +183,6 @@ class EncoderLayer(nn.Module):
    def __init__(self, config: BartConfig):
        super().__init__()
        self.embed_dim = config.d_model
-        self.output_attentions = config.output_attentions
        self.self_attn = SelfAttention(
            self.embed_dim, config.encoder_attention_heads, dropout=config.attention_dropout,
        )
@@ -195,7 +195,7 @@ class EncoderLayer(nn.Module):
        self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
        self.final_layer_norm = LayerNorm(self.embed_dim)

-    def forward(self, x, encoder_padding_mask):
+    def forward(self, x, encoder_padding_mask, output_attentions=False):
        """
        Args:
            x (Tensor): input to the layer of shape `(seq_len, batch, embed_dim)`
@@ -211,7 +211,7 @@ class EncoderLayer(nn.Module):
        if self.normalize_before:
            x = self.self_attn_layer_norm(x)
        x, attn_weights = self.self_attn(
-            query=x, key=x, key_padding_mask=encoder_padding_mask, need_weights=self.output_attentions
+            query=x, key=x, key_padding_mask=encoder_padding_mask, output_attentions=output_attentions
        )
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
@@ -245,7 +245,6 @@ class BartEncoder(nn.Module):

        self.dropout = config.dropout
        self.layerdrop = config.encoder_layerdrop
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

        embed_dim = embed_tokens.embedding_dim
@@ -267,9 +266,7 @@ class BartEncoder(nn.Module):
        # mbart has one extra layer_norm
        self.layer_norm = LayerNorm(config.d_model) if config.normalize_before else None

-    def forward(
-        self, input_ids, attention_mask=None,
-    ):
+    def forward(self, input_ids, attention_mask=None, output_attentions=False):
        """
        Args:
            input_ids (LongTensor): tokens in the source language of shape
@@ -307,9 +304,9 @@ class BartEncoder(nn.Module):
            if self.training and (dropout_probability < self.layerdrop):  # skip the layer
                attn = None
            else:
-                x, attn = encoder_layer(x, attention_mask)
+                x, attn = encoder_layer(x, attention_mask, output_attentions=output_attentions)

-            if self.output_attentions:
+            if output_attentions:
                all_attentions.append(attn)

        if self.layer_norm:
@@ -328,7 +325,6 @@ class DecoderLayer(nn.Module):
    def __init__(self, config: BartConfig):
        super().__init__()
        self.embed_dim = config.d_model
-        self.output_attentions = config.output_attentions
        self.self_attn = SelfAttention(
            embed_dim=self.embed_dim, num_heads=config.decoder_attention_heads, dropout=config.attention_dropout,
        )
@@ -357,6 +353,7 @@ class DecoderLayer(nn.Module):
        layer_state=None,
        causal_mask=None,
        decoder_padding_mask=None,
+        output_attentions=False,
    ):
        residual = x

@@ -372,7 +369,7 @@ class DecoderLayer(nn.Module):
            layer_state=layer_state,  # adds keys to layer state
            key_padding_mask=decoder_padding_mask,
            attn_mask=causal_mask,
-            need_weights=self.output_attentions,
+            output_attentions=output_attentions,
        )
        x = F.dropout(x, p=self.dropout, training=self.training)
        x = residual + x
@@ -424,7 +421,6 @@ class BartDecoder(nn.Module):

    def __init__(self, config: BartConfig, embed_tokens: nn.Embedding):
        super().__init__()
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.dropout = config.dropout
        self.layerdrop = config.decoder_layerdrop
@@ -455,7 +451,8 @@ class BartDecoder(nn.Module):
        decoder_causal_mask,
        decoder_cached_states=None,
        use_cache=False,
-        **unused
+        output_attentions=False,
+        **unused,
    ):
        """
        Includes several features from "Jointly Learning to Align and
@@ -517,6 +514,7 @@ class BartDecoder(nn.Module):
                decoder_padding_mask=decoder_padding_mask,
                layer_state=layer_state,
                causal_mask=decoder_causal_mask,
+                output_attentions=output_attentions,
            )

            if use_cache:
@@ -524,7 +522,7 @@ class BartDecoder(nn.Module):

            if self.layer_norm and (idx == len(self.layers) - 1):  # last layer of mbart
                x = self.layer_norm(x)
-            if self.output_attentions:
+            if output_attentions:
                all_self_attns += (layer_self_attn,)

        # Convert to standard output format: (seq_len, BS, model_dim) -> (BS, seq_len, model_dim)
@@ -582,7 +580,7 @@ class SelfAttention(nn.Module):
        key_padding_mask: Optional[Tensor] = None,
        layer_state: Optional[Dict[str, Optional[Tensor]]] = None,
        attn_mask: Optional[Tensor] = None,
-        need_weights=False,
+        output_attentions=False,
    ) -> Tuple[Tensor, Optional[Tensor]]:
        """Input shape: Time(SeqLen) x Batch x Channel"""
        static_kv: bool = self.encoder_decoder_attention
@@ -654,7 +652,7 @@ class SelfAttention(nn.Module):
        assert attn_output.size() == (bsz * self.num_heads, tgt_len, self.head_dim)
        attn_output = attn_output.transpose(0, 1).contiguous().view(tgt_len, bsz, embed_dim)
        attn_output = self.out_proj(attn_output)
-        if need_weights:
+        if output_attentions:
            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len)
        else:
            attn_weights = None
@@ -796,7 +794,6 @@ def _get_shape(t):
 class BartModel(PretrainedBartModel):
    def __init__(self, config: BartConfig):
        super().__init__(config)
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

        padding_idx, vocab_size = config.pad_token_id, config.vocab_size
@@ -817,7 +814,9 @@ class BartModel(PretrainedBartModel):
        decoder_attention_mask=None,
        decoder_cached_states=None,
        use_cache=False,
+        output_attentions=None,
    ):
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions

        # make masks if user doesn't supply
        if not use_cache:
@@ -832,8 +831,11 @@ class BartModel(PretrainedBartModel):
            decoder_padding_mask, causal_mask = None, None

        assert decoder_input_ids is not None
+
        if encoder_outputs is None:
-            encoder_outputs = self.encoder(input_ids=input_ids, attention_mask=attention_mask)
+            encoder_outputs = self.encoder(
+                input_ids=input_ids, attention_mask=attention_mask, output_attentions=output_attentions,
+            )
        assert isinstance(encoder_outputs, tuple)
        # decoder outputs consists of (dec_features, layer_state, dec_hidden, dec_attn)
        decoder_outputs = self.decoder(
@@ -843,8 +845,10 @@ class BartModel(PretrainedBartModel):
            decoder_padding_mask,
            decoder_causal_mask=causal_mask,
            decoder_cached_states=decoder_cached_states,
+            output_attentions=output_attentions,
            use_cache=use_cache,
        )
+
        # Attention and hidden_states will be [] or None if they aren't needed
        decoder_outputs: Tuple = _filter_out_falsey_values(decoder_outputs)
        assert isinstance(decoder_outputs[0], torch.Tensor)
@@ -900,12 +904,13 @@ class BartForConditionalGeneration(PretrainedBartModel):
        decoder_input_ids=None,
        decoder_attention_mask=None,
        decoder_cached_states=None,
-        lm_labels=None,
+        labels=None,
        use_cache=False,
-        **unused
+        output_attentions=None,
+        **unused,
    ):
        r"""
-        lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should either be in ``[0, ..., config.vocab_size]`` or -100 (see ``input_ids`` docstring).
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens
@@ -914,7 +919,7 @@ class BartForConditionalGeneration(PretrainedBartModel):

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
-        masked_lm_loss (`optional`, returned when ``lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        masked_lm_loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -923,7 +928,7 @@ class BartForConditionalGeneration(PretrainedBartModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -945,6 +950,13 @@ class BartForConditionalGeneration(PretrainedBartModel):
            tokenizer.decode(predictions).split()
            # ['good', 'great', 'all', 'really', 'very']
        """
+        if "lm_labels" in unused:
+            warnings.warn(
+                "The `lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = unused.pop("lm_labels")
+
        outputs = self.model(
            input_ids,
            attention_mask=attention_mask,
@@ -953,13 +965,14 @@ class BartForConditionalGeneration(PretrainedBartModel):
            decoder_attention_mask=decoder_attention_mask,
            decoder_cached_states=decoder_cached_states,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )
        lm_logits = F.linear(outputs[0], self.model.shared.weight, bias=self.final_logits_bias)
        outputs = (lm_logits,) + outputs[1:]  # Add cache, hidden states and attention if they are here
-        if lm_labels is not None:
+        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()
-            # TODO(SS): do we need to ignore pad tokens in lm_labels?
-            masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), lm_labels.view(-1))
+            # TODO(SS): do we need to ignore pad tokens in labels?
+            masked_lm_loss = loss_fct(lm_logits.view(-1, self.config.vocab_size), labels.view(-1))
            outputs = (masked_lm_loss,) + outputs

        return outputs
@@ -1047,6 +1060,7 @@ class BartForSequenceClassification(PretrainedBartModel):
        decoder_input_ids=None,
        decoder_attention_mask=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -1064,7 +1078,7 @@ class BartForSequenceClassification(PretrainedBartModel):
                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
                Attentions weights after the attention softmax, used to compute the weighted average in the
                self-attention
@@ -1090,6 +1104,7 @@ class BartForSequenceClassification(PretrainedBartModel):
            decoder_input_ids=decoder_input_ids,
            decoder_attention_mask=decoder_attention_mask,
            encoder_outputs=encoder_outputs,
+            output_attentions=output_attentions,
        )
        x = outputs[0]  # last hidden state
        eos_mask = input_ids.eq(self.config.eos_token_id)
@@ -19,6 +19,7 @@
 import logging
 import math
 import os
+import warnings

 import torch
 from torch import nn
@@ -27,7 +28,7 @@ from torch.nn import CrossEntropyLoss, MSELoss
 from .activations import gelu, gelu_new, swish
 from .configuration_bert import BertConfig
 from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_utils import PreTrainedModel, prune_linear_layer
+from .modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer


 logger = logging.getLogger(__name__)
@@ -189,7 +190,6 @@ class BertSelfAttention(nn.Module):
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )
-        self.output_attentions = config.output_attentions

        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
@@ -213,6 +213,7 @@ class BertSelfAttention(nn.Module):
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
+        output_attentions=False,
    ):
        mixed_query_layer = self.query(hidden_states)

@@ -255,7 +256,7 @@ class BertSelfAttention(nn.Module):
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)

-        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
        return outputs


@@ -283,14 +284,9 @@ class BertAttention(nn.Module):
    def prune_heads(self, heads):
        if len(heads) == 0:
            return
-        mask = torch.ones(self.self.num_attention_heads, self.self.attention_head_size)
-        heads = set(heads) - self.pruned_heads  # Convert to set and remove already pruned heads
-        for head in heads:
-            # Compute how many pruned heads are before the head and move the index accordingly
-            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
-            mask[head] = 0
-        mask = mask.view(-1).contiguous().eq(1)
-        index = torch.arange(len(mask))[mask].long()
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
+        )

        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
@@ -310,9 +306,10 @@ class BertAttention(nn.Module):
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
+        output_attentions=False,
    ):
        self_outputs = self.self(
-            hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask
+            hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask, output_attentions,
        )
        attention_output = self.output(self_outputs[0], hidden_states)
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
@@ -365,14 +362,22 @@ class BertLayer(nn.Module):
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
+        output_attentions=False,
    ):
-        self_attention_outputs = self.attention(hidden_states, attention_mask, head_mask)
+        self_attention_outputs = self.attention(
+            hidden_states, attention_mask, head_mask, output_attentions=output_attentions,
+        )
        attention_output = self_attention_outputs[0]
        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights

        if self.is_decoder and encoder_hidden_states is not None:
            cross_attention_outputs = self.crossattention(
-                attention_output, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask
+                attention_output,
+                attention_mask,
+                head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions,
            )
            attention_output = cross_attention_outputs[0]
            outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights
@@ -386,7 +391,6 @@ class BertLayer(nn.Module):
 class BertEncoder(nn.Module):
    def __init__(self, config):
        super().__init__()
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])

@@ -397,6 +401,7 @@ class BertEncoder(nn.Module):
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
+        output_attentions=False,
    ):
        all_hidden_states = ()
        all_attentions = ()
@@ -405,11 +410,16 @@ class BertEncoder(nn.Module):
                all_hidden_states = all_hidden_states + (hidden_states,)

            layer_outputs = layer_module(
-                hidden_states, attention_mask, head_mask[i], encoder_hidden_states, encoder_attention_mask
+                hidden_states,
+                attention_mask,
+                head_mask[i],
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions,
            )
            hidden_states = layer_outputs[0]

-            if self.output_attentions:
+            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[1],)

        # Add last layer
@@ -419,7 +429,7 @@ class BertEncoder(nn.Module):
        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)

@@ -643,6 +653,7 @@ class BertModel(BertPreTrainedModel):
        inputs_embeds=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -663,7 +674,7 @@ class BertModel(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -684,6 +695,7 @@ class BertModel(BertPreTrainedModel):
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -732,6 +744,7 @@ class BertModel(BertPreTrainedModel):
            head_mask=head_mask,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_extended_attention_mask,
+            output_attentions=output_attentions,
        )
        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output)
@@ -768,11 +781,13 @@ class BertForPreTraining(BertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
-        masked_lm_labels=None,
+        labels=None,
        next_sentence_label=None,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
-        masked_lm_labels (``torch.LongTensor`` of shape ``(batch_size, sequence_length)``, `optional`, defaults to :obj:`None`):
+        labels (``torch.LongTensor`` of shape ``(batch_size, sequence_length)``, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
@@ -782,10 +797,12 @@ class BertForPreTraining(BertPreTrainedModel):
            Indices should be in ``[0, 1]``.
            ``0`` indicates sequence B is a continuation of sequence A,
            ``1`` indicates sequence B is a random sequence.
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
-        loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -797,7 +814,7 @@ class BertForPreTraining(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -819,6 +836,13 @@ class BertForPreTraining(BertPreTrainedModel):
        prediction_scores, seq_relationship_scores = outputs[:2]

        """
+        if "masked_lm_labels" in kwargs:
+            warnings.warn(
+                "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("masked_lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."

        outputs = self.bert(
            input_ids,
@@ -827,6 +851,7 @@ class BertForPreTraining(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output, pooled_output = outputs[:2]
@@ -836,9 +861,9 @@ class BertForPreTraining(BertPreTrainedModel):
            2:
        ]  # add hidden states and attention if they are here

-        if masked_lm_labels is not None and next_sentence_label is not None:
+        if labels is not None and next_sentence_label is not None:
            loss_fct = CrossEntropyLoss()
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
            total_loss = masked_lm_loss + next_sentence_loss
            outputs = (total_loss,) + outputs
@@ -846,6 +871,7 @@ class BertForPreTraining(BertPreTrainedModel):
        return outputs  # (loss), prediction_scores, seq_relationship_score, (hidden_states), (attentions)


+# TODO: Split with a different BertWithLMHead to get rid of `lm_labels` here and in encoder_decoder.
@add_start_docstrings("""Bert Model with a `language modeling` head on top. """, BERT_START_DOCSTRING)
 class BertForMaskedLM(BertPreTrainedModel):
    def __init__(self, config):
@@ -868,13 +894,15 @@ class BertForMaskedLM(BertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
-        masked_lm_labels=None,
+        labels=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
        lm_labels=None,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
-        masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
@@ -884,10 +912,12 @@ class BertForMaskedLM(BertPreTrainedModel):
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
-        masked_lm_loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        masked_lm_loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        ltr_lm_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`lm_labels` is provided):
                Next token prediction loss.
@@ -898,7 +928,7 @@ class BertForMaskedLM(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -914,11 +944,18 @@ class BertForMaskedLM(BertPreTrainedModel):
            model = BertForMaskedLM.from_pretrained('bert-base-uncased')

            input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
-            outputs = model(input_ids, masked_lm_labels=input_ids)
+            outputs = model(input_ids, labels=input_ids)

            loss, prediction_scores = outputs[:2]

        """
+        if "masked_lm_labels" in kwargs:
+            warnings.warn(
+                "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("masked_lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."

        outputs = self.bert(
            input_ids,
@@ -929,6 +966,7 @@ class BertForMaskedLM(BertPreTrainedModel):
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=encoder_hidden_states,
            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -942,9 +980,9 @@ class BertForMaskedLM(BertPreTrainedModel):
        #    of predictions for masked words.
        # 2. If `lm_labels` is provided we are in a causal scenario where we
        #    try to predict the next token for each input in the decoder.
-        if masked_lm_labels is not None:
+        if labels is not None:
            loss_fct = CrossEntropyLoss()  # -100 index = padding token
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            outputs = (masked_lm_loss,) + outputs

        if lm_labels is not None:
@@ -1002,6 +1040,7 @@ class BertForNextSentencePrediction(BertPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        next_sentence_label=None,
+        output_attentions=None,
    ):
        r"""
        next_sentence_label (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -1021,7 +1060,7 @@ class BertForNextSentencePrediction(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1051,6 +1090,7 @@ class BertForNextSentencePrediction(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        pooled_output = outputs[1]
@@ -1092,6 +1132,7 @@ class BertForSequenceClassification(BertPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -1111,7 +1152,7 @@ class BertForSequenceClassification(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1141,6 +1182,7 @@ class BertForSequenceClassification(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        pooled_output = outputs[1]
@@ -1188,6 +1230,7 @@ class BertForMultipleChoice(BertPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -1208,7 +1251,7 @@ class BertForMultipleChoice(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1248,6 +1291,7 @@ class BertForMultipleChoice(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        pooled_output = outputs[1]
@@ -1292,6 +1336,7 @@ class BertForTokenClassification(BertPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
@@ -1309,7 +1354,7 @@ class BertForTokenClassification(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1339,6 +1384,7 @@ class BertForTokenClassification(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -1390,6 +1436,7 @@ class BertForQuestionAnswering(BertPreTrainedModel):
        inputs_embeds=None,
        start_positions=None,
        end_positions=None,
+        output_attentions=None,
    ):
        r"""
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -1414,7 +1461,7 @@ class BertForQuestionAnswering(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1448,6 +1495,7 @@ class BertForQuestionAnswering(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -83,9 +83,8 @@ def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=N


 class MultiHeadAttention(torch.nn.Module):
-    def __init__(self, d_model_size, num_heads, output_attentions=False):
+    def __init__(self, d_model_size, num_heads):
        super().__init__()
-        self.output_attentions = output_attentions
        self.num_heads = num_heads
        self.d_model_size = d_model_size

@@ -101,7 +100,18 @@ class MultiHeadAttention(torch.nn.Module):
        x = x.reshape(batch_size, -1, self.num_heads, self.depth)
        return x.permute([0, 2, 1, 3])

-    def forward(self, v, k, q, mask, layer_past=None, attention_mask=None, head_mask=None, use_cache=False):
+    def forward(
+        self,
+        v,
+        k,
+        q,
+        mask,
+        layer_past=None,
+        attention_mask=None,
+        head_mask=None,
+        use_cache=False,
+        output_attentions=False,
+    ):
        batch_size = q.shape[0]

        q = self.Wq(q)
@@ -128,7 +138,7 @@ class MultiHeadAttention(torch.nn.Module):
        output = self.dense(original_size_attention)

        outputs = (output, present)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (attn,)
        return outputs

@@ -138,10 +148,10 @@ def point_wise_feed_forward_network(d_model_size, dff):


 class EncoderLayer(torch.nn.Module):
-    def __init__(self, d_model_size, num_heads, dff, rate=0.1, output_attentions=False):
+    def __init__(self, d_model_size, num_heads, dff, rate=0.1):
        super().__init__()

-        self.multi_head_attention = MultiHeadAttention(d_model_size, num_heads, output_attentions)
+        self.multi_head_attention = MultiHeadAttention(d_model_size, num_heads)
        self.ffn = point_wise_feed_forward_network(d_model_size, dff)

        self.layernorm1 = torch.nn.LayerNorm(d_model_size, eps=1e-6)
@@ -150,7 +160,9 @@ class EncoderLayer(torch.nn.Module):
        self.dropout1 = torch.nn.Dropout(rate)
        self.dropout2 = torch.nn.Dropout(rate)

-    def forward(self, x, mask, layer_past=None, attention_mask=None, head_mask=None, use_cache=False):
+    def forward(
+        self, x, mask, layer_past=None, attention_mask=None, head_mask=None, use_cache=False, output_attentions=False
+    ):
        normed = self.layernorm1(x)
        attn_outputs = self.multi_head_attention(
            normed,
@@ -161,6 +173,7 @@ class EncoderLayer(torch.nn.Module):
            attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )
        attn_output = attn_outputs[0]
        attn_output = self.dropout1(attn_output)
@@ -264,7 +277,6 @@ class CTRLModel(CTRLPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.output_hidden_states = config.output_hidden_states
-        self.output_attentions = config.output_attentions

        self.d_model_size = config.n_embd
        self.num_layers = config.n_layer
@@ -275,10 +287,7 @@ class CTRLModel(CTRLPreTrainedModel):

        self.dropout = nn.Dropout(config.embd_pdrop)
        self.h = nn.ModuleList(
-            [
-                EncoderLayer(config.n_embd, config.n_head, config.dff, config.resid_pdrop, config.output_attentions)
-                for _ in range(config.n_layer)
-            ]
+            [EncoderLayer(config.n_embd, config.n_head, config.dff, config.resid_pdrop) for _ in range(config.n_layer)]
        )
        self.layernorm = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)

@@ -308,6 +317,7 @@ class CTRLModel(CTRLPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        use_cache=True,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -322,7 +332,7 @@ class CTRLModel(CTRLPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -343,6 +353,7 @@ class CTRLModel(CTRLPreTrainedModel):
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -424,12 +435,13 @@ class CTRLModel(CTRLPreTrainedModel):
                attention_mask=attention_mask,
                head_mask=head_mask[i],
                use_cache=use_cache,
+                output_attentions=output_attentions,
            )
            hidden_states, present = outputs[:2]
            if use_cache is True:
                presents = presents + (present,)

-            if self.output_attentions:
+            if output_attentions:
                all_attentions.append(outputs[2])

        hidden_states = self.layernorm(hidden_states)
@@ -442,7 +454,7 @@ class CTRLModel(CTRLPreTrainedModel):
            outputs = outputs + (presents,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            # let the number of heads free (-1) so we can extract attention even after head pruning
            attention_output_shape = input_shape[:-1] + (-1,) + all_attentions[0].shape[-2:]
            all_attentions = tuple(t.view(*attention_output_shape) for t in all_attentions)
@@ -485,11 +497,12 @@ class CTRLLMHeadModel(CTRLPreTrainedModel):
        inputs_embeds=None,
        labels=None,
        use_cache=True,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for language modeling.
-            Note that the labels **are shifted** inside the model, i.e. you can set ``lm_labels = input_ids``
+            Note that the labels **are shifted** inside the model, i.e. you can set ``labels = input_ids``
            Indices are selected in ``[-100, 0, ..., config.vocab_size]``
            All labels set to ``-100`` are ignored (masked), the loss is only
            computed for labels in ``[0, ..., config.vocab_size]``
@@ -508,7 +521,7 @@ class CTRLLMHeadModel(CTRLPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -537,6 +550,7 @@ class CTRLLMHeadModel(CTRLPreTrainedModel):
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )

        hidden_states = transformer_outputs[0]
@@ -21,6 +21,7 @@
 import copy
 import logging
 import math
+import warnings

 import numpy as np
 import torch
@@ -30,7 +31,7 @@ from torch.nn import CrossEntropyLoss
 from .activations import gelu
 from .configuration_distilbert import DistilBertConfig
 from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_utils import PreTrainedModel, prune_linear_layer
+from .modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer


 logger = logging.getLogger(__name__)
@@ -104,7 +105,6 @@ class MultiHeadSelfAttention(nn.Module):
        self.n_heads = config.n_heads
        self.dim = config.dim
        self.dropout = nn.Dropout(p=config.attention_dropout)
-        self.output_attentions = config.output_attentions

        assert self.dim % self.n_heads == 0

@@ -119,13 +119,7 @@ class MultiHeadSelfAttention(nn.Module):
        attention_head_size = self.dim // self.n_heads
        if len(heads) == 0:
            return
-        mask = torch.ones(self.n_heads, attention_head_size)
-        heads = set(heads) - self.pruned_heads
-        for head in heads:
-            head -= sum(1 if h < head else 0 for h in self.pruned_heads)
-            mask[head] = 0
-        mask = mask.view(-1).contiguous().eq(1)
-        index = torch.arange(len(mask))[mask].long()
+        heads, index = find_pruneable_heads_and_indices(heads, self.n_heads, attention_head_size, self.pruned_heads)
        # Prune linear layers
        self.q_lin = prune_linear_layer(self.q_lin, index)
        self.k_lin = prune_linear_layer(self.k_lin, index)
@@ -136,7 +130,7 @@ class MultiHeadSelfAttention(nn.Module):
        self.dim = attention_head_size * self.n_heads
        self.pruned_heads = self.pruned_heads.union(heads)

-    def forward(self, query, key, value, mask, head_mask=None):
+    def forward(self, query, key, value, mask, head_mask=None, output_attentions=False):
        """
        Parameters
        ----------
@@ -189,7 +183,7 @@ class MultiHeadSelfAttention(nn.Module):
        context = unshape(context)  # (bs, q_length, dim)
        context = self.out_lin(context)  # (bs, q_length, dim)

-        if self.output_attentions:
+        if output_attentions:
            return (context, weights)
        else:
            return (context,)
@@ -218,8 +212,6 @@ class TransformerBlock(nn.Module):
    def __init__(self, config):
        super().__init__()

-        self.output_attentions = config.output_attentions
-
        assert config.dim % config.n_heads == 0

        self.attention = MultiHeadSelfAttention(config)
@@ -228,7 +220,7 @@ class TransformerBlock(nn.Module):
        self.ffn = FFN(config)
        self.output_layer_norm = nn.LayerNorm(normalized_shape=config.dim, eps=1e-12)

-    def forward(self, x, attn_mask=None, head_mask=None):
+    def forward(self, x, attn_mask=None, head_mask=None, output_attentions=False):
        """
        Parameters
        ----------
@@ -243,8 +235,10 @@ class TransformerBlock(nn.Module):
            The output of the transformer block contextualization.
        """
        # Self-Attention
-        sa_output = self.attention(query=x, key=x, value=x, mask=attn_mask, head_mask=head_mask)
-        if self.output_attentions:
+        sa_output = self.attention(
+            query=x, key=x, value=x, mask=attn_mask, head_mask=head_mask, output_attentions=output_attentions,
+        )
+        if output_attentions:
            sa_output, sa_weights = sa_output  # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)
        else:  # To handle these `output_attention` or `output_hidden_states` cases returning tuples
            assert type(sa_output) == tuple
@@ -256,7 +250,7 @@ class TransformerBlock(nn.Module):
        ffn_output = self.output_layer_norm(ffn_output + sa_output)  # (bs, seq_length, dim)

        output = (ffn_output,)
-        if self.output_attentions:
+        if output_attentions:
            output = (sa_weights,) + output
        return output

@@ -265,13 +259,12 @@ class Transformer(nn.Module):
    def __init__(self, config):
        super().__init__()
        self.n_layers = config.n_layers
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

        layer = TransformerBlock(config)
        self.layer = nn.ModuleList([copy.deepcopy(layer) for _ in range(config.n_layers)])

-    def forward(self, x, attn_mask=None, head_mask=None):
+    def forward(self, x, attn_mask=None, head_mask=None, output_attentions=False):
        """
        Parameters
        ----------
@@ -299,10 +292,12 @@ class Transformer(nn.Module):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_state,)

-            layer_outputs = layer_module(x=hidden_state, attn_mask=attn_mask, head_mask=head_mask[i])
+            layer_outputs = layer_module(
+                x=hidden_state, attn_mask=attn_mask, head_mask=head_mask[i], output_attentions=output_attentions
+            )
            hidden_state = layer_outputs[-1]

-            if self.output_attentions:
+            if output_attentions:
                assert len(layer_outputs) == 2
                attentions = layer_outputs[0]
                all_attentions = all_attentions + (attentions,)
@@ -316,7 +311,7 @@ class Transformer(nn.Module):
        outputs = (hidden_state,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)

@@ -413,7 +408,9 @@ class DistilBertModel(DistilBertPreTrainedModel):
            self.transformer.layer[layer].attention.prune_heads(heads)

    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
-    def forward(self, input_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None):
+    def forward(
+        self, input_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None, output_attentions=None,
+    ):
        r"""
    Return:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.DistilBertConfig`) and inputs:
@@ -424,7 +421,7 @@ class DistilBertModel(DistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -445,6 +442,8 @@ class DistilBertModel(DistilBertPreTrainedModel):
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -464,7 +463,9 @@ class DistilBertModel(DistilBertPreTrainedModel):

        if inputs_embeds is None:
            inputs_embeds = self.embeddings(input_ids)  # (bs, seq_length, dim)
-        tfmr_output = self.transformer(x=inputs_embeds, attn_mask=attention_mask, head_mask=head_mask)
+        tfmr_output = self.transformer(
+            x=inputs_embeds, attn_mask=attention_mask, head_mask=head_mask, output_attentions=output_attentions,
+        )
        hidden_state = tfmr_output[0]
        output = (hidden_state,) + tfmr_output[1:]

@@ -477,7 +478,6 @@ class DistilBertModel(DistilBertPreTrainedModel):
 class DistilBertForMaskedLM(DistilBertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

        self.distilbert = DistilBertModel(config)
@@ -493,17 +493,28 @@ class DistilBertForMaskedLM(DistilBertPreTrainedModel):
        return self.vocab_projector

    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
-    def forward(self, input_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None, masked_lm_labels=None):
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        **kwargs
+    ):
        r"""
-        masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.DistilBertConfig`) and inputs:
-        loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -512,7 +523,7 @@ class DistilBertForMaskedLM(DistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -527,12 +538,24 @@ class DistilBertForMaskedLM(DistilBertPreTrainedModel):
        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-cased')
        model = DistilBertForMaskedLM.from_pretrained('distilbert-base-cased')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
-        outputs = model(input_ids, masked_lm_labels=input_ids)
+        outputs = model(input_ids, labels=input_ids)
        loss, prediction_scores = outputs[:2]

        """
+        if "masked_lm_labels" in kwargs:
+            warnings.warn(
+                "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("masked_lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
+
        dlbrt_output = self.distilbert(
-            input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        hidden_states = dlbrt_output[0]  # (bs, seq_length, dim)
        prediction_logits = self.vocab_transform(hidden_states)  # (bs, seq_length, dim)
@@ -541,10 +564,8 @@ class DistilBertForMaskedLM(DistilBertPreTrainedModel):
        prediction_logits = self.vocab_projector(prediction_logits)  # (bs, seq_length, vocab_size)

        outputs = (prediction_logits,) + dlbrt_output[1:]
-        if masked_lm_labels is not None:
-            mlm_loss = self.mlm_loss_fct(
-                prediction_logits.view(-1, prediction_logits.size(-1)), masked_lm_labels.view(-1)
-            )
+        if labels is not None:
+            mlm_loss = self.mlm_loss_fct(prediction_logits.view(-1, prediction_logits.size(-1)), labels.view(-1))
            outputs = (mlm_loss,) + outputs

        return outputs  # (mlm_loss), prediction_logits, (all hidden_states), (all attentions)
@@ -568,7 +589,15 @@ class DistilBertForSequenceClassification(DistilBertPreTrainedModel):
        self.init_weights()

    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
-    def forward(self, input_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None, labels=None):
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for computing the sequence classification/regression loss.
@@ -587,7 +616,7 @@ class DistilBertForSequenceClassification(DistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -608,7 +637,11 @@ class DistilBertForSequenceClassification(DistilBertPreTrainedModel):

        """
        distilbert_output = self.distilbert(
-            input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        hidden_state = distilbert_output[0]  # (bs, seq_len, dim)
        pooled_output = hidden_state[:, 0]  # (bs, dim)
@@ -655,6 +688,7 @@ class DistilBertForQuestionAnswering(DistilBertPreTrainedModel):
        inputs_embeds=None,
        start_positions=None,
        end_positions=None,
+        output_attentions=None,
    ):
        r"""
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -679,7 +713,7 @@ class DistilBertForQuestionAnswering(DistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -701,7 +735,11 @@ class DistilBertForQuestionAnswering(DistilBertPreTrainedModel):

        """
        distilbert_output = self.distilbert(
-            input_ids=input_ids, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        hidden_states = distilbert_output[0]  # (bs, max_query_len, dim)

@@ -749,7 +787,15 @@ class DistilBertForTokenClassification(DistilBertPreTrainedModel):
        self.init_weights()

    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
-    def forward(self, input_ids=None, attention_mask=None, head_mask=None, inputs_embeds=None, labels=None):
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the token classification loss.
@@ -766,7 +812,7 @@ class DistilBertForTokenClassification(DistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -788,7 +834,11 @@ class DistilBertForTokenClassification(DistilBertPreTrainedModel):
        """

        outputs = self.distilbert(
-            input_ids, attention_mask=attention_mask, head_mask=head_mask, inputs_embeds=inputs_embeds
+            input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -1,5 +1,6 @@
 import logging
 import os
+import warnings

 import torch
 import torch.nn as nn
@@ -268,6 +269,7 @@ class ElectraModel(ElectraPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -279,7 +281,7 @@ class ElectraModel(ElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -300,6 +302,9 @@ class ElectraModel(ElectraPreTrainedModel):
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

        """
+
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -326,7 +331,12 @@ class ElectraModel(ElectraPreTrainedModel):
        if hasattr(self, "embeddings_project"):
            hidden_states = self.embeddings_project(hidden_states)

-        hidden_states = self.encoder(hidden_states, attention_mask=extended_attention_mask, head_mask=head_mask)
+        hidden_states = self.encoder(
+            hidden_states,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+        )

        return hidden_states

@@ -374,6 +384,7 @@ class ElectraForSequenceClassification(ElectraPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -393,7 +404,7 @@ class ElectraForSequenceClassification(ElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -405,8 +416,8 @@ class ElectraForSequenceClassification(ElectraPreTrainedModel):
        from transformers import BertTokenizer, BertForSequenceClassification
        import torch

-        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-        model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+        tokenizer = ElectraTokenizer.from_pretrained('bert-base-uncased')
+        model = ElectraForSequenceClassification.from_pretrained('bert-base-uncased')

        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
@@ -416,7 +427,7 @@ class ElectraForSequenceClassification(ElectraPreTrainedModel):

        """
        discriminator_hidden_states = self.electra(
-            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds
+            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions
        )

        sequence_output = discriminator_hidden_states[0]
@@ -463,6 +474,7 @@ class ElectraForPreTraining(ElectraPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (``torch.LongTensor`` of shape ``(batch_size, sequence_length)``, `optional`, defaults to :obj:`None`):
@@ -482,7 +494,7 @@ class ElectraForPreTraining(ElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -506,7 +518,7 @@ class ElectraForPreTraining(ElectraPreTrainedModel):
        """

        discriminator_hidden_states = self.electra(
-            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds
+            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions,
        )
        discriminator_sequence_output = discriminator_hidden_states[0]

@@ -561,18 +573,22 @@ class ElectraForMaskedLM(ElectraPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
-        masked_lm_labels=None,
+        labels=None,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
-        masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.ElectraConfig`) and inputs:
-        masked_lm_loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        masked_lm_loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -581,7 +597,7 @@ class ElectraForMaskedLM(ElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -597,14 +613,21 @@ class ElectraForMaskedLM(ElectraPreTrainedModel):
            model = ElectraForMaskedLM.from_pretrained('google/electra-small-generator')

            input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
-            outputs = model(input_ids, masked_lm_labels=input_ids)
+            outputs = model(input_ids, labels=input_ids)

            loss, prediction_scores = outputs[:2]

        """
+        if "masked_lm_labels" in kwargs:
+            warnings.warn(
+                "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("masked_lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."

        generator_hidden_states = self.electra(
-            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds
+            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions,
        )
        generator_sequence_output = generator_hidden_states[0]

@@ -614,9 +637,9 @@ class ElectraForMaskedLM(ElectraPreTrainedModel):
        output = (prediction_scores,)

        # Masked language modeling softmax layer
-        if masked_lm_labels is not None:
+        if labels is not None:
            loss_fct = nn.CrossEntropyLoss()  # -100 index = padding token
-            loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+            loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            output = (loss,) + output

        output += generator_hidden_states[1:]
@@ -650,6 +673,7 @@ class ElectraForTokenClassification(ElectraPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
@@ -667,7 +691,7 @@ class ElectraForTokenClassification(ElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -691,7 +715,7 @@ class ElectraForTokenClassification(ElectraPreTrainedModel):
        """

        discriminator_hidden_states = self.electra(
-            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds
+            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, output_attentions,
        )
        discriminator_sequence_output = discriminator_hidden_states[0]

@@ -191,7 +191,7 @@ class EncoderDecoderModel(PreTrainedModel):
        decoder_attention_mask=None,
        decoder_head_mask=None,
        decoder_inputs_embeds=None,
-        masked_lm_labels=None,
+        labels=None,
        lm_labels=None,
        **kwargs,
    ):
@@ -234,7 +234,7 @@ class EncoderDecoderModel(PreTrainedModel):
                Optionally, instead of passing :obj:`decoder_input_ids` you can choose to directly pass an embedded representation.
                This is useful if you want more control over how to convert `decoder_input_ids` indices into associated vectors
                than the model's internal embedding lookup matrix.
-            masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+            labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
                Labels for computing the masked language modeling loss for the decoder.
                Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
                Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
@@ -294,7 +294,7 @@ class EncoderDecoderModel(PreTrainedModel):
            encoder_attention_mask=attention_mask,
            head_mask=decoder_head_mask,
            lm_labels=lm_labels,
-            masked_lm_labels=masked_lm_labels,
+            labels=labels,
            **kwargs_decoder,
        )

@@ -128,6 +128,7 @@ class FlaubertModel(XLMModel):
        cache=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -139,7 +140,7 @@ class FlaubertModel(XLMModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -158,6 +159,8 @@ class FlaubertModel(XLMModel):
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
        # removed: src_enc=None, src_len=None
        if input_ids is not None:
            bs, slen = input_ids.size()
@@ -240,9 +243,11 @@ class FlaubertModel(XLMModel):

            # self attention
            if not self.pre_norm:
-                attn_outputs = self.attentions[i](tensor, attn_mask, cache=cache, head_mask=head_mask[i])
+                attn_outputs = self.attentions[i](
+                    tensor, attn_mask, cache=cache, head_mask=head_mask[i], output_attentions=output_attentions,
+                )
                attn = attn_outputs[0]
-                if self.output_attentions:
+                if output_attentions:
                    attentions = attentions + (attn_outputs[1],)
                attn = F.dropout(attn, p=self.dropout, training=self.training)
                tensor = tensor + attn
@@ -251,7 +256,7 @@ class FlaubertModel(XLMModel):
                tensor_normalized = self.layer_norm1[i](tensor)
                attn_outputs = self.attentions[i](tensor_normalized, attn_mask, cache=cache, head_mask=head_mask[i])
                attn = attn_outputs[0]
-                if self.output_attentions:
+                if output_attentions:
                    attentions = attentions + (attn_outputs[1],)
                attn = F.dropout(attn, p=self.dropout, training=self.training)
                tensor = tensor + attn
@@ -287,7 +292,7 @@ class FlaubertModel(XLMModel):
        outputs = (tensor,)
        if self.output_hidden_states:
            outputs = outputs + (hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (attentions,)
        return outputs  # outputs, (hidden_states), (attentions)

@@ -18,6 +18,7 @@

 import logging
 import os
+import warnings

 import torch
 import torch.nn as nn
@@ -26,7 +27,13 @@ from torch.nn import CrossEntropyLoss
 from .activations import ACT2FN
 from .configuration_gpt2 import GPT2Config
 from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_utils import Conv1D, PreTrainedModel, SequenceSummary, prune_conv1d_layer
+from .modeling_utils import (
+    Conv1D,
+    PreTrainedModel,
+    SequenceSummary,
+    find_pruneable_heads_and_indices,
+    prune_conv1d_layer,
+)


 logger = logging.getLogger(__name__)
@@ -99,7 +106,6 @@ def load_tf_weights_in_gpt2(model, config, gpt2_checkpoint_path):
 class Attention(nn.Module):
    def __init__(self, nx, n_ctx, config, scale=False):
        super().__init__()
-        self.output_attentions = config.output_attentions

        n_state = nx  # in Attention: n_state=768 (nx=n_embd)
        # [switch nx => n_state from Block to Attention to keep identical to TF implem]
@@ -121,14 +127,9 @@ class Attention(nn.Module):
    def prune_heads(self, heads):
        if len(heads) == 0:
            return
-        mask = torch.ones(self.n_head, self.split_size // self.n_head)
-        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
-        for head in heads:
-            # Compute how many pruned heads are before the head and move the index accordingly
-            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
-            mask[head] = 0
-        mask = mask.view(-1).contiguous().eq(1)
-        index = torch.arange(len(mask))[mask].long()
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.n_head, self.split_size // self.n_head, self.pruned_heads
+        )
        index_attn = torch.cat([index, index + self.split_size, index + (2 * self.split_size)])

        # Prune conv1d layers
@@ -140,7 +141,7 @@ class Attention(nn.Module):
        self.n_head = self.n_head - len(heads)
        self.pruned_heads = self.pruned_heads.union(heads)

-    def _attn(self, q, k, v, attention_mask=None, head_mask=None):
+    def _attn(self, q, k, v, attention_mask=None, head_mask=None, output_attentions=False):
        w = torch.matmul(q, k)
        if self.scale:
            w = w / (float(v.size(-1)) ** 0.5)
@@ -160,7 +161,7 @@ class Attention(nn.Module):
            w = w * head_mask

        outputs = [torch.matmul(w, v)]
-        if self.output_attentions:
+        if output_attentions:
            outputs.append(w)
        return outputs

@@ -177,7 +178,9 @@ class Attention(nn.Module):
        else:
            return x.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)

-    def forward(self, x, layer_past=None, attention_mask=None, head_mask=None, use_cache=False):
+    def forward(
+        self, x, layer_past=None, attention_mask=None, head_mask=None, use_cache=False, output_attentions=False
+    ):
        x = self.c_attn(x)
        query, key, value = x.split(self.split_size, dim=2)
        query = self.split_heads(query)
@@ -193,7 +196,7 @@ class Attention(nn.Module):
        else:
            present = (None,)

-        attn_outputs = self._attn(query, key, value, attention_mask, head_mask)
+        attn_outputs = self._attn(query, key, value, attention_mask, head_mask, output_attentions)
        a = attn_outputs[0]

        a = self.merge_heads(a)
@@ -228,13 +231,16 @@ class Block(nn.Module):
        self.ln_2 = nn.LayerNorm(nx, eps=config.layer_norm_epsilon)
        self.mlp = MLP(4 * nx, config)

-    def forward(self, x, layer_past=None, attention_mask=None, head_mask=None, use_cache=False):
+    def forward(
+        self, x, layer_past=None, attention_mask=None, head_mask=None, use_cache=False, output_attentions=False,
+    ):
        output_attn = self.attn(
            self.ln_1(x),
            layer_past=layer_past,
            attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )
        a = output_attn[0]  # output_attn: a, present, (attentions)

@@ -340,7 +346,6 @@ class GPT2Model(GPT2PreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        self.output_hidden_states = config.output_hidden_states
-        self.output_attentions = config.output_attentions

        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
        self.wpe = nn.Embedding(config.n_positions, config.n_embd)
@@ -374,6 +379,7 @@ class GPT2Model(GPT2PreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        use_cache=True,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -389,7 +395,7 @@ class GPT2Model(GPT2PreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -408,6 +414,7 @@ class GPT2Model(GPT2PreTrainedModel):
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -486,13 +493,14 @@ class GPT2Model(GPT2PreTrainedModel):
                attention_mask=attention_mask,
                head_mask=head_mask[i],
                use_cache=use_cache,
+                output_attentions=output_attentions,
            )

            hidden_states, present = outputs[:2]
            if use_cache is True:
                presents = presents + (present,)

-            if self.output_attentions:
+            if output_attentions:
                all_attentions.append(outputs[2])

        hidden_states = self.ln_f(hidden_states)
@@ -507,7 +515,7 @@ class GPT2Model(GPT2PreTrainedModel):
            outputs = outputs + (presents,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            # let the number of heads free (-1) so we can extract attention even after head pruning
            attention_output_shape = input_shape[:-1] + (-1,) + all_attentions[0].shape[-2:]
            all_attentions = tuple(t.view(*attention_output_shape) for t in all_attentions)
@@ -550,6 +558,7 @@ class GPT2LMHeadModel(GPT2PreTrainedModel):
        inputs_embeds=None,
        labels=None,
        use_cache=True,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
@@ -573,7 +582,7 @@ class GPT2LMHeadModel(GPT2PreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -602,6 +611,7 @@ class GPT2LMHeadModel(GPT2PreTrainedModel):
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )
        hidden_states = transformer_outputs[0]

@@ -652,17 +662,19 @@ class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        mc_token_ids=None,
-        lm_labels=None,
+        labels=None,
        mc_labels=None,
        use_cache=True,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
        mc_token_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, num_choices)`, `optional`, default to index of the last token of the input)
            Index of the classification token in each input sequence.
            Selected in the range ``[0, input_ids.size(-1) - 1[``.
-        lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`)
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`)
            Labels for language modeling.
-            Note that the labels **are shifted** inside the model, i.e. you can set ``lm_labels = input_ids``
+            Note that the labels **are shifted** inside the model, i.e. you can set ``labels = input_ids``
            Indices are selected in ``[-1, 0, ..., config.vocab_size]``
            All labels set to ``-100`` are ignored (masked), the loss is only
            computed for labels in ``[0, ..., config.vocab_size]``
@@ -670,12 +682,14 @@ class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
            Labels for computing the multiple choice classification loss.
            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
            of the input tensors. (see `input_ids` above)
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Return:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.GPT2Config`) and inputs:
-        lm_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``lm_labels`` is provided):
+        lm_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided):
            Language modeling loss.
-        mc_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`multiple_choice_labels` is provided):
+        mc_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`mc_labels` is provided):
            Multiple choice classification loss.
        lm_prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -689,7 +703,7 @@ class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -720,6 +734,14 @@ class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
        lm_prediction_scores, mc_prediction_scores = outputs[:2]

        """
+        if "lm_labels" in kwargs:
+            warnings.warn(
+                "The `lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
+
        transformer_outputs = self.transformer(
            input_ids,
            past=past,
@@ -729,6 +751,7 @@ class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )

        hidden_states = transformer_outputs[0]
@@ -741,9 +764,9 @@ class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1))
            outputs = (loss,) + outputs
-        if lm_labels is not None:
+        if labels is not None:
            shift_logits = lm_logits[..., :-1, :].contiguous()
-            shift_labels = lm_labels[..., 1:].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
            outputs = (loss,) + outputs
@@ -16,6 +16,7 @@

 import logging
 import math
+import warnings

 import torch
 import torch.nn as nn
@@ -86,7 +87,6 @@ class LongformerSelfAttention(nn.Module):
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )
-        self.output_attentions = config.output_attentions
        self.num_heads = config.num_attention_heads
        self.head_dim = int(config.hidden_size / config.num_attention_heads)
        self.embed_dim = config.hidden_size
@@ -152,12 +152,11 @@ class LongformerSelfAttention(nn.Module):
        beginning_mask_2d = input_tensor.new_ones(w, w + 1).tril().flip(dims=[0])
        beginning_mask = beginning_mask_2d[None, :, None, :]
        ending_mask = beginning_mask.flip(dims=(1, 3))
-        seqlen = input_tensor.size(1)
        beginning_input = input_tensor[:, :affected_seqlen, :, : w + 1]
-        beginning_mask = beginning_mask[:, :seqlen].expand(beginning_input.size())
+        beginning_mask = beginning_mask.expand(beginning_input.size())
        beginning_input.masked_fill_(beginning_mask == 1, -float("inf"))  # `== 1` converts to bool or uint8
        ending_input = input_tensor[:, -affected_seqlen:, :, -(w + 1) :]
-        ending_mask = ending_mask[:, -seqlen:].expand(ending_input.size())
+        ending_mask = ending_mask.expand(ending_input.size())
        ending_input.masked_fill_(ending_mask == 1, -float("inf"))  # `== 1` converts to bool or uint8

    def _sliding_chunks_matmul_qk(self, q: torch.Tensor, k: torch.Tensor, w: int):
@@ -242,6 +241,7 @@ class LongformerSelfAttention(nn.Module):
        head_mask=None,
        encoder_hidden_states=None,
        encoder_attention_mask=None,
+        output_attentions=False,
    ):
        """
        LongformerSelfAttention expects `len(hidden_states)` to be multiple of `attention_window`.
@@ -300,7 +300,6 @@ class LongformerSelfAttention(nn.Module):
        k = k.view(seqlen, batch_size, self.num_heads, self.head_dim).transpose(0, 1)
        # attn_weights = (batch_size, seqlen, num_heads, window*2+1)
        attn_weights = self._sliding_chunks_matmul_qk(q, k, self.one_sided_attention_window_size)
-        self._mask_invalid_locations(attn_weights, self.one_sided_attention_window_size)
        if remove_from_windowed_attention_mask is not None:
            # This implementation is fast and takes very little memory because num_heads x hidden_size = 1
            # from (batch_size x seqlen) to (batch_size x seqlen x num_heads x hidden_size)
@@ -328,7 +327,7 @@ class LongformerSelfAttention(nn.Module):
            selected_k[selection_padding_mask_nonzeros] = k[extra_attention_mask_nonzeros]
            # (batch_size, seqlen, num_heads, max_num_extra_indices_per_batch)
            selected_attn_weights = torch.einsum("blhd,bshd->blhs", (q, selected_k))
-            selected_attn_weights[selection_padding_mask_zeros[0], :, :, selection_padding_mask_zeros[1]] = -10000
+            selected_attn_weights[selection_padding_mask_zeros[0], :, :, selection_padding_mask_zeros[1]] = -10000.0
            # concat to attn_weights
            # (batch_size, seqlen, num_heads, extra attention count + 2*window+1)
            attn_weights = torch.cat((selected_attn_weights, attn_weights), dim=-1)
@@ -416,7 +415,7 @@ class LongformerSelfAttention(nn.Module):
            )

        context_layer = attn.transpose(0, 1)
-        if self.output_attentions:
+        if output_attentions:
            if extra_attention_mask is not None:
                # With global attention, return global attention probabilities only
                # batch_size x num_heads x max_num_global_attention_tokens x sequence_length
@@ -430,7 +429,7 @@ class LongformerSelfAttention(nn.Module):
                # batch_size x num_heads x sequence_length x window_size
                # which is the attention weights of every token attending to its neighbours
                attn_weights = attn_weights.permute(0, 2, 1, 3)
-        outputs = (context_layer, attn_weights) if self.output_attentions else (context_layer,)
+        outputs = (context_layer, attn_weights) if output_attentions else (context_layer,)
        return outputs


@@ -498,9 +497,10 @@ LONGFORMER_INPUTS_DOCSTRING = r"""
 class LongformerModel(RobertaModel):
    """
    This class overrides :class:`~transformers.RobertaModel` to provide the ability to process
-    long sequences following the selfattention approach described in `Longformer: the Long-Document Transformer`_by
-    Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer selfattention combines a local (sliding window)
-    and global attention to extend to long documents without the O(n^2) increase in memory and compute.
+    long sequences following the selfattention approach described in `Longformer: the Long-Document Transformer
+    <https://arxiv.org/abs/2004.05150>`_ by Iz Beltagy, Matthew E. Peters, and Arman Cohan. Longformer selfattention
+    combines a local (sliding window) and global attention to extend to long documents without the O(n^2) increase in
+    memory and compute.

    The selfattention module `LongformerSelfAttention` implemented here supports the combination of local and
    global attention but it lacks support for autoregressive attention and dilated attention. Autoregressive
@@ -508,9 +508,6 @@ class LongformerModel(RobertaModel):
    tasks. Future release will add support for autoregressive attention, but the support for dilated attention
    requires a custom CUDA kernel to be memory and compute efficient.

-    .. _`Longformer: the Long-Document Transformer`:
-        https://arxiv.org/abs/2004.05150
-
    """

    config_class = LongformerConfig
@@ -587,14 +584,12 @@ class LongformerModel(RobertaModel):
        token_type_ids=None,
        position_ids=None,
        inputs_embeds=None,
-        masked_lm_labels=None,
+        output_attentions=None,
    ):
        r"""

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
-        masked_lm_loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
-            Masked language modeling loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
@@ -602,7 +597,7 @@ class LongformerModel(RobertaModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -629,6 +624,8 @@ class LongformerModel(RobertaModel):
        sequence_output, pooled_output = model(input_ids, attention_mask=attention_mask)
        """

+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
        # padding
        attention_window = (
            self.config.attention_window
@@ -668,6 +665,7 @@ class LongformerModel(RobertaModel):
            inputs_embeds=inputs_embeds,
            encoder_hidden_states=None,
            encoder_attention_mask=None,
+            output_attentions=output_attentions,
        )

        # undo padding
@@ -704,18 +702,22 @@ class LongformerForMaskedLM(BertPreTrainedModel):
        token_type_ids=None,
        position_ids=None,
        inputs_embeds=None,
-        masked_lm_labels=None,
+        labels=None,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
-        masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
-        masked_lm_loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        masked_lm_loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -724,7 +726,7 @@ class LongformerForMaskedLM(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -744,9 +746,17 @@ class LongformerForMaskedLM(BertPreTrainedModel):

        attention_mask = None  # default is local attention everywhere, which is a good choice for MaskedLM
                               # check ``LongformerModel.forward`` for more details how to set `attention_mask`
-        loss, prediction_scores = model(input_ids, attention_mask=attention_mask, masked_lm_labels=input_ids)
+        loss, prediction_scores = model(input_ids, attention_mask=attention_mask, labels=input_ids)
        """

+        if "masked_lm_labels" in kwargs:
+            warnings.warn(
+                "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("masked_lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
+
        outputs = self.longformer(
            input_ids,
            attention_mask=attention_mask,
@@ -754,15 +764,16 @@ class LongformerForMaskedLM(BertPreTrainedModel):
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        sequence_output = outputs[0]
        prediction_scores = self.lm_head(sequence_output)

        outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention if they are here

-        if masked_lm_labels is not None:
+        if labels is not None:
            loss_fct = CrossEntropyLoss()
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            outputs = (masked_lm_loss,) + outputs

        return outputs  # (masked_lm_loss), prediction_scores, (hidden_states), (attentions)
@@ -794,6 +805,7 @@ class LongformerForSequenceClassification(BertPreTrainedModel):
        position_ids=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -813,7 +825,7 @@ class LongformerForSequenceClassification(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -913,6 +925,7 @@ class LongformerForQuestionAnswering(BertPreTrainedModel):
        inputs_embeds=None,
        start_positions=None,
        end_positions=None,
+        output_attentions=None,
    ):
        r"""
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -935,7 +948,7 @@ class LongformerForQuestionAnswering(BertPreTrainedModel):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
@@ -978,6 +991,7 @@ class LongformerForQuestionAnswering(BertPreTrainedModel):
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -1037,6 +1051,7 @@ class LongformerForTokenClassification(BertPreTrainedModel):
        position_ids=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
@@ -1054,7 +1069,7 @@ class LongformerForTokenClassification(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1136,6 +1151,7 @@ class LongformerForMultipleChoice(BertPreTrainedModel):
        labels=None,
        position_ids=None,
        inputs_embeds=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -1145,7 +1161,7 @@ class LongformerForMultipleChoice(BertPreTrainedModel):

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
-        loss (:obj:`torch.FloatTensor`` of shape ``(1,)`, `optional`, returned when :obj:`labels` is provided):
+        loss (:obj:`torch.FloatTensor`` of shape `(1,)`, `optional`, returned when :obj:`labels` is provided):
            Classification loss.
        classification_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices)`):
            `num_choices` is the second dimension of the input tensors. (see `input_ids` above).
@@ -1156,7 +1172,7 @@ class LongformerForMultipleChoice(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -165,7 +165,7 @@ class MMBTModel(nn.Module, ModuleUtilsMixin):
                list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
                of shape ``(batch_size, sequence_length, hidden_size)``:
                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            **attentions**: (`optional`, returned when ``config.output_attentions=True``)
+            **attentions**: (`optional`, returned when ``output_attentions=True``)
                list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

@@ -295,7 +295,7 @@ class MMBTForClassification(nn.Module):
                list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
                of shape ``(batch_size, sequence_length, hidden_size)``:
                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            **attentions**: (`optional`, returned when ``config.output_attentions=True``)
+            **attentions**: (`optional`, returned when ``output_attentions=True``)
                list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.

@@ -20,6 +20,7 @@ import json
 import logging
 import math
 import os
+import warnings

 import torch
 import torch.nn as nn
@@ -28,7 +29,13 @@ from torch.nn import CrossEntropyLoss
 from .activations import gelu_new, swish
 from .configuration_openai import OpenAIGPTConfig
 from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_utils import Conv1D, PreTrainedModel, SequenceSummary, prune_conv1d_layer
+from .modeling_utils import (
+    Conv1D,
+    PreTrainedModel,
+    SequenceSummary,
+    find_pruneable_heads_and_indices,
+    prune_conv1d_layer,
+)


 logger = logging.getLogger(__name__)
@@ -130,8 +137,6 @@ class Attention(nn.Module):
        self.split_size = n_state
        self.scale = scale

-        self.output_attentions = config.output_attentions
-
        self.c_attn = Conv1D(n_state * 3, nx)
        self.c_proj = Conv1D(n_state, nx)
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
@@ -141,13 +146,9 @@ class Attention(nn.Module):
    def prune_heads(self, heads):
        if len(heads) == 0:
            return
-        mask = torch.ones(self.n_head, self.split_size // self.n_head)
-        heads = set(heads) - self.pruned_heads
-        for head in heads:
-            head -= sum(1 if h < head else 0 for h in self.pruned_heads)
-            mask[head] = 0
-        mask = mask.view(-1).contiguous().eq(1)
-        index = torch.arange(len(mask))[mask].long()
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.n_head, self.split_size // self.n_head, self.pruned_heads
+        )
        index_attn = torch.cat([index, index + self.split_size, index + (2 * self.split_size)])
        # Prune conv1d layers
        self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
@@ -157,7 +158,7 @@ class Attention(nn.Module):
        self.n_head = self.n_head - len(heads)
        self.pruned_heads = self.pruned_heads.union(heads)

-    def _attn(self, q, k, v, attention_mask=None, head_mask=None):
+    def _attn(self, q, k, v, attention_mask=None, head_mask=None, output_attentions=False):
        w = torch.matmul(q, k)
        if self.scale:
            w = w / math.sqrt(v.size(-1))
@@ -178,7 +179,7 @@ class Attention(nn.Module):
            w = w * head_mask

        outputs = [torch.matmul(w, v)]
-        if self.output_attentions:
+        if output_attentions:
            outputs.append(w)
        return outputs

@@ -195,14 +196,14 @@ class Attention(nn.Module):
        else:
            return x.permute(0, 2, 1, 3)

-    def forward(self, x, attention_mask=None, head_mask=None):
+    def forward(self, x, attention_mask=None, head_mask=None, output_attentions=False):
        x = self.c_attn(x)
        query, key, value = x.split(self.split_size, dim=2)
        query = self.split_heads(query)
        key = self.split_heads(key, k=True)
        value = self.split_heads(value)

-        attn_outputs = self._attn(query, key, value, attention_mask, head_mask)
+        attn_outputs = self._attn(query, key, value, attention_mask, head_mask, output_attentions)
        a = attn_outputs[0]

        a = self.merge_heads(a)
@@ -237,8 +238,10 @@ class Block(nn.Module):
        self.mlp = MLP(4 * nx, config)
        self.ln_2 = nn.LayerNorm(nx, eps=config.layer_norm_epsilon)

-    def forward(self, x, attention_mask=None, head_mask=None):
-        attn_outputs = self.attn(x, attention_mask=attention_mask, head_mask=head_mask)
+    def forward(self, x, attention_mask=None, head_mask=None, output_attentions=False):
+        attn_outputs = self.attn(
+            x, attention_mask=attention_mask, head_mask=head_mask, output_attentions=output_attentions,
+        )
        a = attn_outputs[0]

        n = self.ln_1(x + a)
@@ -319,6 +322,8 @@ OPENAI_GPT_INPUTS_DOCSTRING = r"""
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
            than the model's internal embedding lookup matrix.
+        output_attentions (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Should the model returns attentions weights.
 """


@@ -329,7 +334,6 @@ OPENAI_GPT_INPUTS_DOCSTRING = r"""
 class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

        self.tokens_embed = nn.Embedding(config.vocab_size, config.n_embd)
@@ -361,6 +365,7 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -372,7 +377,7 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -391,6 +396,8 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -446,9 +453,9 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)

-            outputs = block(hidden_states, attention_mask, head_mask[i])
+            outputs = block(hidden_states, attention_mask, head_mask[i], output_attentions=output_attentions)
            hidden_states = outputs[0]
-            if self.output_attentions:
+            if output_attentions:
                all_attentions = all_attentions + (outputs[1],)

        # Add last layer
@@ -458,7 +465,7 @@ class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
        outputs = (hidden_states.view(*output_shape),)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last hidden state, (all hidden states), (all attentions)

@@ -489,6 +496,7 @@ class OpenAIGPTLMHeadModel(OpenAIGPTPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
@@ -513,7 +521,7 @@ class OpenAIGPTLMHeadModel(OpenAIGPTPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -539,6 +547,7 @@ class OpenAIGPTLMHeadModel(OpenAIGPTPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        hidden_states = transformer_outputs[0]
        lm_logits = self.lm_head(hidden_states)
@@ -588,16 +597,18 @@ class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        mc_token_ids=None,
-        lm_labels=None,
+        labels=None,
        mc_labels=None,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
        mc_token_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, num_choices)`, `optional`, default to index of the last token of the input)
            Index of the classification token in each input sequence.
            Selected in the range ``[0, input_ids.size(-1) - 1[``.
-        lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`)
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`)
            Labels for language modeling.
-            Note that the labels **are shifted** inside the model, i.e. you can set ``lm_labels = input_ids``
+            Note that the labels **are shifted** inside the model, i.e. you can set ``labels = input_ids``
            Indices are selected in ``[-1, 0, ..., config.vocab_size]``
            All labels set to ``-100`` are ignored (masked), the loss is only
            computed for labels in ``[0, ..., config.vocab_size]``
@@ -605,12 +616,14 @@ class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
            Labels for computing the multiple choice classification loss.
            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
            of the input tensors. (see `input_ids` above)
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Return:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.OpenAIGPTConfig`) and inputs:
-        lm_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``lm_labels`` is provided):
+        lm_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided):
            Language modeling loss.
-        mc_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`multiple_choice_labels` is provided):
+        mc_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`mc_labels` is provided):
            Multiple choice classification loss.
        lm_prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -625,7 +638,7 @@ class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -650,6 +663,14 @@ class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
        lm_prediction_scores, mc_prediction_scores = outputs[:2]

    """
+        if "lm_labels" in kwargs:
+            warnings.warn(
+                "The `lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
+
        transformer_outputs = self.transformer(
            input_ids,
            attention_mask=attention_mask,
@@ -657,6 +678,7 @@ class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        hidden_states = transformer_outputs[0]

@@ -668,9 +690,9 @@ class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1))
            outputs = (loss,) + outputs
-        if lm_labels is not None:
+        if labels is not None:
            shift_logits = lm_logits[..., :-1, :].contiguous()
-            shift_labels = lm_labels[..., 1:].contiguous()
+            shift_labels = labels[..., 1:].contiguous()
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
            outputs = (loss,) + outputs
@@ -318,7 +318,7 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
        attention_mask=None,
        head_mask=None,
        num_hashes=None,
-        do_output_attentions=False,
+        output_attentions=False,
        buckets=None,
        **kwargs
    ):
@@ -351,41 +351,48 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
            value_vectors.shape[-1], self.attention_head_size
        )

-        # set `num_buckets` on the fly, recommended way to do it
-        if self.num_buckets is None:
-            self._set_num_buckets(sequence_length)
+        # LSH attention only makes sense if chunked attention should be performed
+        if self.chunk_length < sequence_length:
+            # set `num_buckets` on the fly, recommended way to do it
+            if self.num_buckets is None:
+                self._set_num_buckets(sequence_length)

-        # use cached buckets for backprop only
-        if buckets is None:
-            # hash query key vectors into buckets
-            buckets = self._hash_vectors(query_key_vectors, num_hashes)
+            # use cached buckets for backprop only
+            if buckets is None:
+                # hash query key vectors into buckets
+                buckets = self._hash_vectors(query_key_vectors, num_hashes)

-        assert (
-            int(buckets.shape[-1]) == num_hashes * sequence_length
-        ), "last dim of buckets is {}, but should be {}".format(buckets.shape[-1], num_hashes * sequence_length)
-
-        sorted_bucket_idx, undo_sorted_bucket_idx = self._get_sorted_bucket_idx_and_undo_sorted_bucket_idx(
-            sequence_length, buckets, num_hashes
-        )
-
-        # make sure bucket idx is not longer then sequence length
-        sorted_bucket_idx = sorted_bucket_idx % sequence_length
-
-        # cluster query key value vectors according to hashed buckets
-        query_key_vectors = self._gather_by_expansion(query_key_vectors, sorted_bucket_idx, num_hashes)
-        value_vectors = self._gather_by_expansion(value_vectors, sorted_bucket_idx, num_hashes)
-
-        query_key_vectors = self._split_seq_length_dim_to(
-            query_key_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
-        )
-        value_vectors = self._split_seq_length_dim_to(
-            value_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
-        )
-
-        if self.chunk_length is None:
            assert (
-                self.num_chunks_before == 0 and self.num_chunks_after == 0
-            ), "If `config.chunk_length` is `None`, make sure `config.num_chunks_after` and `config.num_chunks_before` are set to 0."
+                int(buckets.shape[-1]) == num_hashes * sequence_length
+            ), "last dim of buckets is {}, but should be {}".format(buckets.shape[-1], num_hashes * sequence_length)
+
+            sorted_bucket_idx, undo_sorted_bucket_idx = self._get_sorted_bucket_idx_and_undo_sorted_bucket_idx(
+                sequence_length, buckets, num_hashes
+            )
+
+            # make sure bucket idx is not longer then sequence length
+            sorted_bucket_idx = sorted_bucket_idx % sequence_length
+
+            # cluster query key value vectors according to hashed buckets
+            query_key_vectors = self._gather_by_expansion(query_key_vectors, sorted_bucket_idx, num_hashes)
+            value_vectors = self._gather_by_expansion(value_vectors, sorted_bucket_idx, num_hashes)
+
+            query_key_vectors = self._split_seq_length_dim_to(
+                query_key_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
+            )
+            value_vectors = self._split_seq_length_dim_to(
+                value_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
+            )
+
+            if self.chunk_length is None:
+                assert (
+                    self.num_chunks_before == 0 and self.num_chunks_after == 0
+                ), "If `config.chunk_length` is `None`, make sure `config.num_chunks_after` and `config.num_chunks_before` are set to 0."
+        else:
+            # get sequence length indices
+            sorted_bucket_idx = torch.arange(sequence_length, device=query_key_vectors.device).repeat(
+                batch_size, self.num_attention_heads, 1
+            )

        # scale key vectors
        key_vectors = self._len_and_dim_norm(query_key_vectors)
@@ -398,31 +405,35 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
            sorted_bucket_idx=sorted_bucket_idx,
            attention_mask=attention_mask,
            head_mask=head_mask,
+            sequence_length=sequence_length,
        )
+
        # free memory
        del query_key_vectors, key_vectors, value_vectors

-        # sort clusters back to correct ordering
-        out_vectors, logits = ReverseSort.apply(
-            out_vectors, logits, sorted_bucket_idx, undo_sorted_bucket_idx, self.num_hashes
-        )
-
-        # sum up all hash rounds
-        if num_hashes > 1:
-            out_vectors = self._split_seq_length_dim_to(
-                out_vectors, num_hashes, sequence_length, self.num_attention_heads, self.attention_head_size,
+        # re-order out_vectors and logits
+        if self.chunk_length < sequence_length:
+            # sort clusters back to correct ordering
+            out_vectors, logits = ReverseSort.apply(
+                out_vectors, logits, sorted_bucket_idx, undo_sorted_bucket_idx, self.num_hashes
            )
-            logits = self._split_seq_length_dim_to(
-                logits, num_hashes, sequence_length, self.num_attention_heads, self.attention_head_size,
-            ).unsqueeze(-1)

-            probs_vectors = torch.exp(logits - torch.logsumexp(logits, dim=2, keepdim=True))
-            out_vectors = torch.sum(out_vectors * probs_vectors, dim=2)
+            # sum up all hash rounds
+            if num_hashes > 1:
+                out_vectors = self._split_seq_length_dim_to(
+                    out_vectors, num_hashes, sequence_length, self.num_attention_heads, self.attention_head_size,
+                )
+                logits = self._split_seq_length_dim_to(
+                    logits, num_hashes, sequence_length, self.num_attention_heads, self.attention_head_size,
+                ).unsqueeze(-1)
+
+                probs_vectors = torch.exp(logits - torch.logsumexp(logits, dim=2, keepdim=True))
+                out_vectors = torch.sum(out_vectors * probs_vectors, dim=2)
+                # free memory
+                del probs_vectors
+
            # free memory
-            del probs_vectors
-
-        # free memory
-        del logits
+            del logits

        assert out_vectors.shape == (
            batch_size,
@@ -433,7 +444,7 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):

        out_vectors = self._merge_hidden_size_dims(out_vectors, self.num_attention_heads, self.attention_head_size)

-        if do_output_attentions is False:
+        if output_attentions is False:
            attention_probs = ()

        return LSHSelfAttentionOutput(hidden_states=out_vectors, attention_probs=attention_probs, buckets=buckets)
@@ -554,10 +565,13 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
        self.num_buckets = num_buckets

    def _attend(
-        self, query_vectors, key_vectors, value_vectors, sorted_bucket_idx, attention_mask, head_mask,
+        self, query_vectors, key_vectors, value_vectors, sorted_bucket_idx, attention_mask, head_mask, sequence_length
    ):
-        key_vectors = self._look_adjacent(key_vectors, self.num_chunks_before, self.num_chunks_after)
-        value_vectors = self._look_adjacent(value_vectors, self.num_chunks_before, self.num_chunks_after)
+
+        # look at previous and following chunks if chunked attention
+        if self.chunk_length < sequence_length:
+            key_vectors = self._look_adjacent(key_vectors, self.num_chunks_before, self.num_chunks_after)
+            value_vectors = self._look_adjacent(value_vectors, self.num_chunks_before, self.num_chunks_after)

        # get logits and dots
        query_key_dots = torch.matmul(query_vectors, key_vectors.transpose(-1, -2))
@@ -565,10 +579,14 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
        # free memory
        del query_vectors, key_vectors

-        query_bucket_idx = self._split_seq_length_dim_to(
-            sorted_bucket_idx, -1, self.chunk_length, self.num_attention_heads
-        )
-        key_value_bucket_idx = self._look_adjacent(query_bucket_idx, self.num_chunks_before, self.num_chunks_after)
+        # if chunked attention split bucket idxs to query and key
+        if self.chunk_length < sequence_length:
+            query_bucket_idx = self._split_seq_length_dim_to(
+                sorted_bucket_idx, -1, self.chunk_length, self.num_attention_heads
+            )
+            key_value_bucket_idx = self._look_adjacent(query_bucket_idx, self.num_chunks_before, self.num_chunks_after)
+        else:
+            query_bucket_idx = key_value_bucket_idx = sorted_bucket_idx

        # get correct mask values depending on precision
        if query_key_dots.dtype == torch.float16:
@@ -578,7 +596,7 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
            self_mask_value = self.self_mask_value_float32
            mask_value = self.mask_value_float32

-        mask = self._compute_attn_mask(query_bucket_idx, key_value_bucket_idx, attention_mask)
+        mask = self._compute_attn_mask(query_bucket_idx, key_value_bucket_idx, attention_mask, sequence_length)

        if mask is not None:
            query_key_dots = torch.where(mask, query_key_dots, mask_value)
@@ -627,12 +645,13 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
        del value_vectors

        # merge chunk length
-        logits = logits.flatten(start_dim=2, end_dim=3).squeeze(-1)
-        out_vectors = out_vectors.flatten(start_dim=2, end_dim=3)
+        if self.chunk_length < sequence_length:
+            logits = logits.flatten(start_dim=2, end_dim=3).squeeze(-1)
+            out_vectors = out_vectors.flatten(start_dim=2, end_dim=3)

        return out_vectors, logits, attention_probs

-    def _compute_attn_mask(self, query_indices, key_indices, attention_mask):
+    def _compute_attn_mask(self, query_indices, key_indices, attention_mask, sequence_length):
        mask = None

        # Causal mask
@@ -642,15 +661,27 @@ class LSHSelfAttention(nn.Module, EfficientAttentionMixin):
        # Attention mask: chunk, look up correct mask value from key_value_bucket_idx
        # IMPORTANT: official trax code does not use a mask for LSH Atttention. Not sure why.
        if attention_mask is not None:
-            attention_mask = attention_mask.to(torch.uint8)[:, None, None, :]
-            # expand attn_mask to fit with key_value_bucket_idx shape
-            attention_mask = attention_mask.expand(query_indices.shape[:-1] + (-1,))
-            key_attn_mask = torch.gather(attention_mask, -1, key_indices)
-            query_attn_mask = torch.gather(attention_mask, -1, query_indices)
-            # expand to query_key_dots shape: duplicate along query axis since key sorting is the same for each query position in chunk
-            attn_mask = query_attn_mask.unsqueeze(-1) * key_attn_mask.unsqueeze(-2)
+            # if chunked attention, the attention mask has to correspond to LSH order
+            if sequence_length > self.chunk_length:
+                attention_mask = attention_mask.to(torch.uint8)[:, None, None, :]
+                # expand attn_mask to fit with key_value_bucket_idx shape
+                attention_mask = attention_mask.expand(query_indices.shape[:-1] + (-1,))
+                key_attn_mask = torch.gather(attention_mask, -1, key_indices)
+                query_attn_mask = torch.gather(attention_mask, -1, query_indices)
+                # expand to query_key_dots shape: duplicate along query axis since key sorting is the same for each query position in chunk
+                attn_mask = query_attn_mask.unsqueeze(-1) * key_attn_mask.unsqueeze(-2)
+
+                # free memory
+                del query_attn_mask, key_attn_mask
+            else:
+                # usual attention mask creation
+                attention_mask = attention_mask.to(torch.uint8)[:, None, :]
+                attn_mask = (attention_mask.unsqueeze(-1) * attention_mask.unsqueeze(-2)).expand(
+                    query_indices.shape + attention_mask.shape[-1:]
+                )
+
            # free memory
-            del query_attn_mask, key_attn_mask, attention_mask
+            del attention_mask

            # multiply by casaul mask if necessary
            if mask is not None:
@@ -770,7 +801,7 @@ class LocalSelfAttention(nn.Module, EfficientAttentionMixin):
        self.register_buffer("mask_value_float16", torch.tensor(-1e4))
        self.register_buffer("mask_value_float32", torch.tensor(-1e9))

-    def forward(self, hidden_states, attention_mask=None, head_mask=None, do_output_attentions=False, **kwargs):
+    def forward(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, **kwargs):
        sequence_length = hidden_states.shape[1]
        batch_size = hidden_states.shape[0]

@@ -810,36 +841,45 @@ class LocalSelfAttention(nn.Module, EfficientAttentionMixin):
            torch.tensor(self.attention_head_size, device=key_vectors.device, dtype=key_vectors.dtype)
        )

-        # chunk vectors
-        # B x Num_Attn_Head x Seq_Len // chunk_len x chunk_len  x  attn_head_size
-        query_vectors = self._split_seq_length_dim_to(
-            query_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
-        )
-        key_vectors = self._split_seq_length_dim_to(
-            key_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
-        )
-        value_vectors = self._split_seq_length_dim_to(
-            value_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
-        )
-
-        # chunk indices
+        # get sequence length indices
        indices = torch.arange(sequence_length, device=query_vectors.device).repeat(
            batch_size, self.num_attention_heads, 1
        )
-        query_indices = self._split_seq_length_dim_to(indices, -1, self.chunk_length, self.num_attention_heads)
-        key_indices = self._split_seq_length_dim_to(indices, -1, self.chunk_length, self.num_attention_heads)

-        # append chunks before and after
-        key_vectors = self._look_adjacent(key_vectors, self.num_chunks_before, self.num_chunks_after)
-        value_vectors = self._look_adjacent(value_vectors, self.num_chunks_before, self.num_chunks_after)
-        key_indices = self._look_adjacent(key_indices, self.num_chunks_before, self.num_chunks_after)
+        # if input should be chunked
+        if self.chunk_length < sequence_length:
+            # chunk vectors
+            # B x Num_Attn_Head x Seq_Len // chunk_len x chunk_len  x  attn_head_size
+            query_vectors = self._split_seq_length_dim_to(
+                query_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
+            )
+            key_vectors = self._split_seq_length_dim_to(
+                key_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
+            )
+            value_vectors = self._split_seq_length_dim_to(
+                value_vectors, -1, self.chunk_length, self.num_attention_heads, self.attention_head_size,
+            )

+            # chunk indices
+            query_indices = self._split_seq_length_dim_to(indices, -1, self.chunk_length, self.num_attention_heads)
+            key_indices = self._split_seq_length_dim_to(indices, -1, self.chunk_length, self.num_attention_heads)
+
+            # append chunks before and after
+            key_vectors = self._look_adjacent(key_vectors, self.num_chunks_before, self.num_chunks_after)
+            value_vectors = self._look_adjacent(value_vectors, self.num_chunks_before, self.num_chunks_after)
+            key_indices = self._look_adjacent(key_indices, self.num_chunks_before, self.num_chunks_after)
+        else:
+            query_indices = key_indices = indices
+
+        # query-key matmul: QK^T
        query_key_dots = torch.matmul(query_vectors, key_vectors.transpose(-1, -2))

        # free memory
        del query_vectors, key_vectors

-        mask = self._compute_attn_mask(query_indices, key_indices, attention_mask, query_key_dots.shape)
+        mask = self._compute_attn_mask(
+            query_indices, key_indices, attention_mask, query_key_dots.shape, sequence_length
+        )

        if mask is not None:
            # get mask tensor depending on half precision or not
@@ -874,25 +914,30 @@ class LocalSelfAttention(nn.Module, EfficientAttentionMixin):
        del value_vectors

        # merge chunk length
-        out_vectors = out_vectors.flatten(start_dim=2, end_dim=3)
+        if self.chunk_length < sequence_length:
+            out_vectors = out_vectors.flatten(start_dim=2, end_dim=3)

        assert out_vectors.shape == (batch_size, self.num_attention_heads, sequence_length, self.attention_head_size,)

        out_vectors = self._merge_hidden_size_dims(out_vectors, self.num_attention_heads, self.attention_head_size)

-        if do_output_attentions is False:
+        if output_attentions is False:
            attention_probs = ()

        return LocalSelfAttentionOutput(hidden_states=out_vectors, attention_probs=attention_probs)

-    def _compute_attn_mask(self, query_indices, key_indices, attention_mask, query_key_dots_shape):
+    def _compute_attn_mask(self, query_indices, key_indices, attention_mask, query_key_dots_shape, sequence_length):
        mask = None

        # chunk attention mask and look before and after
        if attention_mask is not None:
            attention_mask = attention_mask.to(torch.uint8)[:, None, :]
-            attention_mask = self._split_seq_length_dim_to(attention_mask, -1, self.chunk_length, 1)
-            attention_mask_key = self._look_adjacent(attention_mask, self.num_chunks_before, self.num_chunks_after)
+
+            if self.chunk_length < sequence_length:
+                attention_mask = self._split_seq_length_dim_to(attention_mask, -1, self.chunk_length, 1)
+                attention_mask_key = self._look_adjacent(attention_mask, self.num_chunks_before, self.num_chunks_after)
+            else:
+                attention_mask_key = attention_mask

        # Causal mask
        if self.is_decoder is True:
@@ -956,7 +1001,7 @@ class ReformerAttention(nn.Module):
        attention_mask=None,
        head_mask=None,
        num_hashes=None,
-        do_output_attentions=False,
+        output_attentions=False,
        buckets=None,
    ):
        hidden_states = self.layer_norm(hidden_states)
@@ -967,7 +1012,7 @@ class ReformerAttention(nn.Module):
            head_mask=head_mask,
            attention_mask=attention_mask,
            num_hashes=num_hashes,
-            do_output_attentions=do_output_attentions,
+            output_attentions=output_attentions,
            buckets=buckets,
        )
        attention_output = self.output(self_attention_outputs.hidden_states)
@@ -1094,7 +1139,7 @@ class ReformerLayer(nn.Module):
        attention_mask=None,
        head_mask=None,
        num_hashes=None,
-        do_output_attentions=False,
+        output_attentions=False,
    ):
        with torch.no_grad():
            # every forward pass we sample a different seed
@@ -1106,7 +1151,7 @@ class ReformerLayer(nn.Module):
                head_mask=head_mask,
                attention_mask=attention_mask,
                num_hashes=num_hashes,
-                do_output_attentions=do_output_attentions,
+                output_attentions=output_attentions,
            )
            attn_output = attn_outputs.hidden_states

@@ -1212,7 +1257,7 @@ class _ReversibleFunction(Function):
        all_hidden_states,
        all_attentions,
        do_output_hidden_states,
-        do_output_attentions,
+        output_attentions,
    ):
        all_buckets = ()

@@ -1229,13 +1274,13 @@ class _ReversibleFunction(Function):
                attention_mask=attention_mask,
                head_mask=layer_head_mask,
                num_hashes=num_hashes,
-                do_output_attentions=do_output_attentions,
+                output_attentions=output_attentions,
            )
            attn_output = layer_outputs.attn_output
            hidden_states = layer_outputs.hidden_states
            all_buckets = all_buckets + (layer_outputs.buckets,)

-            if do_output_attentions:
+            if output_attentions:
                all_attentions.append(layer_outputs.attention_probs)

        # Add last layer
@@ -1316,7 +1361,7 @@ class ReformerEncoder(nn.Module):
        head_mask=None,
        num_hashes=None,
        do_output_hidden_states=False,
-        do_output_attentions=False,
+        output_attentions=False,
    ):
        # hidden_states and attention lists to be filled if wished
        all_hidden_states = []
@@ -1333,7 +1378,7 @@ class ReformerEncoder(nn.Module):
            all_hidden_states,
            all_attentions,
            do_output_hidden_states,
-            do_output_attentions,
+            output_attentions,
        )

        # Apply layer norm to concatenated hidden states
@@ -1504,7 +1549,7 @@ class ReformerModel(ReformerPreTrainedModel):
        inputs_embeds=None,
        num_hashes=None,
        do_output_hidden_states=False,
-        do_output_attentions=False,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -1516,7 +1561,7 @@ class ReformerModel(ReformerPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        all_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``do_output_attentions=True``):
+        all_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1537,8 +1582,7 @@ class ReformerModel(ReformerPreTrainedModel):
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple
        """

-        # TODO(PVP): delete when PR to change output_attentions is made
-        do_output_attentions = self.config.output_attentions
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
        do_output_hidden_states = self.config.output_hidden_states

        if input_ids is not None and inputs_embeds is not None:
@@ -1564,7 +1608,9 @@ class ReformerModel(ReformerPreTrainedModel):

        # if needs padding
        least_common_mult_chunk_length = _get_least_common_mult_chunk_len(self.config)
-        must_pad_to_match_chunk_length = input_shape[-1] % least_common_mult_chunk_length != 0
+        must_pad_to_match_chunk_length = (
+            input_shape[-1] % least_common_mult_chunk_length != 0 and input_shape[-1] > least_common_mult_chunk_length
+        )

        if must_pad_to_match_chunk_length:
            padding_length = least_common_mult_chunk_length - input_shape[-1] % least_common_mult_chunk_length
@@ -1596,7 +1642,7 @@ class ReformerModel(ReformerPreTrainedModel):
            attention_mask=attention_mask,
            num_hashes=num_hashes,
            do_output_hidden_states=do_output_hidden_states,
-            do_output_attentions=do_output_attentions,
+            output_attentions=output_attentions,
        )
        sequence_output = encoder_outputs.hidden_states

@@ -1608,7 +1654,7 @@ class ReformerModel(ReformerPreTrainedModel):
        # TODO(PVP): Replace by named tuple after namedtuples are introduced in the library.
        if do_output_hidden_states is True:
            outputs = outputs + (encoder_outputs.all_hidden_states,)
-        if do_output_attentions is True:
+        if output_attentions is True:
            outputs = outputs + (encoder_outputs.all_attentions,)
        return outputs

@@ -1697,7 +1743,7 @@ class ReformerModelWithLMHead(ReformerPreTrainedModel):
        num_hashes=None,
        labels=None,
        do_output_hidden_states=False,
-        do_output_attentions=False,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -1708,7 +1754,7 @@ class ReformerModelWithLMHead(ReformerPreTrainedModel):

    Return:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
-        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`lm_label` is provided):
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
            Classification loss (cross entropy).
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -1717,7 +1763,7 @@ class ReformerModelWithLMHead(ReformerPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        all_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``do_output_attentions=True``):
+        all_attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -1746,7 +1792,7 @@ class ReformerModelWithLMHead(ReformerPreTrainedModel):
            inputs_embeds=inputs_embeds,
            num_hashes=num_hashes,
            do_output_hidden_states=do_output_hidden_states,
-            do_output_attentions=do_output_attentions,
+            output_attentions=output_attentions,
        )

        sequence_output = reformer_outputs[0]
@@ -17,6 +17,7 @@


 import logging
+import warnings

 import torch
 import torch.nn as nn
@@ -183,18 +184,22 @@ class RobertaForMaskedLM(BertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
-        masked_lm_labels=None,
+        labels=None,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
-        masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
            in ``[0, ..., config.vocab_size]``
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+            Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
-        masked_lm_loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+        masked_lm_loss (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -203,7 +208,7 @@ class RobertaForMaskedLM(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -218,10 +223,18 @@ class RobertaForMaskedLM(BertPreTrainedModel):
        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
        model = RobertaForMaskedLM.from_pretrained('roberta-base')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
-        outputs = model(input_ids, masked_lm_labels=input_ids)
+        outputs = model(input_ids, labels=input_ids)
        loss, prediction_scores = outputs[:2]

        """
+        if "masked_lm_labels" in kwargs:
+            warnings.warn(
+                "The `masked_lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("masked_lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
+
        outputs = self.roberta(
            input_ids,
            attention_mask=attention_mask,
@@ -229,15 +242,16 @@ class RobertaForMaskedLM(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        sequence_output = outputs[0]
        prediction_scores = self.lm_head(sequence_output)

        outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention if they are here

-        if masked_lm_labels is not None:
+        if labels is not None:
            loss_fct = CrossEntropyLoss()
-            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), labels.view(-1))
            outputs = (masked_lm_loss,) + outputs

        return outputs  # (masked_lm_loss), prediction_scores, (hidden_states), (attentions)
@@ -294,6 +308,7 @@ class RobertaForSequenceClassification(BertPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -313,7 +328,7 @@ class RobertaForSequenceClassification(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -340,6 +355,7 @@ class RobertaForSequenceClassification(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )
        sequence_output = outputs[0]
        logits = self.classifier(sequence_output)
@@ -386,6 +402,7 @@ class RobertaForMultipleChoice(BertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -395,7 +412,7 @@ class RobertaForMultipleChoice(BertPreTrainedModel):

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
-        loss (:obj:`torch.FloatTensor`` of shape ``(1,)`, `optional`, returned when :obj:`labels` is provided):
+        loss (:obj:`torch.FloatTensor`` of shape `(1,)`, `optional`, returned when :obj:`labels` is provided):
            Classification loss.
        classification_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices)`):
            `num_choices` is the second dimension of the input tensors. (see `input_ids` above).
@@ -406,7 +423,7 @@ class RobertaForMultipleChoice(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -439,6 +456,7 @@ class RobertaForMultipleChoice(BertPreTrainedModel):
            token_type_ids=flat_token_type_ids,
            attention_mask=flat_attention_mask,
            head_mask=head_mask,
+            output_attentions=output_attentions,
        )
        pooled_output = outputs[1]

@@ -485,6 +503,7 @@ class RobertaForTokenClassification(BertPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        labels=None,
+        output_attentions=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
@@ -502,7 +521,7 @@ class RobertaForTokenClassification(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -530,6 +549,7 @@ class RobertaForTokenClassification(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -604,6 +624,7 @@ class RobertaForQuestionAnswering(BertPreTrainedModel):
        inputs_embeds=None,
        start_positions=None,
        end_positions=None,
+        output_attentions=None,
    ):
        r"""
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
@@ -628,7 +649,7 @@ class RobertaForQuestionAnswering(BertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -662,6 +683,7 @@ class RobertaForQuestionAnswering(BertPreTrainedModel):
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
        )

        sequence_output = outputs[0]
@@ -19,6 +19,7 @@ import copy
 import logging
 import math
 import os
+import warnings

 import torch
 import torch.nn.functional as F
@@ -27,7 +28,7 @@ from torch.nn import CrossEntropyLoss

 from .configuration_t5 import T5Config
 from .file_utils import DUMMY_INPUTS, DUMMY_MASK, add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_utils import PreTrainedModel, prune_linear_layer
+from .modeling_utils import PreTrainedModel, find_pruneable_heads_and_indices, prune_linear_layer


 logger = logging.getLogger(__name__)
@@ -194,7 +195,6 @@ class T5Attention(nn.Module):
        self.is_decoder = config.is_decoder
        self.has_relative_attention_bias = has_relative_attention_bias

-        self.output_attentions = config.output_attentions
        self.relative_attention_num_buckets = config.relative_attention_num_buckets
        self.d_model = config.d_model
        self.d_kv = config.d_kv
@@ -215,13 +215,7 @@ class T5Attention(nn.Module):
    def prune_heads(self, heads):
        if len(heads) == 0:
            return
-        mask = torch.ones(self.n_heads, self.d_kv)
-        heads = set(heads) - self.pruned_heads
-        for head in heads:
-            head -= sum(1 if h < head else 0 for h in self.pruned_heads)
-            mask[head] = 0
-        mask = mask.view(-1).contiguous().eq(1)
-        index = torch.arange(len(mask))[mask].long()
+        heads, index = find_pruneable_heads_and_indices(heads, self.n_heads, self.d_kv, self.pruned_heads)
        # Prune linear layers
        self.q = prune_linear_layer(self.q, index)
        self.k = prune_linear_layer(self.k, index)
@@ -305,6 +299,7 @@ class T5Attention(nn.Module):
        head_mask=None,
        query_length=None,
        use_cache=False,
+        output_attentions=False,
    ):
        """
        Self-attention (if kv is None) or attention over source sentence (provided by kv).
@@ -391,7 +386,7 @@ class T5Attention(nn.Module):

        outputs = (context,) + present_key_value_state

-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (weights,)
        if self.has_relative_attention_bias:
            outputs = outputs + (position_bias,)
@@ -413,6 +408,7 @@ class T5LayerSelfAttention(nn.Module):
        head_mask=None,
        past_key_value_state=None,
        use_cache=False,
+        output_attentions=False,
    ):
        norm_x = self.layer_norm(hidden_states)
        attention_output = self.SelfAttention(
@@ -422,6 +418,7 @@ class T5LayerSelfAttention(nn.Module):
            head_mask=head_mask,
            past_key_value_state=past_key_value_state,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )
        y = attention_output[0]
        layer_output = hidden_states + self.dropout(y)
@@ -446,6 +443,7 @@ class T5LayerCrossAttention(nn.Module):
        past_key_value_state=None,
        use_cache=False,
        query_length=None,
+        output_attentions=False,
    ):
        norm_x = self.layer_norm(hidden_states)
        attention_output = self.EncDecAttention(
@@ -457,6 +455,7 @@ class T5LayerCrossAttention(nn.Module):
            past_key_value_state=past_key_value_state,
            use_cache=use_cache,
            query_length=query_length,
+            output_attentions=output_attentions,
        )
        y = attention_output[0]
        layer_output = hidden_states + self.dropout(y)
@@ -486,6 +485,7 @@ class T5Block(nn.Module):
        head_mask=None,
        past_key_value_state=None,
        use_cache=False,
+        output_attentions=False,
    ):

        if past_key_value_state is not None:
@@ -511,6 +511,7 @@ class T5Block(nn.Module):
            head_mask=head_mask,
            past_key_value_state=self_attn_past_key_value_state,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )
        hidden_states, present_key_value_state = self_attention_outputs[:2]
        attention_outputs = self_attention_outputs[2:]  # Keep self-attention outputs and relative position weights
@@ -532,6 +533,7 @@ class T5Block(nn.Module):
                past_key_value_state=cross_attn_past_key_value_state,
                query_length=query_length,
                use_cache=use_cache,
+                output_attentions=output_attentions,
            )
            hidden_states = cross_attention_outputs[0]
            # Combine self attn and cross attn key value states
@@ -616,10 +618,10 @@ class T5PreTrainedModel(PreTrainedModel):
        shifted_input_ids[..., 0] = decoder_start_token_id

        assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."
-        # replace possible -100 values in lm_labels by `pad_token_id`
+        # replace possible -100 values in labels by `pad_token_id`
        shifted_input_ids.masked_fill_(shifted_input_ids == -100, pad_token_id)

-        assert torch.all(shifted_input_ids >= 0).item(), "Verify that `lm_labels` has only positive values and -100"
+        assert torch.all(shifted_input_ids >= 0).item(), "Verify that `labels` has only positive values and -100"

        return shifted_input_ids

@@ -627,7 +629,6 @@ class T5PreTrainedModel(PreTrainedModel):
 class T5Stack(T5PreTrainedModel):
    def __init__(self, config, embed_tokens=None):
        super().__init__(config)
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

        self.embed_tokens = embed_tokens
@@ -660,8 +661,11 @@ class T5Stack(T5PreTrainedModel):
        head_mask=None,
        past_key_value_states=None,
        use_cache=False,
+        output_attentions=None,
    ):

+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -735,6 +739,7 @@ class T5Stack(T5PreTrainedModel):
                head_mask=head_mask[i],
                past_key_value_state=past_key_value_state,
                use_cache=use_cache,
+                output_attentions=output_attentions,
            )
            # layer_outputs is a tuple with:
            # hidden-states, key-value-states, (self-attention weights), (self-attention position bias), (cross-attention weights), (cross-attention position bias)
@@ -743,13 +748,13 @@ class T5Stack(T5PreTrainedModel):
            if i == 0:
                # We share the position biases between the layers - the first layer store them
                # layer_outputs = hidden-states, key-value-states (self-attention weights), (self-attention position bias), (cross-attention weights), (cross-attention position bias)
-                position_bias = layer_outputs[3 if self.output_attentions else 2]
+                position_bias = layer_outputs[3 if output_attentions else 2]
                if self.is_decoder and encoder_hidden_states is not None:
-                    encoder_decoder_position_bias = layer_outputs[5 if self.output_attentions else 3]
+                    encoder_decoder_position_bias = layer_outputs[5 if output_attentions else 3]
            # append next layer key value states
            present_key_value_states = present_key_value_states + (present_key_value_state,)

-            if self.output_attentions:
+            if output_attentions:
                all_attentions = all_attentions + (layer_outputs[2],)  # We keep only self-attention weights for now

        hidden_states = self.final_layer_norm(hidden_states)
@@ -765,7 +770,7 @@ class T5Stack(T5PreTrainedModel):
            outputs = outputs + (present_key_value_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (presents,) (all hidden states), (all attentions)

@@ -892,6 +897,7 @@ class T5Model(T5PreTrainedModel):
        inputs_embeds=None,
        decoder_inputs_embeds=None,
        head_mask=None,
+        output_attentions=None,
    ):
        r"""
    Return:
@@ -908,7 +914,7 @@ class T5Model(T5PreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -930,7 +936,11 @@ class T5Model(T5PreTrainedModel):
        # Encode if needed (training, first prediction pass)
        if encoder_outputs is None:
            encoder_outputs = self.encoder(
-                input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, head_mask=head_mask
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                head_mask=head_mask,
+                output_attentions=output_attentions,
            )

        hidden_states = encoder_outputs[0]
@@ -953,6 +963,7 @@ class T5Model(T5PreTrainedModel):
            encoder_attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )

        if use_cache is True:
@@ -1008,21 +1019,25 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
        decoder_attention_mask=None,
        decoder_past_key_value_states=None,
        use_cache=True,
-        lm_labels=None,
+        labels=None,
        inputs_embeds=None,
        decoder_inputs_embeds=None,
        head_mask=None,
+        output_attentions=None,
+        **kwargs
    ):
        r"""
-        lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
                Labels for computing the sequence classification/regression loss.
                Indices should be in :obj:`[-100, 0, ..., config.vocab_size - 1]`.
                All labels set to ``-100`` are ignored (masked), the loss is only
                computed for labels in ``[0, ..., config.vocab_size]``
+        kwargs (:obj:`Dict[str, any]`, optional, defaults to `{}`):
+                Used to hide legacy arguments that have been deprecated.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.T5Config`) and inputs.
-        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`lm_label` is provided):
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
            Classification loss (cross entropy).
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
@@ -1035,7 +1050,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention.
@@ -1047,7 +1062,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
        tokenizer = T5Tokenizer.from_pretrained('t5-small')
        model = T5ForConditionalGeneration.from_pretrained('t5-small')
        input_ids = tokenizer.encode("Hello, my dog is cute", return_tensors="pt")  # Batch size 1
-        outputs = model(input_ids=input_ids, decoder_input_ids=input_ids, lm_labels=input_ids)
+        outputs = model(input_ids=input_ids, decoder_input_ids=input_ids, labels=input_ids)
        loss, prediction_scores = outputs[:2]

        tokenizer = T5Tokenizer.from_pretrained('t5-small')
@@ -1056,23 +1071,35 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
        outputs = model.generate(input_ids)
        """

+        if "lm_labels" in kwargs:
+            warnings.warn(
+                "The `lm_labels` argument is deprecated and will be removed in a future version, use `labels` instead.",
+                DeprecationWarning,
+            )
+            labels = kwargs.pop("lm_labels")
+        assert kwargs == {}, f"Unexpected keyword arguments: {list(kwargs.keys())}."
+
        # Encode if needed (training, first prediction pass)
        if encoder_outputs is None:
            # Convert encoder inputs in embeddings if needed
            encoder_outputs = self.encoder(
-                input_ids=input_ids, attention_mask=attention_mask, inputs_embeds=inputs_embeds, head_mask=head_mask
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                inputs_embeds=inputs_embeds,
+                head_mask=head_mask,
+                output_attentions=output_attentions,
            )

        hidden_states = encoder_outputs[0]

-        if lm_labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
+        if labels is not None and decoder_input_ids is None and decoder_inputs_embeds is None:
            # get decoder inputs from shifting lm labels to the right
-            decoder_input_ids = self._shift_right(lm_labels)
+            decoder_input_ids = self._shift_right(labels)

        # If decoding with past key value states, only the last tokens
        # should be given as an input
        if decoder_past_key_value_states is not None:
-            assert lm_labels is None, "Decoder should not use cached key value states when training."
+            assert labels is None, "Decoder should not use cached key value states when training."
            if decoder_input_ids is not None:
                decoder_input_ids = decoder_input_ids[:, -1:]
            if decoder_inputs_embeds is not None:
@@ -1088,6 +1115,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
            encoder_attention_mask=attention_mask,
            head_mask=head_mask,
            use_cache=use_cache,
+            output_attentions=output_attentions,
        )

        # insert decoder past at right place
@@ -1103,9 +1131,9 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
        lm_logits = self.lm_head(sequence_output)

        decoder_outputs = (lm_logits,) + decoder_outputs[1:]  # Add hidden states and attention if they are here
-        if lm_labels is not None:
+        if labels is not None:
            loss_fct = CrossEntropyLoss(ignore_index=-100)
-            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), lm_labels.view(-1))
+            loss = loss_fct(lm_logits.view(-1, lm_logits.size(-1)), labels.view(-1))
            # TODO(thom): Add z_loss https://github.com/tensorflow/mesh/blob/fa19d69eafc9a482aff0b59ddd96b025c0cb207d/mesh_tensorflow/layers.py#L666
            decoder_outputs = (loss,) + decoder_outputs

@@ -23,7 +23,17 @@ import tensorflow as tf
 from .configuration_albert import AlbertConfig
 from .file_utils import MULTIPLE_CHOICE_DUMMY_INPUTS, add_start_docstrings, add_start_docstrings_to_callable
 from .modeling_tf_bert import ACT2FN, TFBertSelfAttention
-from .modeling_tf_utils import TFPreTrainedModel, get_initializer, keras_serializable, shape_list
+from .modeling_tf_utils import (
+    TFMultipleChoiceLoss,
+    TFPreTrainedModel,
+    TFQuestionAnsweringLoss,
+    TFSequenceClassificationLoss,
+    TFTokenClassificationLoss,
+    cast_bool_to_primitive,
+    get_initializer,
+    keras_serializable,
+    shape_list,
+)
 from .tokenization_utils import BatchEncoding


@@ -149,7 +159,6 @@ class TFAlbertSelfAttention(tf.keras.layers.Layer):
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )
-        self.output_attentions = config.output_attentions

        self.num_attention_heads = config.num_attention_heads
        assert config.hidden_size % config.num_attention_heads == 0
@@ -173,7 +182,7 @@ class TFAlbertSelfAttention(tf.keras.layers.Layer):
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, inputs, training=False):
-        hidden_states, attention_mask, head_mask = inputs
+        hidden_states, attention_mask, head_mask, output_attentions = inputs

        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
@@ -213,7 +222,9 @@ class TFAlbertSelfAttention(tf.keras.layers.Layer):
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

-        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
+        outputs = (
+            (context_layer, attention_probs) if cast_bool_to_primitive(output_attentions) is True else (context_layer,)
+        )
        return outputs


@@ -250,7 +261,7 @@ class TFAlbertAttention(TFBertSelfAttention):
        raise NotImplementedError

    def call(self, inputs, training=False):
-        input_tensor, attention_mask, head_mask = inputs
+        input_tensor, attention_mask, head_mask, output_attentions = inputs

        batch_size = shape_list(input_tensor)[0]
        mixed_query_layer = self.query(input_tensor)
@@ -290,7 +301,9 @@ class TFAlbertAttention(TFBertSelfAttention):
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

-        self_outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
+        self_outputs = (
+            (context_layer, attention_probs) if cast_bool_to_primitive(output_attentions) is True else (context_layer,)
+        )

        hidden_states = self_outputs[0]

@@ -326,9 +339,11 @@ class TFAlbertLayer(tf.keras.layers.Layer):
        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)

    def call(self, inputs, training=False):
-        hidden_states, attention_mask, head_mask = inputs
+        hidden_states, attention_mask, head_mask, output_attentions = inputs

-        attention_outputs = self.attention([hidden_states, attention_mask, head_mask], training=training)
+        attention_outputs = self.attention(
+            [hidden_states, attention_mask, head_mask, output_attentions], training=training
+        )
        ffn_output = self.ffn(attention_outputs[0])
        ffn_output = self.activation(ffn_output)
        ffn_output = self.ffn_output(ffn_output)
@@ -345,23 +360,24 @@ class TFAlbertLayerGroup(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)

-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.albert_layers = [
            TFAlbertLayer(config, name="albert_layers_._{}".format(i)) for i in range(config.inner_group_num)
        ]

    def call(self, inputs, training=False):
-        hidden_states, attention_mask, head_mask = inputs
+        hidden_states, attention_mask, head_mask, output_attentions = inputs

        layer_hidden_states = ()
        layer_attentions = ()

        for layer_index, albert_layer in enumerate(self.albert_layers):
-            layer_output = albert_layer([hidden_states, attention_mask, head_mask[layer_index]], training=training)
+            layer_output = albert_layer(
+                [hidden_states, attention_mask, head_mask[layer_index], output_attentions], training=training
+            )
            hidden_states = layer_output[0]

-            if self.output_attentions:
+            if cast_bool_to_primitive(output_attentions) is True:
                layer_attentions = layer_attentions + (layer_output[1],)

            if self.output_hidden_states:
@@ -370,7 +386,7 @@ class TFAlbertLayerGroup(tf.keras.layers.Layer):
        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (layer_hidden_states,)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            outputs = outputs + (layer_attentions,)
        # last-layer hidden state, (layer hidden states), (layer attentions)
        return outputs
@@ -381,7 +397,6 @@ class TFAlbertTransformer(tf.keras.layers.Layer):
        super().__init__(**kwargs)

        self.config = config
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.embedding_hidden_mapping_in = tf.keras.layers.Dense(
            config.hidden_size,
@@ -394,7 +409,7 @@ class TFAlbertTransformer(tf.keras.layers.Layer):
        ]

    def call(self, inputs, training=False):
-        hidden_states, attention_mask, head_mask = inputs
+        hidden_states, attention_mask, head_mask, output_attentions = inputs

        hidden_states = self.embedding_hidden_mapping_in(hidden_states)
        all_attentions = ()
@@ -414,12 +429,13 @@ class TFAlbertTransformer(tf.keras.layers.Layer):
                    hidden_states,
                    attention_mask,
                    head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group],
+                    output_attentions,
                ],
                training=training,
            )
            hidden_states = layer_group_output[0]

-            if self.output_attentions:
+            if cast_bool_to_primitive(output_attentions) is True:
                all_attentions = all_attentions + layer_group_output[-1]

            if self.output_hidden_states:
@@ -428,7 +444,7 @@ class TFAlbertTransformer(tf.keras.layers.Layer):
        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            outputs = outputs + (all_attentions,)

        # last-layer hidden state, (all hidden states), (all attentions)
@@ -485,6 +501,7 @@ class TFAlbertMainLayer(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)
        self.num_hidden_layers = config.num_hidden_layers
+        self.output_attentions = config.output_attentions

        self.embeddings = TFAlbertEmbeddings(config, name="embeddings")
        self.encoder = TFAlbertTransformer(config, name="encoder")
@@ -516,6 +533,7 @@ class TFAlbertMainLayer(tf.keras.layers.Layer):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
        training=False,
    ):
        if isinstance(inputs, (tuple, list)):
@@ -525,7 +543,8 @@ class TFAlbertMainLayer(tf.keras.layers.Layer):
            position_ids = inputs[3] if len(inputs) > 3 else position_ids
            head_mask = inputs[4] if len(inputs) > 4 else head_mask
            inputs_embeds = inputs[5] if len(inputs) > 5 else inputs_embeds
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs[6] if len(inputs) > 6 else output_attentions
+            assert len(inputs) <= 7, "Too many inputs."
        elif isinstance(inputs, (dict, BatchEncoding)):
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
@@ -533,10 +552,13 @@ class TFAlbertMainLayer(tf.keras.layers.Layer):
            position_ids = inputs.get("position_ids", position_ids)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 7, "Too many inputs."
        else:
            input_ids = inputs

+        output_attentions = output_attentions if output_attentions is not None else self.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -579,7 +601,9 @@ class TFAlbertMainLayer(tf.keras.layers.Layer):
            # head_mask = tf.constant([0] * self.num_hidden_layers)

        embedding_output = self.embeddings([input_ids, position_ids, token_type_ids, inputs_embeds], training=training)
-        encoder_outputs = self.encoder([embedding_output, extended_attention_mask, head_mask], training=training)
+        encoder_outputs = self.encoder(
+            [embedding_output, extended_attention_mask, head_mask, output_attentions], training=training
+        )

        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output[:, 0])
@@ -695,7 +719,7 @@ class TFAlbertModel(TFAlbertPreTrainedModel):
                of shape :obj:`(batch_size, sequence_length, hidden_size)`.

                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+            attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
                tuple of :obj:`tf.Tensor` (one for each layer) of shape
                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -747,7 +771,7 @@ class TFAlbertForPreTraining(TFAlbertPreTrainedModel):
            tuple of :obj:`tf.Tensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
@@ -807,7 +831,7 @@ class TFAlbertForMaskedLM(TFAlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -841,7 +865,7 @@ class TFAlbertForMaskedLM(TFAlbertPreTrainedModel):
    the pooled output) e.g. for GLUE tasks. """,
    ALBERT_START_DOCSTRING,
 )
-class TFAlbertForSequenceClassification(TFAlbertPreTrainedModel):
+class TFAlbertForSequenceClassification(TFAlbertPreTrainedModel, TFSequenceClassificationLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -852,9 +876,26 @@ class TFAlbertForSequenceClassification(TFAlbertPreTrainedModel):
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
        )

-    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
-    def call(self, inputs, **kwargs):
+    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the sequence classification/regression loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
+            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
+
    Returns:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
        logits (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, config.num_labels)`)
@@ -864,7 +905,7 @@ class TFAlbertForSequenceClassification(TFAlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -878,27 +919,129 @@ class TFAlbertForSequenceClassification(TFAlbertPreTrainedModel):
        tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
        model = TFAlbertForSequenceClassification.from_pretrained('albert-base-v2')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        logits = outputs[0]
+        labels = tf.reshape(tf.constant(1), (-1, 1)) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, logits = outputs[:2]

        """
-        outputs = self.albert(inputs, **kwargs)
+
+        outputs = self.albert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        pooled_output = outputs[1]

-        pooled_output = self.dropout(pooled_output, training=kwargs.get("training", False))
+        pooled_output = self.dropout(pooled_output, training=training)
        logits = self.classifier(pooled_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

-        return outputs  # logits, (hidden_states), (attentions)
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), logits, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+    """Albert Model with a token classification head on top (a linear layer on top of
+    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
+    ALBERT_START_DOCSTRING,
+)
+class TFAlbertForTokenClassification(TFAlbertPreTrainedModel, TFTokenClassificationLoss):
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        self.num_labels = config.num_labels
+
+        self.albert = TFAlbertMainLayer(config, name="albert")
+        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
+        self.classifier = tf.keras.layers.Dense(
+            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
+        )
+
+    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
+        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the token classification loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+
+    Return:
+        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`):
+            Classification scores (before SoftMax).
+        hidden_states (:obj:`tuple(tf.Tensor)`, `optional`, returned when :obj:`config.output_hidden_states=True`):
+            tuple of :obj:`tf.Tensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
+            tuple of :obj:`tf.Tensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
+
+    Examples::
+
+        import tensorflow as tf
+        from transformers import AlbertTokenizer, TFAlbertForTokenClassification
+
+        tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
+        model = TFAlbertForTokenClassification.from_pretrained('albert-base-v2')
+        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
+        labels = tf.reshape(tf.constant([1] * tf.size(input_ids).numpy()), (-1, tf.size(input_ids))) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, scores = outputs[:2]
+
+        """
+        outputs = self.albert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output, training=training)
+        logits = self.classifier(sequence_output)
+
+        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
+
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), logits, (hidden_states), (attentions)


@add_start_docstrings(
    """Albert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`). """,
    ALBERT_START_DOCSTRING,
 )
-class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel):
+class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel, TFQuestionAnsweringLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -908,9 +1051,33 @@ class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel):
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
        )

-    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
-    def call(self, inputs, **kwargs):
+    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        cls_index=None,
+        p_mask=None,
+        is_impossible=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        start_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
        start_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length,)`):
@@ -922,7 +1089,7 @@ class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -938,14 +1105,24 @@ class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel):

        tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
        model = TFAlbertForQuestionAnswering.from_pretrained('albert-base-v2')
-        input_ids = tokenizer.encode("Who was Jim Henson?", "Jim Henson was a nice puppet")
-        start_scores, end_scores = model(tf.constant(input_ids)[None, :]) # Batch size 1
+        question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+        input_dict = tokenizer.encode_plus(question, text, return_tensors='tf')
+        start_scores, end_scores = model(input_dict)

-        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
+        all_tokens = tokenizer.convert_ids_to_tokens(input_dict["input_ids"].numpy()[0])
        answer = ' '.join(all_tokens[tf.math.argmax(start_scores, 1)[0] : tf.math.argmax(end_scores, 1)[0]+1])

        """
-        outputs = self.albert(inputs, **kwargs)
+        outputs = self.albert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        sequence_output = outputs[0]

@@ -956,7 +1133,13 @@ class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel):

        outputs = (start_logits, end_logits,) + outputs[2:]

-        return outputs  # start_logits, end_logits, (hidden_states), (attentions)
+        if start_positions is not None and end_positions is not None:
+            labels = {"start_position": start_positions}
+            labels["end_position"] = end_positions
+            loss = self.compute_loss(labels, outputs[:2])
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)


@add_start_docstrings(
@@ -964,7 +1147,7 @@ class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel):
    the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
    ALBERT_START_DOCSTRING,
 )
-class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel):
+class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel, TFMultipleChoiceLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -992,9 +1175,16 @@ class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
        training=False,
    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the multiple choice classification loss.
+            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
+            of the input tensors. (see `input_ids` above)
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
        classification_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, num_choices)`:
@@ -1006,7 +1196,7 @@ class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -1019,12 +1209,13 @@ class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel):

        tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
        model = TFAlbertForMultipleChoice.from_pretrained('albert-base-v2')
+        choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]

-        example1 = ["This is a context", "Is it a context? Yes"]
-        example2 = ["This is a context", "Is it a context? No"]
-        encoding = tokenizer.batch_encode_plus([example1, example2], return_tensors='tf', truncation_strategy="only_first", pad_to_max_length=True, max_length=128)
-        outputs = model(encoding["input_ids"][None, :])
-        logits = outputs[0]
+        input_ids = tf.constant([tokenizer.encode(s, add_special_tokens=True) for s in choices])[None, :] # Batch size 1, 2 choices
+        labels = tf.reshape(tf.constant(1), (-1, 1))
+        outputs = model(input_ids, labels=labels)
+
+        loss, classification_scores = outputs[:2]

        """
        if isinstance(inputs, (tuple, list)):
@@ -1034,18 +1225,17 @@ class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel):
            position_ids = inputs[3] if len(inputs) > 3 else position_ids
            head_mask = inputs[4] if len(inputs) > 4 else head_mask
            inputs_embeds = inputs[5] if len(inputs) > 5 else inputs_embeds
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs[6] if len(inputs) > 6 else output_attentions
+            assert len(inputs) <= 7, "Too many inputs."
        elif isinstance(inputs, dict):
-            print("isdict(1)")
            input_ids = inputs.get("input_ids")
-            print(input_ids)
-
            attention_mask = inputs.get("attention_mask", attention_mask)
            token_type_ids = inputs.get("token_type_ids", token_type_ids)
            position_ids = inputs.get("position_ids", position_ids)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 7, "Too many inputs."
        else:
            input_ids = inputs

@@ -1068,6 +1258,7 @@ class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel):
            flat_position_ids,
            head_mask,
            inputs_embeds,
+            output_attentions,
        ]

        outputs = self.albert(flat_inputs, training=training)
@@ -1080,4 +1271,8 @@ class TFAlbertForMultipleChoice(TFAlbertPreTrainedModel):

        outputs = (reshaped_logits,) + outputs[2:]  # add hidden states and attention if they are here

-        return outputs  # reshaped_logits, (hidden_states), (attentions)
+        if labels is not None:
+            loss = self.compute_loss(labels, reshaped_logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)
@@ -22,14 +22,18 @@ from .configuration_auto import (
    AlbertConfig,
    AutoConfig,
    BertConfig,
+    CamembertConfig,
    CTRLConfig,
    DistilBertConfig,
+    ElectraConfig,
+    FlaubertConfig,
    GPT2Config,
    OpenAIGPTConfig,
    RobertaConfig,
    T5Config,
    TransfoXLConfig,
    XLMConfig,
+    XLMRobertaConfig,
    XLNetConfig,
 )
 from .configuration_utils import PretrainedConfig
@@ -39,6 +43,7 @@ from .modeling_tf_albert import (
    TFAlbertForPreTraining,
    TFAlbertForQuestionAnswering,
    TFAlbertForSequenceClassification,
+    TFAlbertForTokenClassification,
    TFAlbertModel,
 )
 from .modeling_tf_bert import (
@@ -50,18 +55,43 @@ from .modeling_tf_bert import (
    TFBertForTokenClassification,
    TFBertModel,
 )
+from .modeling_tf_camembert import (
+    TFCamembertForMaskedLM,
+    TFCamembertForMultipleChoice,
+    TFCamembertForQuestionAnswering,
+    TFCamembertForSequenceClassification,
+    TFCamembertForTokenClassification,
+    TFCamembertModel,
+)
 from .modeling_tf_ctrl import TFCTRLLMHeadModel, TFCTRLModel
 from .modeling_tf_distilbert import (
    TFDistilBertForMaskedLM,
+    TFDistilBertForMultipleChoice,
    TFDistilBertForQuestionAnswering,
    TFDistilBertForSequenceClassification,
    TFDistilBertForTokenClassification,
    TFDistilBertModel,
 )
+from .modeling_tf_electra import (
+    TFElectraForMaskedLM,
+    TFElectraForPreTraining,
+    TFElectraForQuestionAnswering,
+    TFElectraForTokenClassification,
+    TFElectraModel,
+)
+from .modeling_tf_flaubert import (
+    TFFlaubertForMultipleChoice,
+    TFFlaubertForQuestionAnsweringSimple,
+    TFFlaubertForSequenceClassification,
+    TFFlaubertForTokenClassification,
+    TFFlaubertModel,
+    TFFlaubertWithLMHeadModel,
+)
 from .modeling_tf_gpt2 import TFGPT2LMHeadModel, TFGPT2Model
 from .modeling_tf_openai import TFOpenAIGPTLMHeadModel, TFOpenAIGPTModel
 from .modeling_tf_roberta import (
    TFRobertaForMaskedLM,
+    TFRobertaForMultipleChoice,
    TFRobertaForQuestionAnswering,
    TFRobertaForSequenceClassification,
    TFRobertaForTokenClassification,
@@ -70,12 +100,23 @@ from .modeling_tf_roberta import (
 from .modeling_tf_t5 import TFT5ForConditionalGeneration, TFT5Model
 from .modeling_tf_transfo_xl import TFTransfoXLLMHeadModel, TFTransfoXLModel
 from .modeling_tf_xlm import (
+    TFXLMForMultipleChoice,
    TFXLMForQuestionAnsweringSimple,
    TFXLMForSequenceClassification,
+    TFXLMForTokenClassification,
    TFXLMModel,
    TFXLMWithLMHeadModel,
 )
+from .modeling_tf_xlm_roberta import (
+    TFXLMRobertaForMaskedLM,
+    TFXLMRobertaForMultipleChoice,
+    TFXLMRobertaForQuestionAnswering,
+    TFXLMRobertaForSequenceClassification,
+    TFXLMRobertaForTokenClassification,
+    TFXLMRobertaModel,
+)
 from .modeling_tf_xlnet import (
+    TFXLNetForMultipleChoice,
    TFXLNetForQuestionAnsweringSimple,
    TFXLNetForSequenceClassification,
    TFXLNetForTokenClassification,
@@ -89,83 +130,118 @@ logger = logging.getLogger(__name__)

 TF_MODEL_MAPPING = OrderedDict(
    [
-        (T5Config, TFT5Model),
-        (DistilBertConfig, TFDistilBertModel),
        (AlbertConfig, TFAlbertModel),
-        (RobertaConfig, TFRobertaModel),
        (BertConfig, TFBertModel),
-        (OpenAIGPTConfig, TFOpenAIGPTModel),
-        (GPT2Config, TFGPT2Model),
-        (TransfoXLConfig, TFTransfoXLModel),
-        (XLNetConfig, TFXLNetModel),
-        (XLMConfig, TFXLMModel),
+        (CamembertConfig, TFCamembertModel),
        (CTRLConfig, TFCTRLModel),
+        (DistilBertConfig, TFDistilBertModel),
+        (ElectraConfig, TFElectraModel),
+        (FlaubertConfig, TFFlaubertModel),
+        (GPT2Config, TFGPT2Model),
+        (OpenAIGPTConfig, TFOpenAIGPTModel),
+        (RobertaConfig, TFRobertaModel),
+        (T5Config, TFT5Model),
+        (TransfoXLConfig, TFTransfoXLModel),
+        (XLMConfig, TFXLMModel),
+        (XLMRobertaConfig, TFXLMRobertaModel),
+        (XLNetConfig, TFXLNetModel),
    ]
 )

 TF_MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(
    [
-        (T5Config, TFT5ForConditionalGeneration),
-        (DistilBertConfig, TFDistilBertForMaskedLM),
        (AlbertConfig, TFAlbertForPreTraining),
-        (RobertaConfig, TFRobertaForMaskedLM),
        (BertConfig, TFBertForPreTraining),
-        (OpenAIGPTConfig, TFOpenAIGPTLMHeadModel),
-        (GPT2Config, TFGPT2LMHeadModel),
-        (TransfoXLConfig, TFTransfoXLLMHeadModel),
-        (XLNetConfig, TFXLNetLMHeadModel),
-        (XLMConfig, TFXLMWithLMHeadModel),
+        (CamembertConfig, TFCamembertForMaskedLM),
        (CTRLConfig, TFCTRLLMHeadModel),
+        (DistilBertConfig, TFDistilBertForMaskedLM),
+        (ElectraConfig, TFElectraForPreTraining),
+        (FlaubertConfig, TFFlaubertWithLMHeadModel),
+        (GPT2Config, TFGPT2LMHeadModel),
+        (OpenAIGPTConfig, TFOpenAIGPTLMHeadModel),
+        (RobertaConfig, TFRobertaForMaskedLM),
+        (T5Config, TFT5ForConditionalGeneration),
+        (TransfoXLConfig, TFTransfoXLLMHeadModel),
+        (XLMConfig, TFXLMWithLMHeadModel),
+        (XLMRobertaConfig, TFXLMRobertaForMaskedLM),
+        (XLNetConfig, TFXLNetLMHeadModel),
    ]
 )

 TF_MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(
    [
-        (T5Config, TFT5ForConditionalGeneration),
-        (DistilBertConfig, TFDistilBertForMaskedLM),
        (AlbertConfig, TFAlbertForMaskedLM),
-        (RobertaConfig, TFRobertaForMaskedLM),
        (BertConfig, TFBertForMaskedLM),
-        (OpenAIGPTConfig, TFOpenAIGPTLMHeadModel),
-        (GPT2Config, TFGPT2LMHeadModel),
-        (TransfoXLConfig, TFTransfoXLLMHeadModel),
-        (XLNetConfig, TFXLNetLMHeadModel),
-        (XLMConfig, TFXLMWithLMHeadModel),
+        (CamembertConfig, TFCamembertForMaskedLM),
        (CTRLConfig, TFCTRLLMHeadModel),
+        (DistilBertConfig, TFDistilBertForMaskedLM),
+        (ElectraConfig, TFElectraForMaskedLM),
+        (FlaubertConfig, TFFlaubertWithLMHeadModel),
+        (GPT2Config, TFGPT2LMHeadModel),
+        (OpenAIGPTConfig, TFOpenAIGPTLMHeadModel),
+        (RobertaConfig, TFRobertaForMaskedLM),
+        (T5Config, TFT5ForConditionalGeneration),
+        (TransfoXLConfig, TFTransfoXLLMHeadModel),
+        (XLMConfig, TFXLMWithLMHeadModel),
+        (XLMRobertaConfig, TFXLMRobertaForMaskedLM),
+        (XLNetConfig, TFXLNetLMHeadModel),
+    ]
+)
+
+TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING = OrderedDict(
+    [
+        (AlbertConfig, TFAlbertForMultipleChoice),
+        (BertConfig, TFBertForMultipleChoice),
+        (CamembertConfig, TFCamembertForMultipleChoice),
+        (DistilBertConfig, TFDistilBertForMultipleChoice),
+        (FlaubertConfig, TFFlaubertForMultipleChoice),
+        (RobertaConfig, TFRobertaForMultipleChoice),
+        (XLMConfig, TFXLMForMultipleChoice),
+        (XLMRobertaConfig, TFXLMRobertaForMultipleChoice),
+        (XLNetConfig, TFXLNetForMultipleChoice),
+    ]
+)
+
+TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
+    [
+        (AlbertConfig, TFAlbertForQuestionAnswering),
+        (BertConfig, TFBertForQuestionAnswering),
+        (CamembertConfig, TFCamembertForQuestionAnswering),
+        (DistilBertConfig, TFDistilBertForQuestionAnswering),
+        (ElectraConfig, TFElectraForQuestionAnswering),
+        (FlaubertConfig, TFFlaubertForQuestionAnsweringSimple),
+        (RobertaConfig, TFRobertaForQuestionAnswering),
+        (XLMConfig, TFXLMForQuestionAnsweringSimple),
+        (XLMRobertaConfig, TFXLMRobertaForQuestionAnswering),
+        (XLNetConfig, TFXLNetForQuestionAnsweringSimple),
    ]
 )

 TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = OrderedDict(
    [
-        (DistilBertConfig, TFDistilBertForSequenceClassification),
        (AlbertConfig, TFAlbertForSequenceClassification),
-        (RobertaConfig, TFRobertaForSequenceClassification),
        (BertConfig, TFBertForSequenceClassification),
-        (XLNetConfig, TFXLNetForSequenceClassification),
+        (CamembertConfig, TFCamembertForSequenceClassification),
+        (DistilBertConfig, TFDistilBertForSequenceClassification),
+        (FlaubertConfig, TFFlaubertForSequenceClassification),
+        (RobertaConfig, TFRobertaForSequenceClassification),
        (XLMConfig, TFXLMForSequenceClassification),
-    ]
-)
-
-TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING = OrderedDict(
-    [(BertConfig, TFBertForMultipleChoice), (AlbertConfig, TFAlbertForMultipleChoice)]
-)
-
-TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
-    [
-        (DistilBertConfig, TFDistilBertForQuestionAnswering),
-        (AlbertConfig, TFAlbertForQuestionAnswering),
-        (RobertaConfig, TFRobertaForQuestionAnswering),
-        (BertConfig, TFBertForQuestionAnswering),
-        (XLNetConfig, TFXLNetForQuestionAnsweringSimple),
-        (XLMConfig, TFXLMForQuestionAnsweringSimple),
+        (XLMRobertaConfig, TFXLMRobertaForSequenceClassification),
+        (XLNetConfig, TFXLNetForSequenceClassification),
    ]
 )

 TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = OrderedDict(
    [
-        (DistilBertConfig, TFDistilBertForTokenClassification),
-        (RobertaConfig, TFRobertaForTokenClassification),
+        (AlbertConfig, TFAlbertForTokenClassification),
        (BertConfig, TFBertForTokenClassification),
+        (CamembertConfig, TFCamembertForTokenClassification),
+        (DistilBertConfig, TFDistilBertForTokenClassification),
+        (ElectraConfig, TFElectraForTokenClassification),
+        (FlaubertConfig, TFFlaubertForTokenClassification),
+        (RobertaConfig, TFRobertaForTokenClassification),
+        (XLMConfig, TFXLMForTokenClassification),
+        (XLMRobertaConfig, TFXLMRobertaForTokenClassification),
        (XLNetConfig, TFXLNetForTokenClassification),
    ]
 )
@@ -632,11 +708,13 @@ class TFAutoModelWithLMHead(object):

        """
        config = kwargs.pop("config", None)
+
        if not isinstance(config, PretrainedConfig):
            config = AutoConfig.from_pretrained(pretrained_model_name_or_path, **kwargs)

        for config_class, model_class in TF_MODEL_WITH_LM_HEAD_MAPPING.items():
-            if isinstance(config, config_class):
+            # Not using isinstance() here to do not take into account inheritance
+            if config_class == type(config):
                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
@@ -23,7 +23,17 @@ import tensorflow as tf

 from .configuration_bert import BertConfig
 from .file_utils import MULTIPLE_CHOICE_DUMMY_INPUTS, add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_tf_utils import TFPreTrainedModel, get_initializer, keras_serializable, shape_list
+from .modeling_tf_utils import (
+    TFMultipleChoiceLoss,
+    TFPreTrainedModel,
+    TFQuestionAnsweringLoss,
+    TFSequenceClassificationLoss,
+    TFTokenClassificationLoss,
+    cast_bool_to_primitive,
+    get_initializer,
+    keras_serializable,
+    shape_list,
+)
 from .tokenization_utils import BatchEncoding


@@ -202,7 +212,6 @@ class TFBertSelfAttention(tf.keras.layers.Layer):
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
            )
-        self.output_attentions = config.output_attentions

        self.num_attention_heads = config.num_attention_heads
        assert config.hidden_size % config.num_attention_heads == 0
@@ -226,7 +235,7 @@ class TFBertSelfAttention(tf.keras.layers.Layer):
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, inputs, training=False):
-        hidden_states, attention_mask, head_mask = inputs
+        hidden_states, attention_mask, head_mask, output_attentions = inputs

        batch_size = shape_list(hidden_states)[0]
        mixed_query_layer = self.query(hidden_states)
@@ -266,7 +275,10 @@ class TFBertSelfAttention(tf.keras.layers.Layer):
            context_layer, (batch_size, -1, self.all_head_size)
        )  # (batch_size, seq_len_q, all_head_size)

-        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
+        outputs = (
+            (context_layer, attention_probs) if cast_bool_to_primitive(output_attentions) is True else (context_layer,)
+        )
+
        return outputs


@@ -298,9 +310,11 @@ class TFBertAttention(tf.keras.layers.Layer):
        raise NotImplementedError

    def call(self, inputs, training=False):
-        input_tensor, attention_mask, head_mask = inputs
+        input_tensor, attention_mask, head_mask, output_attentions = inputs

-        self_outputs = self.self_attention([input_tensor, attention_mask, head_mask], training=training)
+        self_outputs = self.self_attention(
+            [input_tensor, attention_mask, head_mask, output_attentions], training=training
+        )
        attention_output = self.dense_output([self_outputs[0], input_tensor], training=training)
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
        return outputs
@@ -349,9 +363,11 @@ class TFBertLayer(tf.keras.layers.Layer):
        self.bert_output = TFBertOutput(config, name="output")

    def call(self, inputs, training=False):
-        hidden_states, attention_mask, head_mask = inputs
+        hidden_states, attention_mask, head_mask, output_attentions = inputs

-        attention_outputs = self.attention([hidden_states, attention_mask, head_mask], training=training)
+        attention_outputs = self.attention(
+            [hidden_states, attention_mask, head_mask, output_attentions], training=training
+        )
        attention_output = attention_outputs[0]
        intermediate_output = self.intermediate(attention_output)
        layer_output = self.bert_output([intermediate_output, attention_output], training=training)
@@ -362,12 +378,11 @@ class TFBertLayer(tf.keras.layers.Layer):
 class TFBertEncoder(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.layer = [TFBertLayer(config, name="layer_._{}".format(i)) for i in range(config.num_hidden_layers)]

    def call(self, inputs, training=False):
-        hidden_states, attention_mask, head_mask = inputs
+        hidden_states, attention_mask, head_mask, output_attentions = inputs

        all_hidden_states = ()
        all_attentions = ()
@@ -375,10 +390,12 @@ class TFBertEncoder(tf.keras.layers.Layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

-            layer_outputs = layer_module([hidden_states, attention_mask, head_mask[i]], training=training)
+            layer_outputs = layer_module(
+                [hidden_states, attention_mask, head_mask[i], output_attentions], training=training
+            )
            hidden_states = layer_outputs[0]

-            if self.output_attentions:
+            if cast_bool_to_primitive(output_attentions) is True:
                all_attentions = all_attentions + (layer_outputs[1],)

        # Add last layer
@@ -388,7 +405,7 @@ class TFBertEncoder(tf.keras.layers.Layer):
        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            outputs = outputs + (all_attentions,)
        return outputs  # outputs, (hidden states), (attentions)

@@ -480,6 +497,7 @@ class TFBertMainLayer(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)
        self.num_hidden_layers = config.num_hidden_layers
+        self.output_attentions = config.output_attentions

        self.embeddings = TFBertEmbeddings(config, name="embeddings")
        self.encoder = TFBertEncoder(config, name="encoder")
@@ -506,6 +524,7 @@ class TFBertMainLayer(tf.keras.layers.Layer):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
        training=False,
    ):
        if isinstance(inputs, (tuple, list)):
@@ -515,7 +534,8 @@ class TFBertMainLayer(tf.keras.layers.Layer):
            position_ids = inputs[3] if len(inputs) > 3 else position_ids
            head_mask = inputs[4] if len(inputs) > 4 else head_mask
            inputs_embeds = inputs[5] if len(inputs) > 5 else inputs_embeds
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs[6] if len(inputs) > 6 else output_attentions
+            assert len(inputs) <= 7, "Too many inputs."
        elif isinstance(inputs, (dict, BatchEncoding)):
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
@@ -523,10 +543,13 @@ class TFBertMainLayer(tf.keras.layers.Layer):
            position_ids = inputs.get("position_ids", position_ids)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 7, "Too many inputs."
        else:
            input_ids = inputs

+        output_attentions = output_attentions if output_attentions is not None else self.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -569,7 +592,9 @@ class TFBertMainLayer(tf.keras.layers.Layer):
            # head_mask = tf.constant([0] * self.num_hidden_layers)

        embedding_output = self.embeddings([input_ids, position_ids, token_type_ids, inputs_embeds], training=training)
-        encoder_outputs = self.encoder([embedding_output, extended_attention_mask, head_mask], training=training)
+        encoder_outputs = self.encoder(
+            [embedding_output, extended_attention_mask, head_mask, output_attentions], training=training
+        )

        sequence_output = encoder_outputs[0]
        pooled_output = self.pooler(sequence_output)
@@ -688,7 +713,7 @@ class TFBertModel(TFBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -740,7 +765,7 @@ class TFBertForPreTraining(TFBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -794,7 +819,7 @@ class TFBertForMaskedLM(TFBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -844,7 +869,7 @@ class TFBertForNextSentencePrediction(TFBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -880,7 +905,7 @@ class TFBertForNextSentencePrediction(TFBertPreTrainedModel):
    the pooled output) e.g. for GLUE tasks. """,
    BERT_START_DOCSTRING,
 )
-class TFBertForSequenceClassification(TFBertPreTrainedModel):
+class TFBertForSequenceClassification(TFBertPreTrainedModel, TFSequenceClassificationLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -891,9 +916,26 @@ class TFBertForSequenceClassification(TFBertPreTrainedModel):
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
        )

-    @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
-    def call(self, inputs, **kwargs):
+    @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the sequence classification/regression loss.
+            Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
+            If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
        logits (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, config.num_labels)`):
@@ -903,7 +945,7 @@ class TFBertForSequenceClassification(TFBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -916,21 +958,36 @@ class TFBertForSequenceClassification(TFBertPreTrainedModel):

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = TFBertForSequenceClassification.from_pretrained('bert-base-uncased')
-        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        logits = outputs[0]
+        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
+        labels = tf.reshape(tf.constant(1), (-1, 1)) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, logits = outputs[:2]

        """
-        outputs = self.bert(inputs, **kwargs)
+
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        pooled_output = outputs[1]

-        pooled_output = self.dropout(pooled_output, training=kwargs.get("training", False))
+        pooled_output = self.dropout(pooled_output, training=training)
        logits = self.classifier(pooled_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

-        return outputs  # logits, (hidden_states), (attentions)
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), logits, (hidden_states), (attentions)


@add_start_docstrings(
@@ -938,7 +995,7 @@ class TFBertForSequenceClassification(TFBertPreTrainedModel):
    the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
    BERT_START_DOCSTRING,
 )
-class TFBertForMultipleChoice(TFBertPreTrainedModel):
+class TFBertForMultipleChoice(TFBertPreTrainedModel, TFMultipleChoiceLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -966,9 +1023,16 @@ class TFBertForMultipleChoice(TFBertPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
        training=False,
    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the multiple choice classification loss.
+            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
+            of the input tensors. (see `input_ids` above)
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
        classification_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, num_choices)`:
@@ -980,7 +1044,7 @@ class TFBertForMultipleChoice(TFBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -993,15 +1057,14 @@ class TFBertForMultipleChoice(TFBertPreTrainedModel):

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = TFBertForMultipleChoice.from_pretrained('bert-base-uncased')
+        choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]

-        prompt = "In Italy, pizza served in formal settings, such as at a restaurant, is presented unsliced."
-        choice0 = "It is eaten with a fork and a knife."
-        choice1 = "It is eaten while held in the hand."
-        encoding = tokenizer.batch_encode_plus([[prompt, choice0], [prompt, choice1]], return_tensors='tf', pad_to_max_length=True)
+        input_ids = tf.constant([tokenizer.encode(s, add_special_tokens=True) for s in choices])[None, :] # Batch size 1, 2 choices
+        labels = tf.reshape(tf.constant(1), (-1, 1))
+        outputs = model(input_ids, labels=labels)
+
+        loss, classification_scores = outputs[:2]

-        # linear classifier on the output is not yet trained
-        outputs = model(encoding['input_ids'][None, :])
-        logits = outputs[0]
        """
        if isinstance(inputs, (tuple, list)):
            input_ids = inputs[0]
@@ -1010,15 +1073,17 @@ class TFBertForMultipleChoice(TFBertPreTrainedModel):
            position_ids = inputs[3] if len(inputs) > 3 else position_ids
            head_mask = inputs[4] if len(inputs) > 4 else head_mask
            inputs_embeds = inputs[5] if len(inputs) > 5 else inputs_embeds
-            assert len(inputs) <= 6, "Too many inputs."
-        elif isinstance(inputs, dict):
+            output_attentions = inputs[6] if len(inputs) > 6 else output_attentions
+            assert len(inputs) <= 7, "Too many inputs."
+        elif isinstance(inputs, (dict, BatchEncoding)):
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
            token_type_ids = inputs.get("token_type_ids", token_type_ids)
            position_ids = inputs.get("position_ids", position_ids)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 7, "Too many inputs."
        else:
            input_ids = inputs

@@ -1041,6 +1106,7 @@ class TFBertForMultipleChoice(TFBertPreTrainedModel):
            flat_position_ids,
            head_mask,
            inputs_embeds,
+            output_attentions,
        ]

        outputs = self.bert(flat_inputs, training=training)
@@ -1053,7 +1119,11 @@ class TFBertForMultipleChoice(TFBertPreTrainedModel):

        outputs = (reshaped_logits,) + outputs[2:]  # add hidden states and attention if they are here

-        return outputs  # reshaped_logits, (hidden_states), (attentions)
+        if labels is not None:
+            loss = self.compute_loss(labels, reshaped_logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)


@add_start_docstrings(
@@ -1061,7 +1131,7 @@ class TFBertForMultipleChoice(TFBertPreTrainedModel):
    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
    BERT_START_DOCSTRING,
 )
-class TFBertForTokenClassification(TFBertPreTrainedModel):
+class TFBertForTokenClassification(TFBertPreTrainedModel, TFTokenClassificationLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -1072,9 +1142,24 @@ class TFBertForTokenClassification(TFBertPreTrainedModel):
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
        )

-    @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
-    def call(self, inputs, **kwargs):
+    @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the token classification loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
        scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`):
@@ -1084,7 +1169,7 @@ class TFBertForTokenClassification(TFBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -1098,20 +1183,34 @@ class TFBertForTokenClassification(TFBertPreTrainedModel):
        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = TFBertForTokenClassification.from_pretrained('bert-base-uncased')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        scores = outputs[0]
+        labels = tf.reshape(tf.constant([1] * tf.size(input_ids).numpy()), (-1, tf.size(input_ids))) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, scores = outputs[:2]

        """
-        outputs = self.bert(inputs, **kwargs)
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        sequence_output = outputs[0]

-        sequence_output = self.dropout(sequence_output, training=kwargs.get("training", False))
+        sequence_output = self.dropout(sequence_output, training=training)
        logits = self.classifier(sequence_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

-        return outputs  # scores, (hidden_states), (attentions)
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), logits, (hidden_states), (attentions)


@add_start_docstrings(
@@ -1119,7 +1218,7 @@ class TFBertForTokenClassification(TFBertPreTrainedModel):
    the hidden-states output to compute `span start logits` and `span end logits`). """,
    BERT_START_DOCSTRING,
 )
-class TFBertForQuestionAnswering(TFBertPreTrainedModel):
+class TFBertForQuestionAnswering(TFBertPreTrainedModel, TFQuestionAnsweringLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -1129,9 +1228,33 @@ class TFBertForQuestionAnswering(TFBertPreTrainedModel):
            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
        )

-    @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
-    def call(self, inputs, **kwargs):
+    @add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        cls_index=None,
+        p_mask=None,
+        is_impossible=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        start_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
        start_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length,)`):
@@ -1143,7 +1266,7 @@ class TFBertForQuestionAnswering(TFBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -1156,18 +1279,25 @@ class TFBertForQuestionAnswering(TFBertPreTrainedModel):

        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
        model = TFBertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
-
        question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
-        encoding = tokenizer.encode_plus(question, text)
-        input_ids, token_type_ids = encoding["input_ids"], encoding["token_type_ids"]
-        start_scores, end_scores = model(tf.constant(input_ids)[None, :], token_type_ids=tf.constant(token_type_ids)[None, :])
+        input_dict = tokenizer.encode_plus(question, text, return_tensors='tf')
+        start_scores, end_scores = model(input_dict)

-        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
-        answer = ' '.join(all_tokens[tf.math.argmax(tf.squeeze(start_scores)) : tf.math.argmax(tf.squeeze(end_scores))+1])
+        all_tokens = tokenizer.convert_ids_to_tokens(input_dict["input_ids"].numpy()[0])
+        answer = ' '.join(all_tokens[tf.math.argmax(start_scores, 1)[0] : tf.math.argmax(end_scores, 1)[0]+1])
        assert answer == "a nice puppet"

        """
-        outputs = self.bert(inputs, **kwargs)
+        outputs = self.bert(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        sequence_output = outputs[0]

@@ -1178,4 +1308,10 @@ class TFBertForQuestionAnswering(TFBertPreTrainedModel):

        outputs = (start_logits, end_logits,) + outputs[2:]

-        return outputs  # start_logits, end_logits, (hidden_states), (attentions)
+        if start_positions is not None and end_positions is not None:
+            labels = {"start_position": start_positions}
+            labels["end_position"] = end_positions
+            loss = self.compute_loss(labels, outputs[:2])
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)
@@ -22,6 +22,8 @@ from .configuration_camembert import CamembertConfig
 from .file_utils import add_start_docstrings
 from .modeling_tf_roberta import (
    TFRobertaForMaskedLM,
+    TFRobertaForMultipleChoice,
+    TFRobertaForQuestionAnswering,
    TFRobertaForSequenceClassification,
    TFRobertaForTokenClassification,
    TFRobertaModel,
@@ -114,3 +116,30 @@ class TFCamembertForTokenClassification(TFRobertaForTokenClassification):
    """

    config_class = CamembertConfig
+
+
+@add_start_docstrings(
+    """CamemBERT Model with a multiple choice classification head on top (a linear layer on top of
+    the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
+    CAMEMBERT_START_DOCSTRING,
+)
+class TFCamembertForMultipleChoice(TFRobertaForMultipleChoice):
+    """
+    This class overrides :class:`~transformers.TFRobertaForMultipleChoice`. Please check the
+    superclass for the appropriate documentation alongside usage examples.
+    """
+
+    config_class = CamembertConfig
+
+
+@add_start_docstrings(
+    """CamemBERT Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`). """,
+    CAMEMBERT_START_DOCSTRING,
+)
+class TFCamembertForQuestionAnswering(TFRobertaForQuestionAnswering):
+    """
+    This class overrides :class:`~transformers.TFRobertaForQuestionAnswering`. Please check the
+    superclass for the appropriate documentation alongside usage examples.
+    """
+
+    config_class = CamembertConfig
@@ -23,7 +23,13 @@ import tensorflow as tf

 from .configuration_ctrl import CTRLConfig
 from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_tf_utils import TFPreTrainedModel, TFSharedEmbeddings, keras_serializable, shape_list
+from .modeling_tf_utils import (
+    TFPreTrainedModel,
+    TFSharedEmbeddings,
+    cast_bool_to_primitive,
+    keras_serializable,
+    shape_list,
+)
 from .tokenization_utils import BatchEncoding


@@ -78,9 +84,8 @@ def scaled_dot_product_attention(q, k, v, mask, attention_mask=None, head_mask=N


 class TFMultiHeadAttention(tf.keras.layers.Layer):
-    def __init__(self, d_model_size, num_heads, output_attentions=False, **kwargs):
+    def __init__(self, d_model_size, num_heads, **kwargs):
        super().__init__(**kwargs)
-        self.output_attentions = output_attentions
        self.num_heads = num_heads
        self.d_model_size = d_model_size

@@ -97,7 +102,7 @@ class TFMultiHeadAttention(tf.keras.layers.Layer):
        return tf.transpose(x, perm=[0, 2, 1, 3])

    def call(self, inputs, training=False):
-        v, k, q, mask, layer_past, attention_mask, head_mask, use_cache = inputs
+        v, k, q, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions = inputs
        batch_size = shape_list(q)[0]

        q = self.Wq(q)
@@ -114,13 +119,7 @@ class TFMultiHeadAttention(tf.keras.layers.Layer):
            v = tf.concat((past_value, v), axis=-2)

        # to cope with keras serialization
-        # we need to cast `use_cache` to correct bool
-        # if it is a tensor
-        if tf.is_tensor(use_cache):
-            if hasattr(use_cache, "numpy"):
-                use_cache = bool(use_cache.numpy())
-            else:
-                use_cache = True
+        use_cache = cast_bool_to_primitive(use_cache, True)

        if use_cache is True:
            present = tf.stack((k, v), axis=0)
@@ -134,7 +133,7 @@ class TFMultiHeadAttention(tf.keras.layers.Layer):
        output = self.dense(original_size_attention)

        outputs = (output, present)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            outputs = outputs + (attn,)
        return outputs

@@ -147,14 +146,10 @@ def point_wise_feed_forward_network(d_model_size, dff, name=""):


 class TFEncoderLayer(tf.keras.layers.Layer):
-    def __init__(
-        self, d_model_size, num_heads, dff, rate=0.1, layer_norm_epsilon=1e-6, output_attentions=False, **kwargs
-    ):
+    def __init__(self, d_model_size, num_heads, dff, rate=0.1, layer_norm_epsilon=1e-6, **kwargs):
        super().__init__(**kwargs)

-        self.multi_head_attention = TFMultiHeadAttention(
-            d_model_size, num_heads, output_attentions, name="multi_head_attention"
-        )
+        self.multi_head_attention = TFMultiHeadAttention(d_model_size, num_heads, name="multi_head_attention")
        self.ffn = point_wise_feed_forward_network(d_model_size, dff, name="ffn")

        self.layernorm1 = tf.keras.layers.LayerNormalization(epsilon=layer_norm_epsilon, name="layernorm1")
@@ -164,10 +159,11 @@ class TFEncoderLayer(tf.keras.layers.Layer):
        self.dropout2 = tf.keras.layers.Dropout(rate)

    def call(self, inputs, training=False):
-        x, mask, layer_past, attention_mask, head_mask, use_cache = inputs
+        x, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions = inputs
        normed = self.layernorm1(x)
        attn_outputs = self.multi_head_attention(
-            [normed, normed, normed, mask, layer_past, attention_mask, head_mask, use_cache], training=training
+            [normed, normed, normed, mask, layer_past, attention_mask, head_mask, use_cache, output_attentions],
+            training=training,
        )
        attn_output = attn_outputs[0]
        attn_output = self.dropout1(attn_output, training=training)
@@ -208,7 +204,6 @@ class TFCTRLMainLayer(tf.keras.layers.Layer):
                config.dff,
                config.resid_pdrop,
                config.layer_norm_epsilon,
-                config.output_attentions,
                name="h_._{}".format(i),
            )
            for i in range(config.n_layer)
@@ -237,6 +232,7 @@ class TFCTRLMainLayer(tf.keras.layers.Layer):
        head_mask=None,
        inputs_embeds=None,
        use_cache=True,
+        output_attentions=None,
        training=False,
    ):

@@ -249,7 +245,8 @@ class TFCTRLMainLayer(tf.keras.layers.Layer):
            head_mask = inputs[5] if len(inputs) > 5 else head_mask
            inputs_embeds = inputs[6] if len(inputs) > 6 else inputs_embeds
            use_cache = inputs[7] if len(inputs) > 7 else use_cache
-            assert len(inputs) <= 8, "Too many inputs."
+            output_attentions = inputs[8] if len(inputs) > 8 else output_attentions
+            assert len(inputs) <= 9, "Too many inputs."
        elif isinstance(inputs, (dict, BatchEncoding)):
            input_ids = inputs.get("input_ids")
            past = inputs.get("past", past)
@@ -259,10 +256,13 @@ class TFCTRLMainLayer(tf.keras.layers.Layer):
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
            use_cache = inputs.get("use_cache", use_cache)
-            assert len(inputs) <= 8, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 9, "Too many inputs."
        else:
            input_ids = inputs

+        output_attentions = output_attentions if output_attentions is not None else self.output_attentions
+
        # If using past key value states, only the last tokens
        # should be given as an input
        if past is not None:
@@ -349,13 +349,16 @@ class TFCTRLMainLayer(tf.keras.layers.Layer):
        for i, (h, layer_past) in enumerate(zip(self.h, past)):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (tf.reshape(hidden_states, output_shape),)
-            outputs = h([hidden_states, mask, layer_past, attention_mask, head_mask[i], use_cache], training=training)
+            outputs = h(
+                [hidden_states, mask, layer_past, attention_mask, head_mask[i], use_cache, output_attentions],
+                training=training,
+            )
            hidden_states, present = outputs[:2]

            if use_cache is True:
                presents = presents + (present,)

-            if self.output_attentions:
+            if cast_bool_to_primitive(output_attentions) is True:
                all_attentions.append(outputs[2])

        hidden_states = self.layernorm(hidden_states)
@@ -368,7 +371,7 @@ class TFCTRLMainLayer(tf.keras.layers.Layer):
            outputs = outputs + (presents,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            # let the number of heads free (-1) so we can extract attention even after head pruning
            attention_output_shape = input_shape[:-1] + [-1] + shape_list(all_attentions[0])[-2:]
            all_attentions = tuple(tf.reshape(t, attention_output_shape) for t in all_attentions)
@@ -489,7 +492,7 @@ class TFCTRLModel(TFCTRLPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -569,7 +572,7 @@ class TFCTRLLMHeadModel(TFCTRLPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -23,8 +23,19 @@ import numpy as np
 import tensorflow as tf

 from .configuration_distilbert import DistilBertConfig
-from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
-from .modeling_tf_utils import TFPreTrainedModel, TFSharedEmbeddings, get_initializer, shape_list
+from .file_utils import MULTIPLE_CHOICE_DUMMY_INPUTS, add_start_docstrings, add_start_docstrings_to_callable
+from .modeling_tf_utils import (
+    TFMultipleChoiceLoss,
+    TFPreTrainedModel,
+    TFQuestionAnsweringLoss,
+    TFSequenceClassificationLoss,
+    TFSharedEmbeddings,
+    TFTokenClassificationLoss,
+    cast_bool_to_primitive,
+    get_initializer,
+    keras_serializable,
+    shape_list,
+)
 from .tokenization_utils import BatchEncoding


@@ -176,7 +187,6 @@ class TFMultiHeadSelfAttention(tf.keras.layers.Layer):
        self.n_heads = config.n_heads
        self.dim = config.dim
        self.dropout = tf.keras.layers.Dropout(config.attention_dropout)
-        self.output_attentions = config.output_attentions

        assert self.dim % self.n_heads == 0

@@ -214,7 +224,7 @@ class TFMultiHeadSelfAttention(tf.keras.layers.Layer):
        context: tf.Tensor(bs, seq_length, dim)
            Contextualized layer. Optional: only if `output_attentions=True`
        """
-        query, key, value, mask, head_mask = inputs
+        query, key, value, mask, head_mask, output_attentions = inputs
        bs, q_length, dim = shape_list(query)
        k_length = shape_list(key)[1]
        # assert dim == self.dim, 'Dimensions do not match: %s input vs %s configured' % (dim, self.dim)
@@ -253,7 +263,7 @@ class TFMultiHeadSelfAttention(tf.keras.layers.Layer):
        context = unshape(context)  # (bs, q_length, dim)
        context = self.out_lin(context)  # (bs, q_length, dim)

-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            return (context, weights)
        else:
            return (context,)
@@ -293,7 +303,6 @@ class TFTransformerBlock(tf.keras.layers.Layer):
        self.hidden_dim = config.hidden_dim
        self.dropout = tf.keras.layers.Dropout(config.dropout)
        self.activation = config.activation
-        self.output_attentions = config.output_attentions

        assert config.dim % config.n_heads == 0

@@ -317,11 +326,11 @@ class TFTransformerBlock(tf.keras.layers.Layer):
        ffn_output: tf.Tensor(bs, seq_length, dim)
            The output of the transformer block contextualization.
        """
-        x, attn_mask, head_mask = inputs
+        x, attn_mask, head_mask, output_attentions = inputs

        # Self-Attention
-        sa_output = self.attention([x, x, x, attn_mask, head_mask], training=training)
-        if self.output_attentions:
+        sa_output = self.attention([x, x, x, attn_mask, head_mask, output_attentions], training=training)
+        if cast_bool_to_primitive(output_attentions) is True:
            sa_output, sa_weights = sa_output  # (bs, seq_length, dim), (bs, n_heads, seq_length, seq_length)
        else:  # To handle these `output_attention` or `output_hidden_states` cases returning tuples
            # assert type(sa_output) == tuple
@@ -333,7 +342,7 @@ class TFTransformerBlock(tf.keras.layers.Layer):
        ffn_output = self.output_layer_norm(ffn_output + sa_output)  # (bs, seq_length, dim)

        output = (ffn_output,)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            output = (sa_weights,) + output
        return output

@@ -342,7 +351,6 @@ class TFTransformer(tf.keras.layers.Layer):
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)
        self.n_layers = config.n_layers
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

        self.layer = [TFTransformerBlock(config, name="layer_._{}".format(i)) for i in range(config.n_layers)]
@@ -367,7 +375,7 @@ class TFTransformer(tf.keras.layers.Layer):
            Tuple of length n_layers with the attention weights from each layer
            Optional: only if output_attentions=True
        """
-        x, attn_mask, head_mask = inputs
+        x, attn_mask, head_mask, output_attentions = inputs

        all_hidden_states = ()
        all_attentions = ()
@@ -377,10 +385,10 @@ class TFTransformer(tf.keras.layers.Layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_state,)

-            layer_outputs = layer_module([hidden_state, attn_mask, head_mask[i]], training=training)
+            layer_outputs = layer_module([hidden_state, attn_mask, head_mask[i], output_attentions], training=training)
            hidden_state = layer_outputs[-1]

-            if self.output_attentions:
+            if cast_bool_to_primitive(output_attentions) is True:
                assert len(layer_outputs) == 2
                attentions = layer_outputs[0]
                all_attentions = all_attentions + (attentions,)
@@ -394,15 +402,19 @@ class TFTransformer(tf.keras.layers.Layer):
        outputs = (hidden_state,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)


+@keras_serializable
 class TFDistilBertMainLayer(tf.keras.layers.Layer):
+    config_class = DistilBertConfig
+
    def __init__(self, config, **kwargs):
        super().__init__(**kwargs)
        self.num_hidden_layers = config.num_hidden_layers
+        self.output_attentions = config.output_attentions

        self.embeddings = TFEmbeddings(config, name="embeddings")  # Embeddings
        self.transformer = TFTransformer(config, name="transformer")  # Encoder
@@ -416,22 +428,28 @@ class TFDistilBertMainLayer(tf.keras.layers.Layer):
    def _prune_heads(self, heads_to_prune):
        raise NotImplementedError

-    def call(self, inputs, attention_mask=None, head_mask=None, inputs_embeds=None, training=False):
+    def call(
+        self, inputs, attention_mask=None, head_mask=None, inputs_embeds=None, output_attentions=None, training=False
+    ):
        if isinstance(inputs, (tuple, list)):
            input_ids = inputs[0]
            attention_mask = inputs[1] if len(inputs) > 1 else attention_mask
            head_mask = inputs[2] if len(inputs) > 2 else head_mask
            inputs_embeds = inputs[3] if len(inputs) > 3 else inputs_embeds
-            assert len(inputs) <= 4, "Too many inputs."
+            output_attentions = inputs[4] if len(inputs) > 4 else output_attentions
+            assert len(inputs) <= 5, "Too many inputs."
        elif isinstance(inputs, (dict, BatchEncoding)):
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
-            assert len(inputs) <= 4, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 5, "Too many inputs."
        else:
            input_ids = inputs

+        output_attentions = output_attentions if output_attentions is not None else self.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -456,7 +474,9 @@ class TFDistilBertMainLayer(tf.keras.layers.Layer):
            head_mask = [None] * self.num_hidden_layers

        embedding_output = self.embeddings(input_ids, inputs_embeds=inputs_embeds)  # (bs, seq_length, dim)
-        tfmr_output = self.transformer([embedding_output, attention_mask, head_mask], training=training)
+        tfmr_output = self.transformer(
+            [embedding_output, attention_mask, head_mask, output_attentions], training=training
+        )

        return tfmr_output  # last-layer hidden-state, (all hidden_states), (all attentions)

@@ -553,7 +573,7 @@ class TFDistilBertModel(TFDistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -599,7 +619,6 @@ class TFDistilBertLMHead(tf.keras.layers.Layer):
 class TFDistilBertForMaskedLM(TFDistilBertPreTrainedModel):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
-        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.vocab_size = config.vocab_size

@@ -627,7 +646,7 @@ class TFDistilBertForMaskedLM(TFDistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -662,7 +681,7 @@ class TFDistilBertForMaskedLM(TFDistilBertPreTrainedModel):
    the pooled output) e.g. for GLUE tasks. """,
    DISTILBERT_START_DOCSTRING,
 )
-class TFDistilBertForSequenceClassification(TFDistilBertPreTrainedModel):
+class TFDistilBertForSequenceClassification(TFDistilBertPreTrainedModel, TFSequenceClassificationLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -680,8 +699,23 @@ class TFDistilBertForSequenceClassification(TFDistilBertPreTrainedModel):
        self.dropout = tf.keras.layers.Dropout(config.seq_classif_dropout)

    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
-    def call(self, inputs, **kwargs):
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the sequence classification/regression loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
+            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
+
    Returns:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers,DistilBertConfig`) and inputs:
        logits (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, config.num_labels)`):
@@ -691,7 +725,7 @@ class TFDistilBertForSequenceClassification(TFDistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -705,20 +739,33 @@ class TFDistilBertForSequenceClassification(TFDistilBertPreTrainedModel):
        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-cased')
        model = TFDistilBertForSequenceClassification.from_pretrained('distilbert-base-cased')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        logits = outputs[0]
+        labels = tf.reshape(tf.constant(1), (-1, 1)) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, logits = outputs[:2]

        """
-        distilbert_output = self.distilbert(inputs, **kwargs)
+        distilbert_output = self.distilbert(
+            input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        hidden_state = distilbert_output[0]  # (bs, seq_len, dim)
        pooled_output = hidden_state[:, 0]  # (bs, dim)
        pooled_output = self.pre_classifier(pooled_output)  # (bs, dim)
-        pooled_output = self.dropout(pooled_output, training=kwargs.get("training", False))  # (bs, dim)
+        pooled_output = self.dropout(pooled_output, training=training)  # (bs, dim)
        logits = self.classifier(pooled_output)  # (bs, dim)

        outputs = (logits,) + distilbert_output[1:]
-        return outputs  # logits, (hidden_states), (attentions)
+
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), logits, (hidden_states), (attentions)


@add_start_docstrings(
@@ -726,7 +773,7 @@ class TFDistilBertForSequenceClassification(TFDistilBertPreTrainedModel):
    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
    DISTILBERT_START_DOCSTRING,
 )
-class TFDistilBertForTokenClassification(TFDistilBertPreTrainedModel):
+class TFDistilBertForTokenClassification(TFDistilBertPreTrainedModel, TFTokenClassificationLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -738,8 +785,21 @@ class TFDistilBertForTokenClassification(TFDistilBertPreTrainedModel):
        )

    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
-    def call(self, inputs, **kwargs):
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the token classification loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+
    Returns:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers,DistilBertConfig`) and inputs:
        scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`):
@@ -749,7 +809,7 @@ class TFDistilBertForTokenClassification(TFDistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -762,20 +822,162 @@ class TFDistilBertForTokenClassification(TFDistilBertPreTrainedModel):

        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-cased')
        model = TFDistilBertForTokenClassification.from_pretrained('distilbert-base-cased')
-        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        scores = outputs[0]
+        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
+        labels = tf.reshape(tf.constant([1] * tf.size(input_ids).numpy()), (-1, tf.size(input_ids))) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, scores = outputs[:2]
+
        """
-        outputs = self.distilbert(inputs, **kwargs)
+        outputs = self.distilbert(
+            input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        sequence_output = outputs[0]

-        sequence_output = self.dropout(sequence_output, training=kwargs.get("training", False))
+        sequence_output = self.dropout(sequence_output, training=training)
        logits = self.classifier(sequence_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

-        return outputs  # scores, (hidden_states), (attentions)
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), logits, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+    """DistilBert Model with a multiple choice classification head on top (a linear layer on top of
+    the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
+    DISTILBERT_START_DOCSTRING,
+)
+class TFDistilBertForMultipleChoice(TFDistilBertPreTrainedModel, TFMultipleChoiceLoss):
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+
+        self.distilbert = TFDistilBertMainLayer(config, name="distilbert")
+        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
+        self.pre_classifier = tf.keras.layers.Dense(
+            config.dim,
+            kernel_initializer=get_initializer(config.initializer_range),
+            activation="relu",
+            name="pre_classifier",
+        )
+        self.classifier = tf.keras.layers.Dense(
+            1, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
+        )
+
+    @property
+    def dummy_inputs(self):
+        """ Dummy inputs to build the network.
+
+        Returns:
+            tf.Tensor with dummy inputs
+        """
+        return {"input_ids": tf.constant(MULTIPLE_CHOICE_DUMMY_INPUTS)}
+
+    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
+    def call(
+        self,
+        inputs,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
+        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the multiple choice classification loss.
+            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
+            of the input tensors. (see `input_ids` above)
+
+    Return:
+        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        classification_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, num_choices)`:
+            `num_choices` is the size of the second dimension of the input tensors. (see `input_ids` above).
+
+            Classification scores (before SoftMax).
+        hidden_states (:obj:`tuple(tf.Tensor)`, `optional`, returned when :obj:`config.output_hidden_states=True`):
+            tuple of :obj:`tf.Tensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
+            tuple of :obj:`tf.Tensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
+
+    Examples::
+
+        import tensorflow as tf
+        from transformers import DistilBertTokenizer, TFDistilBertForMultipleChoice
+
+        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-uncased')
+        model = TFDistilBertForMultipleChoice.from_pretrained('distilbert-base-uncased')
+        choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]
+
+        input_ids = tf.constant([tokenizer.encode(s, add_special_tokens=True) for s in choices])[None, :] # Batch size 1, 2 choices
+        labels = tf.reshape(tf.constant(1), (-1, 1))
+        outputs = model(input_ids, labels=labels)
+
+        loss, classification_scores = outputs[:2]
+
+        """
+        if isinstance(inputs, (tuple, list)):
+            input_ids = inputs[0]
+            attention_mask = inputs[1] if len(inputs) > 1 else attention_mask
+            head_mask = inputs[2] if len(inputs) > 2 else head_mask
+            inputs_embeds = inputs[3] if len(inputs) > 3 else inputs_embeds
+            assert len(inputs) <= 4, "Too many inputs."
+        elif isinstance(inputs, (dict, BatchEncoding)):
+            input_ids = inputs.get("input_ids")
+            attention_mask = inputs.get("attention_mask", attention_mask)
+            head_mask = inputs.get("head_mask", head_mask)
+            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
+            assert len(inputs) <= 4, "Too many inputs."
+        else:
+            input_ids = inputs
+
+        if input_ids is not None:
+            num_choices = shape_list(input_ids)[1]
+            seq_length = shape_list(input_ids)[2]
+        else:
+            num_choices = shape_list(inputs_embeds)[1]
+            seq_length = shape_list(inputs_embeds)[2]
+
+        flat_input_ids = tf.reshape(input_ids, (-1, seq_length)) if input_ids is not None else None
+        flat_attention_mask = tf.reshape(attention_mask, (-1, seq_length)) if attention_mask is not None else None
+
+        flat_inputs = [
+            flat_input_ids,
+            flat_attention_mask,
+            head_mask,
+            inputs_embeds,
+        ]
+
+        distilbert_output = self.distilbert(flat_inputs, training=training)
+        hidden_state = distilbert_output[0]  # (bs, seq_len, dim)
+        pooled_output = hidden_state[:, 0]  # (bs, dim)
+        pooled_output = self.pre_classifier(pooled_output)  # (bs, dim)
+        pooled_output = self.dropout(pooled_output, training=training)  # (bs, dim)
+        logits = self.classifier(pooled_output)
+        reshaped_logits = tf.reshape(logits, (-1, num_choices))
+
+        outputs = (reshaped_logits,) + distilbert_output[1:]  # add hidden states and attention if they are here
+
+        if labels is not None:
+            loss = self.compute_loss(labels, reshaped_logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)


@add_start_docstrings(
@@ -783,7 +985,7 @@ class TFDistilBertForTokenClassification(TFDistilBertPreTrainedModel):
    the hidden-states output to compute `span start logits` and `span end logits`). """,
    DISTILBERT_START_DOCSTRING,
 )
-class TFDistilBertForQuestionAnswering(TFDistilBertPreTrainedModel):
+class TFDistilBertForQuestionAnswering(TFDistilBertPreTrainedModel, TFQuestionAnsweringLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -795,8 +997,30 @@ class TFDistilBertForQuestionAnswering(TFDistilBertPreTrainedModel):
        self.dropout = tf.keras.layers.Dropout(config.qa_dropout)

    @add_start_docstrings_to_callable(DISTILBERT_INPUTS_DOCSTRING)
-    def call(self, inputs, **kwargs):
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        head_mask=None,
+        inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        cls_index=None,
+        p_mask=None,
+        is_impossible=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        start_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers,DistilBertConfig`) and inputs:
        start_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length,)`):
@@ -808,7 +1032,7 @@ class TFDistilBertForQuestionAnswering(TFDistilBertPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -821,19 +1045,35 @@ class TFDistilBertForQuestionAnswering(TFDistilBertPreTrainedModel):

        tokenizer = DistilBertTokenizer.from_pretrained('distilbert-base-cased')
        model = TFDistilBertForQuestionAnswering.from_pretrained('distilbert-base-cased')
-        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        start_scores, end_scores = outputs[:2]
+        question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+        input_dict = tokenizer.encode_plus(question, text, return_tensors='tf')
+        start_scores, end_scores = model(input_dict)
+
+        all_tokens = tokenizer.convert_ids_to_tokens(input_dict["input_ids"].numpy()[0])
+        answer = ' '.join(all_tokens[tf.math.argmax(start_scores, 1)[0] : tf.math.argmax(end_scores, 1)[0]+1])

        """
-        distilbert_output = self.distilbert(inputs, **kwargs)
+        distilbert_output = self.distilbert(
+            input_ids,
+            attention_mask=attention_mask,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            training=training,
+        )

        hidden_states = distilbert_output[0]  # (bs, max_query_len, dim)
-        hidden_states = self.dropout(hidden_states, training=kwargs.get("training", False))  # (bs, max_query_len, dim)
+        hidden_states = self.dropout(hidden_states, training=training)  # (bs, max_query_len, dim)
        logits = self.qa_outputs(hidden_states)  # (bs, max_query_len, 2)
        start_logits, end_logits = tf.split(logits, 2, axis=-1)
        start_logits = tf.squeeze(start_logits, axis=-1)
        end_logits = tf.squeeze(end_logits, axis=-1)

        outputs = (start_logits, end_logits,) + distilbert_output[1:]
-        return outputs  # start_logits, end_logits, (hidden_states), (attentions)
+
+        if start_positions is not None and end_positions is not None:
+            labels = {"start_position": start_positions}
+            labels["end_position"] = end_positions
+            loss = self.compute_loss(labels, outputs[:2])
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)
@@ -6,7 +6,13 @@ from transformers import ElectraConfig

 from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
 from .modeling_tf_bert import ACT2FN, TFBertEncoder, TFBertPreTrainedModel
-from .modeling_tf_utils import get_initializer, shape_list
+from .modeling_tf_utils import (
+    TFQuestionAnsweringLoss,
+    TFTokenClassificationLoss,
+    get_initializer,
+    keras_serializable,
+    shape_list,
+)
 from .tokenization_utils import BatchEncoding


@@ -194,6 +200,7 @@ class TFElectraPreTrainedModel(TFBertPreTrainedModel):
        return head_mask


+@keras_serializable
 class TFElectraMainLayer(TFElectraPreTrainedModel):

    config_class = ElectraConfig
@@ -228,6 +235,7 @@ class TFElectraMainLayer(TFElectraPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
        training=False,
    ):
        if isinstance(inputs, (tuple, list)):
@@ -237,7 +245,8 @@ class TFElectraMainLayer(TFElectraPreTrainedModel):
            position_ids = inputs[3] if len(inputs) > 3 else position_ids
            head_mask = inputs[4] if len(inputs) > 4 else head_mask
            inputs_embeds = inputs[5] if len(inputs) > 5 else inputs_embeds
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs[6] if len(inputs) > 6 else output_attentions
+            assert len(inputs) <= 7, "Too many inputs."
        elif isinstance(inputs, (dict, BatchEncoding)):
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
@@ -245,10 +254,13 @@ class TFElectraMainLayer(TFElectraPreTrainedModel):
            position_ids = inputs.get("position_ids", position_ids)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 7, "Too many inputs."
        else:
            input_ids = inputs

+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -271,7 +283,9 @@ class TFElectraMainLayer(TFElectraPreTrainedModel):
        if hasattr(self, "embeddings_project"):
            hidden_states = self.embeddings_project(hidden_states, training=training)

-        hidden_states = self.encoder([hidden_states, extended_attention_mask, head_mask], training=training)
+        hidden_states = self.encoder(
+            [hidden_states, extended_attention_mask, head_mask, output_attentions], training=training
+        )

        return hidden_states

@@ -365,7 +379,7 @@ class TFElectraModel(TFElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -414,6 +428,7 @@ class TFElectraForPreTraining(TFElectraPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
        training=False,
    ):
        r"""
@@ -426,7 +441,7 @@ class TFElectraForPreTraining(TFElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -445,7 +460,14 @@ class TFElectraForPreTraining(TFElectraPreTrainedModel):
        """

        discriminator_hidden_states = self.electra(
-            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, training=training
+            input_ids,
+            attention_mask,
+            token_type_ids,
+            position_ids,
+            head_mask,
+            inputs_embeds,
+            output_attentions,
+            training=training,
        )
        discriminator_sequence_output = discriminator_hidden_states[0]
        logits = self.discriminator_predictions(discriminator_sequence_output)
@@ -507,6 +529,7 @@ class TFElectraForMaskedLM(TFElectraPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
        training=False,
    ):
        r"""
@@ -519,7 +542,7 @@ class TFElectraForMaskedLM(TFElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -539,7 +562,14 @@ class TFElectraForMaskedLM(TFElectraPreTrainedModel):
        """

        generator_hidden_states = self.electra(
-            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, training=training
+            input_ids,
+            attention_mask,
+            token_type_ids,
+            position_ids,
+            head_mask,
+            inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
        )
        generator_sequence_output = generator_hidden_states[0]
        prediction_scores = self.generator_predictions(generator_sequence_output, training=training)
@@ -557,13 +587,15 @@ Electra model with a token classification head on top.
 Both the discriminator and generator may be loaded into this model.""",
    ELECTRA_START_DOCSTRING,
 )
-class TFElectraForTokenClassification(TFElectraPreTrainedModel):
+class TFElectraForTokenClassification(TFElectraPreTrainedModel, TFTokenClassificationLoss):
    def __init__(self, config, **kwargs):
        super().__init__(config, **kwargs)

        self.electra = TFElectraMainLayer(config, name="electra")
        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
-        self.classifier = tf.keras.layers.Dense(config.num_labels, name="classifier")
+        self.classifier = tf.keras.layers.Dense(
+            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
+        )

    @add_start_docstrings_to_callable(ELECTRA_INPUTS_DOCSTRING)
    def call(
@@ -574,9 +606,15 @@ class TFElectraForTokenClassification(TFElectraPreTrainedModel):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
        training=False,
    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the token classification loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+
    Returns:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.ElectraConfig`) and inputs:
        scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`):
@@ -586,7 +624,7 @@ class TFElectraForTokenClassification(TFElectraPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -599,18 +637,134 @@ class TFElectraForTokenClassification(TFElectraPreTrainedModel):

        tokenizer = ElectraTokenizer.from_pretrained('google/electra-small-discriminator')
        model = TFElectraForTokenClassification.from_pretrained('google/electra-small-discriminator')
-        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        scores = outputs[0]
+        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
+        labels = tf.reshape(tf.constant([1] * tf.size(input_ids).numpy()), (-1, tf.size(input_ids))) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, scores = outputs[:2]
+
        """

        discriminator_hidden_states = self.electra(
-            input_ids, attention_mask, token_type_ids, position_ids, head_mask, inputs_embeds, training=training
+            input_ids,
+            attention_mask,
+            token_type_ids,
+            position_ids,
+            head_mask,
+            inputs_embeds,
+            output_attentions,
+            training=training,
        )
        discriminator_sequence_output = discriminator_hidden_states[0]
        discriminator_sequence_output = self.dropout(discriminator_sequence_output)
        logits = self.classifier(discriminator_sequence_output)
-        output = (logits,)
-        output += discriminator_hidden_states[1:]

-        return output  # (loss), scores, (hidden_states), (attentions)
+        outputs = (logits,) + discriminator_hidden_states[1:]
+
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), scores, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+    """Electra Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
+    the hidden-states output to compute `span start logits` and `span end logits`). """,
+    ELECTRA_START_DOCSTRING,
+)
+class TFElectraForQuestionAnswering(TFElectraPreTrainedModel, TFQuestionAnsweringLoss):
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        self.num_labels = config.num_labels
+
+        self.electra = TFElectraMainLayer(config, name="electra")
+        self.qa_outputs = tf.keras.layers.Dense(
+            config.num_labels, kernel_initializer=get_initializer(config.initializer_range), name="qa_outputs"
+        )
+
+    @add_start_docstrings_to_callable(ELECTRA_INPUTS_DOCSTRING)
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        cls_index=None,
+        p_mask=None,
+        is_impossible=None,
+        output_attentions=None,
+        training=False,
+    ):
+        r"""
+        start_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+
+    Return:
+        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        start_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length,)`):
+            Span-start scores (before SoftMax).
+        end_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length,)`):
+            Span-end scores (before SoftMax).
+        hidden_states (:obj:`tuple(tf.Tensor)`, `optional`, returned when :obj:`config.output_hidden_states=True`):
+            tuple of :obj:`tf.Tensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
+            tuple of :obj:`tf.Tensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
+
+    Examples::
+
+        import tensorflow as tf
+        from transformers import ElectraTokenizer, TFElectraForQuestionAnswering
+
+        tokenizer = ElectraTokenizer.from_pretrained('google/electra-small-generator')
+        model = TFElectraForQuestionAnswering.from_pretrained('google/electra-small-generator')
+
+        question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+        input_dict = tokenizer.encode_plus(question, text, return_tensors='tf')
+        start_scores, end_scores = model(input_dict)
+
+        all_tokens = tokenizer.convert_ids_to_tokens(input_dict["input_ids"].numpy()[0])
+        answer = ' '.join(all_tokens[tf.math.argmax(start_scores, 1)[0] : tf.math.argmax(end_scores, 1)[0]+1])
+
+        """
+        discriminator_hidden_states = self.electra(
+            input_ids,
+            attention_mask,
+            token_type_ids,
+            position_ids,
+            head_mask,
+            inputs_embeds,
+            output_attentions,
+            training=training,
+        )
+        discriminator_sequence_output = discriminator_hidden_states[0]
+
+        logits = self.qa_outputs(discriminator_sequence_output)
+        start_logits, end_logits = tf.split(logits, 2, axis=-1)
+        start_logits = tf.squeeze(start_logits, axis=-1)
+        end_logits = tf.squeeze(end_logits, axis=-1)
+
+        outputs = (start_logits, end_logits,) + discriminator_hidden_states[1:]
+
+        if start_positions is not None and end_positions is not None:
+            labels = {"start_position": start_positions}
+            labels["end_position"] = end_positions
+            loss = self.compute_loss(labels, outputs[:2])
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)
@@ -22,13 +22,16 @@ import tensorflow as tf

 from .configuration_flaubert import FlaubertConfig
 from .file_utils import add_start_docstrings
+from .modeling_tf_utils import keras_serializable, shape_list
 from .modeling_tf_xlm import (
+    TFXLMForMultipleChoice,
+    TFXLMForQuestionAnsweringSimple,
    TFXLMForSequenceClassification,
+    TFXLMForTokenClassification,
    TFXLMMainLayer,
    TFXLMModel,
    TFXLMWithLMHeadModel,
    get_masks,
-    shape_list,
 )
 from .tokenization_utils import BatchEncoding

@@ -112,6 +115,7 @@ class TFFlaubertModel(TFXLMModel):
        self.transformer = TFFlaubertMainLayer(config, name="transformer")


+@keras_serializable
 class TFFlaubertMainLayer(TFXLMMainLayer):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
@@ -130,6 +134,7 @@ class TFFlaubertMainLayer(TFXLMMainLayer):
        head_mask=None,
        inputs_embeds=None,
        training=False,
+        output_attentions=False,
    ):
        # removed: src_enc=None, src_len=None
        if isinstance(inputs, (tuple, list)):
@@ -251,7 +256,7 @@ class TFFlaubertMainLayer(TFXLMMainLayer):
            if not self.pre_norm:
                attn_outputs = self.attentions[i]([tensor, attn_mask, None, cache, head_mask[i]], training=training)
                attn = attn_outputs[0]
-                if self.output_attentions:
+                if output_attentions:
                    attentions = attentions + (attn_outputs[1],)
                attn = self.dropout(attn, training=training)
                tensor = tensor + attn
@@ -262,7 +267,7 @@ class TFFlaubertMainLayer(TFXLMMainLayer):
                    [tensor_normalized, attn_mask, None, cache, head_mask[i]], training=training
                )
                attn = attn_outputs[0]
-                if self.output_attentions:
+                if output_attentions:
                    attentions = attentions + (attn_outputs[1],)
                attn = self.dropout(attn, training=training)
                tensor = tensor + attn
@@ -298,7 +303,7 @@ class TFFlaubertMainLayer(TFXLMMainLayer):
        outputs = (tensor,)
        if self.output_hidden_states:
            outputs = outputs + (hidden_states,)
-        if self.output_attentions:
+        if output_attentions:
            outputs = outputs + (attentions,)
        return outputs  # outputs, (hidden_states), (attentions)

@@ -327,3 +332,38 @@ class TFFlaubertForSequenceClassification(TFXLMForSequenceClassification):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.transformer = TFFlaubertMainLayer(config, name="transformer")
+
+
+@add_start_docstrings(
+    """Flaubert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
+    the hidden-states output to compute `span start logits` and `span end logits`). """,
+    FLAUBERT_START_DOCSTRING,
+)
+class TFFlaubertForQuestionAnsweringSimple(TFXLMForQuestionAnsweringSimple):
+    config_class = FlaubertConfig
+
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        self.transformer = TFFlaubertMainLayer(config, name="transformer")
+
+
+@add_start_docstrings(
+    """Flaubert Model with a token classification head on top (a linear layer on top of
+    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
+    FLAUBERT_START_DOCSTRING,
+)
+class TFFlaubertForTokenClassification(TFXLMForTokenClassification):
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        self.transformer = TFFlaubertMainLayer(config, name="transformer")
+
+
+@add_start_docstrings(
+    """Flaubert Model with a multiple choice classification head on top (a linear layer on top of
+    the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
+    FLAUBERT_START_DOCSTRING,
+)
+class TFFlaubertForMultipleChoice(TFXLMForMultipleChoice):
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        self.transformer = TFFlaubertMainLayer(config, name="transformer")
@@ -28,6 +28,7 @@ from .modeling_tf_utils import (
    TFPreTrainedModel,
    TFSequenceSummary,
    TFSharedEmbeddings,
+    cast_bool_to_primitive,
    get_initializer,
    keras_serializable,
    shape_list,
@@ -63,7 +64,6 @@ def gelu(x):
 class TFAttention(tf.keras.layers.Layer):
    def __init__(self, nx, n_ctx, config, scale=False, **kwargs):
        super().__init__(**kwargs)
-        self.output_attentions = config.output_attentions

        n_state = nx  # in Attention: n_state=768 (nx=n_embd)
        # [switch nx => n_state from Block to Attention to keep identical to TF implem]
@@ -93,7 +93,7 @@ class TFAttention(tf.keras.layers.Layer):
        return tf.cast(m, dtype)

    def _attn(self, inputs, training=False):
-        q, k, v, attention_mask, head_mask = inputs
+        q, k, v, attention_mask, head_mask, output_attentions = inputs
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        if self.scale:
@@ -118,7 +118,7 @@ class TFAttention(tf.keras.layers.Layer):
            w = w * head_mask

        outputs = [tf.matmul(w, v)]
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            outputs.append(w)
        return outputs

@@ -135,7 +135,7 @@ class TFAttention(tf.keras.layers.Layer):
        return tf.transpose(x, (0, 2, 1, 3))  # (batch, head, seq_length, head_features)

    def call(self, inputs, training=False):
-        x, layer_past, attention_mask, head_mask, use_cache = inputs
+        x, layer_past, attention_mask, head_mask, use_cache, output_attentions = inputs

        x = self.c_attn(x)
        query, key, value = tf.split(x, 3, axis=2)
@@ -148,20 +148,12 @@ class TFAttention(tf.keras.layers.Layer):
            value = tf.concat([past_value, value], axis=-2)

        # to cope with keras serialization
-        # we need to cast `use_cache` to correct bool
-        # if it is a tensor
-        if tf.is_tensor(use_cache):
-            if hasattr(use_cache, "numpy"):
-                use_cache = bool(use_cache.numpy())
-            else:
-                use_cache = True
-
-        if use_cache is True:
+        if cast_bool_to_primitive(use_cache, True) is True:
            present = tf.stack([key, value], axis=0)
        else:
            present = (None,)

-        attn_outputs = self._attn([query, key, value, attention_mask, head_mask], training=training)
+        attn_outputs = self._attn([query, key, value, attention_mask, head_mask, output_attentions], training=training)
        a = attn_outputs[0]

        a = self.merge_heads(a)
@@ -198,10 +190,12 @@ class TFBlock(tf.keras.layers.Layer):
        self.mlp = TFMLP(4 * nx, config, name="mlp")

    def call(self, inputs, training=False):
-        x, layer_past, attention_mask, head_mask, use_cache = inputs
+        x, layer_past, attention_mask, head_mask, use_cache, output_attentions = inputs

        a = self.ln_1(x)
-        output_attn = self.attn([a, layer_past, attention_mask, head_mask, use_cache], training=training)
+        output_attn = self.attn(
+            [a, layer_past, attention_mask, head_mask, use_cache, output_attentions], training=training
+        )
        a = output_attn[0]  # output_attn: a, present, (attentions)
        x = x + a

@@ -219,8 +213,8 @@ class TFGPT2MainLayer(tf.keras.layers.Layer):

    def __init__(self, config, *inputs, **kwargs):
        super().__init__(*inputs, **kwargs)
-        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions
+        self.output_hidden_states = config.output_hidden_states
        self.num_hidden_layers = config.n_layer
        self.vocab_size = config.vocab_size
        self.n_embd = config.n_embd
@@ -261,6 +255,7 @@ class TFGPT2MainLayer(tf.keras.layers.Layer):
        inputs_embeds=None,
        use_cache=True,
        training=False,
+        output_attentions=None,
    ):
        if isinstance(inputs, (tuple, list)):
            input_ids = inputs[0]
@@ -271,7 +266,8 @@ class TFGPT2MainLayer(tf.keras.layers.Layer):
            head_mask = inputs[5] if len(inputs) > 5 else head_mask
            inputs_embeds = inputs[6] if len(inputs) > 6 else inputs_embeds
            use_cache = inputs[7] if len(inputs) > 7 else use_cache
-            assert len(inputs) <= 8, "Too many inputs."
+            output_attentions = inputs[8] if len(inputs) > 7 else output_attentions
+            assert len(inputs) <= 9, "Too many inputs."
        elif isinstance(inputs, (dict, BatchEncoding)):
            input_ids = inputs.get("input_ids")
            past = inputs.get("past", past)
@@ -281,10 +277,13 @@ class TFGPT2MainLayer(tf.keras.layers.Layer):
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
            use_cache = inputs.get("use_cache", use_cache)
-            assert len(inputs) <= 8, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 9, "Too many inputs."
        else:
            input_ids = inputs

+        output_attentions = output_attentions if output_attentions is not None else self.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -355,12 +354,15 @@ class TFGPT2MainLayer(tf.keras.layers.Layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (tf.reshape(hidden_states, output_shape),)

-            outputs = block([hidden_states, layer_past, attention_mask, head_mask[i], use_cache], training=training)
+            outputs = block(
+                [hidden_states, layer_past, attention_mask, head_mask[i], use_cache, output_attentions],
+                training=training,
+            )

            hidden_states, present = outputs[:2]
            presents = presents + (present,)

-            if self.output_attentions:
+            if cast_bool_to_primitive(output_attentions) is True:
                all_attentions.append(outputs[2])

        hidden_states = self.ln_f(hidden_states)
@@ -376,7 +378,7 @@ class TFGPT2MainLayer(tf.keras.layers.Layer):
            outputs = outputs + (presents,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            # let the number of heads free (-1) so we can extract attention even after head pruning
            attention_output_shape = input_shape[:-1] + [-1] + shape_list(all_attentions[0])[-2:]
            all_attentions = tuple(tf.reshape(t, attention_output_shape) for t in all_attentions)
@@ -493,7 +495,7 @@ class TFGPT2Model(TFGPT2PreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -552,7 +554,7 @@ class TFGPT2LMHeadModel(TFGPT2PreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -614,6 +616,7 @@ class TFGPT2DoubleHeadsModel(TFGPT2PreTrainedModel):
        inputs_embeds=None,
        mc_token_ids=None,
        use_cache=True,
+        output_attentions=None,
        training=False,
    ):
        r"""
@@ -636,7 +639,7 @@ class TFGPT2DoubleHeadsModel(TFGPT2PreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -681,7 +684,8 @@ class TFGPT2DoubleHeadsModel(TFGPT2PreTrainedModel):
            inputs_embeds = inputs[6] if len(inputs) > 6 else inputs_embeds
            mc_token_ids = inputs[7] if len(inputs) > 7 else mc_token_ids
            use_cache = inputs[8] if len(inputs) > 8 else use_cache
-            assert len(inputs) <= 9, "Too many inputs."
+            output_attentions = inputs[9] if len(inputs) > 8 else output_attentions
+            assert len(inputs) <= 10, "Too many inputs."
        elif isinstance(inputs, dict):
            input_ids = inputs.get("input_ids")
            past = inputs.get("past", past)
@@ -692,7 +696,8 @@ class TFGPT2DoubleHeadsModel(TFGPT2PreTrainedModel):
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
            mc_token_ids = inputs.get("mc_token_ids", mc_token_ids)
            use_cache = inputs.get("use_cache", use_cache)
-            assert len(inputs) <= 9, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 10, "Too many inputs."
        else:
            input_ids = inputs

@@ -717,6 +722,7 @@ class TFGPT2DoubleHeadsModel(TFGPT2PreTrainedModel):
            head_mask,
            inputs_embeds,
            use_cache,
+            output_attentions,
        ]

        transformer_outputs = self.transformer(flat_inputs, training=training)
@@ -28,7 +28,9 @@ from .modeling_tf_utils import (
    TFPreTrainedModel,
    TFSequenceSummary,
    TFSharedEmbeddings,
+    cast_bool_to_primitive,
    get_initializer,
+    keras_serializable,
    shape_list,
 )
 from .tokenization_utils import BatchEncoding
@@ -69,7 +71,6 @@ ACT_FNS = {
 class TFAttention(tf.keras.layers.Layer):
    def __init__(self, nx, n_ctx, config, scale=False, **kwargs):
        super().__init__(**kwargs)
-        self.output_attentions = config.output_attentions

        n_state = nx  # in Attention: n_state=768 (nx=n_embd)
        # [switch nx => n_state from Block to Attention to keep identical to TF implem]
@@ -99,7 +100,7 @@ class TFAttention(tf.keras.layers.Layer):
        return tf.cast(m, dtype)

    def _attn(self, inputs, training=False):
-        q, k, v, attention_mask, head_mask = inputs
+        q, k, v, attention_mask, head_mask, output_attentions = inputs
        # q, k, v have shape [batch, heads, sequence, features]
        w = tf.matmul(q, k, transpose_b=True)
        if self.scale:
@@ -124,7 +125,7 @@ class TFAttention(tf.keras.layers.Layer):
            w = w * head_mask

        outputs = [tf.matmul(w, v)]
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            outputs.append(w)
        return outputs

@@ -141,7 +142,7 @@ class TFAttention(tf.keras.layers.Layer):
        return tf.transpose(x, (0, 2, 1, 3))  # (batch, head, seq_length, head_features)

    def call(self, inputs, training=False):
-        x, attention_mask, head_mask = inputs
+        x, attention_mask, head_mask, output_attentions = inputs

        x = self.c_attn(x)
        query, key, value = tf.split(x, 3, axis=2)
@@ -149,7 +150,7 @@ class TFAttention(tf.keras.layers.Layer):
        key = self.split_heads(key)
        value = self.split_heads(value)

-        attn_outputs = self._attn([query, key, value, attention_mask, head_mask], training=training)
+        attn_outputs = self._attn([query, key, value, attention_mask, head_mask, output_attentions], training=training)
        a = attn_outputs[0]

        a = self.merge_heads(a)
@@ -186,9 +187,9 @@ class TFBlock(tf.keras.layers.Layer):
        self.ln_2 = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_epsilon, name="ln_2")

    def call(self, inputs, training=False):
-        x, attention_mask, head_mask = inputs
+        x, attention_mask, head_mask, output_attentions = inputs

-        output_attn = self.attn([x, attention_mask, head_mask], training=training)
+        output_attn = self.attn([x, attention_mask, head_mask, output_attentions], training=training)
        a = output_attn[0]  # output_attn: a, (attentions)

        n = self.ln_1(x + a)
@@ -199,7 +200,10 @@ class TFBlock(tf.keras.layers.Layer):
        return outputs  # x, (attentions)


+@keras_serializable
 class TFOpenAIGPTMainLayer(tf.keras.layers.Layer):
+    config_class = OpenAIGPTConfig
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(*inputs, **kwargs)
        self.output_hidden_states = config.output_hidden_states
@@ -240,6 +244,7 @@ class TFOpenAIGPTMainLayer(tf.keras.layers.Layer):
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
+        output_attentions=None,
        training=False,
    ):
        if isinstance(inputs, (tuple, list)):
@@ -249,7 +254,8 @@ class TFOpenAIGPTMainLayer(tf.keras.layers.Layer):
            position_ids = inputs[3] if len(inputs) > 3 else position_ids
            head_mask = inputs[4] if len(inputs) > 4 else head_mask
            inputs_embeds = inputs[5] if len(inputs) > 5 else inputs_embeds
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs[6] if len(inputs) > 6 else output_attentions
+            assert len(inputs) <= 7, "Too many inputs."
        elif isinstance(inputs, (dict, BatchEncoding)):
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
@@ -257,10 +263,13 @@ class TFOpenAIGPTMainLayer(tf.keras.layers.Layer):
            position_ids = inputs.get("position_ids", position_ids)
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
-            assert len(inputs) <= 6, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 7, "Too many inputs."
        else:
            input_ids = inputs

+        output_attentions = output_attentions if output_attentions is not None else self.output_attentions
+
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
@@ -325,9 +334,9 @@ class TFOpenAIGPTMainLayer(tf.keras.layers.Layer):
            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (tf.reshape(hidden_states, output_shape),)

-            outputs = block([hidden_states, attention_mask, head_mask[i]], training=training)
+            outputs = block([hidden_states, attention_mask, head_mask[i], output_attentions], training=training)
            hidden_states = outputs[0]
-            if self.output_attentions:
+            if cast_bool_to_primitive(output_attentions) is True:
                all_attentions.append(outputs[1])

        hidden_states = tf.reshape(hidden_states, output_shape)
@@ -338,7 +347,7 @@ class TFOpenAIGPTMainLayer(tf.keras.layers.Layer):
        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
-        if self.output_attentions:
+        if cast_bool_to_primitive(output_attentions) is True:
            # let the number of heads free (-1) so we can extract attention even after head pruning
            attention_output_shape = input_shape[:-1] + [-1] + shape_list(all_attentions[0])[-2:]
            all_attentions = tuple(tf.reshape(t, attention_output_shape) for t in all_attentions)
@@ -444,7 +453,7 @@ class TFOpenAIGPTModel(TFOpenAIGPTPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -492,7 +501,7 @@ class TFOpenAIGPTLMHeadModel(TFOpenAIGPTPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -551,6 +560,7 @@ class TFOpenAIGPTDoubleHeadsModel(TFOpenAIGPTPreTrainedModel):
        head_mask=None,
        inputs_embeds=None,
        mc_token_ids=None,
+        output_attentions=None,
        training=False,
    ):
        r"""
@@ -573,7 +583,7 @@ class TFOpenAIGPTDoubleHeadsModel(TFOpenAIGPTPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            Tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

@@ -613,7 +623,8 @@ class TFOpenAIGPTDoubleHeadsModel(TFOpenAIGPTPreTrainedModel):
            head_mask = inputs[4] if len(inputs) > 4 else head_mask
            inputs_embeds = inputs[5] if len(inputs) > 5 else inputs_embeds
            mc_token_ids = inputs[6] if len(inputs) > 6 else mc_token_ids
-            assert len(inputs) <= 7, "Too many inputs."
+            output_attentions = inputs[7] if len(inputs) > 7 else output_attentions
+            assert len(inputs) <= 8, "Too many inputs."
        elif isinstance(inputs, dict):
            input_ids = inputs.get("input_ids")
            attention_mask = inputs.get("attention_mask", attention_mask)
@@ -622,7 +633,8 @@ class TFOpenAIGPTDoubleHeadsModel(TFOpenAIGPTPreTrainedModel):
            head_mask = inputs.get("head_mask", head_mask)
            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
            mc_token_ids = inputs.get("mc_token_ids", mc_token_ids)
-            assert len(inputs) <= 7, "Too many inputs."
+            output_attentions = inputs.get("output_attentions", output_attentions)
+            assert len(inputs) <= 8, "Too many inputs."
        else:
            input_ids = inputs

@@ -645,6 +657,7 @@ class TFOpenAIGPTDoubleHeadsModel(TFOpenAIGPTPreTrainedModel):
            flat_position_ids,
            head_mask,
            inputs_embeds,
+            output_attentions,
        ]

        transformer_outputs = self.transformer(flat_inputs, training=training)
@@ -21,9 +21,18 @@ import logging
 import tensorflow as tf

 from .configuration_roberta import RobertaConfig
-from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
+from .file_utils import MULTIPLE_CHOICE_DUMMY_INPUTS, add_start_docstrings, add_start_docstrings_to_callable
 from .modeling_tf_bert import TFBertEmbeddings, TFBertMainLayer, gelu
-from .modeling_tf_utils import TFPreTrainedModel, get_initializer, shape_list
+from .modeling_tf_utils import (
+    TFMultipleChoiceLoss,
+    TFPreTrainedModel,
+    TFQuestionAnsweringLoss,
+    TFSequenceClassificationLoss,
+    TFTokenClassificationLoss,
+    get_initializer,
+    keras_serializable,
+    shape_list,
+)


 logger = logging.getLogger(__name__)
@@ -82,6 +91,7 @@ class TFRobertaEmbeddings(TFBertEmbeddings):
        return super()._embedding([input_ids, position_ids, token_type_ids, inputs_embeds], training=training)


+@keras_serializable
 class TFRobertaMainLayer(TFBertMainLayer):
    """
    Same as TFBertMainLayer but uses TFRobertaEmbeddings.
@@ -203,7 +213,7 @@ class TFRobertaModel(TFRobertaPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -279,7 +289,7 @@ class TFRobertaForMaskedLM(TFRobertaPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -311,7 +321,7 @@ class TFRobertaClassificationHead(tf.keras.layers.Layer):
    """Head for sentence-level classification tasks."""

    def __init__(self, config, **kwargs):
-        super().__init__(config, **kwargs)
+        super().__init__(**kwargs)
        self.dense = tf.keras.layers.Dense(
            config.hidden_size,
            kernel_initializer=get_initializer(config.initializer_range),
@@ -337,7 +347,7 @@ class TFRobertaClassificationHead(tf.keras.layers.Layer):
    on top of the pooled output) e.g. for GLUE tasks. """,
    ROBERTA_START_DOCSTRING,
 )
-class TFRobertaForSequenceClassification(TFRobertaPreTrainedModel):
+class TFRobertaForSequenceClassification(TFRobertaPreTrainedModel, TFSequenceClassificationLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -346,7 +356,18 @@ class TFRobertaForSequenceClassification(TFRobertaPreTrainedModel):
        self.classifier = TFRobertaClassificationHead(config, name="classifier")

    @add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
-    def call(self, inputs, **kwargs):
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
@@ -357,7 +378,7 @@ class TFRobertaForSequenceClassification(TFRobertaPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -370,20 +391,166 @@ class TFRobertaForSequenceClassification(TFRobertaPreTrainedModel):

        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
        model = TFRobertaForSequenceClassification.from_pretrained('roberta-base')
-        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
-        labels = tf.constant([1])[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        logits = outputs[0]
+        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute"))[None, :]  # Batch size 1
+        labels = tf.reshape(tf.constant(1), (-1, 1)) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, logits = outputs[:2]

        """
-        outputs = self.roberta(inputs, **kwargs)
+        outputs = self.roberta(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        sequence_output = outputs[0]
-        logits = self.classifier(sequence_output, training=kwargs.get("training", False))
+        logits = self.classifier(sequence_output, training=training)

        outputs = (logits,) + outputs[2:]

-        return outputs  # logits, (hidden_states), (attentions)
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), logits, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+    """Roberta Model with a multiple choice classification head on top (a linear layer on top of
+    the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
+    ROBERTA_START_DOCSTRING,
+)
+class TFRobertaForMultipleChoice(TFRobertaPreTrainedModel, TFMultipleChoiceLoss):
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+
+        self.roberta = TFBertMainLayer(config, name="roberta")
+        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
+        self.classifier = tf.keras.layers.Dense(
+            1, kernel_initializer=get_initializer(config.initializer_range), name="classifier"
+        )
+
+    @property
+    def dummy_inputs(self):
+        """ Dummy inputs to build the network.
+
+        Returns:
+            tf.Tensor with dummy inputs
+        """
+        return {"input_ids": tf.constant(MULTIPLE_CHOICE_DUMMY_INPUTS)}
+
+    @add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
+    def call(
+        self,
+        inputs,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
+        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the multiple choice classification loss.
+            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
+            of the input tensors. (see `input_ids` above)
+
+    Return:
+        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        classification_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, num_choices)`:
+            `num_choices` is the size of the second dimension of the input tensors. (see `input_ids` above).
+
+            Classification scores (before SoftMax).
+        hidden_states (:obj:`tuple(tf.Tensor)`, `optional`, returned when :obj:`config.output_hidden_states=True`):
+            tuple of :obj:`tf.Tensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
+            tuple of :obj:`tf.Tensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:
+
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
+
+    Examples::
+
+        import tensorflow as tf
+        from transformers import RobertaTokenizer, TFRobertaForMultipleChoice
+
+        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+        model = TFRobertaForMultipleChoice.from_pretrained('roberta-base')
+        choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]
+
+        input_ids = tf.constant([tokenizer.encode(s, add_special_tokens=True) for s in choices])[None, :] # Batch size 1, 2 choices
+        labels = tf.reshape(tf.constant(1), (-1, 1))
+        outputs = model(input_ids, labels=labels)
+
+        loss, classification_scores = outputs[:2]
+
+        """
+        if isinstance(inputs, (tuple, list)):
+            input_ids = inputs[0]
+            attention_mask = inputs[1] if len(inputs) > 1 else attention_mask
+            token_type_ids = inputs[2] if len(inputs) > 2 else token_type_ids
+            position_ids = inputs[3] if len(inputs) > 3 else position_ids
+            head_mask = inputs[4] if len(inputs) > 4 else head_mask
+            inputs_embeds = inputs[5] if len(inputs) > 5 else inputs_embeds
+            assert len(inputs) <= 6, "Too many inputs."
+        elif isinstance(inputs, dict):
+            input_ids = inputs.get("input_ids")
+            attention_mask = inputs.get("attention_mask", attention_mask)
+            token_type_ids = inputs.get("token_type_ids", token_type_ids)
+            position_ids = inputs.get("position_ids", position_ids)
+            head_mask = inputs.get("head_mask", head_mask)
+            inputs_embeds = inputs.get("inputs_embeds", inputs_embeds)
+            assert len(inputs) <= 6, "Too many inputs."
+        else:
+            input_ids = inputs
+
+        if input_ids is not None:
+            num_choices = shape_list(input_ids)[1]
+            seq_length = shape_list(input_ids)[2]
+        else:
+            num_choices = shape_list(inputs_embeds)[1]
+            seq_length = shape_list(inputs_embeds)[2]
+
+        flat_input_ids = tf.reshape(input_ids, (-1, seq_length)) if input_ids is not None else None
+        flat_attention_mask = tf.reshape(attention_mask, (-1, seq_length)) if attention_mask is not None else None
+        flat_token_type_ids = tf.reshape(token_type_ids, (-1, seq_length)) if token_type_ids is not None else None
+        flat_position_ids = tf.reshape(position_ids, (-1, seq_length)) if position_ids is not None else None
+
+        flat_inputs = [
+            flat_input_ids,
+            flat_attention_mask,
+            flat_token_type_ids,
+            flat_position_ids,
+            head_mask,
+            inputs_embeds,
+        ]
+
+        outputs = self.roberta(flat_inputs, training=training)
+
+        pooled_output = outputs[1]
+
+        pooled_output = self.dropout(pooled_output, training=training)
+        logits = self.classifier(pooled_output)
+        reshaped_logits = tf.reshape(logits, (-1, num_choices))
+
+        outputs = (reshaped_logits,) + outputs[2:]  # add hidden states and attention if they are here
+
+        if labels is not None:
+            loss = self.compute_loss(labels, reshaped_logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)


@add_start_docstrings(
@@ -391,7 +558,7 @@ class TFRobertaForSequenceClassification(TFRobertaPreTrainedModel):
    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
    ROBERTA_START_DOCSTRING,
 )
-class TFRobertaForTokenClassification(TFRobertaPreTrainedModel):
+class TFRobertaForTokenClassification(TFRobertaPreTrainedModel, TFTokenClassificationLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -403,8 +570,23 @@ class TFRobertaForTokenClassification(TFRobertaPreTrainedModel):
        )

    @add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
-    def call(self, inputs, **kwargs):
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the token classification loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
        scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`):
@@ -414,7 +596,7 @@ class TFRobertaForTokenClassification(TFRobertaPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -428,27 +610,41 @@ class TFRobertaForTokenClassification(TFRobertaPreTrainedModel):
        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
        model = TFRobertaForTokenClassification.from_pretrained('roberta-base')
        input_ids = tf.constant(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True))[None, :]  # Batch size 1
-        outputs = model(input_ids)
-        scores = outputs[0]
+        labels = tf.reshape(tf.constant([1] * tf.size(input_ids).numpy()), (-1, tf.size(input_ids))) # Batch size 1
+        outputs = model(input_ids, labels=labels)
+        loss, scores = outputs[:2]

        """
-        outputs = self.roberta(inputs, **kwargs)
+        outputs = self.roberta(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        sequence_output = outputs[0]

-        sequence_output = self.dropout(sequence_output, training=kwargs.get("training", False))
+        sequence_output = self.dropout(sequence_output, training=training)
        logits = self.classifier(sequence_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

-        return outputs  # scores, (hidden_states), (attentions)
+        if labels is not None:
+            loss = self.compute_loss(labels, logits)
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), logits, (hidden_states), (attentions)


@add_start_docstrings(
    """RoBERTa Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of the hidden-states output to compute `span start logits` and `span end logits`). """,
    ROBERTA_START_DOCSTRING,
 )
-class TFRobertaForQuestionAnswering(TFRobertaPreTrainedModel):
+class TFRobertaForQuestionAnswering(TFRobertaPreTrainedModel, TFQuestionAnsweringLoss):
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -459,8 +655,32 @@ class TFRobertaForQuestionAnswering(TFRobertaPreTrainedModel):
        )

    @add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
-    def call(self, inputs, **kwargs):
+    def call(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        start_positions=None,
+        end_positions=None,
+        cls_index=None,
+        p_mask=None,
+        is_impossible=None,
+        output_attentions=None,
+        training=False,
+    ):
        r"""
+        start_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`tf.Tensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+
    Return:
        :obj:`tuple(tf.Tensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
        start_scores (:obj:`Numpy array` or :obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length,)`):
@@ -472,7 +692,7 @@ class TFRobertaForQuestionAnswering(TFRobertaPreTrainedModel):
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``config.output_attentions=True``):
+        attentions (:obj:`tuple(tf.Tensor)`, `optional`, returned when ``output_attentions=True``):
            tuple of :obj:`tf.Tensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`:

@@ -488,14 +708,24 @@ class TFRobertaForQuestionAnswering(TFRobertaPreTrainedModel):

        tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
        model = TFRobertaForQuestionAnswering.from_pretrained('roberta-base')
-        input_ids = tokenizer.encode("Who was Jim Henson?", "Jim Henson was a nice puppet")
-        start_scores, end_scores = model(tf.constant(input_ids)[None, :]) # Batch size 1
+        question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+        input_dict = tokenizer.encode_plus(question, text, return_tensors='tf')
+        start_scores, end_scores = model(input_dict)

-        all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
+        all_tokens = tokenizer.convert_ids_to_tokens(input_dict["input_ids"].numpy()[0])
        answer = ' '.join(all_tokens[tf.math.argmax(start_scores, 1)[0] : tf.math.argmax(end_scores, 1)[0]+1])

        """
-        outputs = self.roberta(inputs, **kwargs)
+        outputs = self.roberta(
+            input_ids,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            training=training,
+        )

        sequence_output = outputs[0]

@@ -506,4 +736,10 @@ class TFRobertaForQuestionAnswering(TFRobertaPreTrainedModel):

        outputs = (start_logits, end_logits,) + outputs[2:]

-        return outputs  # start_logits, end_logits, (hidden_states), (attentions)
+        if start_positions is not None and end_positions is not None:
+            labels = {"start_position": start_positions}
+            labels["end_position"] = end_positions
+            loss = self.compute_loss(labels, outputs[:2])
+            outputs = (loss,) + outputs
+
+        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)
Author	SHA1	Message	Date
Lysandre	513ba3b6a0	Better warning	2020-06-09 19:08:27 -04:00
Lysandre	6342e803a9	warn with FutureWarning when using `output_attentions` in the configuration	2020-06-09 19:06:49 -04:00
Lysandre	13aa174112	uninstalled wandb raises AttributeError	2020-06-09 18:50:56 -04:00
Bharat RaghunathanandPatrick von Platen	6e603cb789	[All models] Extend config.output_attentions with output_attentions function arguments (#4538 ) * DOC: Replace instances of ``config.output_attentions`` with function argument ``output_attentions`` * DOC: Apply Black Formatting * Fix errors where output_attentions was undefined * Remove output_attentions in classes per review * Fix regressions on tests having `output_attention` * Fix further regressions in tests relating to `output_attentions` Ensure proper propagation of `output_attentions` as a function parameter to all model subclasses * Fix more regressions in `test_output_attentions` * Fix issues with BertEncoder * Rename related variables to `output_attentions` * fix pytorch tests * fix bert and gpt2 tf * Fix most TF tests for `test_output_attentions` * Fix linter errors and more TF tests * fix conflicts * DOC: Apply Black Formatting * Fix errors where output_attentions was undefined * Remove output_attentions in classes per review * Fix regressions on tests having `output_attention` * fix conflicts * fix conflicts * fix conflicts * fix conflicts * fix pytorch tests * fix conflicts * fix conflicts * Fix linter errors and more TF tests * fix tf tests * make style * fix isort * improve output_attentions * improve tensorflow Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>	2020-06-09 23:39:06 +02:00
Sam Shleifer	f90bc44d9a	[examples] Cleanup summarization docs (#4876 )	2020-06-09 17:38:28 -04:00
Patrick von PlatenandPatrick von Platen	2cfb947f59	[Benchmark] add tpu and torchscipt for benchmark (#4850 ) * add tpu and torchscipt for benchmark * fix name in tests * "fix email" * make style * better log message for tpu * add more print and info for tpu * allow possibility to print tpu metrics * correct cpu usage * fix test for non-install * remove bugus file * include psutil in testing * run a couple of times before tracing in torchscript * do not allow tpu memory tracing for now * make style * add torchscript to env * better name for torch tpu Co-authored-by: Patrick von Platen <patrick@huggingface.co>	2020-06-09 23:12:43 +02:00
Hamza HarkousandHamza Harkous	f0340b3031	Removes from the of the parent of TFRobertaClassificationHead (#4884 ) Co-authored-by: Hamza Harkous <harkous@google.com>	2020-06-09 16:14:01 -04:00
Amil Khare	02e5f79662	[examples] consolidate summarization examples (#4837 )	2020-06-09 11:14:12 -04:00
Julien Plu	9f5d5a531d	Fix the __getattr__ method in BatchEncoding (#4772 )	2020-06-09 09:44:00 +02:00
Sylvain Gugger	41a1d27cde	Add XLMRobertaForQuestionAnswering (#4855 ) * Add XLMRobertaForQuestionAnswering * Formatting * Make test happy	2020-06-08 21:22:37 -04:00
Sam Shleifer	a139d1a160	[cleanup] consolidate some prune_heads logic (#4799 )	2020-06-08 17:08:04 -04:00
ZhuBaohe	4c7f564f9a	fix (#4839 )	2020-06-08 18:28:50 +02:00
Sylvain Gugger	37be3786cf	Clean documentation (#4849 ) * Clean documentation	2020-06-08 11:28:19 -04:00
Lysandre	42860e92a4	Turn off codecov patch for now	2020-06-08 09:47:13 -04:00
Julien Plu	36dfc317b3	TF Checkpoints (#4831 ) * Align checkpoint dir with the PT trainer * Use args for max to keep checkpoints	2020-06-08 09:45:23 -04:00
Patrick von Platen	439f1cab20	[Generate] beam search should generate without replacement (#4845 ) * fix flaky beam search * fix typo	2020-06-08 15:31:32 +02:00
Patrick von Platen	c0554776de	fix PR (#4810 )	2020-06-08 15:31:12 +02:00
Sylvain Gugger	e817747941	Expose classes used in documentation (#4808 ) * Expose classes used in documentation * Format code	2020-06-08 08:14:32 -04:00
daniel-shanandDaniel Shan	b6f365a8ed	Updates args in tf squad example. (#4820 ) Co-authored-by: Daniel Shan <daniel.shan@workday.com>	2020-06-08 05:36:09 -04:00
Bram Vanroy	e33fdc93b4	Export PretrainedBartModel from __init__ (#4819 )	2020-06-07 11:55:10 -04:00
Sam Shleifer	c58e6c129a	[marian tests ] pass device to pipeline (#4815 )	2020-06-06 00:52:17 -04:00
Mr Ruben	ddf9a3dfc7	Updated path "cd examples/text-generation/pplm" (#4778 ) https://github.com/huggingface/transformers/issues/4776	2020-06-05 21:16:48 -04:00
Sylvain Gugger	2d372a990b	Explain how to preview the docs in a PR (#4795 )	2020-06-05 20:47:02 -04:00
Sylvain Gugger	56d5d160cd	Add model and doc badges (#4811 ) * Add badges for models and docs	2020-06-05 18:45:42 -04:00
Sam Shleifer	4ab7424597	[cleanup/marian] pipelines test and new kwarg (#4812 )	2020-06-05 18:45:19 -04:00
Sam Shleifer	875288b344	[isort] add matplotlib to known 3rd party dependencies (#4800 )	2020-06-05 17:27:31 -04:00
Patrick von Platen	8cca875569	[EncoderDecoderConfig] automatically set decoder config to decoder (#4809 ) * automatically set decoder config to decoder * add more tests	2020-06-05 23:16:37 +02:00
Sylvain Gugger	f1fe18465d	Use labels to remove deprecation warnings (#4807 )	2020-06-05 16:41:46 -04:00
Sylvain Gugger	5c0cfc2cf0	Add link to community models (#4804 )	2020-06-05 15:29:20 -04:00
Sylvain Gugger	4dd5cf2207	Fix argument label (#4792 ) * Fix argument label * Fix test	2020-06-05 15:20:29 -04:00
Sam Shleifer	3723f30a18	[cleanup] MarianTokenizer: delete unused constants (#4802 )	2020-06-05 14:57:24 -04:00
Sylvain Gugger	acaa2e6267	Clean-up code (#4790 )	2020-06-05 12:36:22 -04:00
Sylvain Gugger	fa661ce749	Add model summary (#4789 ) * Add model summary * Add link to pretrained models	2020-06-05 12:22:50 -04:00
Lysandre DebutandJulien Chaumond	79ab881eb1	No silent error when d_head already in the configuration (#4747 ) * No silent error when d_head already in the configuration * Update src/transformers/configuration_xlnet.py Co-authored-by: Julien Chaumond <chaumond@gmail.com> Co-authored-by: Julien Chaumond <chaumond@gmail.com>	2020-06-05 12:01:43 -04:00
Julien Chaumond	b9109f2de1	[doc] Make it clearer that `text-generation` does not involve training	2020-06-05 14:59:22 +02:00
Sylvain Gugger	ceaab8dd22	Add .vs to gitignore (#4774 )	2020-06-05 07:56:11 -04:00
Julien PluandJulien Chaumond	f9414f7553	Tensorflow improvements (#4530 ) * Better None gradients handling * Apply Style * Apply Style * Create a loss class per task to compute its respective loss * Add loss classes to the ALBERT TF models * Add loss classes to the BERT TF models * Add question answering and multiple choice to TF Camembert * Remove prints * Add multiple choice model to TF DistilBERT + loss computation * Add question answering model to TF Electra + loss computation * Add token classification, question answering and multiple choice models to TF Flaubert * Add multiple choice model to TF Roberta + loss computation * Add multiple choice model to TF XLM + loss computation * Add multiple choice and question answering models to TF XLM-Roberta * Add multiple choice model to TF XLNet + loss computation * Remove unused parameters * Add task loss classes * Reorder TF imports + add new model classes * Add new model classes * Bugfix in TF T5 model * Bugfix for TF T5 tests * Bugfix in TF T5 model * Fix TF T5 model tests * Fix T5 tests + some renaming * Fix inheritance issue in the AutoX tests * Add tests for TF Flaubert and TF XLM Roberta * Add tests for TF Flaubert and TF XLM Roberta * Remove unused piece of code in the TF trainer * bugfix and remove unused code * Bugfix for TF 2.2 * Apply Style * Divide TFSequenceClassificationAndMultipleChoiceLoss into their two respective name * Apply style * Mirror the PT Trainer in the TF one: fp16, optimizers and tb_writer as class parameter and better dataset handling * Fix TF optimizations tests and apply style * Remove useless parameter * Bugfix and apply style * Fix TF Trainer prediction * Now the TF models return the loss such as their PyTorch couterparts * Apply Style * Ignore some tests output * Take into account the SQuAD cls_index, p_mask and is_impossible parameters for the QuestionAnswering task models. * Fix names for SQuAD data * Apply Style * Fix conflicts with 2.11 release * Fix conflicts with 2.11 * Fix wrongname * Add better documentation on the new create_optimizer function * Fix isort * logging_dir: use same default as PyTorch Co-authored-by: Julien Chaumond <chaumond@gmail.com>	2020-06-04 19:45:53 -04:00
Théophile Blard	ccd26c2862	Create model card for tblard/allocine (#4775 ) https://huggingface.co/tblard/tf-allocine	2020-06-04 19:15:07 -04:00
Stefan Schweter	2a4b9e09c0	NER: Add new WNUT’17 example (#4681 ) * ner: add preprocessing script for examples that splits longer sentences * ner: example shell scripts use local preprocessing now * ner: add new example section for WNUT’17 NER task. Remove old English CoNLL-03 results * ner: satisfy black and isort	2020-06-04 19:13:17 -04:00
Setu Shah	0e1869cc28	Add drop_last arg for data loader	2020-06-04 18:30:31 -04:00
prajjwal1	48a05026de	removed deprecared use of Variable api from pplm example	2020-06-04 18:07:49 -04:00
Sylvain Gugger	12d0eb5f3e	Don't access pad_token_id if there is no pad_token (#4773 )	2020-06-04 17:57:04 -04:00
Manuel Romero	17a88d3192	Create model card for T5-base fine-tuned for Sentiment Span Extraction (#4737 )	2020-06-04 16:59:56 -04:00
Oren Amsalem	fb52143cf6	Create README.md (#4743 )	2020-06-04 16:59:37 -04:00
Suraj Parmar	5f077a3445	Model Card for RoBERTa trained on Sanskrit (#4763 ) * Model cad for SanBERTa Model Card for RoBERTa trained on Sanskrit * Model card for SanBERTa model card for RoBERTa trained on Sanskrit	2020-06-04 16:58:40 -04:00
Sylvain Gugger	cd4e07a85e	Add note about doc generation (#4770 )	2020-06-04 13:43:14 -04:00
Jason Phang	492b352ab6	Remove unnecessary model_type arg in example (#4771 )	2020-06-04 13:41:24 -04:00
Lysandre Debut	e645b9ab94	Codecov setup (#4768 ) * Codecov setup * Understanding codecov	2020-06-04 11:44:38 -04:00
Sam Shleifer	2b8b6c929e	[cleanup] PretrainedModel.generate: remove unused kwargs (#4761 )	2020-06-04 08:13:52 -04:00
Funtowicz Morgan	5bf9afbf35	Introduce a new tensor type for return_tensors on tokenizer for NumPy (#4585 ) * Refactor tensor creation in tokenizers. * Make sure to convert string to TensorType * Refactor convert_to_tensors_ * Introduce numpy tensor creation * Format * Add unittest for TensorType creation from str * sorting imports * Added unittests for numpy tensor conversion. * Do not use in-place version for squeeze as numpy doesn't provide such feature. * Added extra parameter prepend_batch_axis: bool on prepare_for_model. * Ensure test_np_encode_plus_sent_to_model is not executed if encoder/decoder model. * style. * numpy tests require_torch for now while flax not merged. * Hopefully will make flake8 happy. * One more time 🎶	2020-06-04 06:57:01 +02:00
Funtowicz Morgan	efae154929	never_split on slow tokenizers should not split (#4723 ) * Ensure tokens in never_split are not splitted when using basic tokenizer before wordpiece. * never_split only use membership attempt to use a set() which is 10x faster for this operation. * Use union to concatenate two sets. * Updated docstring for never_split parameter. * Avoid set.union() if never_split is None * Added comments. * Correct docstring format.	2020-06-03 16:48:28 -04:00
Lysandre Debut	2e4de76231	Update encode documentation (#4751 )	2020-06-03 16:30:59 -04:00
Patrick von Platen	ed4df85572	fix beam search bug in tf as well (#4745 )	2020-06-03 12:53:23 -04:00
Sylvain Gugger	1b5820a565	Unify label args (#4722 ) * Deprecate masked_lm_labels argument * Apply to all models * Better error message	2020-06-03 09:36:26 -04:00
Abhishek Kumar MishraandPatrick von Platen	3e5928c57d	Adding notebooks for Fine Tuning [Community Notebook] (#4732 ) * Added links to more community notebooks Added links to 3 more community notebooks from the git repo: https://github.com/abhimishra91/transformers-tutorials Different Transformers models are fine tuned on Dataset using PyTorch * Update README.md * Update README.md * Update README.md Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>	2020-06-03 11:07:26 +02:00
Julien Chaumond	99207bd112	Pipelines: miscellanea of QoL improvements and small features... (#4632 ) * [hf_api] Attach all unknown attributes for future-proof compatibility * [Pipeline] NerPipeline is really a TokenClassificationPipeline * modelcard.py: I don't think we need to force the download * Remove config, tokenizer from SUPPORTED_TASKS as we're moving to one model = one weight + one tokenizer * FillMaskPipeline: also output token in string form * TextClassificationPipeline: option to return all scores, not just the argmax * Update docs/source/main_classes/pipelines.rst	2020-06-03 03:51:31 -04:00
David Mezzetti	8ed47aa10b	bert-small-cord19 model cards (#4730 ) * Create README.md * Create README.md * Create README.md	2020-06-03 03:40:14 -04:00
Patrick von Platen	9ca485734a	[Reformer] Improved memory if input is shorter than chunk length (#4720 ) * improve handling of short inputs for reformer * correct typo in assert statement * fix other tests	2020-06-02 23:08:39 +02:00
Jin Young SohnandJulien Chaumond	b231a413f5	Add cache_dir to save features in GLUE + Differentiate match/mismatch for MNLI metrics (#4621 ) * Glue task cleaup * Enable writing cache to cache_dir in case dataset lives in readOnly filesystem. * Differentiate match vs mismatch for MNLI metrics. * Style * Fix pytype * Fix type * Use cache_dir in mnli mismatch eval dataset * Small Tweaks Co-authored-by: Julien Chaumond <chaumond@gmail.com>	2020-06-02 13:40:14 -04:00
Sam Shleifer	70f7423436	TFRobertaModelIntegrationTest requires tf (#4726 )	2020-06-02 12:59:00 -04:00
Lysandre	d976ef262e	Repin versions	2020-06-02 10:27:15 -04:00