Torch import guard.

fully created from static state.

.circleci/config.yml

@@ -79,6 +79,7 @@ jobs:
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
             - run: pip install .[sklearn,tf-cpu,torch,testing,sentencepiece]
+            - run: pip install tapas torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cpu.html
             - save_cache:
                 key: v0.4-{{ checksum "setup.py" }}
                 paths:
@@ -105,6 +106,7 @@ jobs:
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
             - run: pip install .[sklearn,torch,testing,sentencepiece]
+            - run: pip install tapas torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cpu.html
             - save_cache:
                   key: v0.4-torch-{{ checksum "setup.py" }}
                   paths:
@@ -183,6 +185,7 @@ jobs:
                       - v0.4-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
             - run: pip install .[sklearn,torch,testing,sentencepiece]
+            - run: pip install tapas torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cpu.html
             - save_cache:
                   key: v0.4-torch-{{ checksum "setup.py" }}
                   paths:

.circleci/deploy.sh

@@ -53,4 +53,5 @@ deploy_doc "3ebb1b3" v3.2.0
 deploy_doc "0613f05" v3.3.1
 deploy_doc "eb0e0ce" v3.4.0
 deploy_doc "818878d" v3.5.1
-deploy_doc "c781171" # v4.0.0 Latest stable release
+deploy_doc "c781171" v4.0.0
+deploy_doc "bfa4ccf" # v4.1.1 Latest stable release

.github/workflows/model-templates.yml

@@ -40,6 +40,8 @@ jobs:
           transformers-cli add-new-model --testing --testing_file=templates/adding_a_new_model/tests/pt-encoder-bert-tokenizer.json --path=templates/adding_a_new_model
           transformers-cli add-new-model --testing --testing_file=templates/adding_a_new_model/tests/standalone.json --path=templates/adding_a_new_model
           transformers-cli add-new-model --testing --testing_file=templates/adding_a_new_model/tests/tf-encoder-bert-tokenizer.json --path=templates/adding_a_new_model
+          transformers-cli add-new-model --testing --testing_file=templates/adding_a_new_model/tests/tf-seq-2-seq-bart-tokenizer.json --path=templates/adding_a_new_model
+          transformers-cli add-new-model --testing --testing_file=templates/adding_a_new_model/tests/pt-seq-2-seq-bart-tokenizer.json --path=templates/adding_a_new_model
           make style
           python utils/check_table.py --fix_and_overwrite
           python utils/check_dummies.py --fix_and_overwrite

.github/workflows/self-push.yml

@@ -50,6 +50,7 @@ jobs:
           pip install --upgrade pip
           pip install .[torch,sklearn,testing,onnxruntime,sentencepiece]
           pip install git+https://github.com/huggingface/datasets
+          pip install pandas torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
 
       - name: Are GPUs recognized by our DL frameworks
         run: |
@@ -187,6 +188,7 @@ jobs:
           pip install --upgrade pip
           pip install .[torch,sklearn,testing,onnxruntime,sentencepiece]
           pip install git+https://github.com/huggingface/datasets
+          pip install pandas torch-scatter -f https://pytorch-geometric.com/whl/torch-1.7.0+cu102.html
 
       - name: Are GPUs recognized by our DL frameworks
         run: |

README.md

@@ -222,6 +222,7 @@ Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih.
 ultilingual BERT into [DistilmBERT](https://github.com/huggingface/transformers/tree/master/examples/distillation) and a German version of DistilBERT.
 1. **[SqueezeBert](https://huggingface.co/transformers/model_doc/squeezebert.html)** released with the paper [SqueezeBERT: What can computer vision teach NLP about efficient neural networks?](https://arxiv.org/abs/2006.11316) by Forrest N. Iandola, Albert E. Shaw, Ravi Krishna, and Kurt W. Keutzer.
 1. **[T5](https://huggingface.co/transformers/model_doc/t5.html)** (from Google AI) released with the paper [Exploring the Limits of Transfer Learning with a Unified Text-to-Text Transformer](https://arxiv.org/abs/1910.10683) by Colin Raffel and Noam Shazeer and Adam Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
+1. **[TAPAS](https://huggingface.co/transformers/master/model_doc/tapas.html)** (from Google AI) released with the paper [TAPAS: Weakly Supervised Table Parsing via Pre-training](https://arxiv.org/abs/2004.02349) by Jonathan Herzig, Paweł Krzysztof Nowak, Thomas Müller, Francesco Piccinno and Julian Martin Eisenschlos.
 1. **[Transformer-XL](https://huggingface.co/transformers/model_doc/transformerxl.html)** (from Google/CMU) released with the paper [Transformer-XL: Attentive Language Models Beyond a Fixed-Length Context](https://arxiv.org/abs/1901.02860) by Zihang Dai*, Zhilin Yang*, Yiming Yang, Jaime Carbonell, Quoc V. Le, Ruslan Salakhutdinov.
 1. **[XLM](https://huggingface.co/transformers/model_doc/xlm.html)** (from Facebook) released together with the paper [Cross-lingual Language Model Pretraining](https://arxiv.org/abs/1901.07291) by Guillaume Lample and Alexis Conneau.
 1. **[XLM-ProphetNet](https://huggingface.co/transformers/model_doc/xlmprophetnet.html)** (from Microsoft Research) released with the paper [ProphetNet: Predicting Future N-gram for Sequence-to-Sequence Pre-training](https://arxiv.org/abs/2001.04063) by Yu Yan, Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou.

docs/source/_static/js/custom.js

@@ -1,10 +1,11 @@
 // These two things need to be updated at each release for the version selector.
 // Last stable version
-const stableVersion = "v4.0.0"
+const stableVersion = "v4.1.1"
 // Dictionary doc folder to label. The last stable version should have an empty key.
 const versionMapping = {
     "master": "master",
-    "": "v4.0.0 (stable)",
+    "": "v4.1.1 (stable)",
+    "v4.0.1": "v4.0.0/v4.0.1",
     "v3.5.1": "v3.5.0/v3.5.1",
     "v3.4.0": "v3.4.0",
     "v3.3.1": "v3.3.0/v3.3.1",

docs/source/conf.py

@@ -26,7 +26,7 @@ author = u'huggingface'
 # The short X.Y version
 version = u''
 # The full version, including alpha/beta/rc tags
-release = u'4.0.0'
+release = u'4.1.1'
 
 
 # -- General configuration ---------------------------------------------------

docs/source/index.rst

@@ -176,19 +176,22 @@ and conversion utilities for the following models:
 30. :doc:`T5 <model_doc/t5>` (from Google AI) released with the paper `Exploring the Limits of Transfer Learning with a
     Unified Text-to-Text Transformer <https://arxiv.org/abs/1910.10683>`__ by Colin Raffel and Noam Shazeer and Adam
     Roberts and Katherine Lee and Sharan Narang and Michael Matena and Yanqi Zhou and Wei Li and Peter J. Liu.
-31. :doc:`Transformer-XL <model_doc/transformerxl>` (from Google/CMU) released with the paper `Transformer-XL:
+31. `TAPAS <https://huggingface.co/transformers/master/model_doc/tapas.html>`__ (from Google AI) released with the
+    paper `TAPAS: Weakly Supervised Table Parsing via Pre-training <https://arxiv.org/abs/2004.02349>`__ by Jonathan
+    Herzig, Paweł Krzysztof Nowak, Thomas Müller, Francesco Piccinno and Julian Martin Eisenschlos.
+32. :doc:`Transformer-XL <model_doc/transformerxl>` (from Google/CMU) released with the paper `Transformer-XL:
     Attentive Language Models Beyond a Fixed-Length Context <https://arxiv.org/abs/1901.02860>`__ by Zihang Dai*,
     Zhilin Yang*, Yiming Yang, Jaime Carbonell, Quoc V. Le, Ruslan Salakhutdinov.
-32. :doc:`XLM <model_doc/xlm>` (from Facebook) released together with the paper `Cross-lingual Language Model
+33. :doc:`XLM <model_doc/xlm>` (from Facebook) released together with the paper `Cross-lingual Language Model
     Pretraining <https://arxiv.org/abs/1901.07291>`__ by Guillaume Lample and Alexis Conneau.
-33. :doc:`XLM-ProphetNet <model_doc/xlmprophetnet>` (from Microsoft Research) released with the paper `ProphetNet:
+34. :doc:`XLM-ProphetNet <model_doc/xlmprophetnet>` (from Microsoft Research) released with the paper `ProphetNet:
     Predicting Future N-gram for Sequence-to-Sequence Pre-training <https://arxiv.org/abs/2001.04063>`__ by Yu Yan,
     Weizhen Qi, Yeyun Gong, Dayiheng Liu, Nan Duan, Jiusheng Chen, Ruofei Zhang and Ming Zhou.
-34. :doc:`XLM-RoBERTa <model_doc/xlmroberta>` (from Facebook AI), released together with the paper `Unsupervised
+35. :doc:`XLM-RoBERTa <model_doc/xlmroberta>` (from Facebook AI), released together with the paper `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.
-35. :doc:`XLNet <model_doc/xlnet>` (from Google/CMU) released with the paper `​XLNet: Generalized Autoregressive
+36. :doc:`XLNet <model_doc/xlnet>` (from Google/CMU) released with the paper `​XLNet: Generalized Autoregressive
     Pretraining for Language Understanding <https://arxiv.org/abs/1906.08237>`__ by Zhilin Yang*, Zihang Dai*, Yiming
     Yang, Jaime Carbonell, Ruslan Salakhutdinov, Quoc V. Le.
 
@@ -269,6 +272,8 @@ TensorFlow and/or Flax.
 +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+
 |             T5              |       ✅       |       ✅       |       ✅        |         ✅         |      ❌      |
 +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+
+|            TAPAS            |       ✅       |       ❌       |       ✅        |         ❌         |      ❌      |
++-----------------------------+----------------+----------------+-----------------+--------------------+--------------+
 |       Transformer-XL        |       ✅       |       ❌       |       ✅        |         ✅         |      ❌      |
 +-----------------------------+----------------+----------------+-----------------+--------------------+--------------+
 |             XLM             |       ✅       |       ❌       |       ✅        |         ✅         |      ❌      |
@@ -382,6 +387,7 @@ TensorFlow and/or Flax.
     model_doc/roberta
     model_doc/squeezebert
     model_doc/t5
+    model_doc/tapas
     model_doc/transformerxl
     model_doc/xlm
     model_doc/xlmprophetnet

docs/source/main_classes/model.rst

@@ -13,9 +13,10 @@
 Models
 -----------------------------------------------------------------------------------------------------------------------
 
-The base classes :class:`~transformers.PreTrainedModel` and :class:`~transformers.TFPreTrainedModel` implement the
-common methods for loading/saving a model either from a local file or directory, or from a pretrained model
-configuration provided by the library (downloaded from HuggingFace's AWS S3 repository).
+The base classes :class:`~transformers.PreTrainedModel`, :class:`~transformers.TFPreTrainedModel`, and
+:class:`~transformers.FlaxPreTrainedModel` implement the common methods for loading/saving a model either from a local
+file or directory, or from a pretrained model configuration provided by the library (downloaded from HuggingFace's AWS
+S3 repository).
 
 :class:`~transformers.PreTrainedModel` and :class:`~transformers.TFPreTrainedModel` also implement a few methods which
 are common among all the models to:
@@ -57,6 +58,13 @@ TFModelUtilsMixin
     :members:
 
 
+FlaxPreTrainedModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.FlaxPreTrainedModel
+    :members:
+
+
 Generation
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/main_classes/optimizer_schedules.rst

@@ -43,6 +43,10 @@ Schedules
 Learning Rate Schedules (Pytorch)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 
+.. autoclass:: transformers.SchedulerType
+
+.. autofunction:: transformers.get_scheduler
+
 .. autofunction:: transformers.get_constant_schedule
 
 

docs/source/main_classes/pipelines.rst

@@ -34,6 +34,7 @@ There are two categories of pipeline abstractions to be aware about:
     - :class:`~transformers.TranslationPipeline`
     - :class:`~transformers.ZeroShotClassificationPipeline`
     - :class:`~transformers.Text2TextGenerationPipeline`
+    - :class:`~transformers.TableQuestionAnsweringPipeline`
 
 The pipeline abstraction
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -91,6 +92,13 @@ SummarizationPipeline
     :special-members: __call__
     :members:
 
+TableQuestionAnsweringPipeline
+=======================================================================================================================
+
+.. autoclass:: transformers.TableQuestionAnsweringPipeline
+    :special-members: __call__
+
+
 TextClassificationPipeline
 =======================================================================================================================
 

docs/source/main_classes/trainer.rst

@@ -63,6 +63,13 @@ Trainer
     :members:
 
 
+Seq2SeqTrainer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.Seq2SeqTrainer
+    :members: evaluate, predict
+
+
 TFTrainer
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -77,6 +84,13 @@ TrainingArguments
     :members:
 
 
+Seq2SeqTrainingArguments
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.Seq2SeqTrainingArguments
+    :members:
+
+
 TFTrainingArguments
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/model_doc/auto.rst

@@ -114,6 +114,13 @@ AutoModelForQuestionAnswering
     :members:
 
 
+AutoModelForTableQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.AutoModelForTableQuestionAnswering
+    :members:
+
+
 TFAutoModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/model_doc/blenderbot.rst

@@ -100,6 +100,15 @@ BlenderbotSmallTokenizer
     :members:
 
 
+BlenderbotModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+See :obj:`transformers.BartModel` for arguments to `forward` and `generate`
+
+.. autoclass:: transformers.BlenderbotModel
+    :members:
+
+
 BlenderbotForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/model_doc/ctrl.rst

@@ -97,3 +97,8 @@ TFCTRLLMHeadModel
 .. autoclass:: transformers.TFCTRLLMHeadModel
     :members: call
 
+TFCTRLForSequenceClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFCTRLForSequenceClassification
+    :members: call

docs/source/model_doc/gpt.rst

@@ -138,3 +138,9 @@ TFOpenAIGPTDoubleHeadsModel
 
 .. autoclass:: transformers.TFOpenAIGPTDoubleHeadsModel
     :members: call
+
+TFOpenAIGPTForSequenceClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFOpenAIGPTForSequenceClassification
+    :members: call

docs/source/model_doc/longformer.rst

@@ -34,6 +34,12 @@ contrast to most prior work, we also pretrain Longformer and finetune it on a va
 pretrained Longformer consistently outperforms RoBERTa on long document tasks and sets new state-of-the-art results on
 WikiHop and TriviaQA.*
 
+Tips:
+
+- Since the Longformer is based on RoBERTa, it doesn't have :obj:`token_type_ids`. You don't need to indicate which
+  token belongs to which segment. Just separate your segments with the separation token :obj:`tokenizer.sep_token` (or
+  :obj:`</s>`).
+
 The Authors' code can be found `here <https://github.com/allenai/longformer>`__.
 
 Longformer Self Attention

docs/source/model_doc/mbart.rst

@@ -97,6 +97,13 @@ MBartTokenizerFast
     :members:
 
 
+MBartModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.MBartModel
+    :members:
+
+
 MBartForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/model_doc/pegasus.rst

@@ -119,6 +119,12 @@ PegasusTokenizerFast
     :members:
 
 
+PegasusModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.PegasusModel
+
+
 PegasusForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/model_doc/t5.rst

@@ -44,9 +44,9 @@ Tips:
 
   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
-  :obj:`T5ForConditionalGeneration.generate()``. This 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.
+  :obj:`T5ForConditionalGeneration.generate()`. This 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>`__.
 
@@ -55,7 +55,7 @@ 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 :obj:`input_ids``. The target sequence is shifted to the right, i.e., prepended by a start-sequence
+to the model using :obj:`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 :obj:`decoder_input_ids`. In teacher-forcing style, the target sequence is then
 appended by the EOS token and corresponds to the :obj:`labels`. The PAD token is hereby used as the start-sequence
 token. T5 can be trained / fine-tuned both in a supervised and unsupervised fashion.

docs/source/model_doc/tapas.rst

@@ -0,0 +1,434 @@
+TAPAS
+-----------------------------------------------------------------------------------------------------------------------
+
+.. note::
+
+    This is a recently introduced model so the API hasn't been tested extensively. There may be some bugs or slight
+    breaking changes to fix them in the future.
+
+
+
+Overview
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+The TAPAS model was proposed in `TAPAS: Weakly Supervised Table Parsing via Pre-training
+<https://www.aclweb.org/anthology/2020.acl-main.398>`__ by Jonathan Herzig, Paweł Krzysztof Nowak, Thomas Müller,
+Francesco Piccinno and Julian Martin Eisenschlos. It's a BERT-based model specifically designed (and pre-trained) for
+answering questions about tabular data. Compared to BERT, TAPAS uses relative position embeddings and has 7 token types
+that encode tabular structure. TAPAS is pre-trained on the masked language modeling (MLM) objective on a large dataset
+comprising millions of tables from English Wikipedia and corresponding texts. For question answering, TAPAS has 2 heads
+on top: a cell selection head and an aggregation head, for (optionally) performing aggregations (such as counting or
+summing) among selected cells. TAPAS has been fine-tuned on several datasets: `SQA
+<https://www.microsoft.com/en-us/download/details.aspx?id=54253>`__ (Sequential Question Answering by Microsoft), `WTQ
+<https://github.com/ppasupat/WikiTableQuestions>`__ (Wiki Table Questions by Stanford University) and `WikiSQL
+<https://github.com/salesforce/WikiSQL>`__ (by Salesforce). It achieves state-of-the-art on both SQA and WTQ, while
+having comparable performance to SOTA on WikiSQL, with a much simpler architecture.
+
+The abstract from the paper is the following:
+
+*Answering natural language questions over tables is usually seen as a semantic parsing task. To alleviate the
+collection cost of full logical forms, one popular approach focuses on weak supervision consisting of denotations
+instead of logical forms. However, training semantic parsers from weak supervision poses difficulties, and in addition,
+the generated logical forms are only used as an intermediate step prior to retrieving the denotation. In this paper, we
+present TAPAS, an approach to question answering over tables without generating logical forms. TAPAS trains from weak
+supervision, and predicts the denotation by selecting table cells and optionally applying a corresponding aggregation
+operator to such selection. TAPAS extends BERT's architecture to encode tables as input, initializes from an effective
+joint pre-training of text segments and tables crawled from Wikipedia, and is trained end-to-end. We experiment with
+three different semantic parsing datasets, and find that TAPAS outperforms or rivals semantic parsing models by
+improving state-of-the-art accuracy on SQA from 55.1 to 67.2 and performing on par with the state-of-the-art on WIKISQL
+and WIKITQ, but with a simpler model architecture. We additionally find that transfer learning, which is trivial in our
+setting, from WIKISQL to WIKITQ, yields 48.7 accuracy, 4.2 points above the state-of-the-art.*
+
+In addition, the authors have further pre-trained TAPAS to recognize **table entailment**, by creating a balanced
+dataset of millions of automatically created training examples which are learned in an intermediate step prior to
+fine-tuning. The authors of TAPAS call this further pre-training intermediate pre-training (since TAPAS is first
+pre-trained on MLM, and then on another dataset). They found that intermediate pre-training further improves
+performance on SQA, achieving a new state-of-the-art as well as state-of-the-art on `TabFact
+<https://github.com/wenhuchen/Table-Fact-Checking>`__, a large-scale dataset with 16k Wikipedia tables for table
+entailment (a binary classification task). For more details, see their follow-up paper: `Understanding tables with
+intermediate pre-training <https://www.aclweb.org/anthology/2020.findings-emnlp.27/>`__ by Julian Martin Eisenschlos,
+Syrine Krichene and Thomas Müller.
+
+The original code can be found `here <https://github.com/google-research/tapas>`__.
+
+Tips:
+
+- TAPAS is a model that uses relative position embeddings by default (restarting the position embeddings at every cell
+  of the table). Note that this is something that was added after the publication of the original TAPAS paper.
+  According to the authors, this usually results in a slightly better performance, and allows you to encode longer
+  sequences without running out of embeddings. This is reflected in the ``reset_position_index_per_cell`` parameter of
+  :class:`~transformers.TapasConfig`, which is set to ``True`` by default. The default versions of the models available
+  in the `model hub <https://huggingface.co/models?search=tapas>`_ all use relative position embeddings. You can still
+  use the ones with absolute position embeddings by passing in an additional argument ``revision="no_reset"`` when
+  calling the ``.from_pretrained()`` method. Note that it's usually advised to pad the inputs on the right rather than
+  the left.
+- TAPAS is based on BERT, so ``TAPAS-base`` for example corresponds to a ``BERT-base`` architecture. Of course,
+  TAPAS-large will result in the best performance (the results reported in the paper are from TAPAS-large). Results of
+  the various sized models are shown on the `original Github repository <https://github.com/google-research/tapas>`_.
+- TAPAS has checkpoints fine-tuned on SQA, which are capable of answering questions related to a table in a
+  conversational set-up. This means that you can ask follow-up questions such as "what is his age?" related to the
+  previous question. Note that the forward pass of TAPAS is a bit different in case of a conversational set-up: in that
+  case, you have to feed every table-question pair one by one to the model, such that the `prev_labels` token type ids
+  can be overwritten by the predicted `labels` of the model to the previous question. See "Usage" section for more
+  info.
+- TAPAS is similar to BERT and therefore relies on the masked language modeling (MLM) objective. It is therefore
+  efficient at predicting masked tokens and at NLU in general, but is not optimal for text generation. Models trained
+  with a causal language modeling (CLM) objective are better in that regard.
+
+
+Usage: fine-tuning
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Here we explain how you can fine-tune :class:`~transformers.TapasForQuestionAnswering` on your own dataset.
+
+**STEP 1: Choose one of the 3 ways in which you can use TAPAS - or experiment**
+
+Basically, there are 3 different ways in which one can fine-tune :class:`~transformers.TapasForQuestionAnswering`,
+corresponding to the different datasets on which Tapas was fine-tuned:
+
+1. SQA: if you're interested in asking follow-up questions related to a table, in a conversational set-up. For example
+   if you first ask "what's the name of the first actor?" then you can ask a follow-up question such as "how old is
+   he?". Here, questions do not involve any aggregation (all questions are cell selection questions).
+2. WTQ: if you're not interested in asking questions in a conversational set-up, but rather just asking questions
+   related to a table, which might involve aggregation, such as counting a number of rows, summing up cell values or
+   averaging cell values. You can then for example ask "what's the total number of goals Cristiano Ronaldo made in his
+   career?". This case is also called **weak supervision**, since the model itself must learn the appropriate
+   aggregation operator (SUM/COUNT/AVERAGE/NONE) given only the answer to the question as supervision.
+3. WikiSQL-supervised: this dataset is based on WikiSQL with the model being given the ground truth aggregation
+   operator during training. This is also called **strong supervision**. Here, learning the appropriate aggregation
+   operator is much easier.
+
+To summarize:
+
++------------------------------------+----------------------+-------------------------------------------------------------------------------------------------------------------+
+| **Task**                           | **Example dataset**  | **Description**                                                                                                   |
++------------------------------------+----------------------+-------------------------------------------------------------------------------------------------------------------+
+| Conversational                     | SQA                  | Conversational, only cell selection questions                                                                     |
++------------------------------------+----------------------+-------------------------------------------------------------------------------------------------------------------+
+| Weak supervision for aggregation   | WTQ                  | Questions might involve aggregation, and the model must learn this given only the answer as supervision           |
++------------------------------------+----------------------+-------------------------------------------------------------------------------------------------------------------+
+| Strong supervision for aggregation | WikiSQL-supervised   | Questions might involve aggregation, and the model must learn this given the gold aggregation operator            |
++------------------------------------+----------------------+-------------------------------------------------------------------------------------------------------------------+
+
+Initializing a model with a pre-trained base and randomly initialized classification heads from the model hub can be
+done as follows (be sure to have installed the `torch-scatter dependency <https://github.com/rusty1s/pytorch_scatter>`_
+for your environment):
+
+.. code-block::
+
+        >>> from transformers import TapasConfig, TapasForQuestionAnswering
+
+        >>> # for example, the base sized model with default SQA configuration
+        >>> model = TapasForQuestionAnswering.from_pretrained('google/tapas-base')
+
+        >>> # or, the base sized model with WTQ configuration
+        >>> config = TapasConfig.from_pretrained('google/tapas-base-finetuned-wtq')
+        >>> model = TapasForQuestionAnswering.from_pretrained('google/tapas-base', config=config)
+
+        >>> # or, the base sized model with WikiSQL configuration
+        >>> config = TapasConfig('google-base-finetuned-wikisql-supervised')
+        >>> model = TapasForQuestionAnswering.from_pretrained('google/tapas-base', config=config)
+
+
+Of course, you don't necessarily have to follow one of these three ways in which TAPAS was fine-tuned. You can also
+experiment by defining any hyperparameters you want when initializing :class:`~transformers.TapasConfig`, and then
+create a :class:`~transformers.TapasForQuestionAnswering` based on that configuration. For example, if you have a
+dataset that has both conversational questions and questions that might involve aggregation, then you can do it this
+way. Here's an example:
+
+.. code-block::
+
+        >>> from transformers import TapasConfig, TapasForQuestionAnswering
+
+        >>> # you can initialize the classification heads any way you want (see docs of TapasConfig)
+        >>> config = TapasConfig(num_aggregation_labels=3, average_logits_per_cell=True, select_one_column=False)
+        >>> # initializing the pre-trained base sized model with our custom classification heads
+        >>> model = TapasForQuestionAnswering.from_pretrained('google/tapas-base', config=config)
+
+What you can also do is start from an already fine-tuned checkpoint. A note here is that the already fine-tuned
+checkpoint on WTQ has some issues due to the L2-loss which is somewhat brittle. See `here
+<https://github.com/google-research/tapas/issues/91#issuecomment-735719340>`__ for more info.
+
+For a list of all pre-trained and fine-tuned TAPAS checkpoints available in the HuggingFace model hub, see `here
+<https://huggingface.co/models?search=tapas>`__.
+
+**STEP 2: Prepare your data in the SQA format**
+
+Second, no matter what you picked above, you should prepare your dataset in the `SQA format
+<https://www.microsoft.com/en-us/download/details.aspx?id=54253>`__. This format is a TSV/CSV file with the following
+columns:
+
+- ``id``: optional, id of the table-question pair, for bookkeeping purposes.
+- ``annotator``: optional, id of the person who annotated the table-question pair, for bookkeeping purposes.
+- ``position``: integer indicating if the question is the first, second, third,... related to the table. Only required
+  in case of conversational setup (SQA). You don't need this column in case you're going for WTQ/WikiSQL-supervised.
+- ``question``: string
+- ``table_file``: string, name of a csv file containing the tabular data
+- ``answer_coordinates``: list of one or more tuples (each tuple being a cell coordinate, i.e. row, column pair that is
+  part of the answer)
+- ``answer_text``: list of one or more strings (each string being a cell value that is part of the answer)
+- ``aggregation_label``: index of the aggregation operator. Only required in case of strong supervision for aggregation
+  (the WikiSQL-supervised case)
+- ``float_answer``: the float answer to the question, if there is one (np.nan if there isn't). Only required in case of
+  weak supervision for aggregation (such as WTQ and WikiSQL)
+
+The tables themselves should be present in a folder, each table being a separate csv file. Note that the authors of the
+TAPAS algorithm used conversion scripts with some automated logic to convert the other datasets (WTQ, WikiSQL) into the
+SQA format. The author explains this `here
+<https://github.com/google-research/tapas/issues/50#issuecomment-705465960>`__. Interestingly, these conversion scripts
+are not perfect (the ``answer_coordinates`` and ``float_answer`` fields are populated based on the ``answer_text``),
+meaning that WTQ and WikiSQL results could actually be improved.
+
+**STEP 3: Convert your data into PyTorch tensors using TapasTokenizer**
+
+Third, given that you've prepared your data in this TSV/CSV format (and corresponding CSV files containing the tabular
+data), you can then use :class:`~transformers.TapasTokenizer` to convert table-question pairs into :obj:`input_ids`,
+:obj:`attention_mask`, :obj:`token_type_ids` and so on. Again, based on which of the three cases you picked above,
+:class:`~transformers.TapasForQuestionAnswering` requires different inputs to be fine-tuned:
+
++------------------------------------+----------------------------------------------------------------------------------------------+
+| **Task**                           | **Required inputs**                                                                          |
++------------------------------------+----------------------------------------------------------------------------------------------+
+| Conversational                     | ``input_ids``, ``attention_mask``, ``token_type_ids``, ``labels``                            |
++------------------------------------+----------------------------------------------------------------------------------------------+
+| Weak supervision for aggregation   | ``input_ids``, ``attention_mask``, ``token_type_ids``, ``labels``, ``numeric_values``,       |
+|                                    | ``numeric_values_scale``, ``float_answer``                                                   |
++------------------------------------+----------------------------------------------------------------------------------------------+
+| Strong supervision for aggregation | ``input ids``, ``attention mask``, ``token type ids``, ``labels``, ``aggregation_labels``    |
++------------------------------------+----------------------------------------------------------------------------------------------+
+
+:class:`~transformers.TapasTokenizer` creates the ``labels``, ``numeric_values`` and ``numeric_values_scale`` based on
+the ``answer_coordinates`` and ``answer_text`` columns of the TSV file. The ``float_answer`` and ``aggregation_labels``
+are already in the TSV file of step 2. Here's an example:
+
+.. code-block::
+
+        >>> from transformers import TapasTokenizer
+        >>> import pandas as pd
+
+        >>> model_name = 'google/tapas-base'
+        >>> tokenizer = TapasTokenizer.from_pretrained(model_name)
+
+        >>> data = {'Actors': ["Brad Pitt", "Leonardo Di Caprio", "George Clooney"], 'Number of movies': ["87", "53", "69"]}
+        >>> queries = ["What is the name of the first actor?", "How many movies has George Clooney played in?", "What is the total number of movies?"]
+        >>> answer_coordinates = [[(0, 0)], [(2, 1)], [(0, 1), (1, 1), (2, 1)]]
+        >>> answer_text = [["Brad Pitt"], ["69"], ["209"]]
+        >>> table = pd.DataFrame.from_dict(data)
+        >>> inputs = tokenizer(table=table, queries=queries, answer_coordinates=answer_coordinates, answer_text=answer_text, padding='max_length', return_tensors='pt')
+        >>> inputs
+        {'input_ids': tensor([[ ... ]]), 'attention_mask': tensor([[...]]), 'token_type_ids': tensor([[[...]]]),
+        'numeric_values': tensor([[ ... ]]), 'numeric_values_scale: tensor([[ ... ]]), labels: tensor([[ ... ]])}
+
+Note that :class:`~transformers.TapasTokenizer` expects the data of the table to be **text-only**. You can use
+``.astype(str)`` on a dataframe to turn it into text-only data. Of course, this only shows how to encode a single
+training example. It is advised to create a PyTorch dataset and a corresponding dataloader:
+
+.. code-block::
+
+        >>> import torch
+        >>> import pandas as pd
+
+        >>> tsv_path = "your_path_to_the_tsv_file"
+        >>> table_csv_path = "your_path_to_a_directory_containing_all_csv_files"
+
+        >>> class TableDataset(torch.utils.data.Dataset):
+        ...     def __init__(self, data, tokenizer):
+        ...         self.data = data
+        ...         self.tokenizer = tokenizer
+        ...
+        ...     def __getitem__(self, idx):
+        ...         item = data.iloc[idx]
+        ...         table = pd.read_csv(table_csv_path + item.table_file).astype(str) # be sure to make your table data text only
+        ...         encoding = self.tokenizer(table=table, 
+        ...                                   queries=item.question, 
+        ...                                   answer_coordinates=item.answer_coordinates, 
+        ...                                   answer_text=item.answer_text,
+        ...                                   truncation=True,
+        ...                                   padding="max_length",
+        ...                                   return_tensors="pt"
+        ...         )
+        ...         # remove the batch dimension which the tokenizer adds by default
+        ...         encoding = {key: val.squeeze(0) for key, val in encoding.items()}
+        ...         # add the float_answer which is also required (weak supervision for aggregation case)
+        ...         encoding["float_answer"] = torch.tensor(item.float_answer) 
+        ...         return encoding
+        ...
+        ...     def __len__(self):
+        ...        return len(self.data)
+
+        >>> data = pd.read_csv(tsv_path, sep='\t')
+        >>> train_dataset = TableDataset(data, tokenizer)
+        >>> train_dataloader = torch.utils.data.DataLoader(train_dataset, batch_size=32)
+
+Note that here, we encode each table-question pair independently. This is fine as long as your dataset is **not
+conversational**. In case your dataset involves conversational questions (such as in SQA), then you should first group
+together the ``queries``, ``answer_coordinates`` and ``answer_text`` per table (in the order of their ``position``
+index) and batch encode each table with its questions. This will make sure that the ``prev_labels`` token types (see
+docs of :class:`~transformers.TapasTokenizer`) are set correctly. See `this notebook
+<https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Fine_tuning_TapasForQuestionAnswering_on_SQA.ipynb>`__
+for more info.
+
+**STEP 4: Train (fine-tune) TapasForQuestionAnswering**
+
+You can then fine-tune :class:`~transformers.TapasForQuestionAnswering` using native PyTorch as follows (shown here for
+the weak supervision for aggregation case):
+
+.. code-block::
+
+        >>> from transformers import TapasConfig, TapasForQuestionAnswering, AdamW
+
+        >>> # this is the default WTQ configuration
+        >>> config = TapasConfig(
+        ...            num_aggregation_labels = 4,
+        ...            use_answer_as_supervision = True,
+        ...            answer_loss_cutoff = 0.664694,
+        ...            cell_selection_preference = 0.207951,
+        ...            huber_loss_delta = 0.121194,
+        ...            init_cell_selection_weights_to_zero = True,
+        ...            select_one_column = True,
+        ...            allow_empty_column_selection = False,
+        ...            temperature = 0.0352513,
+        ... )
+        >>> model = TapasForQuestionAnswering.from_pretrained("google/tapas-base", config=config)
+
+        >>> optimizer = AdamW(model.parameters(), lr=5e-5)
+
+        >>> for epoch in range(2):  # loop over the dataset multiple times
+        ...    for idx, batch in enumerate(train_dataloader):
+        ...         # get the inputs; 
+        ...         input_ids = batch["input_ids"]
+        ...         attention_mask = batch["attention_mask"]
+        ...         token_type_ids = batch["token_type_ids"]
+        ...         labels = batch["labels"]
+        ...         numeric_values = batch["numeric_values"]
+        ...         numeric_values_scale = batch["numeric_values_scale"]
+        ...         float_answer = batch["float_answer"]
+
+        ...         # zero the parameter gradients
+        ...         optimizer.zero_grad()
+
+        ...         # forward + backward + optimize
+        ...         outputs = model(input_ids=input_ids, attention_mask=attention_mask, token_type_ids=token_type_ids, 
+        ...                        labels=labels, numeric_values=numeric_values, numeric_values_scale=numeric_values_scale, 
+        ...                        float_answer=float_answer)
+        ...         loss = outputs.loss
+        ...         loss.backward()
+        ...         optimizer.step()
+
+Usage: inference
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+Here we explain how you can use :class:`~transformers.TapasForQuestionAnswering` for inference (i.e. making predictions
+on new data). For inference, only ``input_ids``, ``attention_mask`` and ``token_type_ids`` (which you can obtain using
+:class:`~transformers.TapasTokenizer`) have to be provided to the model to obtain the logits. Next, you can use the
+handy ``convert_logits_to_predictions`` method of :class:`~transformers.TapasTokenizer` to convert these into predicted
+coordinates and optional aggregation indices.
+
+However, note that inference is **different** depending on whether or not the setup is conversational. In a
+non-conversational set-up, inference can be done in parallel on all table-question pairs of a batch. Here's an example
+of that:
+
+.. code-block::
+
+        >>> from transformers import TapasTokenizer, TapasForQuestionAnswering
+        >>> import pandas as pd 
+
+        >>> model_name = 'google/tapas-base-finetuned-wtq'
+        >>> model = TapasForQuestionAnswering.from_pretrained(model_name)
+        >>> tokenizer = TapasTokenizer.from_pretrained(model_name)
+
+        >>> data = {'Actors': ["Brad Pitt", "Leonardo Di Caprio", "George Clooney"], 'Number of movies': ["87", "53", "69"]}
+        >>> queries = ["What is the name of the first actor?", "How many movies has George Clooney played in?", "What is the total number of movies?"]
+        >>> table = pd.DataFrame.from_dict(data)
+        >>> inputs = tokenizer(table=table, queries=queries, padding='max_length', return_tensors="pt") 
+        >>> outputs = model(**inputs)
+        >>> predicted_answer_coordinates, predicted_aggregation_indices = tokenizer.convert_logits_to_predictions(
+        ...         inputs, 
+        ...         outputs.logits.detach(), 
+        ...         outputs.logits_aggregation.detach()
+        ... )
+
+        >>> # let's print out the results:
+        >>> id2aggregation = {0: "NONE", 1: "SUM", 2: "AVERAGE", 3:"COUNT"}
+        >>> aggregation_predictions_string = [id2aggregation[x] for x in predicted_aggregation_indices]
+
+        >>> answers = []
+        >>> for coordinates in predicted_answer_coordinates:
+        ...   if len(coordinates) == 1:
+        ...     # only a single cell:
+        ...     answers.append(table.iat[coordinates[0]])
+        ...   else:
+        ...     # multiple cells
+        ...     cell_values = []
+        ...     for coordinate in coordinates:
+        ...        cell_values.append(table.iat[coordinate])
+        ...     answers.append(", ".join(cell_values))
+
+        >>> display(table)
+        >>> print("")
+        >>> for query, answer, predicted_agg in zip(queries, answers, aggregation_predictions_string):
+        ...   print(query)
+        ...   if predicted_agg == "NONE":
+        ...     print("Predicted answer: " + answer)
+        ...   else:
+        ...     print("Predicted answer: " + predicted_agg + " > " + answer)    
+        What is the name of the first actor?
+        Predicted answer: Brad Pitt
+        How many movies has George Clooney played in?
+        Predicted answer: COUNT > 69
+        What is the total number of movies?
+        Predicted answer: SUM > 87, 53, 69
+
+In case of a conversational set-up, then each table-question pair must be provided **sequentially** to the model, such
+that the ``prev_labels`` token types can be overwritten by the predicted ``labels`` of the previous table-question
+pair. Again, more info can be found in `this notebook
+<https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Fine_tuning_TapasForQuestionAnswering_on_SQA.ipynb>`__.
+
+
+Tapas specific outputs
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.models.tapas.modeling_tapas.TableQuestionAnsweringOutput
+    :members:
+
+
+TapasConfig
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasConfig
+    :members:
+
+
+TapasTokenizer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasTokenizer
+    :members: __call__, convert_logits_to_predictions, save_vocabulary
+
+
+TapasModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasModel
+    :members: forward
+
+
+TapasForMaskedLM
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasForMaskedLM
+    :members: forward
+
+
+TapasForSequenceClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasForSequenceClassification
+    :members: forward
+
+
+TapasForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasForQuestionAnswering
+    :members: forward

docs/source/model_doc/transformerxl.rst

@@ -87,12 +87,14 @@ TransfoXLLMHeadModel
 .. autoclass:: transformers.TransfoXLLMHeadModel
     :members: forward
 
+
 TransfoXLForSequenceClassification
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: transformers.TransfoXLForSequenceClassification
     :members: forward
 
+
 TFTransfoXLModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -107,6 +109,13 @@ TFTransfoXLLMHeadModel
     :members: call
 
 
+TFTransfoXLForSequenceClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFTransfoXLForSequenceClassification
+    :members: call
+
+
 Internal Layers
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/preprocessing.rst

@@ -10,7 +10,7 @@
     an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. See the License for the
     specific language governing permissions and limitations under the License.
 
-reprocessing data
+Preprocessing data
 =======================================================================================================================
 
 In this tutorial, we'll explore how to preprocess your data using 🤗 Transformers. The main tool for this is what we

examples/README.md

@@ -54,12 +54,12 @@ Coming soon!
 | Task | Example datasets | Trainer support | TFTrainer support | 🤗 Datasets | Colab
 |---|---|:---:|:---:|:---:|:---:|
 | [**`language-modeling`**](https://github.com/huggingface/transformers/tree/master/examples/language-modeling)       | Raw text        | ✅ | -  | ✅ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)
-| [**`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           | ✅ | ✅ | ✅ | -
+| [**`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           | ✅ | ✅ | ✅ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/huggingface/notebooks/blob/master/examples/question_answering.ipynb)
 | [**`summarization`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)                     | CNN/Daily Mail  | ✅  | - | - | -
 | [**`text-classification`**](https://github.com/huggingface/transformers/tree/master/examples/text-classification)   | GLUE, XNLI      | ✅ | ✅ | ✅ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/huggingface/notebooks/blob/master/examples/text_classification.ipynb)
 | [**`text-generation`**](https://github.com/huggingface/transformers/tree/master/examples/text-generation)           | -               | 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)
-| [**`token-classification`**](https://github.com/huggingface/transformers/tree/master/examples/token-classification) | CoNLL NER       | ✅ | ✅ | ✅ | -
+| [**`token-classification`**](https://github.com/huggingface/transformers/tree/master/examples/token-classification) | CoNLL NER       | ✅ | ✅ | ✅ | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://github.com/huggingface/notebooks/blob/master/examples/token_classification.ipynb)
 | [**`translation`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)                       | WMT             | ✅  | - | - | -
 
 
@@ -69,6 +69,43 @@ Coming soon!
 **Coming soon!**
 -->
 
+## Distributed training and mixed precision
+
+All the PyTorch scripts mentioned above work out of the box with distributed training and mixed precision, thanks to
+the [Trainer API](https://huggingface.co/transformers/main_classes/trainer.html). To launch one of them on _n_ GPUS,
+use the following command:
+
+```bash
+python -m torch.distributed.launch \
+    --nproc_per_node number_of_gpu_you_have path_to_script.py \
+	--all_arguments_of_the_script 
+```
+
+As an example, here is how you would fine-tune the BERT large model (with whole word masking) on the text
+classification MNLI task using the `run_glue` script, with 8 GPUs:
+
+```bash
+python -m torch.distributed.launch \
+    --nproc_per_node 8 text-classification/run_glue.py \
+    --model_name_or_path bert-large-uncased-whole-word-masking \
+    --task_name mnli \
+    --do_train \
+    --do_eval \
+    --max_seq_length 128 \
+    --per_device_train_batch_size 8 \
+    --learning_rate 2e-5 \
+    --num_train_epochs 3.0 \
+    --output_dir /tmp/mnli_output/
+```
+
+If you have a GPU with mixed precision capabilities (architecture Pascal or more recent), you can use mixed precision
+training with PyTorch 1.6.0 or latest, or by installing the [Apex](https://github.com/NVIDIA/apex) library for previous
+versions. Just add the flag `--fp16` to your command launching one of the scripts mentioned above!
+
+Using mixed precision training usually results in 2x-speedup for training with the same final results (as shown in
+[this table](https://github.com/huggingface/transformers/tree/master/examples/text-classification#mixed-precision-training)
+for text classification).
+
 ## Running on TPUs
 
 When using Tensorflow, TPUs are supported out of the box as a `tf.distribute.Strategy`.
@@ -76,27 +113,34 @@ When using Tensorflow, TPUs are supported out of the box as a `tf.distribute.Str
 When using PyTorch, we support TPUs thanks to `pytorch/xla`. For more context and information on how to setup your TPU environment refer to Google's documentation and to the
 very detailed [pytorch/xla README](https://github.com/pytorch/xla/blob/master/README.md).
 
-In this repo, we provide a very simple launcher script named [xla_spawn.py](https://github.com/huggingface/transformers/tree/master/examples/xla_spawn.py) that lets you run our example scripts on multiple TPU cores without any boilerplate.
-Just pass a `--num_cores` flag to this script, then your regular training script with its arguments (this is similar to the `torch.distributed.launch` helper for torch.distributed). 
-Note that this approach does not work for examples that use `pytorch-lightning`.
-
-For example for `run_glue`:
+In this repo, we provide a very simple launcher script named
+[xla_spawn.py](https://github.com/huggingface/transformers/tree/master/examples/xla_spawn.py) that lets you run our
+example scripts on multiple TPU cores without any boilerplate. Just pass a `--num_cores` flag to this script, then your
+regular training script with its arguments (this is similar to the `torch.distributed.launch` helper for
+`torch.distributed`):
 
 ```bash
-python examples/xla_spawn.py --num_cores 8 \
-	examples/text-classification/run_glue.py \
-	--model_name_or_path bert-base-cased \
-	--task_name mnli \
-	--data_dir ./data/glue_data/MNLI \
-	--output_dir ./models/tpu \
-	--overwrite_output_dir \
-	--do_train \
-	--do_eval \
-	--num_train_epochs 1 \
-	--save_steps 20000
+python xla_spawn.py --num_cores num_tpu_you_have \
+    path_to_script.py \
+	--all_arguments_of_the_script 
 ```
 
-Feedback and more use cases and benchmarks involving TPUs are welcome, please share with the community.
+As an example, here is how you would fine-tune the BERT large model (with whole word masking) on the text
+classification MNLI task using the `run_glue` script, with 8 TPUs:
+
+```bash
+python xla_spawn.py --num_cores 8 \
+    text-classification/run_glue.py \
+    --model_name_or_path bert-large-uncased-whole-word-masking \
+    --task_name mnli \
+    --do_train \
+    --do_eval \
+    --max_seq_length 128 \
+    --per_device_train_batch_size 8 \
+    --learning_rate 2e-5 \
+    --num_train_epochs 3.0 \
+    --output_dir /tmp/mnli_output/
+```
 
 ## Logging & Experiment tracking
 

examples/language-modeling/README.md

@@ -25,8 +25,7 @@ objectives in our [model summary](https://huggingface.co/transformers/model_summ
 These scripts leverage the 🤗 Datasets library and the Trainer API. You can easily customize them to your needs if you
 need extra processing on your datasets.
 
-**Note:** The old script `run_language_modeling.py` is still available
-[here](https://github.com/huggingface/transformers/blob/master/examples/contrib/legacy/run_language_modeling.py).
+**Note:** The old script `run_language_modeling.py` is still available [here](https://github.com/huggingface/transformers/blob/master/examples/legacy/run_language_modeling.py).
 
 The following examples, will run on a datasets hosted on our [hub](https://huggingface.co/datasets) or with your own
 text files for training and validation. We give examples of both below.

examples/language-modeling/run_clm.py

@@ -113,6 +113,12 @@ class DataTrainingArguments:
     overwrite_cache: bool = field(
         default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
     )
+    validation_split_percentage: Optional[int] = field(
+        default=5,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
     preprocessing_num_workers: Optional[int] = field(
         default=None,
         metadata={"help": "The number of processes to use for the preprocessing."},
@@ -188,6 +194,17 @@ def main():
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
         datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
+        if "validation" not in datasets.keys():
+            datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+            )
+            datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+            )
     else:
         data_files = {}
         if data_args.train_file is not None:
@@ -324,9 +341,20 @@ def main():
             if (model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path))
             else None
         )
-        trainer.train(model_path=model_path)
+        train_result = trainer.train(model_path=model_path)
         trainer.save_model()  # Saves the tokenizer too for easy upload
 
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
     # Evaluation
     results = {}
     if training_args.do_eval:
@@ -341,7 +369,7 @@ def main():
         if trainer.is_world_process_zero():
             with open(output_eval_file, "w") as writer:
                 logger.info("***** Eval results *****")
-                for key, value in results.items():
+                for key, value in sorted(results.items()):
                     logger.info(f"  {key} = {value}")
                     writer.write(f"{key} = {value}\n")
 

examples/language-modeling/run_mlm.py

@@ -103,6 +103,12 @@ class DataTrainingArguments:
     overwrite_cache: bool = field(
         default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
     )
+    validation_split_percentage: Optional[int] = field(
+        default=5,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
     max_seq_length: Optional[int] = field(
         default=None,
         metadata={
@@ -199,6 +205,17 @@ def main():
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
         datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
+        if "validation" not in datasets.keys():
+            datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+            )
+            datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+            )
     else:
         data_files = {}
         if data_args.train_file is not None:
@@ -359,9 +376,20 @@ def main():
             if (model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path))
             else None
         )
-        trainer.train(model_path=model_path)
+        train_result = trainer.train(model_path=model_path)
         trainer.save_model()  # Saves the tokenizer too for easy upload
 
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
     # Evaluation
     results = {}
     if training_args.do_eval:
@@ -376,7 +404,7 @@ def main():
         if trainer.is_world_process_zero():
             with open(output_eval_file, "w") as writer:
                 logger.info("***** Eval results *****")
-                for key, value in results.items():
+                for key, value in sorted(results.items()):
                     logger.info(f"  {key} = {value}")
                     writer.write(f"{key} = {value}\n")
 

examples/language-modeling/run_mlm_flax.py

@@ -134,6 +134,12 @@ class DataTrainingArguments:
     overwrite_cache: bool = field(
         default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
     )
+    validation_split_percentage: Optional[int] = field(
+        default=5,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
     max_seq_length: Optional[int] = field(
         default=None,
         metadata={
@@ -379,7 +385,7 @@ def training_step(optimizer, batch, dropout_rng):
         # Hide away tokens which doesn't participate in the optimization
         token_mask = jnp.where(targets > 0, 1.0, 0.0)
 
-        pooled, logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)
+        logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
         loss, weight_sum = cross_entropy(logits, targets, token_mask)
         return loss / weight_sum
 
@@ -401,7 +407,7 @@ def eval_step(params, batch):
 
     # Hide away tokens which doesn't participate in the optimization
     token_mask = jnp.where(targets > 0, 1.0, 0.0)
-    _, logits = model(**batch, params=params, train=False)
+    logits = model(**batch, params=params, train=False)[0]
 
     return compute_metrics(logits, targets, token_mask)
 
@@ -413,7 +419,7 @@ def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndar
     if samples_to_remove != 0:
         samples_idx = samples_idx[:-samples_to_remove]
     sections_split = nb_samples // batch_size
-    batch_idx = jnp.split(samples_idx, sections_split)
+    batch_idx = np.split(samples_idx, sections_split)
     return batch_idx
 
 
@@ -473,6 +479,17 @@ if __name__ == "__main__":
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
         datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
+        if "validation" not in datasets.keys():
+            datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+            )
+            datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+            )
     else:
         data_files = {}
         if data_args.train_file is not None:
@@ -525,9 +542,9 @@ if __name__ == "__main__":
 
     def tokenize_function(examples):
         # Remove empty lines
-        examples["text"] = [line for line in examples["text"] if len(line) > 0 and not line.isspace()]
+        examples = [line for line in examples if len(line) > 0 and not line.isspace()]
         return tokenizer(
-            examples["text"],
+            examples,
             return_special_tokens_mask=True,
             padding=padding,
             truncation=True,
@@ -536,9 +553,10 @@ if __name__ == "__main__":
 
     tokenized_datasets = datasets.map(
         tokenize_function,
+        input_columns=[text_column_name],
         batched=True,
         num_proc=data_args.preprocessing_num_workers,
-        remove_columns=[text_column_name],
+        remove_columns=column_names,
         load_from_cache_file=not data_args.overwrite_cache,
     )
 
@@ -554,8 +572,13 @@ if __name__ == "__main__":
     rng = jax.random.PRNGKey(training_args.seed)
     dropout_rngs = jax.random.split(rng, jax.local_device_count())
 
-    model = FlaxBertForMaskedLM.from_pretrained("bert-base-cased", dtype=jnp.float32, dropout_rate=0.1)
-    model.init(jax.random.PRNGKey(training_args.seed), (training_args.train_batch_size, model.config.max_length))
+    model = FlaxBertForMaskedLM.from_pretrained(
+        "bert-base-cased",
+        dtype=jnp.float32,
+        input_shape=(training_args.train_batch_size, config.max_position_embeddings),
+        seed=training_args.seed,
+        dropout_rate=0.1,
+    )
 
     # Setup optimizer
     optimizer = Adam(
@@ -566,8 +589,9 @@ if __name__ == "__main__":
     ).create(model.params)
 
     # Create learning rate scheduler
+    # warmup_steps = 0 causes the Flax optimizer to return NaNs; warmup_steps = 1 is functionally equivalent.
     lr_scheduler_fn = create_learning_rate_scheduler(
-        base_learning_rate=training_args.learning_rate, warmup_steps=training_args.warmup_steps
+        base_learning_rate=training_args.learning_rate, warmup_steps=min(training_args.warmup_steps, 1)
     )
 
     # Create parallel version of the training and evaluation steps
@@ -606,13 +630,13 @@ if __name__ == "__main__":
         epochs.write(f"Loss: {loss}")
 
         # ======================== Evaluating ==============================
-        nb_eval_samples = len(tokenized_datasets["test"])
+        nb_eval_samples = len(tokenized_datasets["validation"])
         eval_samples_idx = jnp.arange(nb_eval_samples)
         eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)
 
         eval_metrics = []
         for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
-            samples = [tokenized_datasets["test"][int(idx)] for idx in batch_idx]
+            samples = [tokenized_datasets["validation"][int(idx)] for idx in batch_idx]
             model_inputs = data_collator(samples, pad_to_multiple_of=16)
 
             # Model forward

examples/language-modeling/run_mlm_wwm.py

@@ -91,6 +91,12 @@ class DataTrainingArguments:
     Arguments pertaining to what data we are going to input our model for training and eval.
     """
 
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
     train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
     validation_file: Optional[str] = field(
         default=None,
@@ -107,6 +113,12 @@ class DataTrainingArguments:
     overwrite_cache: bool = field(
         default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
     )
+    validation_split_percentage: Optional[int] = field(
+        default=5,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
     max_seq_length: Optional[int] = field(
         default=None,
         metadata={
@@ -203,15 +215,30 @@ def main():
     #
     # In distributed training, the load_dataset function guarantee that only one local process can concurrently
     # download the dataset.
-    data_files = {}
-    if data_args.train_file is not None:
-        data_files["train"] = data_args.train_file
-    if data_args.validation_file is not None:
-        data_files["validation"] = data_args.validation_file
-    extension = data_args.train_file.split(".")[-1]
-    if extension == "txt":
-        extension = "text"
-    datasets = load_dataset(extension, data_files=data_files)
+    if data_args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
+        if "validation" not in datasets.keys():
+            datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+            )
+            datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+            )
+    else:
+        data_files = {}
+        if data_args.train_file is not None:
+            data_files["train"] = data_args.train_file
+        if data_args.validation_file is not None:
+            data_files["validation"] = data_args.validation_file
+        extension = data_args.train_file.split(".")[-1]
+        if extension == "txt":
+            extension = "text"
+        datasets = load_dataset(extension, data_files=data_files)
     # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
     # https://huggingface.co/docs/datasets/loading_datasets.html.
 
@@ -307,9 +334,20 @@ def main():
             if (model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path))
             else None
         )
-        trainer.train(model_path=model_path)
+        train_result = trainer.train(model_path=model_path)
         trainer.save_model()  # Saves the tokenizer too for easy upload
 
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
     # Evaluation
     results = {}
     if training_args.do_eval:
@@ -324,7 +362,7 @@ def main():
         if trainer.is_world_process_zero():
             with open(output_eval_file, "w") as writer:
                 logger.info("***** Eval results *****")
-                for key, value in results.items():
+                for key, value in sorted(results.items()):
                     logger.info(f"  {key} = {value}")
                     writer.write(f"{key} = {value}\n")
 

examples/language-modeling/run_plm.py

@@ -93,6 +93,12 @@ class DataTrainingArguments:
     overwrite_cache: bool = field(
         default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
     )
+    validation_split_percentage: Optional[int] = field(
+        default=5,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
     max_seq_length: int = field(
         default=512,
         metadata={
@@ -196,6 +202,17 @@ def main():
     if data_args.dataset_name is not None:
         # Downloading and loading a dataset from the hub.
         datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
+        if "validation" not in datasets.keys():
+            datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+            )
+            datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+            )
     else:
         data_files = {}
         if data_args.train_file is not None:
@@ -346,9 +363,20 @@ def main():
             if (model_args.model_name_or_path is not None and os.path.isdir(model_args.model_name_or_path))
             else None
         )
-        trainer.train(model_path=model_path)
+        train_result = trainer.train(model_path=model_path)
         trainer.save_model()  # Saves the tokenizer too for easy upload
 
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
     # Evaluation
     results = {}
     if training_args.do_eval:
@@ -363,7 +391,7 @@ def main():
         if trainer.is_world_process_zero():
             with open(output_eval_file, "w") as writer:
                 logger.info("***** Eval results *****")
-                for key, value in results.items():
+                for key, value in sorted(results.items()):
                     logger.info(f"  {key} = {value}")
                     writer.write(f"{key} = {value}\n")
 

examples/legacy/multiple_choice/utils_multiple_choice.py

@@ -0,0 +1,579 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Multiple choice fine-tuning: utilities to work with multiple choice tasks of reading comprehension """
+
+
+import csv
+import glob
+import json
+import logging
+import os
+from dataclasses import dataclass
+from enum import Enum
+from typing import List, Optional
+
+import tqdm
+
+from filelock import FileLock
+from transformers import PreTrainedTokenizer, is_tf_available, is_torch_available
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass(frozen=True)
+class InputExample:
+    """
+    A single training/test example for multiple choice
+
+    Args:
+        example_id: Unique id for the example.
+        question: string. The untokenized text of the second sequence (question).
+        contexts: list of str. The untokenized text of the first sequence (context of corresponding question).
+        endings: list of str. multiple choice's options. Its length must be equal to contexts' length.
+        label: (Optional) string. The label of the example. This should be
+        specified for train and dev examples, but not for test examples.
+    """
+
+    example_id: str
+    question: str
+    contexts: List[str]
+    endings: List[str]
+    label: Optional[str]
+
+
+@dataclass(frozen=True)
+class InputFeatures:
+    """
+    A single set of features of data.
+    Property names are the same names as the corresponding inputs to a model.
+    """
+
+    example_id: str
+    input_ids: List[List[int]]
+    attention_mask: Optional[List[List[int]]]
+    token_type_ids: Optional[List[List[int]]]
+    label: Optional[int]
+
+
+class Split(Enum):
+    train = "train"
+    dev = "dev"
+    test = "test"
+
+
+if is_torch_available():
+    import torch
+    from torch.utils.data.dataset import Dataset
+
+    class MultipleChoiceDataset(Dataset):
+        """
+        This will be superseded by a framework-agnostic approach
+        soon.
+        """
+
+        features: List[InputFeatures]
+
+        def __init__(
+            self,
+            data_dir: str,
+            tokenizer: PreTrainedTokenizer,
+            task: str,
+            max_seq_length: Optional[int] = None,
+            overwrite_cache=False,
+            mode: Split = Split.train,
+        ):
+            processor = processors[task]()
+
+            cached_features_file = os.path.join(
+                data_dir,
+                "cached_{}_{}_{}_{}".format(
+                    mode.value,
+                    tokenizer.__class__.__name__,
+                    str(max_seq_length),
+                    task,
+                ),
+            )
+
+            # Make sure only the first process in distributed training processes the dataset,
+            # and the others will use the cache.
+            lock_path = cached_features_file + ".lock"
+            with FileLock(lock_path):
+
+                if os.path.exists(cached_features_file) and not overwrite_cache:
+                    logger.info(f"Loading features from cached file {cached_features_file}")
+                    self.features = torch.load(cached_features_file)
+                else:
+                    logger.info(f"Creating features from dataset file at {data_dir}")
+                    label_list = processor.get_labels()
+                    if mode == Split.dev:
+                        examples = processor.get_dev_examples(data_dir)
+                    elif mode == Split.test:
+                        examples = processor.get_test_examples(data_dir)
+                    else:
+                        examples = processor.get_train_examples(data_dir)
+                    logger.info("Training examples: %s", len(examples))
+                    self.features = convert_examples_to_features(
+                        examples,
+                        label_list,
+                        max_seq_length,
+                        tokenizer,
+                    )
+                    logger.info("Saving features into cached file %s", cached_features_file)
+                    torch.save(self.features, cached_features_file)
+
+        def __len__(self):
+            return len(self.features)
+
+        def __getitem__(self, i) -> InputFeatures:
+            return self.features[i]
+
+
+if is_tf_available():
+    import tensorflow as tf
+
+    class TFMultipleChoiceDataset:
+        """
+        This will be superseded by a framework-agnostic approach
+        soon.
+        """
+
+        features: List[InputFeatures]
+
+        def __init__(
+            self,
+            data_dir: str,
+            tokenizer: PreTrainedTokenizer,
+            task: str,
+            max_seq_length: Optional[int] = 128,
+            overwrite_cache=False,
+            mode: Split = Split.train,
+        ):
+            processor = processors[task]()
+
+            logger.info(f"Creating features from dataset file at {data_dir}")
+            label_list = processor.get_labels()
+            if mode == Split.dev:
+                examples = processor.get_dev_examples(data_dir)
+            elif mode == Split.test:
+                examples = processor.get_test_examples(data_dir)
+            else:
+                examples = processor.get_train_examples(data_dir)
+            logger.info("Training examples: %s", len(examples))
+
+            self.features = convert_examples_to_features(
+                examples,
+                label_list,
+                max_seq_length,
+                tokenizer,
+            )
+
+            def gen():
+                for (ex_index, ex) in tqdm.tqdm(enumerate(self.features), desc="convert examples to features"):
+                    if ex_index % 10000 == 0:
+                        logger.info("Writing example %d of %d" % (ex_index, len(examples)))
+
+                    yield (
+                        {
+                            "example_id": 0,
+                            "input_ids": ex.input_ids,
+                            "attention_mask": ex.attention_mask,
+                            "token_type_ids": ex.token_type_ids,
+                        },
+                        ex.label,
+                    )
+
+            self.dataset = tf.data.Dataset.from_generator(
+                gen,
+                (
+                    {
+                        "example_id": tf.int32,
+                        "input_ids": tf.int32,
+                        "attention_mask": tf.int32,
+                        "token_type_ids": tf.int32,
+                    },
+                    tf.int64,
+                ),
+                (
+                    {
+                        "example_id": tf.TensorShape([]),
+                        "input_ids": tf.TensorShape([None, None]),
+                        "attention_mask": tf.TensorShape([None, None]),
+                        "token_type_ids": tf.TensorShape([None, None]),
+                    },
+                    tf.TensorShape([]),
+                ),
+            )
+
+        def get_dataset(self):
+            self.dataset = self.dataset.apply(tf.data.experimental.assert_cardinality(len(self.features)))
+
+            return self.dataset
+
+        def __len__(self):
+            return len(self.features)
+
+        def __getitem__(self, i) -> InputFeatures:
+            return self.features[i]
+
+
+class DataProcessor:
+    """Base class for data converters for multiple choice data sets."""
+
+    def get_train_examples(self, data_dir):
+        """Gets a collection of `InputExample`s for the train set."""
+        raise NotImplementedError()
+
+    def get_dev_examples(self, data_dir):
+        """Gets a collection of `InputExample`s for the dev set."""
+        raise NotImplementedError()
+
+    def get_test_examples(self, data_dir):
+        """Gets a collection of `InputExample`s for the test set."""
+        raise NotImplementedError()
+
+    def get_labels(self):
+        """Gets the list of labels for this data set."""
+        raise NotImplementedError()
+
+
+class RaceProcessor(DataProcessor):
+    """Processor for the RACE data set."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} train".format(data_dir))
+        high = os.path.join(data_dir, "train/high")
+        middle = os.path.join(data_dir, "train/middle")
+        high = self._read_txt(high)
+        middle = self._read_txt(middle)
+        return self._create_examples(high + middle, "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} dev".format(data_dir))
+        high = os.path.join(data_dir, "dev/high")
+        middle = os.path.join(data_dir, "dev/middle")
+        high = self._read_txt(high)
+        middle = self._read_txt(middle)
+        return self._create_examples(high + middle, "dev")
+
+    def get_test_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} test".format(data_dir))
+        high = os.path.join(data_dir, "test/high")
+        middle = os.path.join(data_dir, "test/middle")
+        high = self._read_txt(high)
+        middle = self._read_txt(middle)
+        return self._create_examples(high + middle, "test")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1", "2", "3"]
+
+    def _read_txt(self, input_dir):
+        lines = []
+        files = glob.glob(input_dir + "/*txt")
+        for file in tqdm.tqdm(files, desc="read files"):
+            with open(file, "r", encoding="utf-8") as fin:
+                data_raw = json.load(fin)
+                data_raw["race_id"] = file
+                lines.append(data_raw)
+        return lines
+
+    def _create_examples(self, lines, set_type):
+        """Creates examples for the training and dev sets."""
+        examples = []
+        for (_, data_raw) in enumerate(lines):
+            race_id = "%s-%s" % (set_type, data_raw["race_id"])
+            article = data_raw["article"]
+            for i in range(len(data_raw["answers"])):
+                truth = str(ord(data_raw["answers"][i]) - ord("A"))
+                question = data_raw["questions"][i]
+                options = data_raw["options"][i]
+
+                examples.append(
+                    InputExample(
+                        example_id=race_id,
+                        question=question,
+                        contexts=[article, article, article, article],  # this is not efficient but convenient
+                        endings=[options[0], options[1], options[2], options[3]],
+                        label=truth,
+                    )
+                )
+        return examples
+
+
+class SynonymProcessor(DataProcessor):
+    """Processor for the Synonym data set."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} train".format(data_dir))
+        return self._create_examples(self._read_csv(os.path.join(data_dir, "mctrain.csv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} dev".format(data_dir))
+        return self._create_examples(self._read_csv(os.path.join(data_dir, "mchp.csv")), "dev")
+
+    def get_test_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} dev".format(data_dir))
+
+        return self._create_examples(self._read_csv(os.path.join(data_dir, "mctest.csv")), "test")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1", "2", "3", "4"]
+
+    def _read_csv(self, input_file):
+        with open(input_file, "r", encoding="utf-8") as f:
+            return list(csv.reader(f))
+
+    def _create_examples(self, lines: List[List[str]], type: str):
+        """Creates examples for the training and dev sets."""
+
+        examples = [
+            InputExample(
+                example_id=line[0],
+                question="",  # in the swag dataset, the
+                # common beginning of each
+                # choice is stored in "sent2".
+                contexts=[line[1], line[1], line[1], line[1], line[1]],
+                endings=[line[2], line[3], line[4], line[5], line[6]],
+                label=line[7],
+            )
+            for line in lines  # we skip the line with the column names
+        ]
+
+        return examples
+
+
+class SwagProcessor(DataProcessor):
+    """Processor for the SWAG data set."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} train".format(data_dir))
+        return self._create_examples(self._read_csv(os.path.join(data_dir, "train.csv")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} dev".format(data_dir))
+        return self._create_examples(self._read_csv(os.path.join(data_dir, "val.csv")), "dev")
+
+    def get_test_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} dev".format(data_dir))
+        raise ValueError(
+            "For swag testing, the input file does not contain a label column. It can not be tested in current code"
+            "setting!"
+        )
+        return self._create_examples(self._read_csv(os.path.join(data_dir, "test.csv")), "test")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1", "2", "3"]
+
+    def _read_csv(self, input_file):
+        with open(input_file, "r", encoding="utf-8") as f:
+            return list(csv.reader(f))
+
+    def _create_examples(self, lines: List[List[str]], type: str):
+        """Creates examples for the training and dev sets."""
+        if type == "train" and lines[0][-1] != "label":
+            raise ValueError("For training, the input file must contain a label column.")
+
+        examples = [
+            InputExample(
+                example_id=line[2],
+                question=line[5],  # in the swag dataset, the
+                # common beginning of each
+                # choice is stored in "sent2".
+                contexts=[line[4], line[4], line[4], line[4]],
+                endings=[line[7], line[8], line[9], line[10]],
+                label=line[11],
+            )
+            for line in lines[1:]  # we skip the line with the column names
+        ]
+
+        return examples
+
+
+class ArcProcessor(DataProcessor):
+    """Processor for the ARC data set (request from allennlp)."""
+
+    def get_train_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} train".format(data_dir))
+        return self._create_examples(self._read_json(os.path.join(data_dir, "train.jsonl")), "train")
+
+    def get_dev_examples(self, data_dir):
+        """See base class."""
+        logger.info("LOOKING AT {} dev".format(data_dir))
+        return self._create_examples(self._read_json(os.path.join(data_dir, "dev.jsonl")), "dev")
+
+    def get_test_examples(self, data_dir):
+        logger.info("LOOKING AT {} test".format(data_dir))
+        return self._create_examples(self._read_json(os.path.join(data_dir, "test.jsonl")), "test")
+
+    def get_labels(self):
+        """See base class."""
+        return ["0", "1", "2", "3"]
+
+    def _read_json(self, input_file):
+        with open(input_file, "r", encoding="utf-8") as fin:
+            lines = fin.readlines()
+            return lines
+
+    def _create_examples(self, lines, type):
+        """Creates examples for the training and dev sets."""
+
+        # There are two types of labels. They should be normalized
+        def normalize(truth):
+            if truth in "ABCD":
+                return ord(truth) - ord("A")
+            elif truth in "1234":
+                return int(truth) - 1
+            else:
+                logger.info("truth ERROR! %s", str(truth))
+                return None
+
+        examples = []
+        three_choice = 0
+        four_choice = 0
+        five_choice = 0
+        other_choices = 0
+        # we deleted example which has more than or less than four choices
+        for line in tqdm.tqdm(lines, desc="read arc data"):
+            data_raw = json.loads(line.strip("\n"))
+            if len(data_raw["question"]["choices"]) == 3:
+                three_choice += 1
+                continue
+            elif len(data_raw["question"]["choices"]) == 5:
+                five_choice += 1
+                continue
+            elif len(data_raw["question"]["choices"]) != 4:
+                other_choices += 1
+                continue
+            four_choice += 1
+            truth = str(normalize(data_raw["answerKey"]))
+            assert truth != "None"
+            question_choices = data_raw["question"]
+            question = question_choices["stem"]
+            id = data_raw["id"]
+            options = question_choices["choices"]
+            if len(options) == 4:
+                examples.append(
+                    InputExample(
+                        example_id=id,
+                        question=question,
+                        contexts=[
+                            options[0]["para"].replace("_", ""),
+                            options[1]["para"].replace("_", ""),
+                            options[2]["para"].replace("_", ""),
+                            options[3]["para"].replace("_", ""),
+                        ],
+                        endings=[options[0]["text"], options[1]["text"], options[2]["text"], options[3]["text"]],
+                        label=truth,
+                    )
+                )
+
+        if type == "train":
+            assert len(examples) > 1
+            assert examples[0].label is not None
+        logger.info("len examples: %s}", str(len(examples)))
+        logger.info("Three choices: %s", str(three_choice))
+        logger.info("Five choices: %s", str(five_choice))
+        logger.info("Other choices: %s", str(other_choices))
+        logger.info("four choices: %s", str(four_choice))
+
+        return examples
+
+
+def convert_examples_to_features(
+    examples: List[InputExample],
+    label_list: List[str],
+    max_length: int,
+    tokenizer: PreTrainedTokenizer,
+) -> List[InputFeatures]:
+    """
+    Loads a data file into a list of `InputFeatures`
+    """
+
+    label_map = {label: i for i, label in enumerate(label_list)}
+
+    features = []
+    for (ex_index, example) in tqdm.tqdm(enumerate(examples), desc="convert examples to features"):
+        if ex_index % 10000 == 0:
+            logger.info("Writing example %d of %d" % (ex_index, len(examples)))
+        choices_inputs = []
+        for ending_idx, (context, ending) in enumerate(zip(example.contexts, example.endings)):
+            text_a = context
+            if example.question.find("_") != -1:
+                # this is for cloze question
+                text_b = example.question.replace("_", ending)
+            else:
+                text_b = example.question + " " + ending
+
+            inputs = tokenizer(
+                text_a,
+                text_b,
+                add_special_tokens=True,
+                max_length=max_length,
+                padding="max_length",
+                truncation=True,
+                return_overflowing_tokens=True,
+            )
+            if "num_truncated_tokens" in inputs and inputs["num_truncated_tokens"] > 0:
+                logger.info(
+                    "Attention! you are cropping tokens (swag task is ok). "
+                    "If you are training ARC and RACE and you are poping question + options,"
+                    "you need to try to use a bigger max seq length!"
+                )
+
+            choices_inputs.append(inputs)
+
+        label = label_map[example.label]
+
+        input_ids = [x["input_ids"] for x in choices_inputs]
+        attention_mask = (
+            [x["attention_mask"] for x in choices_inputs] if "attention_mask" in choices_inputs[0] else None
+        )
+        token_type_ids = (
+            [x["token_type_ids"] for x in choices_inputs] if "token_type_ids" in choices_inputs[0] else None
+        )
+
+        features.append(
+            InputFeatures(
+                example_id=example.example_id,
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                token_type_ids=token_type_ids,
+                label=label,
+            )
+        )
+
+    for f in features[:2]:
+        logger.info("*** Example ***")
+        logger.info("feature: %s" % f)
+
+    return features
+
+
+processors = {"race": RaceProcessor, "swag": SwagProcessor, "arc": ArcProcessor, "syn": SynonymProcessor}
+MULTIPLE_CHOICE_TASKS_NUM_LABELS = {"race", 4, "swag", 4, "arc", 4, "syn", 5}

examples/legacy/pytorch-lightning/requirements.txt

@@ -19,3 +19,4 @@ pytest
 conllu
 sentencepiece != 0.1.92
 protobuf
+ray

examples/multiple-choice/README.md

@@ -16,27 +16,20 @@ limitations under the License.
 
 ## Multiple Choice
 
-Based on the script [`run_multiple_choice.py`]().
+Based on the script [`run_swag.py`]().
 
 #### Fine-tuning on SWAG
-Download [swag](https://github.com/rowanz/swagaf/tree/master/data) data
 
 ```bash
-#training on 4 tesla V100(16GB) GPUS
-export SWAG_DIR=/path/to/swag_data_dir
-python ./examples/multiple-choice/run_multiple_choice.py \
---task_name swag \
+python examples/multiple-choice/run_swag.py \
 --model_name_or_path roberta-base \
 --do_train \
 --do_eval \
---data_dir $SWAG_DIR \
 --learning_rate 5e-5 \
 --num_train_epochs 3 \
---max_seq_length 80 \
---output_dir models_bert/swag_base \
+--output_dir /tmp/swag_base \
 --per_gpu_eval_batch_size=16 \
 --per_device_train_batch_size=16 \
---gradient_accumulation_steps 2 \
 --overwrite_output
 ```
 Training with the defined hyper-parameters yields the following results:

examples/multiple-choice/run_swag.py

@@ -0,0 +1,360 @@
+# coding=utf-8
+# Copyright The HuggingFace Team and The HuggingFace Inc. team. All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for multiple choice.
+"""
+# You can also adapt this script on your own multiple choice task. Pointers for this are left as comments.
+
+import logging
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Optional, Union
+
+import numpy as np
+import torch
+from datasets import load_dataset
+
+import transformers
+from transformers import (
+    AutoConfig,
+    AutoModelForMultipleChoice,
+    AutoTokenizer,
+    HfArgumentParser,
+    Trainer,
+    TrainingArguments,
+    default_data_collator,
+    set_seed,
+)
+from transformers.tokenization_utils_base import PaddingStrategy, PreTrainedTokenizerBase
+from transformers.trainer_utils import is_main_process
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None,
+        metadata={"help": "Where do you want to store the pretrained models downloaded from huggingface.co"},
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
+    validation_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    max_seq_length: int = field(
+        default=None,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. If passed, sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to pad all samples to the maximum sentence length. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch. More "
+            "efficient on GPU but very bad for TPU."
+        },
+    )
+
+    def __post_init__(self):
+        if self.train_file is not None:
+            extension = self.train_file.split(".")[-1]
+            assert extension in ["csv", "json"], "`train_file` should be a csv or a json file."
+        if self.validation_file is not None:
+            extension = self.validation_file.split(".")[-1]
+            assert extension in ["csv", "json"], "`validation_file` should be a csv or a json file."
+
+
+@dataclass
+class DataCollatorForMultipleChoice:
+    """
+    Data collator that will dynamically pad the inputs for multiple choice received.
+
+    Args:
+        tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
+            The tokenizer used for encoding the data.
+        padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`True`):
+            Select a strategy to pad the returned sequences (according to the model's padding side and padding index)
+            among:
+
+            * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a single
+              sequence if provided).
+            * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+              maximum acceptable input length for the model if that argument is not provided.
+            * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+              different lengths).
+        max_length (:obj:`int`, `optional`):
+            Maximum length of the returned list and optionally padding length (see above).
+        pad_to_multiple_of (:obj:`int`, `optional`):
+            If set will pad the sequence to a multiple of the provided value.
+
+            This is especially useful to enable the use of Tensor Cores on NVIDIA hardware with compute capability >=
+            7.5 (Volta).
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    padding: Union[bool, str, PaddingStrategy] = True
+    max_length: Optional[int] = None
+    pad_to_multiple_of: Optional[int] = None
+
+    def __call__(self, features):
+        label_name = "label" if "label" in features[0].keys() else "labels"
+        labels = [feature.pop(label_name) for feature in features]
+        batch_size = len(features)
+        num_choices = len(features[0]["input_ids"])
+        flattened_features = [
+            [{k: v[i] for k, v in feature.items()} for i in range(num_choices)] for feature in features
+        ]
+        flattened_features = sum(flattened_features, [])
+
+        batch = self.tokenizer.pad(
+            flattened_features,
+            padding=self.padding,
+            max_length=self.max_length,
+            pad_to_multiple_of=self.pad_to_multiple_of,
+            return_tensors="pt",
+        )
+
+        # Un-flatten
+        batch = {k: v.view(batch_size, num_choices, -1) for k, v in batch.items()}
+        # Add back labels
+        batch["labels"] = torch.tensor(labels, dtype=torch.int64)
+        return batch
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    if (
+        os.path.exists(training_args.output_dir)
+        and os.listdir(training_args.output_dir)
+        and training_args.do_train
+        and not training_args.overwrite_output_dir
+    ):
+        raise ValueError(
+            f"Output directory ({training_args.output_dir}) already exists and is not empty."
+            "Use --overwrite_output_dir to overcome."
+        )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO if is_main_process(training_args.local_rank) else logging.WARN,
+    )
+
+    # Log on each process the small summary:
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+    # Set the verbosity to info of the Transformers logger (on main process only):
+    if is_main_process(training_args.local_rank):
+        transformers.utils.logging.set_verbosity_info()
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+
+    # In distributed training, the load_dataset function guarantee that only one local process can concurrently
+    # download the dataset.
+    if data_args.train_file is not None or data_args.validation_file is not None:
+        data_files = {}
+        if data_args.train_file is not None:
+            data_files["train"] = data_args.train_file
+        if data_args.validation_file is not None:
+            data_files["validation"] = data_args.validation_file
+        extension = data_args.train_file.split(".")[-1]
+        datasets = load_dataset(extension, data_files=data_files)
+    else:
+        # Downloading and loading the swag dataset from the hub.
+        datasets = load_dataset("swag", "regular")
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+    config = AutoConfig.from_pretrained(
+        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+        use_fast=model_args.use_fast_tokenizer,
+    )
+    model = AutoModelForMultipleChoice.from_pretrained(
+        model_args.model_name_or_path,
+        from_tf=bool(".ckpt" in model_args.model_name_or_path),
+        config=config,
+        cache_dir=model_args.cache_dir,
+    )
+
+    # When using your own dataset or a different dataset from swag, you will probably need to change this.
+    ending_names = [f"ending{i}" for i in range(4)]
+    context_name = "sent1"
+    question_header_name = "sent2"
+
+    # Preprocessing the datasets.
+    def preprocess_function(examples):
+        first_sentences = [[context] * 4 for context in examples[context_name]]
+        question_headers = examples[question_header_name]
+        second_sentences = [
+            [f"{header} {examples[end][i]}" for end in ending_names] for i, header in enumerate(question_headers)
+        ]
+
+        # Flatten out
+        first_sentences = sum(first_sentences, [])
+        second_sentences = sum(second_sentences, [])
+
+        # Tokenize
+        tokenized_examples = tokenizer(
+            first_sentences,
+            second_sentences,
+            truncation=True,
+            max_length=data_args.max_seq_length,
+            padding="max_length" if data_args.pad_to_max_length else False,
+        )
+        # Un-flatten
+        return {k: [v[i : i + 4] for i in range(0, len(v), 4)] for k, v in tokenized_examples.items()}
+
+    tokenized_datasets = datasets.map(
+        preprocess_function,
+        batched=True,
+        num_proc=data_args.preprocessing_num_workers,
+        load_from_cache_file=not data_args.overwrite_cache,
+    )
+
+    # Data collator
+    data_collator = (
+        default_data_collator if data_args.pad_to_max_length else DataCollatorForMultipleChoice(tokenizer=tokenizer)
+    )
+
+    # Metric
+    def compute_metrics(eval_predictions):
+        predictions, label_ids = eval_predictions
+        preds = np.argmax(predictions, axis=1)
+        return {"accuracy": (preds == label_ids).astype(np.float32).mean().item()}
+
+    # Initialize our Trainer
+    trainer = Trainer(
+        model=model,
+        args=training_args,
+        train_dataset=tokenized_datasets["train"] if training_args.do_train else None,
+        eval_dataset=tokenized_datasets["validation"] if training_args.do_eval else None,
+        tokenizer=tokenizer,
+        data_collator=data_collator,
+        compute_metrics=compute_metrics,
+    )
+
+    # Training
+    if training_args.do_train:
+        train_result = trainer.train(
+            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
+        )
+        trainer.save_model()  # Saves the tokenizer too for easy upload
+
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
+    # Evaluation
+    results = {}
+    if training_args.do_eval:
+        logger.info("*** Evaluate ***")
+
+        results = trainer.evaluate()
+
+        output_eval_file = os.path.join(training_args.output_dir, "eval_results_swag.txt")
+        if trainer.is_world_process_zero():
+            with open(output_eval_file, "w") as writer:
+                logger.info("***** Eval results *****")
+                for key, value in sorted(results.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+    return results
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()

examples/question-answering/README.md

@@ -23,8 +23,7 @@ uses special features of those tokenizers. You can check if your favorite model
 [this table](https://huggingface.co/transformers/index.html#bigtable), if it doesn't you can still use the old version
 of the script.
 
-The old version of this script can be found [here](https://github.com/huggingface/transformers/blob/master/examples/contrib/legacy/question-answering/run_squad.py).
-
+The old version of this script can be found [here](https://github.com/huggingface/transformers/tree/master/examples/legacy/question-answering).
 #### Fine-tuning BERT on SQuAD1.0
 
 This example code fine-tunes BERT on the SQuAD1.0 dataset. It runs in 24 min (with BERT-base) or 68 min (with BERT-large)

examples/question-answering/run_qa.py

@@ -438,11 +438,22 @@ def main():
 
     # Training
     if training_args.do_train:
-        trainer.train(
+        train_result = trainer.train(
             model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
         )
         trainer.save_model()  # Saves the tokenizer too for easy upload
 
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
     # Evaluation
     results = {}
     if training_args.do_eval:
@@ -453,7 +464,7 @@ def main():
         if trainer.is_world_process_zero():
             with open(output_eval_file, "w") as writer:
                 logger.info("***** Eval results *****")
-                for key, value in results.items():
+                for key, value in sorted(results.items()):
                     logger.info(f"  {key} = {value}")
                     writer.write(f"{key} = {value}\n")
 

examples/question-answering/run_qa_beam_search.py

@@ -481,11 +481,22 @@ def main():
 
     # Training
     if training_args.do_train:
-        trainer.train(
+        train_result = trainer.train(
             model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
         )
         trainer.save_model()  # Saves the tokenizer too for easy upload
 
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
     # Evaluation
     results = {}
     if training_args.do_eval:
@@ -496,7 +507,7 @@ def main():
         if trainer.is_world_process_zero():
             with open(output_eval_file, "w") as writer:
                 logger.info("***** Eval results *****")
-                for key, value in results.items():
+                for key, value in sorted(results.items()):
                     logger.info(f"  {key} = {value}")
                     writer.write(f"{key} = {value}\n")
 

examples/question-answering/utils_qa.py

@@ -206,7 +206,7 @@ def postprocess_qa_predictions(
 
         # Make `predictions` JSON-serializable by casting np.float back to float.
         all_nbest_json[example["id"]] = [
-            {k: (float(v) if isinstance(v, (np.float32, np.float64)) else v) for k, v in pred.items()}
+            {k: (float(v) if isinstance(v, (np.float16, np.float32, np.float64)) else v) for k, v in pred.items()}
             for pred in predictions
         ]
 
@@ -394,7 +394,7 @@ def postprocess_qa_predictions_with_beam_search(
 
         # Make `predictions` JSON-serializable by casting np.float back to float.
         all_nbest_json[example["id"]] = [
-            {k: (float(v) if isinstance(v, (np.float32, np.float64)) else v) for k, v in pred.items()}
+            {k: (float(v) if isinstance(v, (np.float16, np.float32, np.float64)) else v) for k, v in pred.items()}
             for pred in predictions
         ]
 

examples/research_projects/bertology/run_prune_gpt.py

@@ -0,0 +1,388 @@
+#!/usr/bin/env python3
+""" This script is adapted from the Bertology pruning code (https://github.com/huggingface/transformers/blob/783d7d2629e97c5f0c5f9ef01b8c66410275c204/examples/research_projects/bertology/run_bertology.py)
+to prune GPT-like models. The author is @altsoph.
+"""
+
+import argparse
+import logging
+import os
+from datetime import datetime
+
+import numpy as np
+import torch
+from torch.utils.data import DataLoader, RandomSampler, TensorDataset
+from tqdm import tqdm
+
+from transformers import GPT2LMHeadModel
+
+
+logger = logging.getLogger(__name__)
+
+
+def save_model(model, dirpath):
+    # save results
+    if os.path.exists(dirpath):
+        if os.path.exists(os.path.join(dirpath, "config.json")) and os.path.isfile(
+            os.path.join(dirpath, "config.json")
+        ):
+            os.remove(os.path.join(dirpath, "config.json"))
+        if os.path.exists(os.path.join(dirpath, "pytorch_model.bin")) and os.path.isfile(
+            os.path.join(dirpath, "pytorch_model.bin")
+        ):
+            os.remove(os.path.join(dirpath, "pytorch_model.bin"))
+    else:
+        os.makedirs(dirpath)
+    model.save_pretrained(dirpath)
+
+
+def entropy(p, unlogit=False):
+    """ Compute the entropy of a probability distribution """
+    exponent = 2
+    if unlogit:
+        p = torch.pow(p, exponent)
+    plogp = p * torch.log(p)
+    plogp[p == 0] = 0
+    return -plogp.sum(dim=-1)
+
+
+def print_2d_tensor(tensor):
+    """ Print a 2D tensor """
+    logger.info("lv, h >\t" + "\t".join(f"{x + 1}" for x in range(len(tensor))))
+    for row in range(len(tensor)):
+        if tensor.dtype != torch.long:
+            logger.info(f"layer {row + 1}:\t" + "\t".join(f"{x:.5f}" for x in tensor[row].cpu().data))
+        else:
+            logger.info(f"layer {row + 1}:\t" + "\t".join(f"{x:d}" for x in tensor[row].cpu().data))
+
+
+def compute_heads_importance(
+    args, model, eval_dataloader, compute_entropy=True, compute_importance=True, head_mask=None, actually_pruned=False
+):
+    """This method shows how to compute:
+    - head attention entropy
+    - head importance scores according to http://arxiv.org/abs/1905.10650
+    """
+    # Prepare our tensors
+    n_layers, n_heads = model.config.num_hidden_layers, model.config.num_attention_heads
+    head_importance = torch.zeros(n_layers, n_heads).to(args.device)
+    attn_entropy = torch.zeros(n_layers, n_heads).to(args.device)
+
+    if head_mask is None:
+        head_mask = torch.ones(n_layers, n_heads).to(args.device)
+
+    head_mask.requires_grad_(requires_grad=True)
+    # If actually pruned attention multi-head, set head mask to None to avoid shape mismatch
+    if actually_pruned:
+        head_mask = None
+
+    tot_tokens = 0.0
+    total_loss = 0.0
+    for step, inputs in enumerate(tqdm(eval_dataloader, desc="Iteration", disable=args.local_rank not in [-1, 0])):
+        inputs = tuple(t.to(args.device) for t in inputs)
+        (input_ids,) = inputs
+
+        # Do a forward pass (not with torch.no_grad() since we need gradients for importance score - see below)
+        outputs = model(input_ids, labels=input_ids, head_mask=head_mask)
+        #  (loss), lm_logits, presents, (all hidden_states), (attentions)
+        loss, _, all_attentions = (
+            outputs[0],
+            outputs[1],
+            outputs[-1],
+        )  # Loss and logits are the first, attention the last
+        loss.backward()  # Backpropagate to populate the gradients in the head mask
+        total_loss += loss.detach().cpu().numpy()
+        if compute_entropy:
+            for layer, attn in enumerate(all_attentions):
+                masked_entropy = entropy(attn.detach(), True)
+                attn_entropy[layer] += masked_entropy.sum(-1).sum(0).sum(0).detach()
+
+        if compute_importance:
+            head_importance += head_mask.grad.abs().detach()
+        tot_tokens += torch.ones_like(input_ids).float().detach().sum().data
+
+    # Normalize
+    attn_entropy /= tot_tokens
+    head_importance /= tot_tokens
+    # Layerwise importance normalization
+    if not args.dont_normalize_importance_by_layer:
+        exponent = 2
+        norm_by_layer = torch.pow(torch.pow(head_importance, exponent).sum(-1), 1 / exponent)
+        head_importance /= norm_by_layer.unsqueeze(-1) + 1e-20
+
+    if not args.dont_normalize_global_importance:
+        head_importance = (head_importance - head_importance.min()) / (head_importance.max() - head_importance.min())
+
+    # Print matrices
+    if compute_entropy:
+        logger.info("Attention entropies")
+        print_2d_tensor(attn_entropy)
+    if compute_importance:
+        logger.info("Head importance scores")
+        print_2d_tensor(head_importance)
+    logger.info("Head ranked by importance scores")
+    head_ranks = torch.zeros(head_importance.numel(), dtype=torch.long, device=args.device)
+    head_ranks[head_importance.view(-1).sort(descending=True)[1]] = torch.arange(
+        head_importance.numel(), device=args.device
+    )
+    head_ranks = head_ranks.view_as(head_importance)
+    print_2d_tensor(head_ranks)
+    return attn_entropy, head_importance, total_loss
+
+
+def mask_heads(args, model, eval_dataloader):
+    """This method shows how to mask head (set some heads to zero), to test the effect on the network,
+    based on the head importance scores, as described in Michel et al. (http://arxiv.org/abs/1905.10650)
+    """
+    _, head_importance, loss = compute_heads_importance(args, model, eval_dataloader, compute_entropy=False)
+    original_score = 1 / loss  # instead of downsteam score use the LM loss
+    logger.info("Pruning: original score: %f, threshold: %f", original_score, original_score * args.masking_threshold)
+
+    new_head_mask = torch.ones_like(head_importance)
+    num_to_mask = max(1, int(new_head_mask.numel() * args.masking_amount))
+
+    current_score = original_score
+    while current_score >= original_score * args.masking_threshold:
+        head_mask = new_head_mask.clone().detach()  # save current head mask
+        # heads from least important to most - keep only not-masked heads
+        head_importance[head_mask == 0.0] = float("Inf")
+        current_heads_to_mask = head_importance.view(-1).sort()[1]
+
+        if len(current_heads_to_mask) <= num_to_mask:
+            print("BREAK BY num_to_mask")
+            break
+
+        # mask heads
+        current_heads_to_mask = current_heads_to_mask[:num_to_mask]
+        logger.info("Heads to mask: %s", str(current_heads_to_mask.tolist()))
+        new_head_mask = new_head_mask.view(-1)
+        new_head_mask[current_heads_to_mask] = 0.0
+        new_head_mask = new_head_mask.view_as(head_mask)
+        new_head_mask = new_head_mask.clone().detach()
+        print_2d_tensor(new_head_mask)
+
+        # Compute metric and head importance again
+        _, head_importance, loss = compute_heads_importance(
+            args, model, eval_dataloader, compute_entropy=False, head_mask=new_head_mask
+        )
+        current_score = 1 / loss
+        logger.info(
+            "Masking: current score: %f, remaining heads %d (%.1f percents)",
+            current_score,
+            new_head_mask.sum(),
+            new_head_mask.sum() / new_head_mask.numel() * 100,
+        )
+
+    logger.info("Final head mask")
+    print_2d_tensor(head_mask)
+    np.save(os.path.join(args.output_dir, "head_mask.npy"), head_mask.detach().cpu().numpy())
+
+    return head_mask
+
+
+def prune_heads(args, model, eval_dataloader, head_mask):
+    """This method shows how to prune head (remove heads weights) based on
+    the head importance scores as described in Michel et al. (http://arxiv.org/abs/1905.10650)
+    """
+    # Try pruning and test time speedup
+    # Pruning is like masking but we actually remove the masked weights
+    before_time = datetime.now()
+    _, _, loss = compute_heads_importance(
+        args, model, eval_dataloader, compute_entropy=False, compute_importance=False, head_mask=head_mask
+    )
+    score_masking = 1 / loss
+    original_time = datetime.now() - before_time
+
+    original_num_params = sum(p.numel() for p in model.parameters())
+    heads_to_prune = dict(
+        (layer, (1 - head_mask[layer].long()).nonzero().squeeze().tolist()) for layer in range(len(head_mask))
+    )
+
+    for k, v in heads_to_prune.items():
+        if isinstance(v, int):
+            heads_to_prune[k] = [
+                v,
+            ]
+
+    assert sum(len(h) for h in heads_to_prune.values()) == (1 - head_mask.long()).sum().item()
+    model.prune_heads(heads_to_prune)
+    pruned_num_params = sum(p.numel() for p in model.parameters())
+
+    before_time = datetime.now()
+    _, _, loss = compute_heads_importance(
+        args,
+        model,
+        eval_dataloader,
+        compute_entropy=False,
+        compute_importance=False,
+        head_mask=None,
+        actually_pruned=True,
+    )
+
+    score_pruning = 1 / loss
+    new_time = datetime.now() - before_time
+
+    logger.info(
+        "Pruning: original num of params: %.2e, after pruning %.2e (%.1f percents)",
+        original_num_params,
+        pruned_num_params,
+        pruned_num_params / original_num_params * 100,
+    )
+    logger.info("Pruning: score with masking: %f score with pruning: %f", score_masking, score_pruning)
+    logger.info("Pruning: speed ratio (original timing / new timing): %f percents", original_time / new_time * 100)
+    save_model(model, args.output_dir)
+
+
+def main():
+    parser = argparse.ArgumentParser()
+    # Required parameters
+    parser.add_argument(
+        "--data_dir",
+        default=None,
+        type=str,
+        required=True,
+        help="The input data dir. Should contain the .tsv files (or other data files) for the task.",
+    )
+    parser.add_argument(
+        "--model_name_or_path",
+        default=None,
+        type=str,
+        required=True,
+        help="Path to pretrained model or model identifier from huggingface.co/models",
+    )
+    parser.add_argument(
+        "--output_dir",
+        default=None,
+        type=str,
+        required=True,
+        help="The output directory where the model predictions and checkpoints will be written.",
+    )
+
+    # Other parameters
+    parser.add_argument(
+        "--config_name",
+        default="",
+        type=str,
+        help="Pretrained config name or path if not the same as model_name_or_path",
+    )
+    parser.add_argument(
+        "--tokenizer_name",
+        default="",
+        type=str,
+        help="Pretrained tokenizer name or path if not the same as model_name_or_path",
+    )
+    parser.add_argument(
+        "--cache_dir",
+        default=None,
+        type=str,
+        help="Where do you want to store the pre-trained models downloaded from s3",
+    )
+    parser.add_argument(
+        "--data_subset", type=int, default=-1, help="If > 0: limit the data to a subset of data_subset instances."
+    )
+    parser.add_argument(
+        "--overwrite_output_dir", action="store_true", help="Whether to overwrite data in output directory"
+    )
+    parser.add_argument(
+        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets"
+    )
+
+    parser.add_argument(
+        "--dont_normalize_importance_by_layer", action="store_true", help="Don't normalize importance score by layers"
+    )
+    parser.add_argument(
+        "--dont_normalize_global_importance",
+        action="store_true",
+        help="Don't normalize all importance scores between 0 and 1",
+    )
+
+    parser.add_argument(
+        "--try_masking", action="store_true", help="Whether to try to mask head until a threshold of accuracy."
+    )
+    parser.add_argument(
+        "--masking_threshold",
+        default=0.9,
+        type=float,
+        help="masking threshold in term of metrics (stop masking when metric < threshold * original metric value).",
+    )
+    parser.add_argument(
+        "--masking_amount", default=0.1, type=float, help="Amount to heads to masking at each masking step."
+    )
+    parser.add_argument("--metric_name", default="acc", type=str, help="Metric to use for head masking.")
+
+    parser.add_argument(
+        "--max_seq_length",
+        default=128,
+        type=int,
+        help="The maximum total input sequence length after WordPiece tokenization. \n"
+        "Sequences longer than this will be truncated, sequences shorter padded.",
+    )
+    parser.add_argument("--batch_size", default=1, type=int, help="Batch size.")
+
+    parser.add_argument("--seed", type=int, default=42)
+    parser.add_argument("--local_rank", type=int, default=-1, help="local_rank for distributed training on gpus")
+    parser.add_argument("--no_cuda", action="store_true", help="Whether not to use CUDA when available")
+    parser.add_argument("--server_ip", type=str, default="", help="Can be used for distant debugging.")
+    parser.add_argument("--server_port", type=str, default="", help="Can be used for distant debugging.")
+    args = parser.parse_args()
+
+    if args.server_ip and args.server_port:
+        # Distant debugging - see https://code.visualstudio.com/docs/python/debugging#_attach-to-a-local-script
+        import ptvsd
+
+        print("Waiting for debugger attach")
+        ptvsd.enable_attach(address=(args.server_ip, args.server_port), redirect_output=True)
+        ptvsd.wait_for_attach()
+
+    # Setup devices and distributed training
+    if args.local_rank == -1 or args.no_cuda:
+        args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
+        args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
+    else:
+        torch.cuda.set_device(args.local_rank)
+        args.device = torch.device("cuda", args.local_rank)
+        args.n_gpu = 1
+        torch.distributed.init_process_group(backend="nccl")  # Initializes the distributed backend
+
+    # Setup logging
+    logging.basicConfig(level=logging.INFO if args.local_rank in [-1, 0] else logging.WARN)
+    logger.info("device: {} n_gpu: {}, distributed: {}".format(args.device, args.n_gpu, bool(args.local_rank != -1)))
+
+    model = GPT2LMHeadModel.from_pretrained(args.model_name_or_path)
+
+    # Distributed and parallel training
+    model.to(args.device)
+    if args.local_rank != -1:
+        model = torch.nn.parallel.DistributedDataParallel(
+            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True
+        )
+    elif args.n_gpu > 1:
+        model = torch.nn.DataParallel(model)
+
+    # Print/save training arguments
+    os.makedirs(args.output_dir, exist_ok=True)
+    torch.save(args, os.path.join(args.output_dir, "run_args.bin"))
+    logger.info("Training/evaluation parameters %s", args)
+
+    # Prepare dataset
+    numpy_data = np.concatenate(
+        [
+            np.loadtxt(args.data_dir, dtype=np.int64),
+        ]
+    )
+    train_tensor_dataset = (torch.from_numpy(numpy_data),)
+    train_data = TensorDataset(*train_tensor_dataset)
+    train_sampler = RandomSampler(train_data)
+    eval_dataloader = DataLoader(train_data, sampler=train_sampler, batch_size=args.batch_size)
+
+    # Compute head entropy and importance score
+    compute_heads_importance(args, model, eval_dataloader)
+
+    # Try head masking (set heads to zero until the score goes under a threshole)
+    # and head pruning (remove masked heads and see the effect on the network)
+    if args.try_masking and args.masking_threshold > 0.0 and args.masking_threshold < 1.0:
+        head_mask = mask_heads(args, model, eval_dataloader)
+        prune_heads(args, model, eval_dataloader, head_mask)
+
+
+if __name__ == "__main__":
+    main()

examples/research_projects/performer/README.md

@@ -0,0 +1,25 @@
+# Performer fine-tuning
+
+Example authors: @TevenLeScao, @Patrickvonplaten
+
+Paper authors: Krzysztof Choromanski, Valerii Likhosherstov, David Dohan, Xingyou Song, Andreea Gane, Tamas Sarlos, Peter Hawkins, Jared Davis, Afroz Mohiuddin, Lukasz Kaiser, David Belanger, Lucy Colwell, Adrian Weller
+
+## Requirements
+
+`datasets`, `flax` and `jax`. `wandb` integration is built-in if you want to use it.
+
+## Examples
+
+`sanity_script.sh` will launch performer fine-tuning from the bert-base-cased checkpoint on the Simple Wikipedia dataset (a small, easy-language English Wikipedia) from `datasets`.
+`full_script.sh` will launch performer fine-tuning from the bert-large-cased checkpoint on the English Wikipedia dataset from `datasets`.
+
+Here are a few key arguments:
+- Remove the `--performer` argument to use a standard Bert model.
+  
+- Add `--reinitialize` to start from a blank model rather than a Bert checkpoint. 
+  
+- You may change the Bert size by passing a different [checkpoint](https://huggingface.co/transformers/pretrained_models.html) to the `--model_name_or_path` argument.
+
+- Passing your user name to the `--wandb_user_name` argument will trigger weights and biases logging.
+
+- You can choose a dataset with `--dataset_name` and `--dataset_config`. Our [viewer](https://huggingface.co/datasets/viewer/) will help you find what you need.

examples/research_projects/performer/full_script.sh

@@ -0,0 +1 @@
+TOKENIZERS_PARALLELISM=true python run_mlm_performer.py  --output_dir experiments --dataset_name wikipedia --dataset_config_name 20200501.en --model_name_or_path bert-large-cased --tokenizer_name bert-large-cased --do_train --overwrite_output_dir --per_device_train_batch_size 4 --learning_rate 5e-4 --warmup_steps 100 --num_train_epochs 3 --performer

examples/research_projects/performer/modeling_flax_performer.py

@@ -0,0 +1,553 @@
+# coding=utf-8
+# Copyright 2018 The Google Flax Team Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+from typing import Callable, Dict, Tuple
+
+import numpy as np
+
+import flax.linen as nn
+import jax
+import jax.numpy as jnp
+from jax.random import PRNGKey
+from modeling_flax_performer_utils import make_fast_softmax_attention
+from transformers.file_utils import add_start_docstrings
+from transformers.modeling_flax_utils import ACT2FN
+from transformers.models.bert.configuration_bert import BertConfig
+from transformers.models.bert.modeling_flax_bert import FlaxBertOnlyMLMHead, FlaxBertPreTrainedModel
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "BertConfig"
+_TOKENIZER_FOR_DOC = "BertTokenizer"
+
+BERT_START_DOCSTRING = r"""
+
+    This model inherits from :class:`~transformers.PreTrainedModel`. Check the superclass documentation for the generic
+    methods the library implements for all its model (such as downloading or saving, resizing the input embeddings,
+    pruning heads etc.)
+
+    This model is also a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`__
+    subclass. Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to
+    general usage and behavior.
+
+    Parameters:
+        config (:class:`~transformers.BertConfig`): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the
+            configuration. Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model
+            weights.
+"""
+
+BERT_INPUTS_DOCSTRING = r"""
+    Args:
+        input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`):
+            Indices of input sequence tokens in the vocabulary.
+
+            Indices can be obtained using :class:`~transformers.BertTokenizer`. See
+            :meth:`transformers.PreTrainedTokenizer.encode` and :meth:`transformers.PreTrainedTokenizer.__call__` for
+            details.
+
+            `What are input IDs? <../glossary.html#input-ids>`__
+        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
+            Mask to avoid performing attention on padding token indices. Mask values selected in ``[0, 1]``:
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **masked**.
+
+            `What are attention masks? <../glossary.html#attention-mask>`__
+        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
+            Segment token indices to indicate first and second portions of the inputs. Indices are selected in ``[0,
+            1]``:
+
+            - 0 corresponds to a `sentence A` token,
+            - 1 corresponds to a `sentence B` token.
+
+            `What are token type IDs? <../glossary.html#token-type-ids>`_
+        position_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
+            Indices of positions of each input sequence tokens in the position embeddings. Selected in the range ``[0,
+            config.max_position_embeddings - 1]``.
+
+            `What are position IDs? <../glossary.html#position-ids>`_
+        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
+            Mask to nullify selected heads of the self-attention modules. Mask values selected in ``[0, 1]``:
+
+            - 1 indicates the head is **not masked**,
+            - 0 indicates the head is **masked**.
+
+        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`({0}, hidden_size)`, `optional`):
+            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 :obj:`input_ids` indices into associated
+            vectors than the model's internal embedding lookup matrix.
+        output_attentions (:obj:`bool`, `optional`):
+            Whether or not to return the attentions tensors of all attention layers. See ``attentions`` under returned
+            tensors for more detail.
+        output_hidden_states (:obj:`bool`, `optional`):
+            Whether or not to return the hidden states of all layers. See ``hidden_states`` under returned tensors for
+            more detail.
+        return_dict (:obj:`bool`, `optional`):
+            Whether or not to return a :class:`~transformers.file_utils.ModelOutput` instead of a plain tuple.
+"""
+
+
+class FlaxPerformerLayerNorm(nn.Module):
+    """
+    Layer normalization (https://arxiv.org/abs/1607.06450). Operates on the last axis of the input data.
+    """
+
+    epsilon: float = 1e-6
+    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+    bias: bool = True  # If True, bias (beta) is added.
+    scale: bool = True  # If True, multiply by scale (gamma). When the next layer is linear
+    # (also e.g. nn.relu), this can be disabled since the scaling will be
+    # done by the next layer.
+    bias_init: jnp.ndarray = nn.initializers.zeros
+    scale_init: jnp.ndarray = nn.initializers.ones
+
+    @nn.compact
+    def __call__(self, x):
+        """
+        Applies layer normalization on the input. It normalizes the activations of the layer for each given example in
+        a batch independently, rather than across a batch like Batch Normalization. i.e. applies a transformation that
+        maintains the mean activation within each example close to 0 and the activation standard deviation close to 1
+
+        Args:
+          x: the inputs
+
+        Returns:
+          Normalized inputs (the same shape as inputs).
+        """
+        features = x.shape[-1]
+        mean = jnp.mean(x, axis=-1, keepdims=True)
+        mean2 = jnp.mean(jax.lax.square(x), axis=-1, keepdims=True)
+        var = mean2 - jax.lax.square(mean)
+        mul = jax.lax.rsqrt(var + self.epsilon)
+        if self.scale:
+            mul = mul * jnp.asarray(self.param("gamma", self.scale_init, (features,)), self.dtype)
+        y = (x - mean) * mul
+        if self.bias:
+            y = y + jnp.asarray(self.param("beta", self.bias_init, (features,)), self.dtype)
+        return y
+
+
+class FlaxPerformerEmbedding(nn.Module):
+    """
+    Specify a new class for doing the embedding stuff as Flax's one use 'embedding' for the parameter name and PyTorch
+    use 'weight'
+    """
+
+    vocab_size: int
+    hidden_size: int
+    emb_init: Callable[..., np.ndarray] = nn.initializers.normal(stddev=0.1)
+
+    @nn.compact
+    def __call__(self, inputs):
+        embedding = self.param("weight", self.emb_init, (self.vocab_size, self.hidden_size))
+        return jnp.take(embedding, inputs, axis=0)
+
+
+class FlaxPerformerEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings."""
+
+    vocab_size: int
+    hidden_size: int
+    type_vocab_size: int
+    max_length: int
+
+    @nn.compact
+    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):
+        # Embed
+        w_emb = FlaxPerformerEmbedding(self.vocab_size, self.hidden_size, name="word_embeddings")(
+            jnp.atleast_2d(input_ids.astype("i4"))
+        )
+        p_emb = FlaxPerformerEmbedding(self.max_length, self.hidden_size, name="position_embeddings")(
+            jnp.atleast_2d(position_ids.astype("i4"))
+        )
+        t_emb = FlaxPerformerEmbedding(self.type_vocab_size, self.hidden_size, name="token_type_embeddings")(
+            jnp.atleast_2d(token_type_ids.astype("i4"))
+        )
+
+        # Sum all embeddings
+        summed_emb = w_emb + jnp.broadcast_to(p_emb, w_emb.shape) + t_emb
+
+        # Layer Norm
+        layer_norm = FlaxPerformerLayerNorm(name="layer_norm")(summed_emb)
+
+        return layer_norm
+
+
+class FlaxPerformerAttention(nn.Module):
+    num_heads: int
+    head_size: int
+
+    @nn.compact
+    def __call__(self, hidden_state, attention_mask):
+        single_head_dim = self.head_size // self.num_heads
+        fast_softmax_attention = make_fast_softmax_attention(qkv_dim=single_head_dim)
+        self_att = nn.attention.SelfAttention(
+            num_heads=self.num_heads, qkv_features=self.head_size, name="self", attention_fn=fast_softmax_attention
+        )(hidden_state, attention_mask)
+
+        layer_norm = FlaxPerformerLayerNorm(name="layer_norm")(self_att + hidden_state)
+        return layer_norm
+
+
+class FlaxPerformerIntermediate(nn.Module):
+    output_size: int
+    hidden_act: str = "gelu"
+
+    @nn.compact
+    def __call__(self, hidden_state):
+        # TODO: Add ACT2FN reference to change activation function
+        dense = nn.Dense(features=self.output_size, name="dense")(hidden_state)
+        return ACT2FN[self.hidden_act](dense)
+
+
+class FlaxPerformerOutput(nn.Module):
+    @nn.compact
+    def __call__(self, intermediate_output, attention_output):
+        hidden_state = nn.Dense(attention_output.shape[-1], name="dense")(intermediate_output)
+        hidden_state = FlaxPerformerLayerNorm(name="layer_norm")(hidden_state + attention_output)
+        return hidden_state
+
+
+class FlaxPerformerLayer(nn.Module):
+    num_heads: int
+    head_size: int
+    intermediate_size: int
+    hidden_act: str = "gelu"
+
+    @nn.compact
+    def __call__(self, hidden_state, attention_mask):
+        attention = FlaxPerformerAttention(self.num_heads, self.head_size, name="attention")(
+            hidden_state, attention_mask
+        )
+        intermediate = FlaxPerformerIntermediate(
+            self.intermediate_size, name="intermediate", hidden_act=self.hidden_act
+        )(attention)
+        output = FlaxPerformerOutput(name="output")(intermediate, attention)
+
+        return output
+
+
+class FlaxPerformerLayerCollection(nn.Module):
+    """
+    Stores N BertLayer(s)
+    """
+
+    num_layers: int
+    num_heads: int
+    head_size: int
+    intermediate_size: int
+    hidden_act: str = "gelu"
+
+    @nn.compact
+    def __call__(self, inputs, attention_mask):
+        assert self.num_layers > 0, f"num_layers should be >= 1, got ({self.num_layers})"
+
+        # Initialize input / output
+        input_i = inputs
+
+        # Forward over all encoders
+        for i in range(self.num_layers):
+            layer = FlaxPerformerLayer(
+                self.num_heads, self.head_size, self.intermediate_size, hidden_act=self.hidden_act, name=f"{i}"
+            )
+            input_i = layer(input_i, attention_mask)
+        return input_i
+
+
+class FlaxPerformerEncoder(nn.Module):
+    num_layers: int
+    num_heads: int
+    head_size: int
+    intermediate_size: int
+    hidden_act: str = "gelu"
+
+    @nn.compact
+    def __call__(self, hidden_state, attention_mask):
+        layer = FlaxPerformerLayerCollection(
+            self.num_layers,
+            self.num_heads,
+            self.head_size,
+            self.intermediate_size,
+            name="layer",
+            hidden_act=self.hidden_act,
+        )(hidden_state, attention_mask)
+        return layer
+
+
+class FlaxPerformerPooler(nn.Module):
+    @nn.compact
+    def __call__(self, hidden_state):
+        cls_token = hidden_state[:, 0]
+        out = nn.Dense(hidden_state.shape[-1], name="dense")(cls_token)
+        return jax.lax.tanh(out)
+
+
+class FlaxPerformerModule(nn.Module):
+    vocab_size: int
+    hidden_size: int
+    type_vocab_size: int
+    max_length: int
+    num_encoder_layers: int
+    num_heads: int
+    head_size: int
+    intermediate_size: int
+    hidden_act: str = "gelu"
+    add_pooling_layer: bool = True
+
+    @nn.compact
+    def __call__(self, input_ids, token_type_ids, position_ids, attention_mask):
+        # Embedding
+        embeddings = FlaxPerformerEmbeddings(
+            self.vocab_size, self.hidden_size, self.type_vocab_size, self.max_length, name="embeddings"
+        )(input_ids, token_type_ids, position_ids, attention_mask)
+
+        # N stacked encoding layers
+        encoder = FlaxPerformerEncoder(
+            self.num_encoder_layers,
+            self.num_heads,
+            self.head_size,
+            self.intermediate_size,
+            hidden_act=self.hidden_act,
+            name="encoder",
+        )(embeddings, attention_mask)
+
+        if not self.add_pooling_layer:
+            return encoder
+
+        pooled = FlaxPerformerPooler(name="pooler")(encoder)
+        return encoder, pooled
+
+
+@add_start_docstrings(
+    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
+    BERT_START_DOCSTRING,
+)
+class FlaxPerformerModel(FlaxBertPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    """
+
+    model_class = FlaxPerformerModule
+    config_class = BertConfig
+    base_model_prefix = "bert"
+
+    @staticmethod
+    def convert_from_pytorch(pt_state: Dict, config: BertConfig) -> Dict:
+        jax_state = dict(pt_state)
+
+        # Need to change some parameters name to match Flax names so that we don't have to fork any layer
+        for key, tensor in pt_state.items():
+            # Key parts
+            key_parts = set(key.split("."))
+
+            # Every dense layer has "kernel" parameters instead of "weight"
+            if "dense.weight" in key:
+                del jax_state[key]
+                key = key.replace("weight", "kernel")
+                jax_state[key] = tensor
+
+            # SelfAttention needs also to replace "weight" by "kernel"
+            if {"query", "key", "value"} & key_parts:
+
+                # Flax SelfAttention decomposes the heads (num_head, size // num_heads)
+                if "bias" in key:
+                    jax_state[key] = tensor.reshape((config.num_attention_heads, -1))
+                elif "weight":
+                    del jax_state[key]
+                    key = key.replace("weight", "kernel")
+                    tensor = tensor.reshape((config.num_attention_heads, -1, config.hidden_size)).transpose((2, 0, 1))
+                    jax_state[key] = tensor
+
+            # SelfAttention output is not a separate layer, remove one nesting
+            if "attention.output.dense" in key:
+                del jax_state[key]
+                key = key.replace("attention.output.dense", "attention.self.out")
+                jax_state[key] = tensor
+
+            # SelfAttention output is not a separate layer, remove nesting on layer norm
+            if "attention.output.LayerNorm" in key:
+                del jax_state[key]
+                key = key.replace("attention.output.LayerNorm", "attention.LayerNorm")
+                jax_state[key] = tensor
+
+            # There are some transposed parameters w.r.t their PyTorch counterpart
+            if "intermediate.dense.kernel" in key or "output.dense.kernel" in key:
+                jax_state[key] = tensor.T
+
+            # Self Attention output projection needs to be transposed
+            if "out.kernel" in key:
+                jax_state[key] = tensor.reshape((config.hidden_size, config.num_attention_heads, -1)).transpose(
+                    1, 2, 0
+                )
+
+            # Pooler needs to transpose its kernel
+            if "pooler.dense.kernel" in key:
+                jax_state[key] = tensor.T
+
+            # Handle LayerNorm conversion
+            if "LayerNorm" in key:
+                del jax_state[key]
+
+                # Replace LayerNorm by layer_norm
+                new_key = key.replace("LayerNorm", "layer_norm")
+
+                if "weight" in key:
+                    new_key = new_key.replace("weight", "gamma")
+                elif "bias" in key:
+                    new_key = new_key.replace("bias", "beta")
+
+                jax_state[new_key] = tensor
+
+        return jax_state
+
+    def __init__(
+        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
+    ):
+        module = FlaxPerformerModule(
+            vocab_size=config.vocab_size,
+            hidden_size=config.hidden_size,
+            type_vocab_size=config.type_vocab_size,
+            max_length=config.max_position_embeddings,
+            num_encoder_layers=config.num_hidden_layers,
+            num_heads=config.num_attention_heads,
+            head_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            dropout_rate=config.hidden_dropout_prob,
+            hidden_act=config.hidden_act,
+        )
+
+        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
+
+    @property
+    def module(self) -> nn.Module:
+        return self._module
+
+    def __call__(
+        self, input_ids, token_type_ids=None, position_ids=None, dropout_rng: PRNGKey = None, attention_mask=None
+    ):
+
+        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
+            input_ids, attention_mask, token_type_ids, position_ids
+        )
+
+        # Handle any PRNG if needed
+        rngs = {}
+        if dropout_rng is not None:
+            rngs["dropout"] = dropout_rng
+
+        return self.module.apply(
+            {"params": self.params},
+            jnp.array(input_ids, dtype="i4"),
+            jnp.array(token_type_ids, dtype="i4"),
+            jnp.array(position_ids, dtype="i4"),
+            jnp.array(attention_mask, dtype="i4"),
+            rng=rngs,
+        )
+
+
+class FlaxPerformerForMaskedLM(FlaxBertPreTrainedModel):
+    def __init__(
+        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
+    ):
+        module = FlaxPerformerForMaskedLMModule(
+            vocab_size=config.vocab_size,
+            type_vocab_size=config.type_vocab_size,
+            hidden_size=config.hidden_size,
+            intermediate_size=config.intermediate_size,
+            head_size=config.hidden_size,
+            num_heads=config.num_attention_heads,
+            num_encoder_layers=config.num_hidden_layers,
+            max_length=config.max_position_embeddings,
+            hidden_act=config.hidden_act,
+            **kwargs,
+        )
+
+        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
+
+    def __call__(
+        self,
+        input_ids,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        params: dict = None,
+        train: bool = False,
+        dropout_rng: PRNGKey = None,
+    ):
+        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
+            input_ids, attention_mask, token_type_ids, position_ids
+        )
+
+        # Handle any PRNG if needed
+        rngs = {}
+        if dropout_rng is not None:
+            rngs["dropout"] = dropout_rng
+
+        return self.module.apply(
+            {"params": params or self.params},
+            jnp.array(input_ids, dtype="i4"),
+            jnp.array(attention_mask, dtype="i4"),
+            jnp.array(token_type_ids, dtype="i4"),
+            jnp.array(position_ids, dtype="i4"),
+            not train,
+            rngs=rngs,
+        )
+
+
+class FlaxPerformerForMaskedLMModule(nn.Module):
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    head_size: int
+    num_heads: int
+    num_encoder_layers: int
+    type_vocab_size: int
+    max_length: int
+    hidden_act: str
+    dropout_rate: float = 0.0
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(
+        self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, deterministic: bool = True
+    ):
+        # Model
+        encoder = FlaxPerformerModule(
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            type_vocab_size=self.type_vocab_size,
+            max_length=self.max_length,
+            num_encoder_layers=self.num_encoder_layers,
+            num_heads=self.num_heads,
+            head_size=self.hidden_size,
+            intermediate_size=self.intermediate_size,
+            hidden_act=self.hidden_act,
+            add_pooling_layer=False,
+            name="bert",
+        )(input_ids, attention_mask, token_type_ids, position_ids)
+
+        # Compute the prediction scores
+        encoder = nn.Dropout(rate=self.dropout_rate)(encoder, deterministic=deterministic)
+        logits = FlaxBertOnlyMLMHead(
+            vocab_size=self.vocab_size, hidden_act=self.hidden_act, name="cls", dtype=self.dtype
+        )(encoder)
+
+        return (logits,)

examples/research_projects/performer/modeling_flax_performer_utils.py

@@ -0,0 +1,660 @@
+# coding=utf-8
+# Copyright 2020 The Google Research Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+"""
+IMPORTANT:
+
+This code was copied from
+https://github.com/google-research/google-research/blob/master/performer/fast_self_attention/fast_self_attention.py on
+6/11/2020. This is very new code, so it might be prone to change soon -> make sure to check the original code and
+update accordingly
+
+Core Fast Attention Module for Flax. Implementation of the approximate fast softmax and generalized attention mechanism
+leveraging structured random feature maps [RFM] techniques and low rank decomposition of the attention matrix.
+"""
+# pylint: disable=invalid-name, missing-function-docstring, line-too-long
+
+import abc
+import functools
+from collections.abc import Iterable  # pylint: disable=g-importing-member
+
+import numpy as onp
+from absl import logging
+
+import jax
+import jax.numpy as jnp
+from jax import lax, random
+
+
+def nonnegative_softmax_kernel_feature_creator(
+    data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=True, eps=0.0001
+):
+    """
+    Constructs nonnegative kernel features for fast softmax attention
+
+    Args:
+      data: input for which features are computes
+      projection_matrix: random matrix used to compute features
+      attention_dims_t: tuple of attention dimensions
+      batch_dims_t: tuple of batch dimensions
+      precision: precision parameter
+      is_query: predicate indicating whether input data corresponds to queries or
+        keys
+      normalize_data: predicate indicating whether data should be normalized,
+      eps: numerical stabilizer
+
+    Returns:
+      Random features for fast softmax attention.
+    """
+    del attention_dims_t
+    if normalize_data:
+        # We have e^{qk^T/sqrt{d}} = e^{q_norm k_norm^T}, where
+        # w_norm = w * data_normalizer for w in {q,k}.
+        data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1])))
+    else:
+        data_normalizer = 1.0
+    ratio = 1.0 / jnp.sqrt(projection_matrix.shape[0])
+    data_mod_shape = data.shape[0 : len(batch_dims_t)] + projection_matrix.shape
+    data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix
+
+    data_dash = lax.dot_general(
+        data_normalizer * data,
+        data_thick_random_matrix,
+        (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), (batch_dims_t, batch_dims_t)),
+        precision=precision,
+    )
+
+    diag_data = jnp.square(data)
+    diag_data = jnp.sum(diag_data, axis=data.ndim - 1)
+    diag_data = (diag_data / 2.0) * data_normalizer * data_normalizer
+    diag_data = jnp.expand_dims(diag_data, axis=data.ndim - 1)
+
+    if is_query:
+        last_dims_t = (len(data_dash.shape) - 1,)
+        data_dash = ratio * (
+            jnp.exp(data_dash - diag_data - jnp.max(data_dash, axis=last_dims_t, keepdims=True)) + eps
+        )
+    else:
+        data_dash = ratio * (jnp.exp(data_dash - diag_data - jnp.max(data_dash)) + eps)
+
+    return data_dash
+
+
+def sincos_softmax_kernel_feature_creator(
+    data, projection_matrix, attention_dims_t, batch_dims_t, precision, normalize_data=True
+):
+    """
+    Constructs kernel sin-cos features for fast softmax attention
+
+    Args:
+      data: input for which features are computes
+      projection_matrix: random matrix used to compute features
+      attention_dims_t: tuple of attention dimensions
+      batch_dims_t: tuple of batch dimensions
+      precision: precision parameter
+      normalize_data: predicate indicating whether data should be normalized
+
+    Returns:
+      Random features for fast softmax attention.
+    """
+    if normalize_data:
+        # We have: exp(qk^T/sqrt{d}) = exp(|q|^2/2sqrt{d}) * exp(|k|^2/2sqrt{d}) *
+        # exp(-(|q*c-k*c|^2)/2), where c = 1.0 / sqrt{sqrt{d}}.
+        data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1])))
+    else:
+        data_normalizer = 1.0
+    ratio = 1.0 / jnp.sqrt(projection_matrix.shape[0])
+    data_mod_shape = data.shape[0 : len(batch_dims_t)] + projection_matrix.shape
+    data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix
+
+    data_dash = lax.dot_general(
+        data_normalizer * data,
+        data_thick_random_matrix,
+        (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), (batch_dims_t, batch_dims_t)),
+        precision=precision,
+    )
+    data_dash_cos = ratio * jnp.cos(data_dash)
+    data_dash_sin = ratio * jnp.sin(data_dash)
+    data_dash = jnp.concatenate((data_dash_cos, data_dash_sin), axis=-1)
+
+    # Constructing D_data and data^{'}
+    diag_data = jnp.square(data)
+    diag_data = jnp.sum(diag_data, axis=data.ndim - 1)
+    diag_data = (diag_data / 2.0) * data_normalizer * data_normalizer
+    diag_data = jnp.expand_dims(diag_data, axis=data.ndim - 1)
+    # Additional renormalization for numerical stability
+    data_renormalizer = jnp.max(diag_data, attention_dims_t, keepdims=True)
+    diag_data -= data_renormalizer
+    diag_data = jnp.exp(diag_data)
+    data_prime = data_dash * diag_data
+    return data_prime
+
+
+def generalized_kernel_feature_creator(
+    data, projection_matrix, batch_dims_t, precision, kernel_fn, kernel_epsilon, normalize_data
+):
+    """
+    Constructs kernel features for fast generalized attention
+
+    Args:
+      data: input for which features are computes
+      projection_matrix: matrix used to compute features
+      batch_dims_t: tuple of batch dimensions
+      precision: precision parameter
+      kernel_fn: kernel function used
+      kernel_epsilon: additive positive term added to every feature for numerical
+        stability
+      normalize_data: predicate indicating whether data should be normalized
+
+    Returns:
+      Random features for fast generalized attention.
+    """
+    if normalize_data:
+        data_normalizer = 1.0 / (jnp.sqrt(jnp.sqrt(data.shape[-1])))
+    else:
+        data_normalizer = 1.0
+    if projection_matrix is None:
+        return kernel_fn(data_normalizer * data) + kernel_epsilon
+    else:
+        data_mod_shape = data.shape[0 : len(batch_dims_t)] + projection_matrix.shape
+        data_thick_random_matrix = jnp.zeros(data_mod_shape) + projection_matrix
+        data_dash = lax.dot_general(
+            data_normalizer * data,
+            data_thick_random_matrix,
+            (((data.ndim - 1,), (data_thick_random_matrix.ndim - 1,)), (batch_dims_t, batch_dims_t)),
+            precision=precision,
+        )
+    data_prime = kernel_fn(data_dash) + kernel_epsilon
+    return data_prime
+
+
+def make_fast_softmax_attention(
+    qkv_dim,
+    renormalize_attention=True,
+    numerical_stabilizer=0.000001,
+    nb_features=256,
+    ortho_features=True,
+    ortho_scaling=0.0,
+    redraw_features=True,
+    unidirectional=False,
+    nonnegative_features=True,
+    lax_scan_unroll=1,
+):
+    """Construct a fast softmax attention method."""
+    logging.info(
+        "Fast softmax attention: %s features and orthogonal=%s, renormalize=%s",
+        nb_features,
+        ortho_features,
+        renormalize_attention,
+    )
+    if ortho_features:
+        matrix_creator = functools.partial(GaussianOrthogonalRandomMatrix, nb_features, qkv_dim, scaling=ortho_scaling)
+    else:
+        matrix_creator = functools.partial(GaussianUnstructuredRandomMatrix, nb_features, qkv_dim)
+    if nonnegative_features:
+
+        def kernel_feature_creator(
+            data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=True
+        ):
+            return nonnegative_softmax_kernel_feature_creator(
+                data,
+                projection_matrix,
+                attention_dims_t,
+                batch_dims_t,
+                precision,
+                is_query,
+                normalize_data,
+                numerical_stabilizer,
+            )
+
+    else:
+
+        def kernel_feature_creator(
+            data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=True
+        ):
+            del is_query
+            return sincos_softmax_kernel_feature_creator(
+                data, projection_matrix, attention_dims_t, batch_dims_t, precision, normalize_data
+            )
+
+    attention_fn = FastAttentionviaLowRankDecomposition(
+        matrix_creator,
+        kernel_feature_creator,
+        renormalize_attention=renormalize_attention,
+        numerical_stabilizer=numerical_stabilizer,
+        redraw_features=redraw_features,
+        unidirectional=unidirectional,
+        lax_scan_unroll=lax_scan_unroll,
+    ).dot_product_attention
+    return attention_fn
+
+
+def make_fast_generalized_attention(
+    qkv_dim,
+    renormalize_attention=True,
+    numerical_stabilizer=0.0,
+    nb_features=256,
+    features_type="deterministic",
+    kernel_fn=jax.nn.relu,
+    kernel_epsilon=0.001,
+    redraw_features=False,
+    unidirectional=False,
+    lax_scan_unroll=1,
+):
+    """Construct a fast generalized attention menthod."""
+    logging.info("Fast generalized attention.: %s features and renormalize=%s", nb_features, renormalize_attention)
+    if features_type == "ortho":
+        matrix_creator = functools.partial(GaussianOrthogonalRandomMatrix, nb_features, qkv_dim, scaling=False)
+    elif features_type == "iid":
+        matrix_creator = functools.partial(GaussianUnstructuredRandomMatrix, nb_features, qkv_dim)
+    elif features_type == "deterministic":
+        matrix_creator = None
+    else:
+        raise ValueError("Unknown feature value type")
+
+    def kernel_feature_creator(
+        data, projection_matrix, attention_dims_t, batch_dims_t, precision, is_query, normalize_data=False
+    ):
+        del attention_dims_t
+        del is_query
+        return generalized_kernel_feature_creator(
+            data, projection_matrix, batch_dims_t, precision, kernel_fn, kernel_epsilon, normalize_data
+        )
+
+    attention_fn = FastAttentionviaLowRankDecomposition(
+        matrix_creator,
+        kernel_feature_creator,
+        renormalize_attention=renormalize_attention,
+        numerical_stabilizer=numerical_stabilizer,
+        redraw_features=redraw_features,
+        unidirectional=unidirectional,
+        lax_scan_unroll=lax_scan_unroll,
+    ).dot_product_attention
+    return attention_fn
+
+
+class RandomMatrix(object):
+    r"""
+    Abstract class providing a method for constructing 2D random arrays. Class is responsible for constructing 2D
+    random arrays.
+    """
+
+    __metaclass__ = abc.ABCMeta
+
+    @abc.abstractmethod
+    def get_2d_array(self):
+        raise NotImplementedError("Abstract method")
+
+
+class GaussianUnstructuredRandomMatrix(RandomMatrix):
+    def __init__(self, nb_rows, nb_columns, key):
+        self.nb_rows = nb_rows
+        self.nb_columns = nb_columns
+        self.key = key
+
+    def get_2d_array(self):
+        return random.normal(self.key, (self.nb_rows, self.nb_columns))
+
+
+class GaussianOrthogonalRandomMatrix(RandomMatrix):
+    r"""
+    Class providing a method to create Gaussian orthogonal matrix. Class is responsible for constructing 2D Gaussian
+    orthogonal arrays.
+    """
+
+    def __init__(self, nb_rows, nb_columns, key, scaling=0):
+        self.nb_rows = nb_rows
+        self.nb_columns = nb_columns
+        self.key = key
+        self.scaling = scaling
+
+    def get_2d_array(self):
+        nb_full_blocks = int(self.nb_rows / self.nb_columns)
+        block_list = []
+        rng = self.key
+        for _ in range(nb_full_blocks):
+            rng, rng_input = jax.random.split(rng)
+            unstructured_block = random.normal(rng_input, (self.nb_columns, self.nb_columns))
+            q, _ = jnp.linalg.qr(unstructured_block)
+            q = jnp.transpose(q)
+            block_list.append(q)
+        remaining_rows = self.nb_rows - nb_full_blocks * self.nb_columns
+        if remaining_rows > 0:
+            rng, rng_input = jax.random.split(rng)
+            unstructured_block = random.normal(rng_input, (self.nb_columns, self.nb_columns))
+            q, _ = jnp.linalg.qr(unstructured_block)
+            q = jnp.transpose(q)
+            block_list.append(q[0:remaining_rows])
+        final_matrix = jnp.vstack(block_list)
+
+        if self.scaling == 0:
+            multiplier = jnp.linalg.norm(random.normal(self.key, (self.nb_rows, self.nb_columns)), axis=1)
+        elif self.scaling == 1:
+            multiplier = jnp.sqrt(float(self.nb_columns)) * jnp.ones((self.nb_rows))
+        else:
+            raise ValueError("Scaling must be one of {0, 1}. Was %s" % self._scaling)
+
+        return jnp.matmul(jnp.diag(multiplier), final_matrix)
+
+
+class FastAttention(object):
+    r"""
+    Abstract class providing a method for fast attention. Class is responsible for providing a method
+    <dot_product_attention> for fast approximate attention.
+    """
+
+    __metaclass__ = abc.ABCMeta
+
+    @abc.abstractmethod
+    def dot_product_attention(
+        self,
+        query,
+        key,
+        value,
+        dtype=jnp.float32,
+        bias=None,
+        axis=None,
+        broadcast_dropout=True,
+        dropout_rng=None,
+        dropout_rate=0.0,
+        deterministic=False,
+        precision=None,
+    ):
+        """
+        Computes dot-product attention given query, key, and value. This is the core function for applying fast
+        approximate dot-product attention. It calculates the attention weights given query and key and combines the
+        values using the attention weights. This function supports multi-dimensional inputs
+
+        Args:
+          query: queries for calculating attention with shape of [batch_size, dim1,
+            dim2, ..., dimN, num_heads, mem_channels].
+          key: keys for calculating attention with shape of [batch_size, dim1, dim2,
+            ..., dimN, num_heads, mem_channels].
+          value: values to be used in attention with shape of [batch_size, dim1,
+            dim2,..., dimN, num_heads, value_channels].
+          dtype: the dtype of the computation (default: float32)
+          bias: bias for the attention weights. This can be used for incorporating
+            autoregressive mask, padding mask, proximity bias.
+          axis: axises over which the attention is applied.
+          broadcast_dropout: bool: use a broadcasted dropout along batch dims.
+          dropout_rng: JAX PRNGKey: to be used for dropout.
+          dropout_rate: dropout rate.
+          deterministic: bool, deterministic or not (to apply dropout).
+          precision: numerical precision of the computation see `jax.lax.Precision`
+            for details
+
+        Returns:
+          Output of shape [bs, dim1, dim2, ..., dimN,, num_heads, value_channels].
+        """
+        raise NotImplementedError("Abstract method")
+
+
+def _numerator(z_slice_shape, precision, unroll=1):
+    def fwd(qs, ks, vs):
+        def body(p, qkv):
+            (q, k, v) = qkv
+            p += jnp.einsum("...m,...d->...md", k, v, precision=precision)
+            X_slice = jnp.einsum("...m,...md->...d", q, p, precision=precision)
+            return p, X_slice
+
+        init_value = jnp.zeros(z_slice_shape)
+        p, W = lax.scan(body, init_value, (qs, ks, vs), unroll=unroll)
+        return W, (p, qs, ks, vs)
+
+    def bwd(pqkv, W_ct):
+        def body(carry, qkv_xct):
+            p, p_ct = carry
+            q, k, v, x_ct = qkv_xct
+            q_ct = jnp.einsum("...d,...md->...m", x_ct, p, precision=precision)
+            p_ct += jnp.einsum("...d,...m->...md", x_ct, q, precision=precision)
+            k_ct = jnp.einsum("...md,...d->...m", p_ct, v, precision=precision)
+            v_ct = jnp.einsum("...md,...m->...d", p_ct, k, precision=precision)
+            p -= jnp.einsum("...m,...d->...md", k, v, precision=precision)
+            return (p, p_ct), (q_ct, k_ct, v_ct)
+
+        p, qs, ks, vs = pqkv
+        _, (qs_ct, ks_ct, vs_ct) = lax.scan(
+            body, (p, jnp.zeros_like(p)), (qs, ks, vs, W_ct), reverse=True, unroll=unroll
+        )
+        return qs_ct, ks_ct, vs_ct
+
+    @jax.custom_vjp
+    def _numerator_impl(qs, ks, vs):
+        W, _ = fwd(qs, ks, vs)
+        return W
+
+    _numerator_impl.defvjp(fwd, bwd)
+
+    return _numerator_impl
+
+
+def _denominator(t_slice_shape, precision, unroll=1):
+    def fwd(qs, ks):
+        def body(p, qk):
+            q, k = qk
+            p += k
+            x = jnp.einsum("...m,...m->...", q, p, precision=precision)
+            return p, x
+
+        p = jnp.zeros(t_slice_shape)
+        p, R = lax.scan(body, p, (qs, ks), unroll=unroll)
+        return R, (qs, ks, p)
+
+    def bwd(qkp, R_ct):
+        def body(carry, qkx):
+            p, p_ct = carry
+            q, k, x_ct = qkx
+            q_ct = jnp.einsum("...,...m->...m", x_ct, p, precision=precision)
+            p_ct += jnp.einsum("...,...m->...m", x_ct, q, precision=precision)
+            k_ct = p_ct
+            p -= k
+            return (p, p_ct), (q_ct, k_ct)
+
+        qs, ks, p = qkp
+        _, (qs_ct, ks_ct) = lax.scan(body, (p, jnp.zeros_like(p)), (qs, ks, R_ct), reverse=True, unroll=unroll)
+        return (qs_ct, ks_ct)
+
+    @jax.custom_vjp
+    def _denominator_impl(qs, ks):
+        R, _ = fwd(qs, ks)
+        return R
+
+    _denominator_impl.defvjp(fwd, bwd)
+
+    return _denominator_impl
+
+
+class FastAttentionviaLowRankDecomposition(FastAttention):
+    r"""
+    Class providing a method for fast attention via low rank decomposition. Class is responsible for providing a method
+    <dot_product_attention> for fast dot-product attention with the use of low rank decomposition (e.g. with random
+    feature maps).
+    """
+
+    def __init__(
+        self,
+        matrix_creator,
+        kernel_feature_creator,
+        renormalize_attention,
+        numerical_stabilizer,
+        redraw_features,
+        unidirectional,
+        lax_scan_unroll=1,
+    ):  # For optimal GPU performance, set to 16.
+        rng = random.PRNGKey(0)
+        self.matrix_creator = matrix_creator
+        self.projection_matrix = self.draw_weights(rng)
+        self.kernel_feature_creator = kernel_feature_creator
+        self.renormalize_attention = renormalize_attention
+        self.numerical_stabilizer = numerical_stabilizer
+        self.redraw_features = redraw_features
+        self.unidirectional = unidirectional
+        self.lax_scan_unroll = lax_scan_unroll
+
+    def draw_weights(self, key):
+        if self.matrix_creator is None:
+            return None
+        matrixrng, _ = random.split(key)
+        projection_matrix = self.matrix_creator(key=matrixrng).get_2d_array()
+        return projection_matrix
+
+    def dot_product_attention(
+        self,
+        query,
+        key,
+        value,
+        dtype=jnp.float32,
+        bias=None,
+        axis=None,
+        broadcast_dropout=True,
+        dropout_rng=None,
+        dropout_rate=0.0,
+        deterministic=False,
+        precision=None,
+    ):
+
+        assert key.shape[:-1] == value.shape[:-1]
+        assert query.shape[0:1] == key.shape[0:1] and query.shape[-1] == key.shape[-1]
+        if axis is None:
+            axis = tuple(range(1, key.ndim - 2))
+        if not isinstance(axis, Iterable):
+            axis = (axis,)
+        assert key.ndim == query.ndim
+        assert key.ndim == value.ndim
+        for ax in axis:
+            if not (query.ndim >= 3 and 1 <= ax < query.ndim - 2):
+                raise ValueError("Attention axis must be between the batch " "axis and the last-two axes.")
+        n = key.ndim
+
+        # Constructing projection tensor.
+        if self.redraw_features:
+            # TODO(kchoro): Get rid of the constant below.
+            query_seed = lax.convert_element_type(jnp.ceil(jnp.sum(query) * 10000000.0), jnp.int32)
+            rng = random.PRNGKey(query_seed)
+            self.projection_matrix = self.draw_weights(rng)
+
+        # batch_dims is  <bs, <non-attention dims>, num_heads>
+        batch_dims = tuple(onp.delete(range(n), axis + (n - 1,)))
+        # q & k -> (bs, <non-attention dims>, num_heads, <attention dims>, channels)
+        qk_perm = batch_dims + axis + (n - 1,)
+        k_extra_perm = axis + batch_dims + (n - 1,)
+        key_extra = key.transpose(k_extra_perm)
+        key = key.transpose(qk_perm)
+        query = query.transpose(qk_perm)
+        # v -> (bs, <non-attention dims>, num_heads, <attention dims>, channels)
+        v_perm = batch_dims + axis + (n - 1,)
+        value = value.transpose(v_perm)
+        batch_dims_t = tuple(range(len(batch_dims)))
+        attention_dims_t = tuple(range(len(batch_dims), len(batch_dims) + len(axis)))
+
+        # Constructing tensors Q^{'} and K^{'}.
+        query_prime = self.kernel_feature_creator(
+            query, self.projection_matrix, attention_dims_t, batch_dims_t, precision, True
+        )
+        key_prime = self.kernel_feature_creator(
+            key, self.projection_matrix, attention_dims_t, batch_dims_t, precision, False
+        )
+
+        if self.unidirectional:
+            index = attention_dims_t[0]
+            z_slice_shape = key_prime.shape[0 : len(batch_dims_t)] + (key_prime.shape[-1],) + (value.shape[-1],)
+
+            numerator_fn = _numerator(z_slice_shape, precision, self.lax_scan_unroll)
+            W = numerator_fn(
+                jnp.moveaxis(query_prime, index, 0), jnp.moveaxis(key_prime, index, 0), jnp.moveaxis(value, index, 0)
+            )
+
+            # Constructing W = (Q^{'}(K^{'})^{T})_{masked}V
+            W = jnp.moveaxis(W, 0, index)
+
+            if not self.renormalize_attention:
+                # Unidirectional, not-normalized attention.
+                perm_inv = _invert_perm(qk_perm)
+                result = W.transpose(perm_inv)
+                return result
+            else:
+                # Unidirectional, normalized attention.
+                thick_all_ones = jnp.zeros(key.shape[0:-1]) + jnp.ones(key_extra.shape[0 : len(axis)])
+
+                index = attention_dims_t[0]
+                t_slice_shape = key_prime.shape[0 : len(batch_dims_t)] + (key_prime.shape[-1],)
+                denominator_fn = _denominator(t_slice_shape, precision, self.lax_scan_unroll)
+                R = denominator_fn(jnp.moveaxis(query_prime, index, 0), jnp.moveaxis(key_prime, index, 0))
+
+                R = jnp.moveaxis(R, 0, index)
+        else:
+            contract_query = tuple(range(len(batch_dims) + len(axis), len(batch_dims) + len(axis) + 1))
+            contract_z = tuple(range(len(batch_dims), len(batch_dims) + 1))
+            # Constructing Z = (K^{'})^{T}V
+            # Z (bs, <non-attention dims>, num_heads, channels_m, channels_v)
+            Z = lax.dot_general(
+                key_prime,
+                value,
+                ((attention_dims_t, attention_dims_t), (batch_dims_t, batch_dims_t)),
+                precision=precision,
+            )
+            # Constructing W = Q^{'}Z = Q^{'}(K^{'})^{T}V
+            # q (bs, <non-attention dims>, num_heads, <attention dims>, channels_m)
+            # Z (bs, <non-attention dims>, num_heads, channels_m, channels_v)
+            # W (bs,  <non-attention dims>, num_heads, <attention dims>, channels_v)
+            W = lax.dot_general(
+                query_prime, Z, ((contract_query, contract_z), (batch_dims_t, batch_dims_t)), precision=precision
+            )
+            if not self.renormalize_attention:
+                # Bidirectional, not-normalized attention.
+                perm_inv = _invert_perm(qk_perm)
+                result = W.transpose(perm_inv)
+                return result
+            else:
+                # Bidirectional, normalized attention.
+                thick_all_ones = jnp.zeros(key.shape[0:-1]) + jnp.ones(key_extra.shape[0 : len(axis)])
+                contract_key = tuple(range(len(batch_dims), len(batch_dims) + len(axis)))
+                contract_thick_all_ones = tuple(range(thick_all_ones.ndim - len(axis), thick_all_ones.ndim))
+                # Construct T = (K^{'})^{T} 1_L
+                # k (bs, <non-attention dims>, num_heads, <attention dims>, channels)
+                T = lax.dot_general(
+                    key_prime,
+                    thick_all_ones,
+                    ((contract_key, contract_thick_all_ones), (batch_dims_t, batch_dims_t)),
+                    precision=precision,
+                )
+
+                # Construct partition function: R = Q^{'} T = Q^{'}(K^{'})^{T} 1_L
+                # q_p (bs, <non-attention dims>, num_heads, <attention dims>, channs_m)
+                # T   (bs, <non-attention dims>, num_heads, channels_m)
+                R = lax.dot_general(
+                    query_prime,
+                    T,
+                    (((query_prime.ndim - 1,), (T.ndim - 1,)), (batch_dims_t, range(0, len(T.shape) - 1))),
+                    precision=precision,
+                )
+
+        R = R + 2 * self.numerical_stabilizer * (jnp.abs(R) <= self.numerical_stabilizer)
+        R = jnp.reciprocal(R)
+        R = jnp.expand_dims(R, len(R.shape))
+        # W (bs, <non-attention dims>, num_heads, <attention dims>, channels_v)
+        # R (bs, <non-attention dims>, num_heads, <attention dims>, extra_channel)
+        result = W * R
+        # back to (bs, dim1, dim2, ..., dimN, num_heads, channels)
+        perm_inv = _invert_perm(qk_perm)
+        result = result.transpose(perm_inv)
+        return result
+
+
+def _invert_perm(perm):
+    perm_inv = [0] * len(perm)
+    for i, j in enumerate(perm):
+        perm_inv[j] = i
+    return tuple(perm_inv)

examples/research_projects/performer/run_mlm_performer.py

@@ -0,0 +1,685 @@
+# coding=utf-8
+# Copyright 2020 The HuggingFace Team All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""
+Fine-tuning the library models for masked language modeling (BERT, ALBERT, RoBERTa...) with whole word masking on a
+text file or a dataset.
+
+Here is the full list of checkpoints on the hub that can be fine-tuned by this script:
+https://huggingface.co/models?filter=masked-lm
+"""
+import logging
+import os
+import sys
+from dataclasses import dataclass, field
+
+# You can also adapt this script on your own masked language modeling task. Pointers for this are left as comments.
+from pathlib import Path
+from typing import Dict, List, Optional, Tuple
+
+import numpy as np
+from datasets import load_dataset
+from tqdm import tqdm
+
+import jax
+import jax.numpy as jnp
+from flax import jax_utils
+from flax.optim import Adam
+from flax.training import common_utils
+from flax.training.common_utils import get_metrics
+from jax.nn import log_softmax
+from modeling_flax_performer import FlaxPerformerForMaskedLM
+from transformers import (
+    MODEL_FOR_MASKED_LM_MAPPING,
+    AutoTokenizer,
+    BertConfig,
+    FlaxBertForMaskedLM,
+    HfArgumentParser,
+    PreTrainedTokenizerBase,
+    TensorType,
+    TrainingArguments,
+    is_tensorboard_available,
+    set_seed,
+)
+
+
+# Cache the result
+has_tensorboard = is_tensorboard_available()
+if has_tensorboard:
+    try:
+        from flax.metrics.tensorboard import SummaryWriter
+    except ImportError as ie:
+        has_tensorboard = False
+        print(f"Unable to display metrics through TensorBoard because some package are not installed: {ie}")
+
+else:
+    print(
+        "Unable to display metrics through TensorBoard because the package is not installed: "
+        "Please run pip install tensorboard to enable."
+    )
+
+MODEL_CONFIG_CLASSES = list(MODEL_FOR_MASKED_LM_MAPPING.keys())
+MODEL_TYPES = tuple(conf.model_type for conf in MODEL_CONFIG_CLASSES)
+
+
+@dataclass
+class WandbArguments:
+    """
+    Arguments for logging
+    """
+
+    wandb_user_name: Optional[str] = field(
+        default=None,
+        metadata={"help": "The WandB user name for potential logging. If left None, no logging"},
+    )
+    wandb_project_name: Optional[str] = field(
+        default="performer-experiments",
+        metadata={"help": "The WandB project name for potential logging"},
+    )
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
+    """
+
+    model_name_or_path: Optional[str] = field(
+        default=None,
+        metadata={
+            "help": "The model checkpoint for weights initialization."
+            "Don't set if you want to train a model from scratch."
+        },
+    )
+    performer: bool = field(
+        default=False,
+        metadata={"help": "Whether to use FAVOR+ attention"},
+    )
+    reinitialize: bool = field(
+        default=False,
+        metadata={"help": "Whether to use a blank model without pretraining"},
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    use_fast_tokenizer: bool = field(
+        default=True,
+        metadata={"help": "Whether to use one of the fast tokenizer (backed by the tokenizers library) or not."},
+    )
+    cache_dir: Optional[str] = field(
+        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
+    )
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    dataset_name: Optional[str] = field(
+        default=None, metadata={"help": "The name of the dataset to use (via the datasets library)."}
+    )
+    dataset_config_name: Optional[str] = field(
+        default=None, metadata={"help": "The configuration name of the dataset to use (via the datasets library)."}
+    )
+    train_file: Optional[str] = field(default=None, metadata={"help": "The input training data file (a text file)."})
+    validation_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input evaluation data file to evaluate the perplexity on (a text file)."},
+    )
+    train_ref_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input train ref data file for whole word masking in Chinese."},
+    )
+    validation_ref_file: Optional[str] = field(
+        default=None,
+        metadata={"help": "An optional input validation ref data file for whole word masking in Chinese."},
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+    validation_split_percentage: Optional[int] = field(
+        default=5,
+        metadata={
+            "help": "The percentage of the train set used as validation set in case there's no validation split"
+        },
+    )
+    max_seq_length: Optional[int] = field(
+        default=None,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated. Default to the max input length of the model."
+        },
+    )
+    preprocessing_num_workers: Optional[int] = field(
+        default=None,
+        metadata={"help": "The number of processes to use for the preprocessing."},
+    )
+    mlm_probability: float = field(
+        default=0.15, metadata={"help": "Ratio of tokens to mask for masked language modeling loss"}
+    )
+    pad_to_max_length: bool = field(
+        default=False,
+        metadata={
+            "help": "Whether to pad all samples to `max_seq_length`. "
+            "If False, will pad the samples dynamically when batching to the maximum length in the batch."
+        },
+    )
+
+    def __post_init__(self):
+        if self.dataset_name is None and self.train_file is None and self.validation_file is None:
+            raise ValueError("Need either a dataset name or a training/validation file.")
+        else:
+            if self.train_file is not None:
+                extension = self.train_file.split(".")[-1]
+                assert extension in ["csv", "json", "txt"], "`train_file` should be a csv, a json or a txt file."
+            if self.validation_file is not None:
+                extension = self.validation_file.split(".")[-1]
+                assert extension in ["csv", "json", "txt"], "`validation_file` should be a csv, a json or a txt file."
+
+
+# Adapted from transformers/data/data_collator.py
+# Letting here for now, let's discuss where it should live
+@dataclass
+class FlaxDataCollatorForLanguageModeling:
+    """
+    Data collator used for language modeling. Inputs are dynamically padded to the maximum length of a batch if they
+    are not all of the same length.
+
+    Args:
+        tokenizer (:class:`~transformers.PreTrainedTokenizer` or :class:`~transformers.PreTrainedTokenizerFast`):
+            The tokenizer used for encoding the data.
+        mlm (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether or not to use masked language modeling. If set to :obj:`False`, the labels are the same as the
+            inputs with the padding tokens ignored (by setting them to -100). Otherwise, the labels are -100 for
+            non-masked tokens and the value to predict for the masked token.
+        mlm_probability (:obj:`float`, `optional`, defaults to 0.15):
+            The probability with which to (randomly) mask tokens in the input, when :obj:`mlm` is set to :obj:`True`.
+
+    .. note::
+
+        For best performance, this data collator should be used with a dataset having items that are dictionaries or
+        BatchEncoding, with the :obj:`"special_tokens_mask"` key, as returned by a
+        :class:`~transformers.PreTrainedTokenizer` or a :class:`~transformers.PreTrainedTokenizerFast` with the
+        argument :obj:`return_special_tokens_mask=True`.
+    """
+
+    tokenizer: PreTrainedTokenizerBase
+    mlm: bool = True
+    mlm_probability: float = 0.15
+
+    def __post_init__(self):
+        if self.mlm and self.tokenizer.mask_token is None:
+            raise ValueError(
+                "This tokenizer does not have a mask token which is necessary for masked language modeling. "
+                "You should pass `mlm=False` to train on causal language modeling instead."
+            )
+
+    def __call__(self, examples: List[Dict[str, np.ndarray]], pad_to_multiple_of: int) -> Dict[str, np.ndarray]:
+        # Handle dict or lists with proper padding and conversion to tensor.
+        batch = self.tokenizer.pad(examples, pad_to_multiple_of=pad_to_multiple_of, return_tensors=TensorType.NUMPY)
+
+        # If special token mask has been preprocessed, pop it from the dict.
+        special_tokens_mask = batch.pop("special_tokens_mask", None)
+        if self.mlm:
+            batch["input_ids"], batch["labels"] = self.mask_tokens(
+                batch["input_ids"], special_tokens_mask=special_tokens_mask
+            )
+        else:
+            labels = batch["input_ids"].copy()
+            if self.tokenizer.pad_token_id is not None:
+                labels[labels == self.tokenizer.pad_token_id] = -100
+            batch["labels"] = labels
+        return batch
+
+    def mask_tokens(
+        self, inputs: np.ndarray, special_tokens_mask: Optional[np.ndarray]
+    ) -> Tuple[jnp.ndarray, jnp.ndarray]:
+        """
+        Prepare masked tokens inputs/labels for masked language modeling: 80% MASK, 10% random, 10% original.
+        """
+        labels = inputs.copy()
+        # We sample a few tokens in each sequence for MLM training (with probability `self.mlm_probability`)
+        probability_matrix = np.full(labels.shape, self.mlm_probability)
+        special_tokens_mask = special_tokens_mask.astype("bool")
+
+        probability_matrix[special_tokens_mask] = 0.0
+        masked_indices = np.random.binomial(1, probability_matrix).astype("bool")
+        labels[~masked_indices] = -100  # We only compute loss on masked tokens
+
+        # 80% of the time, we replace masked input tokens with tokenizer.mask_token ([MASK])
+        indices_replaced = np.random.binomial(1, np.full(labels.shape, 0.8)).astype("bool") & masked_indices
+        inputs[indices_replaced] = self.tokenizer.convert_tokens_to_ids(self.tokenizer.mask_token)
+
+        # 10% of the time, we replace masked input tokens with random word
+        indices_random = np.random.binomial(1, np.full(labels.shape, 0.5)).astype("bool")
+        indices_random &= masked_indices & ~indices_replaced
+
+        random_words = np.random.randint(self.tokenizer.vocab_size, size=labels.shape, dtype="i4")
+        inputs[indices_random] = random_words[indices_random]
+
+        # The rest of the time (10% of the time) we keep the masked input tokens unchanged
+        return inputs, labels
+
+
+def create_learning_rate_scheduler(
+    factors="constant * linear_warmup * rsqrt_decay",
+    base_learning_rate=0.5,
+    warmup_steps=1000,
+    decay_factor=0.5,
+    steps_per_decay=20000,
+    steps_per_cycle=100000,
+):
+    """Creates learning rate schedule.
+    Interprets factors in the factors string which can consist of:
+    * constant: interpreted as the constant value,
+    * linear_warmup: interpreted as linear warmup until warmup_steps,
+    * rsqrt_decay: divide by square root of max(step, warmup_steps)
+    * rsqrt_normalized_decay: divide by square root of max(step/warmup_steps, 1)
+    * decay_every: Every k steps decay the learning rate by decay_factor.
+    * cosine_decay: Cyclic cosine decay, uses steps_per_cycle parameter.
+    Args:
+      factors: string, factors separated by "*" that defines the schedule.
+      base_learning_rate: float, the starting constant for the lr schedule.
+      warmup_steps: int, how many steps to warm up for in the warmup schedule.
+      decay_factor: float, the amount to decay the learning rate by.
+      steps_per_decay: int, how often to decay the learning rate.
+      steps_per_cycle: int, steps per cycle when using cosine decay.
+    Returns:
+      a function learning_rate(step): float -> {"learning_rate": float}, the
+      step-dependent lr.
+    """
+    factors = [n.strip() for n in factors.split("*")]
+
+    def step_fn(step):
+        """Step to learning rate function."""
+        ret = 1.0
+        for name in factors:
+            if name == "constant":
+                ret *= base_learning_rate
+            elif name == "linear_warmup":
+                ret *= jnp.minimum(1.0, step / warmup_steps)
+            elif name == "rsqrt_decay":
+                ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
+            elif name == "rsqrt_normalized_decay":
+                ret *= jnp.sqrt(warmup_steps)
+                ret /= jnp.sqrt(jnp.maximum(step, warmup_steps))
+            elif name == "decay_every":
+                ret *= decay_factor ** (step // steps_per_decay)
+            elif name == "cosine_decay":
+                progress = jnp.maximum(0.0, (step - warmup_steps) / float(steps_per_cycle))
+                ret *= jnp.maximum(0.0, 0.5 * (1.0 + jnp.cos(jnp.pi * (progress % 1.0))))
+            else:
+                raise ValueError("Unknown factor %s." % name)
+        return jnp.asarray(ret, dtype=jnp.float32)
+
+    return step_fn
+
+
+def compute_metrics(logits, labels, weights, label_smoothing=0.0):
+    """Compute summary metrics."""
+    loss, normalizer = cross_entropy(logits, labels, weights, label_smoothing)
+    acc, _ = accuracy(logits, labels, weights)
+    metrics = {"loss": loss, "accuracy": acc, "normalizer": normalizer}
+    metrics = jax.lax.psum(metrics, axis_name="batch")
+    return metrics
+
+
+def accuracy(logits, targets, weights=None):
+    """Compute weighted accuracy for log probs and targets.
+    Args:
+     logits: [batch, length, num_classes] float array.
+     targets: categorical targets [batch, length] int array.
+     weights: None or array of shape [batch, length]
+    Returns:
+      Tuple of scalar loss and batch normalizing factor.
+    """
+    if logits.ndim != targets.ndim + 1:
+        raise ValueError(
+            "Incorrect shapes. Got shape %s logits and %s targets" % (str(logits.shape), str(targets.shape))
+        )
+
+    loss = jnp.equal(jnp.argmax(logits, axis=-1), targets)
+    loss *= weights
+
+    return loss.sum(), weights.sum()
+
+
+def cross_entropy(logits, targets, weights=None, label_smoothing=0.0):
+    """Compute cross entropy and entropy for log probs and targets.
+    Args:
+     logits: [batch, length, num_classes] float array.
+     targets: categorical targets [batch, length] int array.
+     weights: None or array of shape [batch, length]
+     label_smoothing: label smoothing constant, used to determine the on and off values.
+    Returns:
+      Tuple of scalar loss and batch normalizing factor.
+    """
+    if logits.ndim != targets.ndim + 1:
+        raise ValueError(
+            "Incorrect shapes. Got shape %s logits and %s targets" % (str(logits.shape), str(targets.shape))
+        )
+
+    vocab_size = logits.shape[-1]
+    confidence = 1.0 - label_smoothing
+    low_confidence = (1.0 - confidence) / (vocab_size - 1)
+    normalizing_constant = -(
+        confidence * jnp.log(confidence) + (vocab_size - 1) * low_confidence * jnp.log(low_confidence + 1e-20)
+    )
+    soft_targets = common_utils.onehot(targets, vocab_size, on_value=confidence, off_value=low_confidence)
+
+    loss = -jnp.sum(soft_targets * log_softmax(logits), axis=-1)
+    loss = loss - normalizing_constant
+
+    if weights is not None:
+        loss = loss * weights
+        normalizing_factor = weights.sum()
+    else:
+        normalizing_factor = np.prod(targets.shape)
+
+    return loss.sum(), normalizing_factor
+
+
+def training_step(optimizer, batch, dropout_rng):
+    dropout_rng, new_dropout_rng = jax.random.split(dropout_rng)
+
+    def loss_fn(params):
+        targets = batch.pop("labels")
+
+        # Hide away tokens which doesn't participate in the optimization
+        token_mask = jnp.where(targets > 0, 1.0, 0.0)
+
+        logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
+        loss, weight_sum = cross_entropy(logits, targets, token_mask)
+        return loss / weight_sum
+
+    step = optimizer.state.step
+    lr = lr_scheduler_fn(step)
+    grad_fn = jax.value_and_grad(loss_fn)
+    loss, grad = grad_fn(optimizer.target)
+    grad = jax.lax.pmean(grad, "batch")
+    optimizer = optimizer.apply_gradient(grad, learning_rate=lr)
+
+    return loss, optimizer, new_dropout_rng
+
+
+def eval_step(params, batch):
+    """
+    Calculate evaluation metrics on a batch.
+    """
+    targets = batch.pop("labels")
+
+    # Hide away tokens which doesn't participate in the optimization
+    token_mask = jnp.where(targets > 0, 1.0, 0.0)
+    logits = model(**batch, params=params, train=False)[0]
+
+    return compute_metrics(logits, targets, token_mask)
+
+
+def generate_batch_splits(samples_idx: jnp.ndarray, batch_size: int) -> jnp.ndarray:
+    nb_samples = len(samples_idx)
+    samples_to_remove = nb_samples % batch_size
+
+    if samples_to_remove != 0:
+        samples_idx = samples_idx[:-samples_to_remove]
+    sections_split = nb_samples // batch_size
+    batch_idx = np.split(samples_idx, sections_split)
+    return batch_idx
+
+
+if __name__ == "__main__":
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, TrainingArguments, WandbArguments))
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args, wandb_args = parser.parse_json_file(
+            json_file=os.path.abspath(sys.argv[1])
+        )
+    else:
+        model_args, data_args, training_args, wandb_args = parser.parse_args_into_dataclasses()
+
+    if (
+        os.path.exists(training_args.output_dir)
+        and os.listdir(training_args.output_dir)
+        and training_args.do_train
+        and not training_args.overwrite_output_dir
+    ):
+        raise ValueError(
+            f"Output directory ({training_args.output_dir}) already exists and is not empty."
+            "Use --overwrite_output_dir to overcome."
+        )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        level="NOTSET",
+        datefmt="[%X]",
+    )
+
+    # Log on each process the small summary:
+    logger = logging.getLogger(__name__)
+    logger.warning(
+        f"Process rank: {training_args.local_rank}, device: {training_args.device}, n_gpu: {training_args.n_gpu}"
+        + f"distributed training: {bool(training_args.local_rank != -1)}, 16-bits training: {training_args.fp16}"
+    )
+
+    # Set the verbosity to info of the Transformers logger (on main process only):
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set seed before initializing model.
+    set_seed(training_args.seed)
+
+    # Get the datasets: you can either provide your own CSV/JSON/TXT training and evaluation files (see below)
+    # or just provide the name of one of the public datasets available on the hub at https://huggingface.co/datasets/
+    # (the dataset will be downloaded automatically from the datasets Hub).
+    #
+    # For CSV/JSON files, this script will use the column called 'text' or the first column if no column called
+    # 'text' is found. You can easily tweak this behavior (see below).
+    #
+    # In distributed training, the load_dataset function guarantees that only one local process can concurrently
+    # download the dataset.
+    if data_args.dataset_name is not None:
+        # Downloading and loading a dataset from the hub.
+        datasets = load_dataset(data_args.dataset_name, data_args.dataset_config_name)
+        if "validation" not in datasets.keys():
+            datasets["validation"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[:{data_args.validation_split_percentage}%]",
+            )
+            datasets["train"] = load_dataset(
+                data_args.dataset_name,
+                data_args.dataset_config_name,
+                split=f"train[{data_args.validation_split_percentage}%:]",
+            )
+    else:
+        data_files = {}
+        if data_args.train_file is not None:
+            data_files["train"] = data_args.train_file
+        if data_args.validation_file is not None:
+            data_files["validation"] = data_args.validation_file
+        extension = data_args.train_file.split(".")[-1]
+        if extension == "txt":
+            extension = "text"
+        datasets = load_dataset(extension, data_files=data_files)
+    # See more about loading any type of standard or custom dataset (from files, python dict, pandas DataFrame, etc) at
+    # https://huggingface.co/docs/datasets/loading_datasets.html.
+
+    # Load pretrained model and tokenizer
+
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+
+    rng = jax.random.PRNGKey(training_args.seed)
+    dropout_rngs = jax.random.split(rng, jax.local_device_count())
+
+    config = BertConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
+    lm_class = FlaxPerformerForMaskedLM if model_args.performer else FlaxBertForMaskedLM
+    if model_args.reinitialize:
+        model = lm_class(config=BertConfig.from_pretrained(model_args.model_name_or_path))
+    else:
+        model = lm_class.from_pretrained(
+            model_args.model_name_or_path,
+            dtype=jnp.float32,
+            input_shape=(training_args.train_batch_size, config.max_position_embeddings),
+            seed=training_args.seed,
+            dropout_rate=0.1,
+        )
+
+    if model_args.tokenizer_name:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.tokenizer_name, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+        )
+    elif model_args.model_name_or_path:
+        tokenizer = AutoTokenizer.from_pretrained(
+            model_args.model_name_or_path, cache_dir=model_args.cache_dir, use_fast=model_args.use_fast_tokenizer
+        )
+    else:
+        raise ValueError(
+            "You are instantiating a new tokenizer from scratch. This is not supported by this script."
+            "You can do it from another script, save it, and load it from here, using --tokenizer_name."
+        )
+
+    # Preprocessing the datasets.
+    # First we tokenize all the texts.
+    if training_args.do_train:
+        column_names = datasets["train"].column_names
+    else:
+        column_names = datasets["validation"].column_names
+    text_column_name = "text" if "text" in column_names else column_names[0]
+
+    padding = "max_length" if data_args.pad_to_max_length else False
+
+    def tokenize_function(examples):
+        # Remove empty lines
+        examples = [line for line in examples if len(line) > 0 and not line.isspace()]
+        return tokenizer(
+            examples,
+            return_special_tokens_mask=True,
+            padding=padding,
+            truncation=True,
+            max_length=data_args.max_seq_length,
+        )
+
+    tokenized_datasets = datasets.map(
+        tokenize_function,
+        input_columns=[text_column_name],
+        batched=True,
+        num_proc=data_args.preprocessing_num_workers,
+        remove_columns=column_names,
+        load_from_cache_file=not data_args.overwrite_cache,
+    )
+
+    # Enable tensorboard only on the master node
+    if has_tensorboard and jax.host_id() == 0:
+        summary_writer = SummaryWriter(log_dir=Path(training_args.output_dir).joinpath("logs").as_posix())
+
+    # Data collator
+    # This one will take care of randomly masking the tokens.
+    data_collator = FlaxDataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=data_args.mlm_probability)
+
+    # Setup optimizer
+    optimizer = Adam(
+        learning_rate=training_args.learning_rate,
+        weight_decay=training_args.weight_decay,
+        beta1=training_args.adam_beta1,
+        beta2=training_args.adam_beta2,
+    ).create(model.params)
+
+    # Create learning rate scheduler
+    lr_scheduler_fn = create_learning_rate_scheduler(
+        base_learning_rate=training_args.learning_rate, warmup_steps=max(training_args.warmup_steps, 1)
+    )
+
+    # Create parallel version of the training and evaluation steps
+    p_training_step = jax.pmap(training_step, "batch", donate_argnums=(0,))
+    p_eval_step = jax.pmap(eval_step, "batch", donate_argnums=(0,))
+
+    # Replicate the optimizer on each device
+    optimizer = jax_utils.replicate(optimizer)
+
+    # Store some constant
+    nb_epochs = int(training_args.num_train_epochs)
+    batch_size = int(training_args.train_batch_size)
+    eval_batch_size = int(training_args.eval_batch_size)
+
+    if wandb_args.wandb_user_name is not None:
+        import wandb
+
+        wandb.init(project=wandb_args.wandb_project_name, entity=wandb_args.wandb_user_name)
+
+    epochs = tqdm(range(nb_epochs), desc=f"Epoch ... (1/{nb_epochs})", position=0)
+    for epoch in epochs:
+
+        # ======================== Training ================================
+        # Create sampling rng
+        rng, training_rng, eval_rng = jax.random.split(rng, 3)
+
+        # Generate an epoch by shuffling sampling indices from the train dataset
+        nb_training_samples = len(tokenized_datasets["train"])
+        training_samples_idx = jax.random.permutation(training_rng, jnp.arange(nb_training_samples))
+        training_batch_idx = generate_batch_splits(training_samples_idx, batch_size)
+
+        # Gather the indexes for creating the batch and do a training step
+        for batch_idx in tqdm(training_batch_idx, desc="Training...", position=1):
+            samples = [tokenized_datasets["train"][int(idx)] for idx in batch_idx]
+            model_inputs = data_collator(samples, pad_to_multiple_of=16)
+
+            # Model forward
+            model_inputs = common_utils.shard(model_inputs.data)
+            loss, optimizer, dropout_rngs = p_training_step(optimizer, model_inputs, dropout_rngs)
+
+            if wandb_args.wandb_user_name is not None:
+                wandb.log({"Training loss": np.array(loss).mean()})
+
+        epochs.write(f"Loss: {loss}")
+
+        # ======================== Evaluating ==============================
+        nb_eval_samples = len(tokenized_datasets["validation"])
+        eval_samples_idx = jnp.arange(nb_eval_samples)
+        eval_batch_idx = generate_batch_splits(eval_samples_idx, eval_batch_size)
+
+        eval_metrics = []
+        for i, batch_idx in enumerate(tqdm(eval_batch_idx, desc="Evaluating ...", position=2)):
+            samples = [tokenized_datasets["validation"][int(idx)] for idx in batch_idx]
+            model_inputs = data_collator(samples, pad_to_multiple_of=16)
+
+            # Model forward
+            model_inputs = common_utils.shard(model_inputs.data)
+            metrics = p_eval_step(optimizer.target, model_inputs)
+            eval_metrics.append(metrics)
+
+        eval_metrics_np = get_metrics(eval_metrics)
+        eval_metrics_np = jax.tree_map(jnp.sum, eval_metrics_np)
+        eval_normalizer = eval_metrics_np.pop("normalizer")
+        eval_summary = jax.tree_map(lambda x: x / eval_normalizer, eval_metrics_np)
+
+        # Update progress bar
+        epochs.desc = (
+            f"Epoch... ({epoch + 1}/{nb_epochs} | Loss: {eval_summary['loss']}, Acc: {eval_summary['accuracy']})"
+        )
+
+        if wandb_args.wandb_user_name is not None:
+            wandb.log({"Eval loss": np.array(eval_summary["loss"]).mean()})
+
+        # Save metrics
+        if has_tensorboard and jax.host_id() == 0:
+            for name, value in eval_summary.items():
+                summary_writer.scalar(name, value, epoch)

examples/research_projects/performer/sanity_script.sh

@@ -0,0 +1 @@
+TOKENIZERS_PARALLELISM=true python run_mlm_performer.py  --output_dir experiments --dataset_name wikipedia --dataset_config_name 20200501.simple --model_name_or_path bert-base-cased --tokenizer_name bert-base-cased --do_train --overwrite_output_dir --per_device_train_batch_size 4 --learning_rate 5e-4 --warmup_steps 100 --num_train_epochs 3 --performer

examples/research_projects/rag/README.md

@@ -50,6 +50,44 @@ python examples/rag/consolidate_rag_checkpoint.py \
 ```
 You will then be able to pass `path/to/checkpoint` as `model_name_or_path` to the `finetune_rag.py` script.
 
+## Document Retrieval
+When running distributed fine-tuning, each training worker needs to retrieve contextual documents
+for its input by querying a index loaded into memory. RAG provides two implementations for document retrieval, 
+one with [`torch.distributed`](https://pytorch.org/docs/stable/distributed.html) communication package and the other 
+with [`Ray`](https://docs.ray.io/en/master/).
+
+This option can be configured with the `--distributed_retriever` flag which can either be set to `pytorch` or `ray`.
+By default this flag is set to `pytorch`.
+
+For the Pytorch implementation, only training worker 0 loads the index into CPU memory, and a gather/scatter pattern is used
+to collect the inputs from the other training workers and send back the corresponding document embeddings.
+
+For the Ray implementation, the index is loaded in *separate* process(es). The training workers randomly select which 
+retriever worker to query. To use Ray for distributed retrieval, you have to set the `--distributed_retriever` arg to `ray`.
+To configure the number of retrieval workers (the number of processes that load the index), you can set the `num_retrieval_workers` flag.
+Also make sure to start the Ray cluster before running fine-tuning.
+
+```bash
+# Start a single-node Ray cluster.
+ray start --head
+
+python examples/rag/finetune_rag.py \
+    --data_dir $DATA_DIR \
+    --output_dir $OUTPUT_DIR \
+    --model_name_or_path $MODEL_NAME_OR_PATH \
+    --model_type rag_sequence \
+    --fp16 \
+    --gpus 8
+    --distributed_retriever ray \
+    --num_retrieval_workers 4
+
+# Stop the ray cluster once fine-tuning has finished.
+ray stop
+```
+
+Using Ray can lead to retrieval speedups on multi-GPU settings since multiple processes load the index rather than
+just the rank 0 training worker. Using Ray also allows you to load the index on GPU since the index is loaded on a separate
+processes than the model, while with pytorch distributed retrieval, both are loaded in the same process potentially leading to GPU OOM.
 
 # Evaluation
 Our evaluation script enables two modes of evaluation (controlled by the `eval_mode` argument): `e2e` - end2end evaluation, returns EM (exact match) and F1 scores calculated for the downstream task and `retrieval` - which returns precision@k of the documents retrieved for provided inputs.

examples/research_projects/rag/_test_finetune_rag.py

@@ -9,6 +9,7 @@ from transformers.file_utils import is_apex_available
 from transformers.testing_utils import (
     TestCasePlus,
     execute_subprocess_async,
+    require_ray,
     require_torch_gpu,
     require_torch_multi_gpu,
 )
@@ -29,7 +30,7 @@ class RagFinetuneExampleTests(TestCasePlus):
                 with open(os.path.join(data_dir, f"{split}.{field}"), "w") as f:
                     f.write(content)
 
-    def _run_finetune(self, gpus: int):
+    def _run_finetune(self, gpus: int, distributed_retriever: str = "pytorch"):
         stream_handler = logging.StreamHandler(sys.stdout)
         logger.addHandler(stream_handler)
 
@@ -66,6 +67,7 @@ class RagFinetuneExampleTests(TestCasePlus):
                 --gradient_accumulation_steps 1 \
                 --distributed-port 8787 \
                 --use_dummy_dataset 1 \
+                --distributed_retriever {distributed_retriever} \
             """.split()
 
         if gpus > 0:
@@ -94,3 +96,15 @@ class RagFinetuneExampleTests(TestCasePlus):
     def test_finetune_multigpu(self):
         result = self._run_finetune(gpus=2)
         self.assertGreaterEqual(result["test"][0]["test_avg_em"], 0.2)
+
+    @require_torch_gpu
+    @require_ray
+    def test_finetune_gpu_ray_retrieval(self):
+        result = self._run_finetune(gpus=1, distributed_retriever="ray")
+        self.assertGreaterEqual(result["test"][0]["test_avg_em"], 0.2)
+
+    @require_torch_multi_gpu
+    @require_ray
+    def test_finetune_multigpu_ray_retrieval(self):
+        result = self._run_finetune(gpus=1, distributed_retriever="ray")
+        self.assertGreaterEqual(result["test"][0]["test_avg_em"], 0.2)

examples/research_projects/rag/distributed_pytorch_retriever.py

@@ -31,14 +31,13 @@ class RagPyTorchDistributedRetriever(RagRetriever):
             If specified, use this index instead of the one built using the configuration
     """
 
-    _init_retrieval = False
-
     def __init__(self, config, question_encoder_tokenizer, generator_tokenizer, index=None):
         super().__init__(
             config,
             question_encoder_tokenizer=question_encoder_tokenizer,
             generator_tokenizer=generator_tokenizer,
             index=index,
+            init_retrieval=False,
         )
         self.process_group = None
 

examples/research_projects/rag/distributed_ray_retriever.py

@@ -0,0 +1,154 @@
+import logging
+import random
+
+import ray
+from transformers import RagConfig, RagRetriever, RagTokenizer
+from transformers.file_utils import requires_datasets, requires_faiss
+from transformers.models.rag.retrieval_rag import CustomHFIndex
+
+
+logger = logging.getLogger(__name__)
+
+
+class RayRetriever:
+    def __init__(self):
+        self.initialized = False
+
+    def create_rag_retriever(self, config, question_encoder_tokenizer, generator_tokenizer, index):
+        if not self.initialized:
+            self.retriever = RagRetriever(
+                config,
+                question_encoder_tokenizer=question_encoder_tokenizer,
+                generator_tokenizer=generator_tokenizer,
+                index=index,
+                init_retrieval=False,
+            )
+            self.initialized = True
+
+    def init_retrieval(self):
+        self.retriever.index.init_index()
+
+    def retrieve(self, question_hidden_states, n_docs):
+        doc_ids, retrieved_doc_embeds = self.retriever._main_retrieve(question_hidden_states, n_docs)
+        return doc_ids, retrieved_doc_embeds
+
+
+class RagRayDistributedRetriever(RagRetriever):
+    """
+    A distributed retriever built on top of the ``Ray`` API, a library
+    for building distributed applications (https://docs.ray.io/en/master/).
+    package. During training, all training workers initialize their own
+    instance of a `RagRayDistributedRetriever`, and each instance of
+    this distributed retriever shares a common set of Retrieval Ray
+    Actors (https://docs.ray.io/en/master/walkthrough.html#remote
+    -classes-actors) that load the index on separate processes. Ray
+    handles the communication between the `RagRayDistributedRetriever`
+    instances and the remote Ray actors. If training is done in a
+    non-distributed setup, the index will simply be loaded in the same
+    process as the training worker and Ray will not be used.
+
+    Args:
+        config (:class:`~transformers.RagConfig`):
+            The configuration of the RAG model this Retriever is used with. Contains parameters indicating which ``Index`` to build.
+        question_encoder_tokenizer (:class:`~transformers.PretrainedTokenizer`):
+            The tokenizer that was used to tokenize the question.
+            It is used to decode the question and then use the generator_tokenizer.
+        generator_tokenizer (:class:`~transformers.PretrainedTokenizer`):
+            The tokenizer used for the generator part of the RagModel.
+        retrieval_workers (:obj:`List[ray.ActorClass(RayRetriever)]`): A list of already initialized `RayRetriever` actors.
+            These actor classes run on remote processes and are responsible for performing the index lookup.
+        index (:class:`~transformers.retrieval_rag.Index`, optional, defaults to the one defined by the configuration):
+            If specified, use this index instead of the one built using the configuration
+    """
+
+    def __init__(self, config, question_encoder_tokenizer, generator_tokenizer, retrieval_workers, index=None):
+        if index is not None and index.is_initialized() and len(retrieval_workers) > 0:
+            raise ValueError(
+                "When using Ray for distributed fine-tuning, "
+                "you'll need to provide the paths instead, "
+                "as the dataset and the index are loaded "
+                "separately. More info in examples/rag/use_own_knowledge_dataset.py "
+            )
+        super().__init__(
+            config,
+            question_encoder_tokenizer=question_encoder_tokenizer,
+            generator_tokenizer=generator_tokenizer,
+            index=index,
+            init_retrieval=False,
+        )
+        self.retrieval_workers = retrieval_workers
+        if len(self.retrieval_workers) > 0:
+            ray.get(
+                [
+                    worker.create_rag_retriever.remote(config, question_encoder_tokenizer, generator_tokenizer, index)
+                    for worker in self.retrieval_workers
+                ]
+            )
+
+    def init_retrieval(self):
+        """
+        Retriever initialization function, needs to be called from the
+        training process. This function triggers retrieval initialization
+        for all retrieval actors if using distributed setting, or loads
+        index into current process if training is not distributed.
+        """
+        logger.info("initializing retrieval")
+
+        if len(self.retrieval_workers) > 0:
+            ray.get([worker.init_retrieval.remote() for worker in self.retrieval_workers])
+        else:
+            # Non-distributed training. Load index into this same process.
+            self.index.init_index()
+
+    def retrieve(self, question_hidden_states, n_docs):
+        """
+        Retrieves documents for specified ``question_hidden_states``. If
+        running training with multiple workers, a random retrieval actor is
+        selected to perform the index lookup and return the result.
+
+        Args:
+            question_hidden_states (:obj:`np.ndarray` of shape :obj:`(batch_size, vector_size)`):
+                A batch of query vectors to retrieve with.
+            n_docs (:obj:`int`):
+                The number of docs retrieved per query.
+
+        Output:
+            retrieved_doc_embeds (:obj:`np.ndarray` of shape :obj:`(batch_size, n_docs, dim)`
+                The retrieval embeddings of the retrieved docs per query.
+            doc_ids (:obj:`np.ndarray` of shape :obj:`batch_size, n_docs`)
+                The ids of the documents in the index
+            doc_dicts (:obj:`List[dict]`):
+                The retrieved_doc_embeds examples per query.
+        """
+        if len(self.retrieval_workers) > 0:
+            # Select a random retrieval actor.
+            random_worker = self.retrieval_workers[random.randint(0, len(self.retrieval_workers) - 1)]
+            doc_ids, retrieved_doc_embeds = ray.get(random_worker.retrieve.remote(question_hidden_states, n_docs))
+        else:
+            doc_ids, retrieved_doc_embeds = self._main_retrieve(question_hidden_states, n_docs)
+        return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(doc_ids)
+
+    @classmethod
+    def get_tokenizers(cls, retriever_name_or_path, indexed_dataset=None, **kwargs):
+        return super(RagRayDistributedRetriever, cls).get_tokenizers(retriever_name_or_path, indexed_dataset, **kwargs)
+
+    @classmethod
+    def from_pretrained(cls, retriever_name_or_path, actor_handles, indexed_dataset=None, **kwargs):
+        requires_datasets(cls)
+        requires_faiss(cls)
+        config = kwargs.pop("config", None) or RagConfig.from_pretrained(retriever_name_or_path, **kwargs)
+        rag_tokenizer = RagTokenizer.from_pretrained(retriever_name_or_path, config=config)
+        question_encoder_tokenizer = rag_tokenizer.question_encoder
+        generator_tokenizer = rag_tokenizer.generator
+        if indexed_dataset is not None:
+            config.index_name = "custom"
+            index = CustomHFIndex(config.retrieval_vector_size, indexed_dataset)
+        else:
+            index = cls._build_index(config)
+        return cls(
+            config,
+            question_encoder_tokenizer=question_encoder_tokenizer,
+            generator_tokenizer=generator_tokenizer,
+            retrieval_workers=actor_handles,
+            index=index,
+        )

examples/research_projects/rag/eval_rag.py

@@ -96,7 +96,7 @@ def evaluate_batch_retrieval(args, rag_model, questions):
     )["input_ids"].to(args.device)
 
     question_enc_outputs = rag_model.rag.question_encoder(retriever_input_ids)
-    question_enc_pool_output = question_enc_outputs.pooler_output
+    question_enc_pool_output = question_enc_outputs[0]
 
     result = rag_model.retriever(
         retriever_input_ids,

examples/research_projects/rag/finetune_rag.py

@@ -29,6 +29,12 @@ from transformers import (
     T5ForConditionalGeneration,
 )
 from transformers import logging as transformers_logging
+from transformers.integrations import is_ray_available
+
+
+if is_ray_available():
+    import ray
+    from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever
 
 
 from callbacks_rag import (  # noqa: E402 # isort:skipq
@@ -36,7 +42,8 @@ from callbacks_rag import (  # noqa: E402 # isort:skipq
     get_early_stopping_callback,
     Seq2SeqLoggingCallback,
 )
-from distributed_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
+
+from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
 from utils_rag import (  # noqa: E402 # isort:skip
     calculate_exact_match,
     flatten_list,
@@ -88,7 +95,12 @@ class CustomAccel(DDPAccelerator):
         os.environ["MASTER_PORT"] = str(self.distributed_port)
         super().init_ddp_connection(global_rank, world_size, is_slurm_managing_tasks)
         if module.is_rag_model:
-            module.model.rag.retriever.init_retrieval(self.distributed_port)
+            if module.distributed_retriever == "pytorch":
+                module.model.rag.retriever.init_retrieval(self.distributed_port)
+            elif module.distributed_retriever == "ray" and global_rank == 0:
+                # For the Ray retriever, only initialize it once when global
+                # rank is 0.
+                module.model.rag.retriever.init_retrieval()
 
 
 class GenerativeQAModule(BaseTransformer):
@@ -127,7 +139,13 @@ class GenerativeQAModule(BaseTransformer):
                 config.generator.prefix = hparams.prefix
             config.label_smoothing = hparams.label_smoothing
             hparams, config.generator = set_extra_model_params(extra_model_params, hparams, config.generator)
-            retriever = RagPyTorchDistributedRetriever.from_pretrained(hparams.model_name_or_path, config=config)
+            if hparams.distributed_retriever == "pytorch":
+                retriever = RagPyTorchDistributedRetriever.from_pretrained(hparams.model_name_or_path, config=config)
+            elif hparams.distributed_retriever == "ray":
+                # The Ray retriever needs the handles to the retriever actors.
+                retriever = RagRayDistributedRetriever.from_pretrained(
+                    hparams.model_name_or_path, hparams.actor_handles, config=config
+                )
             model = self.model_class.from_pretrained(hparams.model_name_or_path, config=config, retriever=retriever)
             prefix = config.question_encoder.prefix
         else:
@@ -180,7 +198,12 @@ class GenerativeQAModule(BaseTransformer):
         # For single GPU training, init_ddp_connection is not called.
         # So we need to initialize the retrievers here.
         if hparams.gpus <= 1:
-            self.model.retriever.init_retrieval(self.distributed_port)
+            if hparams.distributed_retriever == "ray":
+                self.model.retriever.init_retrieval()
+            elif hparams.distributed_retriever == "pytorch":
+                self.model.retriever.init_retrieval(self.distributed_port)
+
+        self.distributed_retriever = hparams.distributed_retriever
 
     def forward(self, input_ids, **kwargs):
         return self.model(input_ids, **kwargs)
@@ -420,6 +443,7 @@ class GenerativeQAModule(BaseTransformer):
             type=str,
             help="RAG model type: sequence or token, if none specified, the type is inferred from the model_name_or_path",
         )
+
         return parser
 
     @staticmethod
@@ -442,12 +466,58 @@ class GenerativeQAModule(BaseTransformer):
             default=None,
             help="Path to the faiss index for custom index. More info about custom indexes in the RagRetriever documentation as well as in `examples/rag/use_own_knowledge_dataset.py`",
         )
+        parser.add_argument(
+            "--distributed_retriever",
+            choices=["ray", "pytorch"],
+            type=str,
+            default="pytorch",
+            help="What implementation to use for distributed retriever? If "
+            "pytorch is selected, the index is loaded on training "
+            "worker 0, and torch.distributed is used to handle "
+            "communication between training worker 0, and the other "
+            "training workers. If ray is selected, the Ray library is "
+            "used to create load the index on separate processes, "
+            "and Ray handles the communication between the training "
+            "workers and the retrieval actors.",
+        )
         parser.add_argument(
             "--use_dummy_dataset",
             type=bool,
             default=False,
             help="Whether to use the dummy version of the dataset index. More info about custom indexes in the RagRetriever documentation as well as in `examples/rag/use_own_knowledge_dataset.py`",
         )
+
+        parser.add_argument(
+            "--num_retrieval_workers",
+            type=int,
+            default=1,
+            help="The number of retrieval actors to use when Ray is selected"
+            "for the distributed retriever. Has no effect when "
+            "distributed_retriever is set to pytorch.",
+        )
+
+    @staticmethod
+    def add_ray_specific_args(parser):
+        parser.add_argument(
+            "--num_retrieval_workers",
+            type=int,
+            default=1,
+            help="The number of retrieval actors to use when Ray is selected"
+            "for the distributed retriever. Has no effect when "
+            "distributed_retriever is set to pytorch.",
+        )
+
+        # Ray cluster address.
+        parser.add_argument(
+            "--ray-address",
+            default="auto",
+            type=str,
+            help="The address of the Ray cluster to connect to. If not "
+            "specified, Ray will attempt to automatically detect the "
+            "cluster. Has no effect if pytorch is used as the distributed "
+            "retriever.",
+        )
+
         return parser
 
 
@@ -461,6 +531,46 @@ def main(args=None, model=None) -> GenerativeQAModule:
     args = args or parser.parse_args()
 
     Path(args.output_dir).mkdir(exist_ok=True)
+
+    named_actors = []
+    if args.distributed_retriever == "ray" and args.gpus > 1:
+        if not is_ray_available():
+            raise RuntimeError("Please install Ray to use the Ray " "distributed retriever.")
+        # Connect to an existing Ray cluster.
+        try:
+            ray.init(address=args.ray_address)
+        except (ConnectionError, ValueError):
+            logger.warning(
+                "Connection to Ray cluster failed. Make sure a Ray"
+                "cluster is running by either using Ray's cluster "
+                "launcher (`ray up`) or by manually starting Ray on "
+                "each node via `ray start --head` for the head node "
+                "and `ray start --address='<ip address>:6379'` for "
+                "additional nodes. See "
+                "https://docs.ray.io/en/master/cluster/index.html "
+                "for more info."
+            )
+            raise
+
+        # Create Ray actors only for rank 0.
+        if ("LOCAL_RANK" not in os.environ or os.environ["LOCAL_RANK"] == 0) and (
+            "NODE_RANK" not in os.environ or os.environ["NODE_RANK"] == 0
+        ):
+            remote_cls = ray.remote(RayRetriever)
+            named_actors = [
+                remote_cls.options(name="retrieval_worker_{}".format(i)).remote()
+                for i in range(args.num_retrieval_workers)
+            ]
+        else:
+            logger.info(
+                "Getting named actors for NODE_RANK {}, LOCAL_RANK {}".format(
+                    os.environ["NODE_RANK"], os.environ["LOCAL_RANK"]
+                )
+            )
+            named_actors = [ray.get_actor("retrieval_worker_{}".format(i)) for i in range(args.num_retrieval_workers)]
+    args.actor_handles = named_actors
+    assert args.actor_handles == named_actors
+
     if model is None:
         model: GenerativeQAModule = GenerativeQAModule(args)
 
@@ -471,17 +581,17 @@ def main(args=None, model=None) -> GenerativeQAModule:
         or str(args.output_dir).startswith("/tmp")
         or str(args.output_dir).startswith("/var")
     ):
-        logger = True  # don't pollute wandb logs unnecessarily
+        training_logger = True  # don't pollute wandb logs unnecessarily
     elif args.logger_name == "wandb":
         from pytorch_lightning.loggers import WandbLogger
 
         project = os.environ.get("WANDB_PROJECT", dataset)
-        logger = WandbLogger(name=model.output_dir.name, project=project)
+        training_logger = WandbLogger(name=model.output_dir.name, project=project)
 
     elif args.logger_name == "wandb_shared":
         from pytorch_lightning.loggers import WandbLogger
 
-        logger = WandbLogger(name=model.output_dir.name, project=f"hf_{dataset}")
+        training_logger = WandbLogger(name=model.output_dir.name, project=f"hf_{dataset}")
 
     es_callback = (
         get_early_stopping_callback(model.val_metric, args.early_stopping_patience)
@@ -495,8 +605,9 @@ def main(args=None, model=None) -> GenerativeQAModule:
         logging_callback=Seq2SeqLoggingCallback(),
         checkpoint_callback=get_checkpoint_callback(args.output_dir, model.val_metric),
         early_stopping_callback=es_callback,
-        logger=logger,
+        logger=training_logger,
         accelerator=CustomAccel() if args.gpus > 1 else None,
+        profiler=pl.profiler.AdvancedProfiler() if args.profile else None,
     )
     pickle_save(model.hparams, model.output_dir / "hparams.pkl")
 
@@ -509,4 +620,19 @@ def main(args=None, model=None) -> GenerativeQAModule:
 
 
 if __name__ == "__main__":
-    main()
+    parser = argparse.ArgumentParser()
+    parser = pl.Trainer.add_argparse_args(parser)
+    parser = GenerativeQAModule.add_model_specific_args(parser, os.getcwd())
+    parser = GenerativeQAModule.add_retriever_specific_args(parser)
+    parser = GenerativeQAModule.add_ray_specific_args(parser)
+
+    # Pytorch Lightning Profiler
+    parser.add_argument(
+        "--profile",
+        action="store_true",
+        help="If True, use pytorch_lightning.profiler.AdvancedProfiler to profile the Trainer.",
+    )
+
+    args = parser.parse_args()
+
+    main(args)

examples/research_projects/rag/finetune_rag.sh

@@ -2,7 +2,7 @@
 export PYTHONPATH="../":"${PYTHONPATH}"
 
 # A sample finetuning run, you need to specify data_dir, output_dir and model_name_or_path
-# run ./examples/rag/finetune.sh --help to see all the possible options
+# run ./examples/rag/finetune_rag.sh --help to see all the possible options
 
 python examples/rag/finetune_rag.py \
     --data_dir $DATA_DIR \
@@ -11,10 +11,10 @@ python examples/rag/finetune_rag.py \
     --model_type rag_sequence \
     --fp16 \
     --gpus 8 \
+    --profile \
     --do_train \
     --do_predict \
     --n_val -1 \
-    --val_check_interval 0.25 \
     --train_batch_size 8 \
     --eval_batch_size 1 \
     --max_source_length 128 \
@@ -31,4 +31,4 @@ python examples/rag/finetune_rag.py \
     --learning_rate 3e-05 \
     --num_train_epochs 100 \
     --warmup_steps 500 \
-    --gradient_accumulation_steps 1
+    --gradient_accumulation_steps 1 \

examples/research_projects/rag/finetune_rag_ray.sh

@@ -0,0 +1,44 @@
+# Sample script to finetune RAG using Ray for distributed retrieval.
+
+# Add parent directory to python path to access lightning_base.py
+export PYTHONPATH="../":"${PYTHONPATH}"
+
+# Start a single-node Ray cluster.
+ray start --head
+
+# A sample finetuning run, you need to specify data_dir, output_dir and model_name_or_path
+# run ./examples/rag/finetune_rag_ray.sh --help to see all the possible options
+
+python examples/rag/finetune_rag.py \
+    --data_dir $DATA_DIR \
+    --output_dir $OUTPUT_DIR \
+    --model_name_or_path $MODEL_NAME_OR_PATH \
+    --model_type rag_sequence \
+    --fp16 \
+    --gpus 8 \
+    --profile \
+    --do_train \
+    --do_predict \
+    --n_val -1 \
+    --train_batch_size 8 \
+    --eval_batch_size 1 \
+    --max_source_length 128 \
+    --max_target_length 25 \
+    --val_max_target_length 25 \
+    --test_max_target_length 25 \
+    --label_smoothing 0.1 \
+    --dropout 0.1 \
+    --attention_dropout 0.1 \
+    --weight_decay 0.001 \
+    --adam_epsilon 1e-08 \
+    --max_grad_norm 0.1 \
+    --lr_scheduler polynomial \
+    --learning_rate 3e-05 \
+    --num_train_epochs 100 \
+    --warmup_steps 500 \
+    --gradient_accumulation_steps 1 \
+    --distributed_retriever ray \
+    --num_retrieval_workers 4
+
+# Stop the Ray cluster.
+ray stop

examples/research_projects/rag/test_distributed_retriever.py

@@ -13,15 +13,27 @@ from datasets import Dataset
 import faiss
 from transformers import BartConfig, BartTokenizer, DPRConfig, DPRQuestionEncoderTokenizer, RagConfig
 from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
+from transformers.integrations import is_ray_available
 from transformers.models.bert.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
-from transformers.models.rag.retrieval_rag import CustomHFIndex
+from transformers.models.rag.retrieval_rag import CustomHFIndex, RagRetriever
 from transformers.models.roberta.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
-from transformers.testing_utils import require_torch_non_multi_gpu_but_fix_me
+from transformers.testing_utils import require_ray, require_torch_non_multi_gpu_but_fix_me
 
 
 sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip
 
-from distributed_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
+if is_torch_available():
+    from distributed_pytorch_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
+else:
+    RagPyTorchDistributedRetriever = None
+
+if is_ray_available():
+    import ray  # noqa: E402 # isort:skip
+    from distributed_ray_retriever import RagRayDistributedRetriever, RayRetriever  # noqa: E402 # isort:skip
+else:
+    ray = None
+    RagRayDistributedRetriever = None
+    RayRetriever = None
 
 
 def require_distributed_retrieval(test_case):
@@ -32,8 +44,8 @@ def require_distributed_retrieval(test_case):
     These tests are skipped when respective libraries are not installed.
 
     """
-    if not (is_torch_available() and is_datasets_available() and is_faiss_available() and is_psutil_available()):
-        test_case = unittest.skip("test requires PyTorch, Datasets, Faiss, psutil")(test_case)
+    if not (is_datasets_available() and is_faiss_available() and is_psutil_available()):
+        test_case = unittest.skip("test requires Datasets, Faiss, psutil")(test_case)
     return test_case
 
 
@@ -144,7 +156,31 @@ class RagRetrieverTest(TestCase):
                 retriever.init_retrieval(port)
         return retriever
 
-    def get_dummy_custom_hf_index_retriever(self, init_retrieval: bool, from_disk: bool, port=12345):
+    def get_dummy_ray_distributed_retriever(self, init_retrieval: bool) -> RagRayDistributedRetriever:
+        # Have to run in local mode because sys.path modifications at top of
+        # file are not propogated to remote workers.
+        # https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
+        ray.init(local_mode=True)
+        config = RagConfig(
+            retrieval_vector_size=self.retrieval_vector_size,
+            question_encoder=DPRConfig().to_dict(),
+            generator=BartConfig().to_dict(),
+        )
+        remote_cls = ray.remote(RayRetriever)
+        workers = [remote_cls.remote() for _ in range(1)]
+        with patch("transformers.models.rag.retrieval_rag.load_dataset") as mock_load_dataset:
+            mock_load_dataset.return_value = self.get_dummy_dataset()
+            retriever = RagRayDistributedRetriever(
+                config,
+                question_encoder_tokenizer=self.get_dpr_tokenizer(),
+                generator_tokenizer=self.get_bart_tokenizer(),
+                retrieval_workers=workers,
+            )
+            if init_retrieval:
+                retriever.init_retrieval()
+        return retriever
+
+    def get_dummy_custom_hf_index_pytorch_retriever(self, init_retrieval: bool, from_disk: bool, port=12345):
         dataset = self.get_dummy_dataset()
         config = RagConfig(
             retrieval_vector_size=self.retrieval_vector_size,
@@ -175,13 +211,51 @@ class RagRetrieverTest(TestCase):
             retriever.init_retrieval(port)
         return retriever
 
-    @require_torch_non_multi_gpu_but_fix_me
-    def test_pytorch_distributed_retriever_retrieve(self):
-        n_docs = 1
-        retriever = self.get_dummy_pytorch_distributed_retriever(init_retrieval=True)
-        hidden_states = np.array(
-            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+    def get_dummy_custom_hf_index_ray_retriever(self, init_retrieval: bool, from_disk: bool):
+        # Have to run in local mode because sys.path modifications at top of
+        # file are not propogated to remote workers.
+        # https://stackoverflow.com/questions/54338013/parallel-import-a-python-file-from-sibling-folder
+        ray.init(local_mode=True)
+        dataset = self.get_dummy_dataset()
+        config = RagConfig(
+            retrieval_vector_size=self.retrieval_vector_size,
+            question_encoder=DPRConfig().to_dict(),
+            generator=BartConfig().to_dict(),
+            index_name="custom",
         )
+        remote_cls = ray.remote(RayRetriever)
+        workers = [remote_cls.remote() for _ in range(1)]
+        if from_disk:
+            config.passages_path = os.path.join(self.tmpdirname, "dataset")
+            config.index_path = os.path.join(self.tmpdirname, "index.faiss")
+            dataset.get_index("embeddings").save(os.path.join(self.tmpdirname, "index.faiss"))
+            dataset.drop_index("embeddings")
+            dataset.save_to_disk(os.path.join(self.tmpdirname, "dataset"))
+            del dataset
+            retriever = RagRayDistributedRetriever(
+                config,
+                question_encoder_tokenizer=self.get_dpr_tokenizer(),
+                generator_tokenizer=self.get_bart_tokenizer(),
+                retrieval_workers=workers,
+                index=CustomHFIndex.load_from_disk(
+                    vector_size=config.retrieval_vector_size,
+                    dataset_path=config.passages_path,
+                    index_path=config.index_path,
+                ),
+            )
+        else:
+            retriever = RagRayDistributedRetriever(
+                config,
+                question_encoder_tokenizer=self.get_dpr_tokenizer(),
+                generator_tokenizer=self.get_bart_tokenizer(),
+                retrieval_workers=workers,
+                index=CustomHFIndex(config.retrieval_vector_size, dataset),
+            )
+        if init_retrieval:
+            retriever.init_retrieval()
+        return retriever
+
+    def distributed_retriever_check(self, retriever: RagRetriever, hidden_states: np.array, n_docs: int) -> None:
         retrieved_doc_embeds, doc_ids, doc_dicts = retriever.retrieve(hidden_states, n_docs=n_docs)
         self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
         self.assertEqual(len(doc_dicts), 2)
@@ -192,33 +266,76 @@ class RagRetrieverTest(TestCase):
         self.assertListEqual(doc_ids.tolist(), [[1], [0]])
 
     @require_torch_non_multi_gpu_but_fix_me
-    def test_custom_hf_index_retriever_retrieve(self):
+    def test_pytorch_distributed_retriever_retrieve(self):
         n_docs = 1
-        retriever = self.get_dummy_custom_hf_index_retriever(init_retrieval=True, from_disk=False)
         hidden_states = np.array(
             [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
         )
-        retrieved_doc_embeds, doc_ids, doc_dicts = retriever.retrieve(hidden_states, n_docs=n_docs)
-        self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
-        self.assertEqual(len(doc_dicts), 2)
-        self.assertEqual(sorted(doc_dicts[0]), ["embeddings", "id", "text", "title"])
-        self.assertEqual(len(doc_dicts[0]["id"]), n_docs)
-        self.assertEqual(doc_dicts[0]["id"][0], "1")  # max inner product is reached with second doc
-        self.assertEqual(doc_dicts[1]["id"][0], "0")  # max inner product is reached with first doc
-        self.assertListEqual(doc_ids.tolist(), [[1], [0]])
+
+        self.distributed_retriever_check(
+            self.get_dummy_pytorch_distributed_retriever(init_retrieval=True), hidden_states, n_docs
+        )
+
+    @require_torch_non_multi_gpu_but_fix_me
+    def test_custom_hf_index_pytorch_retriever_retrieve(self):
+        n_docs = 1
+        hidden_states = np.array(
+            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+        )
+
+        self.distributed_retriever_check(
+            self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=True, from_disk=False),
+            hidden_states,
+            n_docs,
+        )
 
     @require_torch_non_multi_gpu_but_fix_me
     def test_custom_pytorch_distributed_retriever_retrieve_from_disk(self):
         n_docs = 1
-        retriever = self.get_dummy_custom_hf_index_retriever(init_retrieval=True, from_disk=True)
         hidden_states = np.array(
             [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
         )
-        retrieved_doc_embeds, doc_ids, doc_dicts = retriever.retrieve(hidden_states, n_docs=n_docs)
-        self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
-        self.assertEqual(len(doc_dicts), 2)
-        self.assertEqual(sorted(doc_dicts[0]), ["embeddings", "id", "text", "title"])
-        self.assertEqual(len(doc_dicts[0]["id"]), n_docs)
-        self.assertEqual(doc_dicts[0]["id"][0], "1")  # max inner product is reached with second doc
-        self.assertEqual(doc_dicts[1]["id"][0], "0")  # max inner product is reached with first doc
-        self.assertListEqual(doc_ids.tolist(), [[1], [0]])
+
+        self.distributed_retriever_check(
+            self.get_dummy_custom_hf_index_pytorch_retriever(init_retrieval=True, from_disk=True),
+            hidden_states,
+            n_docs,
+        )
+
+    @require_ray
+    def test_ray_distributed_retriever_retrieve(self):
+        n_docs = 1
+        hidden_states = np.array(
+            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+        )
+
+        self.distributed_retriever_check(
+            self.get_dummy_ray_distributed_retriever(init_retrieval=True), hidden_states, n_docs
+        )
+        ray.shutdown()
+
+    @require_ray
+    def test_custom_hf_index_ray_retriever_retrieve(self):
+        n_docs = 1
+        hidden_states = np.array(
+            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+        )
+        with self.assertRaises(ValueError):
+            self.distributed_retriever_check(
+                self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=True, from_disk=False),
+                hidden_states,
+                n_docs,
+            )
+        ray.shutdown()
+
+    @require_ray
+    def test_custom_ray_distributed_retriever_retrieve_from_disk(self):
+        n_docs = 1
+        hidden_states = np.array(
+            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+        )
+
+        self.distributed_retriever_check(
+            self.get_dummy_custom_hf_index_ray_retriever(init_retrieval=True, from_disk=True), hidden_states, n_docs
+        )
+        ray.shutdown()

examples/seq2seq/README.md

@@ -18,7 +18,7 @@ limitations under the License.
 
 This directory contains examples for finetuning and evaluating transformers on summarization and translation tasks.
 Please tag @patil-suraj with any issues/unexpected behaviors, or send a PR!
-For deprecated `bertabs` instructions, see [`bertabs/README.md`](bertabs/README.md).
+For deprecated `bertabs` instructions, see [`bertabs/README.md`](https://github.com/huggingface/transformers/blob/master/examples/research_projects/bertabs/README.md).
 
 ### Supported Architectures
 
@@ -97,7 +97,7 @@ The `.source` files are the input, the `.target` files are the desired output.
 
 ### Potential issues
 
-- native AMP (`--fp16` and no apex) may lead to a huge memory leak and require 10x gpu memory. This has been fixed in pytorch-nightly and the minimal official version to have this fix will be pytorch-1.8. Until then if you have to use mixed precision please use AMP only with pytorch-nightly or NVIDIA's apex. Reference: https://github.com/huggingface/transformers/issues/8403
+- native AMP (`--fp16` and no apex) may lead to a huge memory leak and require 10x gpu memory. This has been fixed in pytorch-nightly and the minimal official version to have this fix will be pytorch-1.7.1. Until then if you have to use mixed precision please use AMP only with pytorch-nightly or NVIDIA's apex. Reference: https://github.com/huggingface/transformers/issues/8403
 
 
 ### Tips and Tricks

examples/seq2seq/finetune_trainer.py

@@ -20,9 +20,16 @@ from dataclasses import dataclass, field
 from typing import Optional
 
 import transformers
-from seq2seq_trainer import Seq2SeqTrainer
-from seq2seq_training_args import Seq2SeqTrainingArguments
-from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer, HfArgumentParser, MBartTokenizer, set_seed
+from transformers import (
+    AutoConfig,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    HfArgumentParser,
+    MBartTokenizer,
+    Seq2SeqTrainer,
+    Seq2SeqTrainingArguments,
+    set_seed,
+)
 from transformers.trainer_utils import EvaluationStrategy, is_main_process
 from transformers.training_args import ParallelMode
 from utils import (
@@ -97,7 +104,9 @@ class DataTrainingArguments:
         default=142,
         metadata={
             "help": "The maximum total sequence length for validation target text after tokenization. Sequences longer "
-            "than this will be truncated, sequences shorter will be padded."
+            "than this will be truncated, sequences shorter will be padded. "
+            "This argument is also used to override the ``max_length`` param of ``model.generate``, which is used "
+            "during ``evaluate`` and ``predict``."
         },
     )
     test_max_target_length: Optional[int] = field(
@@ -119,6 +128,22 @@ class DataTrainingArguments:
     )
 
 
+def handle_metrics(split, metrics, output_dir):
+    """
+    Log and save metrics
+
+    Args:
+    - split: one of train, val, test
+    - metrics: metrics dict
+    - output_dir: where to save the metrics
+    """
+
+    logger.info(f"***** {split} metrics *****")
+    for key in sorted(metrics.keys()):
+        logger.info(f"  {key} = {metrics[key]}")
+    save_json(metrics, os.path.join(output_dir, f"{split}_results.json"))
+
+
 def main():
     # See all possible arguments in src/transformers/training_args.py
     # or by passing the --help flag to this script.
@@ -256,54 +281,69 @@ def main():
     )
     trainer = Seq2SeqTrainer(
         model=model,
-        config=config,
         args=training_args,
         train_dataset=train_dataset,
         eval_dataset=eval_dataset,
         data_collator=Seq2SeqDataCollator(tokenizer, data_args, training_args.tpu_num_cores),
         compute_metrics=compute_metrics_fn,
-        data_args=data_args,
+        tokenizer=tokenizer,
     )
 
+    all_metrics = {}
     # Training
     if training_args.do_train:
-        trainer.train(
+        logger.info("*** Train ***")
+
+        train_result = trainer.train(
             model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
         )
-        trainer.save_model()
-        # For convenience, we also re-save the tokenizer to the same directory,
-        # so that you can share your model easily on huggingface.co/models =)
+        metrics = train_result.metrics
+        metrics["train_n_objs"] = data_args.n_train
+
+        trainer.save_model()  # this also saves the tokenizer
+
         if trainer.is_world_process_zero():
+            handle_metrics("train", metrics, training_args.output_dir)
+            all_metrics.update(metrics)
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
             trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
+            # For convenience, we also re-save the tokenizer to the same directory,
+            # so that you can share your model easily on huggingface.co/models =)
             tokenizer.save_pretrained(training_args.output_dir)
 
     # Evaluation
-    eval_results = {}
     if training_args.do_eval:
         logger.info("*** Evaluate ***")
 
-        result = trainer.evaluate()
+        metrics = trainer.evaluate(
+            metric_key_prefix="val", max_length=data_args.val_max_target_length, num_beams=data_args.eval_beams
+        )
+        metrics["val_n_objs"] = data_args.n_val
+        metrics["val_loss"] = round(metrics["val_loss"], 4)
 
         if trainer.is_world_process_zero():
-            logger.info("***** Eval results *****")
-            for key, value in result.items():
-                logger.info("  %s = %s", key, value)
-            save_json(result, os.path.join(training_args.output_dir, "eval_results.json"))
-            eval_results.update(result)
+
+            handle_metrics("val", metrics, training_args.output_dir)
+            all_metrics.update(metrics)
 
     if training_args.do_predict:
-        logging.info("*** Test ***")
+        logger.info("*** Predict ***")
 
-        test_output = trainer.predict(test_dataset=test_dataset)
-        test_metrics = {k.replace("eval", "test"): v for k, v in test_output.metrics.items()}
+        test_output = trainer.predict(
+            test_dataset=test_dataset,
+            metric_key_prefix="test",
+            max_length=data_args.val_max_target_length,
+            num_beams=data_args.eval_beams,
+        )
+        metrics = test_output.metrics
+        metrics["test_n_objs"] = data_args.n_test
 
         if trainer.is_world_process_zero():
-            logger.info("***** Test results *****")
-            for key, value in test_metrics.items():
-                logger.info("  %s = %s", key, value)
-
-            save_json(test_metrics, os.path.join(training_args.output_dir, "test_results.json"))
-            eval_results.update(test_metrics)
+            metrics["test_loss"] = round(metrics["test_loss"], 4)
+            handle_metrics("test", metrics, training_args.output_dir)
+            all_metrics.update(metrics)
 
             if training_args.predict_with_generate:
                 test_preds = tokenizer.batch_decode(
@@ -313,8 +353,9 @@ def main():
                 write_txt_file(test_preds, os.path.join(training_args.output_dir, "test_generations.txt"))
 
     if trainer.is_world_process_zero():
-        save_json(eval_results, "all_results.json")
-    return eval_results
+        save_json(all_metrics, os.path.join(training_args.output_dir, "all_results.json"))
+
+    return all_metrics
 
 
 def _mp_fn(index):

examples/seq2seq/seq2seq_trainer.py

@@ -20,6 +20,7 @@ from torch.utils.data import DistributedSampler, RandomSampler
 
 from transformers import PreTrainedModel, Trainer, logging
 from transformers.file_utils import is_torch_tpu_available
+from transformers.integrations import is_fairscale_available
 from transformers.models.fsmt.configuration_fsmt import FSMTConfig
 from transformers.optimization import (
     Adafactor,
@@ -35,6 +36,10 @@ from transformers.trainer_pt_utils import get_tpu_sampler
 from transformers.training_args import ParallelMode
 
 
+if is_fairscale_available():
+    from fairscale.optim import OSS
+
+
 logger = logging.get_logger(__name__)
 
 arg_to_scheduler = {
@@ -99,18 +104,25 @@ class Seq2SeqTrainer(Trainer):
                     "weight_decay": 0.0,
                 },
             ]
+            optimizer_cls = Adafactor if self.args.adafactor else AdamW
             if self.args.adafactor:
-                self.optimizer = Adafactor(
-                    optimizer_grouped_parameters,
-                    lr=self.args.learning_rate,
-                    scale_parameter=False,
-                    relative_step=False,
-                )
-
+                optimizer_cls = Adafactor
+                optimizer_kwargs = {"scale_parameter": False, "relative_step": False}
             else:
-                self.optimizer = AdamW(
-                    optimizer_grouped_parameters, lr=self.args.learning_rate, eps=self.args.adam_epsilon
+                optimizer_cls = AdamW
+                optimizer_kwargs = {
+                    "betas": (self.args.adam_beta1, self.args.adam_beta2),
+                    "eps": self.args.adam_epsilon,
+                }
+            optimizer_kwargs["lr"] = self.args.learning_rate
+            if self.sharded_dpp:
+                self.optimizer = OSS(
+                    params=optimizer_grouped_parameters,
+                    optim=optimizer_cls,
+                    **optimizer_kwargs,
                 )
+            else:
+                self.optimizer = optimizer_cls(optimizer_grouped_parameters, **optimizer_kwargs)
 
         if self.lr_scheduler is None:
             self.lr_scheduler = self._get_lr_scheduler(num_training_steps)

examples/seq2seq/test_finetune_trainer.py

@@ -14,10 +14,11 @@
 
 import os
 import sys
+import unittest
 from unittest.mock import patch
 
-from transformers import BertTokenizer, EncoderDecoderModel
-from transformers.file_utils import is_datasets_available
+from transformers.file_utils import is_apex_available
+from transformers.integrations import is_fairscale_available
 from transformers.testing_utils import (
     TestCasePlus,
     execute_subprocess_async,
@@ -29,8 +30,7 @@ from transformers.testing_utils import (
 from transformers.trainer_callback import TrainerState
 from transformers.trainer_utils import set_seed
 
-from .finetune_trainer import Seq2SeqTrainingArguments, main
-from .seq2seq_trainer import Seq2SeqTrainer
+from .finetune_trainer import main
 
 
 set_seed(42)
@@ -38,9 +38,31 @@ MARIAN_MODEL = "sshleifer/student_marian_en_ro_6_1"
 MBART_TINY = "sshleifer/tiny-mbart"
 
 
+# a candidate for testing_utils
+def require_fairscale(test_case):
+    """
+    Decorator marking a test that requires fairscale
+    """
+    if not is_fairscale_available():
+        return unittest.skip("test requires fairscale")(test_case)
+    else:
+        return test_case
+
+
+# a candidate for testing_utils
+def require_apex(test_case):
+    """
+    Decorator marking a test that requires apex
+    """
+    if not is_apex_available():
+        return unittest.skip("test requires apex")(test_case)
+    else:
+        return test_case
+
+
 class TestFinetuneTrainer(TestCasePlus):
-    def finetune_trainer_quick(self, distributed=None):
-        output_dir = self.run_trainer(1, "12", MBART_TINY, 1, distributed)
+    def finetune_trainer_quick(self, distributed=None, extra_args_str=None):
+        output_dir = self.run_trainer(1, "12", MBART_TINY, 1, distributed, extra_args_str)
         logs = TrainerState.load_from_json(os.path.join(output_dir, "trainer_state.json")).log_history
         eval_metrics = [log for log in logs if "eval_loss" in log.keys()]
         first_step_stats = eval_metrics[0]
@@ -59,6 +81,21 @@ class TestFinetuneTrainer(TestCasePlus):
     def test_finetune_trainer_ddp(self):
         self.finetune_trainer_quick(distributed=True)
 
+    # it's crucial to test --sharded_ddp w/ and w/o --fp16
+    @require_torch_multi_gpu
+    @require_fairscale
+    def test_finetune_trainer_ddp_sharded_ddp(self):
+        self.finetune_trainer_quick(distributed=True, extra_args_str="--sharded_ddp")
+
+    @require_torch_multi_gpu
+    @require_fairscale
+    def test_finetune_trainer_ddp_sharded_ddp_fp16(self):
+        self.finetune_trainer_quick(distributed=True, extra_args_str="--sharded_ddp --fp16")
+
+    @require_apex
+    def test_finetune_trainer_apex(self):
+        self.finetune_trainer_quick(extra_args_str="--fp16 --fp16_backend=apex")
+
     @slow
     def test_finetune_trainer_slow(self):
         # There is a missing call to __init__process_group somewhere
@@ -81,121 +118,14 @@ class TestFinetuneTrainer(TestCasePlus):
         assert "test_generations.txt" in contents
         assert "test_results.json" in contents
 
-    @slow
-    def test_finetune_bert2bert(self):
-        if not is_datasets_available():
-            return
-
-        import datasets
-
-        bert2bert = EncoderDecoderModel.from_encoder_decoder_pretrained("prajjwal1/bert-tiny", "prajjwal1/bert-tiny")
-        tokenizer = BertTokenizer.from_pretrained("bert-base-uncased")
-
-        bert2bert.config.vocab_size = bert2bert.config.encoder.vocab_size
-        bert2bert.config.eos_token_id = tokenizer.sep_token_id
-        bert2bert.config.decoder_start_token_id = tokenizer.cls_token_id
-        bert2bert.config.max_length = 128
-
-        train_dataset = datasets.load_dataset("cnn_dailymail", "3.0.0", split="train[:1%]")
-        val_dataset = datasets.load_dataset("cnn_dailymail", "3.0.0", split="validation[:1%]")
-
-        train_dataset = train_dataset.select(range(32))
-        val_dataset = val_dataset.select(range(16))
-
-        rouge = datasets.load_metric("rouge")
-
-        batch_size = 4
-
-        def _map_to_encoder_decoder_inputs(batch):
-            # Tokenizer will automatically set [BOS] <text> [EOS]
-            inputs = tokenizer(batch["article"], padding="max_length", truncation=True, max_length=512)
-            outputs = tokenizer(batch["highlights"], padding="max_length", truncation=True, max_length=128)
-            batch["input_ids"] = inputs.input_ids
-            batch["attention_mask"] = inputs.attention_mask
-
-            batch["decoder_input_ids"] = outputs.input_ids
-            batch["labels"] = outputs.input_ids.copy()
-            batch["labels"] = [
-                [-100 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch["labels"]
-            ]
-            batch["decoder_attention_mask"] = outputs.attention_mask
-
-            assert all([len(x) == 512 for x in inputs.input_ids])
-            assert all([len(x) == 128 for x in outputs.input_ids])
-
-            return batch
-
-        def _compute_metrics(pred):
-            labels_ids = pred.label_ids
-            pred_ids = pred.predictions
-
-            # all unnecessary tokens are removed
-            pred_str = tokenizer.batch_decode(pred_ids, skip_special_tokens=True)
-            label_str = tokenizer.batch_decode(labels_ids, skip_special_tokens=True)
-
-            rouge_output = rouge.compute(predictions=pred_str, references=label_str, rouge_types=["rouge2"])[
-                "rouge2"
-            ].mid
-
-            return {
-                "rouge2_precision": round(rouge_output.precision, 4),
-                "rouge2_recall": round(rouge_output.recall, 4),
-                "rouge2_fmeasure": round(rouge_output.fmeasure, 4),
-            }
-
-        # map train dataset
-        train_dataset = train_dataset.map(
-            _map_to_encoder_decoder_inputs,
-            batched=True,
-            batch_size=batch_size,
-            remove_columns=["article", "highlights"],
-        )
-        train_dataset.set_format(
-            type="torch",
-            columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"],
-        )
-
-        # same for validation dataset
-        val_dataset = val_dataset.map(
-            _map_to_encoder_decoder_inputs,
-            batched=True,
-            batch_size=batch_size,
-            remove_columns=["article", "highlights"],
-        )
-        val_dataset.set_format(
-            type="torch",
-            columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"],
-        )
-
-        output_dir = self.get_auto_remove_tmp_dir()
-
-        training_args = Seq2SeqTrainingArguments(
-            output_dir=output_dir,
-            per_device_train_batch_size=batch_size,
-            per_device_eval_batch_size=batch_size,
-            predict_with_generate=True,
-            evaluation_strategy="steps",
-            do_train=True,
-            do_eval=True,
-            warmup_steps=0,
-            eval_steps=2,
-            logging_steps=2,
-        )
-
-        # instantiate trainer
-        trainer = Seq2SeqTrainer(
-            model=bert2bert,
-            args=training_args,
-            compute_metrics=_compute_metrics,
-            train_dataset=train_dataset,
-            eval_dataset=val_dataset,
-        )
-
-        # start training
-        trainer.train()
-
     def run_trainer(
-        self, eval_steps: int, max_len: str, model_name: str, num_train_epochs: int, distributed: bool = False
+        self,
+        eval_steps: int,
+        max_len: str,
+        model_name: str,
+        num_train_epochs: int,
+        distributed: bool = False,
+        extra_args_str: str = None,
     ):
         data_dir = self.examples_dir / "seq2seq/test_data/wmt_en_ro"
         output_dir = self.get_auto_remove_tmp_dir()
@@ -231,6 +161,9 @@ class TestFinetuneTrainer(TestCasePlus):
         """.split()
         # --eval_beams  2
 
+        if extra_args_str is not None:
+            args.extend(extra_args_str.split())
+
         if distributed:
             n_gpu = get_gpu_count()
             distributed_args = f"""

examples/seq2seq/train_distil_marian_enro.sh

@@ -29,7 +29,8 @@ python finetune_trainer.py \
     --freeze_encoder --freeze_embeds \
     --num_train_epochs=6 \
     --save_steps 3000 --eval_steps 3000 \
-    --max_source_length $MAX_LEN --max_target_length $MAX_LEN --val_max_target_length $MAX_LEN --test_max_target_length $MAX_LEN \
+    --max_source_length $MAX_LEN --max_target_length $MAX_LEN \
+    --val_max_target_length $MAX_TGT_LEN --test_max_target_length $MAX_TGT_LEN \
     --do_train --do_eval --do_predict \
     --evaluation_strategy steps \
     --predict_with_generate --logging_first_step \

examples/seq2seq/train_distil_marian_enro_tpu.sh

@@ -30,7 +30,8 @@ python xla_spawn.py --num_cores $TPU_NUM_CORES \
     --num_train_epochs=6 \
     --save_steps 500 --eval_steps 500 \
     --logging_first_step --logging_steps 200 \
-    --max_source_length $MAX_LEN --max_target_length $MAX_LEN --val_max_target_length $MAX_LEN --test_max_target_length $MAX_LEN \
+    --max_source_length $MAX_LEN --max_target_length $MAX_LEN \
+    --val_max_target_length $MAX_TGT_LEN --test_max_target_length $MAX_TGT_LEN \
     --do_train --do_eval \
     --evaluation_strategy steps \
     --prediction_loss_only \

examples/seq2seq/train_distilbart_cnn.sh

@@ -32,7 +32,7 @@ python finetune_trainer.py \
     --num_train_epochs=2 \
     --save_steps 3000 --eval_steps 3000 \
     --logging_first_step \
-    --max_target_length 56 --val_max_target_length $MAX_TGT_LEN --test_max_target_length $MAX_TGT_LEN \
+    --max_target_length 56 --val_max_target_length $MAX_TGT_LEN --test_max_target_length $MAX_TGT_LEN\
     --do_train --do_eval --do_predict \
     --evaluation_strategy steps \
     --predict_with_generate --sortish_sampler \

examples/seq2seq/train_mbart_cc25_enro.sh

@@ -24,8 +24,7 @@ python finetune_trainer.py \
     --src_lang en_XX --tgt_lang ro_RO \
     --freeze_embeds \
     --per_device_train_batch_size=4 --per_device_eval_batch_size=4 \
-    --max_source_length 128 --max_target_length 128 \
-    --val_max_target_length 128 --test_max_target_length 128 \
+    --max_source_length 128 --max_target_length 128 --val_max_target_length 128 --test_max_target_length 128\
     --sortish_sampler \
     --num_train_epochs 6 \
     --save_steps 25000 --eval_steps 25000 --logging_steps 1000 \

examples/seq2seq/utils.py

@@ -434,7 +434,8 @@ def use_task_specific_params(model, task):
 
     if task_specific_params is not None:
         pars = task_specific_params.get(task, {})
-        logger.info(f"using task specific params for {task}: {pars}")
+        logger.info(f"setting model.config to task specific params for {task}:\n {pars}")
+        logger.info("note: command line args may override some of these")
         model.config.update(pars)
 
 
@@ -462,7 +463,7 @@ def save_git_info(folder_path: str) -> None:
 
 def save_json(content, path, indent=4, **json_dump_kwargs):
     with open(path, "w") as f:
-        json.dump(content, f, indent=indent, **json_dump_kwargs)
+        json.dump(content, f, indent=indent, sort_keys=True, **json_dump_kwargs)
 
 
 def load_json(path):

examples/test_examples.py

@@ -33,6 +33,7 @@ SRC_DIRS = [
         "text-classification",
         "token-classification",
         "language-modeling",
+        "multiple-choice",
         "question-answering",
     ]
 ]
@@ -46,6 +47,7 @@ if SRC_DIRS is not None:
     import run_mlm
     import run_ner
     import run_qa as run_squad
+    import run_swag
 
 
 logging.basicConfig(level=logging.DEBUG)
@@ -95,9 +97,7 @@ class ExamplesTests(TestCasePlus):
 
         with patch.object(sys, "argv", testargs):
             result = run_glue.main()
-            del result["eval_loss"]
-            for value in result.values():
-                self.assertGreaterEqual(value, 0.75)
+            self.assertGreaterEqual(result["eval_accuracy"], 0.75)
 
     @require_torch_non_multi_gpu_but_fix_me
     def test_run_clm(self):
@@ -216,6 +216,32 @@ class ExamplesTests(TestCasePlus):
             self.assertGreaterEqual(result["f1"], 30)
             self.assertGreaterEqual(result["exact"], 30)
 
+    @require_torch_non_multi_gpu_but_fix_me
+    def test_run_swag(self):
+        stream_handler = logging.StreamHandler(sys.stdout)
+        logger.addHandler(stream_handler)
+
+        tmp_dir = self.get_auto_remove_tmp_dir()
+        testargs = f"""
+            run_swag.py
+            --model_name_or_path bert-base-uncased
+            --train_file tests/fixtures/tests_samples/swag/sample.json
+            --validation_file tests/fixtures/tests_samples/swag/sample.json
+            --output_dir {tmp_dir}
+            --overwrite_output_dir
+            --max_steps=20
+            --warmup_steps=2
+            --do_train
+            --do_eval
+            --learning_rate=2e-4
+            --per_device_train_batch_size=2
+            --per_device_eval_batch_size=1
+        """.split()
+
+        with patch.object(sys, "argv", testargs):
+            result = run_swag.main()
+            self.assertGreaterEqual(result["eval_accuracy"], 0.8)
+
     @require_torch_non_multi_gpu_but_fix_me
     def test_generation(self):
         stream_handler = logging.StreamHandler(sys.stdout)

examples/text-classification/README.md

@@ -14,7 +14,76 @@ See the License for the specific language governing permissions and
 limitations under the License.
 -->
 
-## GLUE Benchmark
+# Text classification examples
+
+## PyTorch version
+
+Based on the script [`run_glue.py`](https://github.com/huggingface/transformers/blob/master/examples/text-classification/run_glue.py).
+
+Fine-tuning the library models for sequence classification on the GLUE benchmark: [General Language Understanding
+Evaluation](https://gluebenchmark.com/). This script can fine-tune any of the models on the [hub](https://huggingface.co/models)
+and can also be used for your own data in a csv or a JSON file (the script might need some tweaks in that case, refer
+to the comments inside for help).
+
+GLUE is made up of a total of 9 different tasks. Here is how to run the script on one of them:
+
+```bash
+export TASK_NAME=mrpc
+
+python run_glue.py \
+  --model_name_or_path bert-base-cased \
+  --task_name $TASK_NAME \
+  --do_train \
+  --do_eval \
+  --max_seq_length 128 \
+  --per_device_train_batch_size 32 \
+  --learning_rate 2e-5 \
+  --num_train_epochs 3 \
+  --output_dir /tmp/$TASK_NAME/
+```
+
+where task name can be one of cola, sst2, mrpc, stsb, qqp, mnli, qnli, rte, wnli.
+
+We get the following results on the dev set of the benchmark with the previous commands (with an exception for MRPC and
+WNLI which are tiny and where we used 5 epochs isntead of 3). Trainings are seeded so you should obtain the same
+results with PyTorch 1.6.0 (and close results with different versions), training times are given for information (a
+single Titan RTX was used):
+
+| Task  | Metric                       | Result      | Training time |
+|-------|------------------------------|-------------|---------------|
+| CoLA  | Matthew's corr               | 56.53       | 3:17          |
+| SST-2 | Accuracy                     | 92.32       | 26:06         |
+| MRPC  | F1/Accuracy                  | 88.85/84.07 | 2:21          |
+| STS-B | Person/Spearman corr.        | 88.64/88.48 | 2:13          |
+| QQP   | Accuracy/F1                  | 90.71/87.49 | 2:22:26       |
+| MNLI  | Matched acc./Mismatched acc. | 83.91/84.10 | 2:35:23       |
+| QNLI  | Accuracy                     | 90.66       | 40:57         |
+| RTE   | Accuracy                     | 65.70       | 57            |
+| WNLI  | Accuracy                     | 56.34       | 24            |
+
+Some of these results are significantly different from the ones reported on the test set of GLUE benchmark on the
+website. For QQP and WNLI, please refer to [FAQ #12](https://gluebenchmark.com/faq) on the website.
+
+### Mixed precision training
+
+If you have a GPU with mixed precision capabilities (architecture Pascal or more recent), you can use mixed precision
+training with PyTorch 1.6.0 or latest, or by installing the [Apex](https://github.com/NVIDIA/apex) library for previous
+versions. Just add the flag `--fp16` to your command launching one of the scripts mentioned above!
+
+Using mixed precision training usually results in 2x-speedup for training with the same final results:
+
+| Task  | Metric                       | Result      | Training time | Result (FP16) | Training time (FP16) |
+|-------|------------------------------|-------------|---------------|---------------|----------------------|
+| CoLA  | Matthew's corr               | 56.53       | 3:17          | 56.78         | 1:41                 |
+| SST-2 | Accuracy                     | 92.32       | 26:06         | 91.74         | 13:11                |
+| MRPC  | F1/Accuracy                  | 88.85/84.07 | 2:21          | 88.12/83.58   | 1:10                 |
+| STS-B | Person/Spearman corr.        | 88.64/88.48 | 2:13          | 88.71/88.55   | 1:08                 |
+| QQP   | Accuracy/F1                  | 90.71/87.49 | 2:22:26       | 90.67/87.43   | 1:11:54              |
+| MNLI  | Matched acc./Mismatched acc. | 83.91/84.10 | 2:35:23       | 84.04/84.06   | 1:17:06              |
+| QNLI  | Accuracy                     | 90.66       | 40:57         | 90.96         | 20:16                |
+| RTE   | Accuracy                     | 65.70       | 57            | 65.34         | 29                   |
+| WNLI  | Accuracy                     | 56.34       | 24            | 56.34         | 12                   |
+
 
 # Run TensorFlow 2.0 version
 
@@ -65,191 +134,8 @@ python run_tf_text_classification.py \
   --max_seq_length 128
 ```
 
-# Run PyTorch version
 
-Based on the script [`run_glue.py`](https://github.com/huggingface/transformers/blob/master/examples/text-classification/run_glue.py).
-
-Fine-tuning the library models for sequence classification on the GLUE benchmark: [General Language Understanding
-Evaluation](https://gluebenchmark.com/). This script can fine-tune the following models: BERT, XLM, XLNet and RoBERTa.
-
-GLUE is made up of a total of 9 different tasks. We get the following results on the dev set of the benchmark with an
-uncased  BERT base model (the checkpoint `bert-base-uncased`). All experiments ran single V100 GPUs with a total train
-batch sizes between 16 and 64. Some of these tasks have a small dataset and training can lead to high variance in the results
-between different runs. We report the median on 5 runs (with different seeds) for each of the metrics.
-
-| Task  | Metric                       | Result      |
-|-------|------------------------------|-------------|
-| CoLA  | Matthew's corr               | 49.23       |
-| SST-2 | Accuracy                     | 91.97       |
-| MRPC  | F1/Accuracy                  | 89.47/85.29 |
-| STS-B | Person/Spearman corr.        | 83.95/83.70 |
-| QQP   | Accuracy/F1                  | 88.40/84.31 |
-| MNLI  | Matched acc./Mismatched acc. | 80.61/81.08 |
-| QNLI  | Accuracy                     | 87.46       |
-| RTE   | Accuracy                     | 61.73       |
-| WNLI  | Accuracy                     | 45.07       |
-
-Some of these results are significantly different from the ones reported on the test set
-of GLUE benchmark on the website. For QQP and WNLI, please refer to [FAQ #12](https://gluebenchmark.com/faq) on the
-website.
-
-```bash
-export TASK_NAME=MRPC
-
-python run_glue.py \
-  --model_name_or_path bert-base-cased \
-  --task_name $TASK_NAME \
-  --do_train \
-  --do_eval \
-  --max_seq_length 128 \
-  --per_device_train_batch_size 32 \
-  --learning_rate 2e-5 \
-  --num_train_epochs 3.0 \
-  --output_dir /tmp/$TASK_NAME/
-```
-
-where task name can be one of CoLA, SST-2, MRPC, STS-B, QQP, MNLI, QNLI, RTE, WNLI.
-
-The dev set results will be present within the text file `eval_results.txt` in the specified output_dir.
-In case of MNLI, since there are two separate dev sets (matched and mismatched), there will be a separate
-output folder called `/tmp/MNLI-MM/` in addition to `/tmp/MNLI/`.
-
-The code has not been tested with half-precision training with apex on any GLUE task apart from MRPC, MNLI,
-CoLA, SST-2. The following section provides details on how to run half-precision training with MRPC. With that being
-said, there shouldn’t be any issues in running half-precision training with the remaining GLUE tasks as well,
-since the data processor for each task inherits from the base class DataProcessor.
-
-## Running on TPUs in PyTorch
-
-Even when running PyTorch, you can accelerate your workloads on Google's TPUs, using `pytorch/xla`. For information on
-how to setup your TPU environment refer to the
-[pytorch/xla README](https://github.com/pytorch/xla/blob/master/README.md).
-
-For running your GLUE task on MNLI dataset you can run something like the following form the root of the transformers
-repo:
-
-```
-python examples/xla_spawn.py \
-  --num_cores=8 \
-  transformers/examples/text-classification/run_glue.py \
-  --do_train \
-  --do_eval \
-  --task_name=mrpc \
-  --num_train_epochs=3 \
-  --max_seq_length=128 \
-  --learning_rate=5e-5 \
-  --output_dir=/tmp/mrpc \
-  --overwrite_output_dir \
-  --logging_steps=5 \
-  --save_steps=5 \
-  --tpu_metrics_debug \
-  --model_name_or_path=bert-base-cased \
-  --per_device_train_batch_size=64 \
-  --per_device_eval_batch_size=64
-```
-
-
-#### Using Apex and mixed-precision
-
-Using Apex and 16 bit precision, the fine-tuning on MRPC only takes 27 seconds. First install
-[apex](https://github.com/NVIDIA/apex), then run the following example:
-
-```bash
-
-python run_glue.py \
-  --model_name_or_path bert-base-cased \
-  --task_name MRPC \
-  --do_train \
-  --do_eval \
-  --max_seq_length 128 \
-  --per_device_train_batch_size 32 \
-  --learning_rate 2e-5 \
-  --num_train_epochs 3.0 \
-  --output_dir /tmp/mrpc_output/ \
-  --fp16
-```
-
-#### Distributed training
-
-Here is an example using distributed training on 8 V100 GPUs. The model used is the BERT whole-word-masking and it
-reaches F1 > 92 on MRPC.
-
-```bash
-
-python -m torch.distributed.launch \
-    --nproc_per_node 8 run_glue.py \
-    --model_name_or_path bert-base-cased \
-    --task_name mrpc \
-    --do_train \
-    --do_eval \
-    --max_seq_length 128 \
-    --per_device_train_batch_size 8 \
-    --learning_rate 2e-5 \
-    --num_train_epochs 3.0 \
-    --output_dir /tmp/mrpc_output/
-```
-
-Training with these hyper-parameters gave us the following results:
-
-```bash
-acc = 0.8823529411764706
-acc_and_f1 = 0.901702786377709
-eval_loss = 0.3418912578906332
-f1 = 0.9210526315789473
-global_step = 174
-loss = 0.07231863956341798
-```
-
-### MNLI
-
-The following example uses the BERT-large, uncased, whole-word-masking model and fine-tunes it on the MNLI task.
-
-```bash
-export GLUE_DIR=/path/to/glue
-
-python -m torch.distributed.launch \
-    --nproc_per_node 8 run_glue.py \
-    --model_name_or_path bert-base-cased \
-    --task_name mnli \
-    --do_train \
-    --do_eval \
-    --max_seq_length 128 \
-    --per_device_train_batch_size 8 \
-    --learning_rate 2e-5 \
-    --num_train_epochs 3.0 \
-    --output_dir output_dir \
-```
-
-The results  are the following:
-
-```bash
-***** Eval results *****
-  acc = 0.8679706601466992
-  eval_loss = 0.4911287787382479
-  global_step = 18408
-  loss = 0.04755385363816904
-
-***** Eval results *****
-  acc = 0.8747965825874695
-  eval_loss = 0.45516540421714036
-  global_step = 18408
-  loss = 0.04755385363816904
-```
-
-# Run PyTorch version using PyTorch-Lightning
-
-Run `bash run_pl.sh` from the `glue` directory. This will also install `pytorch-lightning` and the requirements in
-`examples/requirements.txt`. It is a shell pipeline that will automatically download, preprocess the data and run the
-specified models. Logs are saved in `lightning_logs` directory.
-
-Pass `--gpus` 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}
-```
-
-
-# XNLI
+## XNLI
 
 Based on the script [`run_xnli.py`](https://github.com/huggingface/transformers/blob/master/examples/text-classification/run_xnli.py).
 

examples/text-classification/run_glue.py

@@ -350,11 +350,24 @@ def main():
 
     # Training
     if training_args.do_train:
-        trainer.train(
+        train_result = trainer.train(
             model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
         )
+        metrics = train_result.metrics
+
         trainer.save_model()  # Saves the tokenizer too for easy upload
 
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
     # Evaluation
     eval_results = {}
     if training_args.do_eval:
@@ -374,7 +387,7 @@ def main():
             if trainer.is_world_process_zero():
                 with open(output_eval_file, "w") as writer:
                     logger.info(f"***** Eval results {task} *****")
-                    for key, value in eval_result.items():
+                    for key, value in sorted(eval_result.items()):
                         logger.info(f"  {key} = {value}")
                         writer.write(f"{key} = {value}\n")
 

examples/token-classification/README.md

@@ -55,7 +55,7 @@ of the script.
 
 ## Old version of the script
 
-You can find the old version of the PyTorch script [here](https://github.com/huggingface/transformers/blob/master/examples/contrib/legacy/token-classification/run_ner_old.py).
+You can find the old version of the PyTorch script [here](https://github.com/huggingface/transformers/blob/master/examples/legacy/token-classification/run_ner.py).
 
 ### TensorFlow version
 

examples/token-classification/run_ner.py

@@ -340,11 +340,22 @@ def main():
 
     # Training
     if training_args.do_train:
-        trainer.train(
+        train_result = trainer.train(
             model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
         )
         trainer.save_model()  # Saves the tokenizer too for easy upload
 
+        output_train_file = os.path.join(training_args.output_dir, "train_results.txt")
+        if trainer.is_world_process_zero():
+            with open(output_train_file, "w") as writer:
+                logger.info("***** Train results *****")
+                for key, value in sorted(train_result.metrics.items()):
+                    logger.info(f"  {key} = {value}")
+                    writer.write(f"{key} = {value}\n")
+
+            # Need to save the state, since Trainer.save_model saves only the tokenizer with the model
+            trainer.state.save_to_json(os.path.join(training_args.output_dir, "trainer_state.json"))
+
     # Evaluation
     results = {}
     if training_args.do_eval:
@@ -377,7 +388,7 @@ def main():
         output_test_results_file = os.path.join(training_args.output_dir, "test_results.txt")
         if trainer.is_world_process_zero():
             with open(output_test_results_file, "w") as writer:
-                for key, value in metrics.items():
+                for key, value in sorted(metrics.items()):
                     logger.info(f"  {key} = {value}")
                     writer.write(f"{key} = {value}\n")
 

model_cards/google/tapas-base/README.md

@@ -0,0 +1,123 @@
+---
+language: en
+tags:
+- tapas
+- masked-lm
+license: apache-2.0
+---
+
+# TAPAS base model 
+
+This model corresponds to the `tapas_inter_masklm_base_reset` checkpoint of the [original Github repository](https://github.com/google-research/tapas). 
+
+Disclaimer: The team releasing TAPAS did not write a model card for this model so this model card has been written by
+the Hugging Face team and contributors.
+
+## Model description
+
+TAPAS is a BERT-like transformers model pretrained on a large corpus of English data from Wikipedia in a self-supervised fashion. 
+This means it was pretrained on the raw tables and associated texts only, with no humans labelling them in any way (which is why it
+can use lots of publicly available data) with an automatic process to generate inputs and labels from those texts. More precisely, it
+was pretrained with two objectives:
+
+- Masked language modeling (MLM): taking a (flattened) table and associated context, the model randomly masks 15% of the words in 
+  the input, then runs the entire (partially masked) sequence through the model. The model then has to predict the masked words. 
+  This is different from traditional recurrent neural networks (RNNs) that usually see the words one after the other, 
+  or from autoregressive models like GPT which internally mask the future tokens. It allows the model to learn a bidirectional 
+  representation of a table and associated text.
+- Intermediate pre-training: to encourage numerical reasoning on tables, the authors additionally pre-trained the model by creating 
+  a balanced dataset of millions of syntactically created training examples. Here, the model must predict (classify) whether a sentence 
+  is supported or refuted by the contents of a table. The training examples are created based on synthetic as well as counterfactual statements.
+
+This way, the model learns an inner representation of the English language used in tables and associated texts, which can then be used 
+to extract features useful for downstream tasks such as answering questions about a table, or determining whether a sentence is entailed
+or refuted by the contents of a table. Fine-tuning is done by adding classification heads on top of the pre-trained model, and then jointly
+train the randomly initialized classification heads with the base model on a labelled dataset. 
+
+## Intended uses & limitations
+
+You can use the raw model for masked language modeling, but it's mostly intended to be fine-tuned on a downstream task. 
+See the [model hub](https://huggingface.co/models?filter=tapas) to look for fine-tuned versions on a task that interests you.
+
+
+Here is how to use this model to get the features of a given table-text pair in PyTorch:
+
+```python
+from transformers import TapasTokenizer, TapasModel
+import pandas as pd
+tokenizer = TapasTokenizer.from_pretrained('tapase-base')
+model = TapasModel.from_pretrained("tapas-base")
+data = {'Actors': ["Brad Pitt", "Leonardo Di Caprio", "George Clooney"],
+         'Age': ["56", "45", "59"],
+         'Number of movies': ["87", "53", "69"]
+}
+table = pd.DataFrame.from_dict(data)
+queries = ["How many movies has George Clooney played in?"]
+text = "Replace me by any text you'd like."
+encoded_input = tokenizer(table=table, queries=queries, return_tensors='pt')
+output = model(**encoded_input)
+```
+
+## Training data
+
+For masked language modeling (MLM), a collection of 6.2 million tables was extracted from English Wikipedia: 3.3M of class [Infobox](https://en.wikipedia.org/wiki/Help:Infobox)
+and 2.9M of class WikiTable. The author only considered tables with at most 500 cells. As a proxy for questions that appear in the 
+downstream tasks, the authros extracted the table caption, article title, article description, segment title and text of the segment 
+the table occurs in as relevant text snippets. In this way, 21.3M snippets were created. For more info, see the original [TAPAS paper](https://www.aclweb.org/anthology/2020.acl-main.398.pdf).
+
+For intermediate pre-training, 2 tasks are introduced: one based on synthetic and the other from counterfactual statements. The first one 
+generates a sentence by sampling from a set of logical expressions that filter, combine and compare the information on the table, which is 
+required in table entailment (e.g., knowing that Gerald Ford is taller than the average president requires summing
+all presidents and dividing by the number of presidents). The second one corrupts sentences about tables appearing on Wikipedia by swapping 
+entities for plausible alternatives. Examples of the two tasks can be seen in Figure 1. The procedure is described in detail in section 3 of 
+the [TAPAS follow-up paper](https://www.aclweb.org/anthology/2020.findings-emnlp.27.pdf).
+
+## Training procedure
+
+### Preprocessing
+
+The texts are lowercased and tokenized using WordPiece and a vocabulary size of 30,000. The inputs of the model are
+then of the form:
+
+```
+[CLS] Context [SEP] Flattened table [SEP]
+```
+
+The details of the masking procedure for each sequence are the following:
+- 15% of the tokens are masked.
+- In 80% of the cases, the masked tokens are replaced by `[MASK]`.
+- In 10% of the cases, the masked tokens are replaced by a random token (different) from the one they replace.
+- In the 10% remaining cases, the masked tokens are left as is.
+
+The details of the creation of the synthetic and counterfactual examples can be found in the [follow-up paper](https://arxiv.org/abs/2010.00571). 
+
+### Pretraining
+
+The model was trained on 32 Cloud TPU v3 cores for one million steps with maximum sequence length 512 and batch size of 512.
+In this setup, pre-training takes around 3 days. The optimizer used is Adam with a learning rate of 5e-5, and a warmup ratio 
+of 0.10. 
+
+
+### BibTeX entry and citation info
+
+```bibtex
+@misc{herzig2020tapas,
+      title={TAPAS: Weakly Supervised Table Parsing via Pre-training}, 
+      author={Jonathan Herzig and Paweł Krzysztof Nowak and Thomas Müller and Francesco Piccinno and Julian Martin Eisenschlos},
+      year={2020},
+      eprint={2004.02349},
+      archivePrefix={arXiv},
+      primaryClass={cs.IR}
+}
+```
+
+```bibtex
+@misc{eisenschlos2020understanding,
+      title={Understanding tables with intermediate pre-training}, 
+      author={Julian Martin Eisenschlos and Syrine Krichene and Thomas Müller},
+      year={2020},
+      eprint={2010.00571},
+      archivePrefix={arXiv},
+      primaryClass={cs.CL}
+}
+```

notebooks/02-transformers.ipynb

@@ -73,12 +73,14 @@
     "\n",
     "The transformers library allows you to benefits from large, pretrained language models without requiring a huge and costly computational\n",
     "infrastructure. Most of the State-of-the-Art models are provided directly by their author and made available in the library \n",
-    "in PyTorch and TensorFlow in a transparent and interchangeable way. "
+    "in PyTorch and TensorFlow in a transparent and interchangeable way. \n",
+    "\n",
+    "If you're executing this notebook in Colab, you will need to install the transformers library. You can do so with this command:"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 2,
    "metadata": {
     "id": "KnT3Jn6fSXai",
     "pycharm": {
@@ -89,13 +91,12 @@
    },
    "outputs": [],
    "source": [
-    "!pip install transformers\n",
-    "!pip install --upgrade tensorflow"
+    "# !pip install transformers"
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 3,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -111,13 +112,11 @@
     {
      "data": {
       "text/plain": [
-       "<torch.autograd.grad_mode.set_grad_enabled at 0x7f9c03e5b3c8>"
+       "<torch.autograd.grad_mode.set_grad_enabled at 0x7ff0cc2a2c50>"
       ]
      },
-     "execution_count": 2,
-     "metadata": {
-      "tags": []
-     },
+     "execution_count": 3,
+     "metadata": {},
      "output_type": "execute_result"
     }
    ],
@@ -130,7 +129,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 4,
    "metadata": {
     "id": "1xMDTHQXSXai",
     "pycharm": {
@@ -159,12 +158,50 @@
    "source": [
     "With only the above two lines of code, you're ready to use a BERT pre-trained model. \n",
     "The tokenizers will allow us to map a raw textual input to a sequence of integers representing our textual input\n",
-    "in a way the model can manipulate."
+    "in a way the model can manipulate. Since we will be using a PyTorch model, we ask the tokenizer to return to us PyTorch tensors."
    ]
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 6,
+   "metadata": {},
+   "outputs": [
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "input_ids:\n",
+      "\ttensor([[ 101, 1188, 1110, 1126, 7758, 1859,  102]])\n",
+      "token_type_ids:\n",
+      "\ttensor([[0, 0, 0, 0, 0, 0, 0]])\n",
+      "attention_mask:\n",
+      "\ttensor([[1, 1, 1, 1, 1, 1, 1]])\n"
+     ]
+    }
+   ],
+   "source": [
+    "tokens_pt = tokenizer(\"This is an input example\", return_tensors=\"pt\")\n",
+    "for key, value in tokens_pt.items():\n",
+    "    print(\"{}:\\n\\t{}\".format(key, value))"
+   ]
+  },
+  {
+   "cell_type": "markdown",
+   "metadata": {},
+   "source": [
+    "The tokenizer automatically converted our input to all the inputs expected by the model. It generated some additional tensors on top of the IDs: \n",
+    "\n",
+    "- token_type_ids: This tensor will map every tokens to their corresponding segment (see below).\n",
+    "- attention_mask: This tensor is used to \"mask\" padded values in a batch of sequence with different lengths (see below).\n",
+    "\n",
+    "You can check our [glossary](https://huggingface.co/transformers/glossary.html) for more information about each of those keys. \n",
+    "\n",
+    "We can just feed this directly into our model:"
+   ]
+  },
+  {
+   "cell_type": "code",
+   "execution_count": 7,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -181,32 +218,12 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "Tokens: ['This', 'is', 'an', 'input', 'example']\n",
-      "Tokens id: [1188, 1110, 1126, 7758, 1859]\n",
-      "Tokens PyTorch: tensor([[ 101, 1188, 1110, 1126, 7758, 1859,  102]])\n",
       "Token wise output: torch.Size([1, 7, 768]), Pooled output: torch.Size([1, 768])\n"
      ]
     }
    ],
    "source": [
-    "# Tokens comes from a process that splits the input into sub-entities with interesting linguistic properties. \n",
-    "tokens = tokenizer.tokenize(\"This is an input example\")\n",
-    "print(\"Tokens: {}\".format(tokens))\n",
-    "\n",
-    "# This is not sufficient for the model, as it requires integers as input, \n",
-    "# not a problem, let's convert tokens to ids.\n",
-    "tokens_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
-    "print(\"Tokens id: {}\".format(tokens_ids))\n",
-    "\n",
-    "# Add the required special tokens\n",
-    "tokens_ids = tokenizer.build_inputs_with_special_tokens(tokens_ids)\n",
-    "\n",
-    "# We need to convert to a Deep Learning framework specific format, let's use PyTorch for now.\n",
-    "tokens_pt = torch.tensor([tokens_ids])\n",
-    "print(\"Tokens PyTorch: {}\".format(tokens_pt))\n",
-    "\n",
-    "# Now we're ready to go through BERT with out input\n",
-    "outputs = model(tokens_pt)\n",
+    "outputs = model(**tokens_pt)\n",
     "last_hidden_state = outputs.last_hidden_state\n",
     "pooler_output = outputs.pooler_output\n",
     "\n",
@@ -233,85 +250,9 @@
     "require a fine-grained token-level. This is the case for Sentiment-Analysis of the sequence or Information Retrieval."
    ]
   },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "DCxuDWH2SXaj",
-    "pycharm": {
-     "name": "#%% md\n"
-    }
-   },
-   "source": [
-    "The code you saw in the previous section introduced all the steps required to do simple model invocation.\n",
-    "For more day-to-day usage, transformers provides you higher-level methods which will makes your NLP journey easier.\n",
-    "Let's improve our previous example"
-   ]
-  },
   {
    "cell_type": "code",
-   "execution_count": null,
-   "metadata": {
-    "colab": {
-     "base_uri": "https://localhost:8080/"
-    },
-    "id": "sgcNCdXUSXaj",
-    "outputId": "af2fb928-7c17-475b-cf81-89cfc4b1d9e5",
-    "pycharm": {
-     "is_executing": false,
-     "name": "#%% code\n"
-    }
-   },
-   "outputs": [
-    {
-     "name": "stdout",
-     "output_type": "stream",
-     "text": [
-      "input_ids:\n",
-      "\ttensor([[ 101, 1188, 1110, 1126, 7758, 1859,  102]])\n",
-      "token_type_ids:\n",
-      "\ttensor([[0, 0, 0, 0, 0, 0, 0]])\n",
-      "attention_mask:\n",
-      "\ttensor([[1, 1, 1, 1, 1, 1, 1]])\n",
-      "Difference with previous code: (0.0, 0.0)\n"
-     ]
-    }
-   ],
-   "source": [
-    "# tokens = tokenizer.tokenize(\"This is an input example\")\n",
-    "# tokens_ids = tokenizer.convert_tokens_to_ids(tokens)\n",
-    "# tokens_pt = torch.tensor([tokens_ids])\n",
-    "\n",
-    "# This code can be factored into one-line as follow\n",
-    "tokens_pt2 = tokenizer(\"This is an input example\", return_tensors=\"pt\")\n",
-    "\n",
-    "for key, value in tokens_pt2.items():\n",
-    "    print(\"{}:\\n\\t{}\".format(key, value))\n",
-    "\n",
-    "outputs2 = model(**tokens_pt2)\n",
-    "last_hidden_state2 = outputs2.last_hidden_state\n",
-    "pooler_output2 = outputs2.pooler_output\n",
-    "\n",
-    "print(\"Difference with previous code: ({}, {})\".format((last_hidden_state2 - last_hidden_state).sum(), (pooler_output2 - pooler_output).sum()))"
-   ]
-  },
-  {
-   "cell_type": "markdown",
-   "metadata": {
-    "id": "gC-7xGYPSXal"
-   },
-   "source": [
-    "As you can see above, calling the tokenizer provides a convenient way to generate all the required parameters\n",
-    "that will go through the model. \n",
-    "\n",
-    "Moreover, you might have noticed it generated some additional tensors: \n",
-    "\n",
-    "- token_type_ids: This tensor will map every tokens to their corresponding segment (see below).\n",
-    "- attention_mask: This tensor is used to \"mask\" padded values in a batch of sequence with different lengths (see below)."
-   ]
-  },
-  {
-   "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 8,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -357,7 +298,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 9,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -414,14 +355,47 @@
   },
   {
    "cell_type": "code",
-   "execution_count": null,
+   "execution_count": 10,
    "metadata": {
     "id": "Kubwm-wJSXan",
     "pycharm": {
      "is_executing": false
     }
    },
-   "outputs": [],
+   "outputs": [
+    {
+     "data": {
+      "application/vnd.jupyter.widget-view+json": {
+       "model_id": "3b971be3639d4fedb02778fb5c6898a0",
+       "version_major": 2,
+       "version_minor": 0
+      },
+      "text/plain": [
+       "HBox(children=(FloatProgress(value=0.0, description='Downloading', max=526681800.0, style=ProgressStyle(descri…"
+      ]
+     },
+     "metadata": {},
+     "output_type": "display_data"
+    },
+    {
+     "name": "stdout",
+     "output_type": "stream",
+     "text": [
+      "\n"
+     ]
+    },
+    {
+     "name": "stderr",
+     "output_type": "stream",
+     "text": [
+      "Some layers from the model checkpoint at bert-base-cased were not used when initializing TFBertModel: ['nsp___cls', 'mlm___cls']\n",
+      "- This IS expected if you are initializing TFBertModel from the checkpoint of a model trained on another task or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n",
+      "- This IS NOT expected if you are initializing TFBertModel from the checkpoint of a model that you expect to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model).\n",
+      "All the layers of TFBertModel were initialized from the model checkpoint at bert-base-cased.\n",
+      "If your task is similar to the task the model of the checkpoint was trained on, you can already use TFBertModel for predictions without further training.\n"
+     ]
+    }
+   ],
    "source": [
     "from transformers import TFBertModel, BertModel\n",
     "\n",
@@ -432,7 +406,7 @@
   },
   {
    "cell_type": "code",
-   "execution_count": 12,
+   "execution_count": 11,
    "metadata": {
     "colab": {
      "base_uri": "https://localhost:8080/"
@@ -448,8 +422,8 @@
      "name": "stdout",
      "output_type": "stream",
      "text": [
-      "last_hidden_state differences: 1.0094e-05\n",
-      "pooler_output differences: 7.2969e-07\n"
+      "last_hidden_state differences: 1.2933e-05\n",
+      "pooler_output differences: 2.9691e-06\n"
      ]
     }
    ],
@@ -482,7 +456,7 @@
     "\n",
     "For example, Google released a few months ago **T5** an Encoder/Decoder architecture based on Transformer and available in `transformers` with no more than 11 billions parameters. Microsoft also recently entered the game with **Turing-NLG** using 17 billions parameters. This kind of model requires tens of gigabytes to store the weights and a tremendous compute infrastructure to run such models which makes it impracticable for the common man !\n",
     "\n",
-    "![transformers-parameters](https://lh5.googleusercontent.com/NRdXzEcgZV3ooykjIaTm9uvbr9QnSjDQHHAHb2kk_Lm9lIF0AhS-PJdXGzpcBDztax922XAp386hyNmWZYsZC1lUN2r4Ip5p9v-PHO19-jevRGg4iQFxgv5Olq4DWaqSA_8ptep7)\n",
+    "![transformers-parameters](https://github.com/huggingface/notebooks/blob/master/examples/images/model_parameters.png?raw=true)\n",
     "\n",
     "With the goal of making Transformer-based NLP accessible to everyone we @huggingface developed models that take advantage of a training process called **Distillation** which allows us to drastically reduce the resources needed to run such models with almost zero drop in performances.\n",
     "\n",
@@ -673,7 +647,7 @@
    "name": "python",
    "nbconvert_exporter": "python",
    "pygments_lexer": "ipython3",
-   "version": "3.7.3"
+   "version": "3.7.9"
   },
   "pycharm": {
    "stem_cell": {

notebooks/README.md

@@ -31,7 +31,11 @@ Pull Request so it can be included under the Community notebooks.
 | [Getting Started Tokenizers](https://github.com/huggingface/transformers/blob/master/notebooks/01-training-tokenizers.ipynb)  | How to train and use your very own tokenizer  |[![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/01-training-tokenizers.ipynb) |
 | [Getting Started Transformers](https://github.com/huggingface/transformers/blob/master/notebooks/02-transformers.ipynb)   | How to easily start using transformers  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/02-transformers.ipynb) |
 | [How to use Pipelines](https://github.com/huggingface/transformers/blob/master/notebooks/03-pipelines.ipynb)  | Simple and efficient way to use State-of-the-Art models on downstream tasks through transformers | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/03-pipelines.ipynb) |
-| [How to train a language model](https://github.com/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)| Highlight all the steps to effectively train Transformer model on custom data | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)|
+| [How to fine-tune a model on text classification](https://github.com/huggingface/notebooks/blob/master/examples/text_classification.ipynb) | Show how to preprocess the data and fine-tune a pretrained model on any GLUE task. | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/text_classification.ipynb)|
+| [How to fine-tune a model on language modeling](https://github.com/huggingface/notebooks/blob/master/examples/language_modeling.ipynb) | Show how to preprocess the data and fine-tune a pretrained model on a causal or masked LM task. | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/language_modeling.ipynb)|
+| [How to fine-tune a model on token classification](https://github.com/huggingface/notebooks/blob/master/examples/token_classification.ipynb) | Show how to preprocess the data and fine-tune a pretrained model on a token classification task (NER, PoS). | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/token_classification.ipynb)|
+| [How to fine-tune a model on question answering](https://github.com/huggingface/notebooks/blob/master/examples/question_answering.ipynb) | Show how to preprocess the data and fine-tune a pretrained model on SQUAD. | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/notebooks/blob/master/examples/question_answering.ipynb)|
+| [How to train a language model from scratch](https://github.com/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)| Highlight all the steps to effectively train Transformer model on custom data | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/01_how_to_train.ipynb)|
 | [How to generate text](https://github.com/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)| How to use different decoding methods for language generation with transformers | [![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)|
 | [How to export model to ONNX](https://github.com/huggingface/transformers/blob/master/notebooks/04-onnx-export.ipynb) | Highlight how to export and run inference workloads through ONNX |
 | [How to use Benchmarks](https://github.com/huggingface/transformers/blob/master/notebooks/05-benchmark.ipynb) | How to benchmark models with transformers | [![Open in Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/transformers/blob/master/notebooks/05-benchmark.ipynb)|
@@ -69,3 +73,5 @@ Pull Request so it can be included under the Community notebooks.
 |[Classify text with DistilBERT and Tensorflow](https://github.com/peterbayerle/huggingface_notebook/blob/main/distilbert_tf.ipynb) | How to fine-tune DistilBERT for text classification in TensorFlow | [Peter Bayerle](https://github.com/peterbayerle) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/peterbayerle/huggingface_notebook/blob/main/distilbert_tf.ipynb)|
 |[Leverage BERT for Encoder-Decoder Summarization on CNN/Dailymail](https://github.com/patrickvonplaten/notebooks/blob/master/BERT2BERT_for_CNN_Dailymail.ipynb) | How to warm-start a *EncoderDecoderModel* with a *bert-base-uncased* checkpoint for summarization on CNN/Dailymail | [Patrick von Platen](https://github.com/patrickvonplaten) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/patrickvonplaten/notebooks/blob/master/BERT2BERT_for_CNN_Dailymail.ipynb)|
 |[Leverage RoBERTa for Encoder-Decoder Summarization on BBC XSum](https://github.com/patrickvonplaten/notebooks/blob/master/RoBERTaShared_for_BBC_XSum.ipynb) | How to warm-start a shared *EncoderDecoderModel* with a *roberta-base* checkpoint for summarization on BBC/XSum | [Patrick von Platen](https://github.com/patrickvonplaten) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/patrickvonplaten/notebooks/blob/master/RoBERTaShared_for_BBC_XSum.ipynb)|
+|[Fine-tuning TAPAS on Sequential Question Answering (SQA)](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Fine_tuning_TapasForQuestionAnswering_on_SQA.ipynb) | How to fine-tune *TapasForQuestionAnswering* with a *tapas-base* checkpoint on the Sequential Question Answering (SQA) dataset | [Niels Rogge](https://github.com/nielsrogge) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NielsRogge/Transformers-Tutorials/blob/master/Fine_tuning_TapasForQuestionAnswering_on_SQAipynb)|
+|[Evaluating TAPAS on Table Fact Checking (TabFact)](https://github.com/NielsRogge/Transformers-Tutorials/blob/master/Evaluating_TAPAS_on_the_Tabfact_test_set.ipynb) | How to evaluate a fine-tuned *TapasForSequenceClassification* with a *tapas-base-finetuned-tabfact* checkpoint using a combination of the 🤗 datasets and 🤗 transformers libraries | [Niels Rogge](https://github.com/nielsrogge) | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/NielsRogge/Transformers-Tutorials/blob/master/Evaluating_TAPAS_on_the_Tabfact_test_set.ipynb)|

setup.py

@@ -97,7 +97,7 @@ _deps = [
     "fastapi",
     "filelock",
     "flake8>=3.8.3",
-    "flax==0.2.2",
+    "flax>=0.2.2",
     "fugashi>=1.0",
     "ipadic>=1.0.0,<2.0",
     "isort>=5.5.4",
@@ -127,8 +127,8 @@ _deps = [
     "sphinx-rtd-theme==0.4.3",  # sphinx-rtd-theme==0.5.0 introduced big changes in the style.
     "sphinx==3.2.1",
     "starlette",
-    "tensorflow-cpu>=2.0,<2.4",
-    "tensorflow>=2.0,<2.4",
+    "tensorflow-cpu>=2.3",
+    "tensorflow>=2.3",
     "timeout-decorator",
     "tokenizers==0.9.4",
     "torch>=1.0",
@@ -175,7 +175,7 @@ class DepsTableUpdateCommand(Command):
             "deps = {",
             entries,
             "}",
-            ""
+            "",
         ]
         target = "src/transformers/dependency_versions_table.py"
         print(f"updating {target}")
@@ -232,19 +232,19 @@ extras["dev"] = (
 # when modifying the following list, make sure to update src/transformers/dependency_versions_check.py
 install_requires = [
     deps["dataclasses"] + ";python_version<'3.7'",  # dataclasses for Python versions that don't have it
-    deps["filelock"],   # filesystem locks, e.g., to prevent parallel downloads
+    deps["filelock"],  # filesystem locks, e.g., to prevent parallel downloads
     deps["numpy"],
     deps["packaging"],  # utilities from PyPA to e.g., compare versions
-    deps["regex"],      # for OpenAI GPT
-    deps["requests"],   # for downloading models over HTTPS
-    deps["sacremoses"], # for XLM
+    deps["regex"],  # for OpenAI GPT
+    deps["requests"],  # for downloading models over HTTPS
+    deps["sacremoses"],  # for XLM
     deps["tokenizers"],
-    deps["tqdm"],       # progress bars in model download and training scripts
+    deps["tqdm"],  # progress bars in model download and training scripts
 ]
 
 setup(
     name="transformers",
-    version="4.1.0.dev0",
+    version="4.2.0.dev0", # expected format is one of x.y.z.dev0, or x.y.z.rc1 or x.y.z (no to dashes, yes to dots)
     author="Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Sam Shleifer, Patrick von Platen, Sylvain Gugger, Google AI Language Team Authors, Open AI team Authors, Facebook AI Authors, Carnegie Mellon University Authors",
     author_email="thomas@huggingface.co",
     description="State-of-the-art Natural Language Processing for TensorFlow 2.0 and PyTorch",

src/transformers/__init__.py

@@ -16,7 +16,7 @@
 # See the License for the specific language governing permissions and
 # limitations under the License.
 
-__version__ = "4.1.0.dev0"
+__version__ = "4.2.0dev0"
 
 # Work around to update TensorFlow's absl.logging threshold which alters the
 # default Python logging output behavior when present.
@@ -164,6 +164,7 @@ from .models.retribert import RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RetriBert
 from .models.roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig, RobertaTokenizer
 from .models.squeezebert import SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, SqueezeBertConfig, SqueezeBertTokenizer
 from .models.t5 import T5_PRETRAINED_CONFIG_ARCHIVE_MAP, T5Config
+from .models.tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig, TapasTokenizer
 from .models.transfo_xl import (
     TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP,
     TransfoXLConfig,
@@ -189,6 +190,7 @@ from .pipelines import (
     PipelineDataFormat,
     QuestionAnsweringPipeline,
     SummarizationPipeline,
+    TableQuestionAnsweringPipeline,
     Text2TextGenerationPipeline,
     TextClassificationPipeline,
     TextGenerationPipeline,
@@ -217,6 +219,7 @@ from .integrations import (  # isort:skip
     is_comet_available,
     is_optuna_available,
     is_ray_available,
+    is_ray_tune_available,
     is_tensorboard_available,
     is_wandb_available,
 )
@@ -284,8 +287,9 @@ from .trainer_callback import (
     TrainerControl,
     TrainerState,
 )
-from .trainer_utils import EvalPrediction, EvaluationStrategy, set_seed
+from .trainer_utils import EvalPrediction, EvaluationStrategy, SchedulerType, set_seed
 from .training_args import TrainingArguments
+from .training_args_seq2seq import Seq2SeqTrainingArguments
 from .training_args_tf import TFTrainingArguments
 from .utils import logging
 
@@ -357,6 +361,7 @@ if is_torch_available():
         MODEL_FOR_QUESTION_ANSWERING_MAPPING,
         MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
         MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
+        MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
         MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
         MODEL_MAPPING,
         MODEL_WITH_LM_HEAD_MAPPING,
@@ -369,6 +374,7 @@ if is_torch_available():
         AutoModelForQuestionAnswering,
         AutoModelForSeq2SeqLM,
         AutoModelForSequenceClassification,
+        AutoModelForTableQuestionAnswering,
         AutoModelForTokenClassification,
         AutoModelWithLMHead,
     )
@@ -401,7 +407,11 @@ if is_torch_available():
         BertGenerationEncoder,
         load_tf_weights_in_bert_generation,
     )
-    from .models.blenderbot import BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST, BlenderbotForConditionalGeneration
+    from .models.blenderbot import (
+        BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST,
+        BlenderbotForConditionalGeneration,
+        BlenderbotModel,
+    )
     from .models.camembert import (
         CAMEMBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
         CamembertForCausalLM,
@@ -517,7 +527,7 @@ if is_torch_available():
         LxmertXLayer,
     )
     from .models.marian import MarianMTModel
-    from .models.mbart import MBartForConditionalGeneration
+    from .models.mbart import MBartForConditionalGeneration, MBartModel
     from .models.mmbt import MMBTForClassification, MMBTModel, ModalEmbeddings
     from .models.mobilebert import (
         MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
@@ -554,7 +564,7 @@ if is_torch_available():
         OpenAIGPTPreTrainedModel,
         load_tf_weights_in_openai_gpt,
     )
-    from .models.pegasus import PegasusForConditionalGeneration
+    from .models.pegasus import PegasusForConditionalGeneration, PegasusModel
     from .models.prophetnet import (
         PROPHETNET_PRETRAINED_MODEL_ARCHIVE_LIST,
         ProphetNetDecoder,
@@ -605,6 +615,13 @@ if is_torch_available():
         T5PreTrainedModel,
         load_tf_weights_in_t5,
     )
+    from .models.tapas import (
+        TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TapasForMaskedLM,
+        TapasForQuestionAnswering,
+        TapasForSequenceClassification,
+        TapasModel,
+    )
     from .models.transfo_xl import (
         TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
         AdaptiveEmbedding,
@@ -666,11 +683,13 @@ if is_torch_available():
         get_cosine_with_hard_restarts_schedule_with_warmup,
         get_linear_schedule_with_warmup,
         get_polynomial_decay_schedule_with_warmup,
+        get_scheduler,
     )
 
     # Trainer
     from .trainer import Trainer
     from .trainer_pt_utils import torch_distributed_zero_first
+    from .trainer_seq2seq import Seq2SeqTrainer
 else:
     from .utils.dummy_pt_objects import *
 
@@ -717,7 +736,7 @@ if is_tf_available():
         TFAutoModelForTokenClassification,
         TFAutoModelWithLMHead,
     )
-    from .models.bart import TFBartForConditionalGeneration, TFBartModel
+    from .models.bart import TFBartForConditionalGeneration, TFBartModel, TFBartPretrainedModel
     from .models.bert import (
         TF_BERT_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFBertEmbeddings,
@@ -745,6 +764,7 @@ if is_tf_available():
     )
     from .models.ctrl import (
         TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFCTRLForSequenceClassification,
         TFCTRLLMHeadModel,
         TFCTRLModel,
         TFCTRLPreTrainedModel,
@@ -859,6 +879,7 @@ if is_tf_available():
     from .models.openai import (
         TF_OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFOpenAIGPTDoubleHeadsModel,
+        TFOpenAIGPTForSequenceClassification,
         TFOpenAIGPTLMHeadModel,
         TFOpenAIGPTMainLayer,
         TFOpenAIGPTModel,
@@ -886,6 +907,7 @@ if is_tf_available():
     from .models.transfo_xl import (
         TF_TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST,
         TFAdaptiveEmbedding,
+        TFTransfoXLForSequenceClassification,
         TFTransfoXLLMHeadModel,
         TFTransfoXLMainLayer,
         TFTransfoXLModel,
@@ -936,6 +958,7 @@ else:
 
 
 if is_flax_available():
+    from .modeling_flax_utils import FlaxPreTrainedModel
     from .models.auto import FLAX_MODEL_MAPPING, FlaxAutoModel
     from .models.bert import FlaxBertForMaskedLM, FlaxBertModel
     from .models.roberta import FlaxRobertaModel

src/transformers/configuration_utils.py

@@ -317,6 +317,9 @@ class PretrainedConfig(object):
             proxies (:obj:`Dict[str, str]`, `optional`):
                 A dictionary of proxy servers to use by protocol or endpoint, e.g., :obj:`{'http': 'foo.bar:3128',
                 'http://hostname': 'foo.bar:4012'}.` The proxies are used on each request.
+            use_auth_token (:obj:`str` or `bool`, `optional`):
+                The token to use as HTTP bearer authorization for remote files. If :obj:`True`, will use the token
+                generated when running :obj:`transformers-cli login` (stored in :obj:`~/.huggingface`).
             revision(:obj:`str`, `optional`, defaults to :obj:`"main"`):
                 The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                 git-based system for storing models and other artifacts on huggingface.co, so ``revision`` can be any
@@ -332,6 +335,10 @@ class PretrainedConfig(object):
                 values. Behavior concerning key/value pairs whose keys are *not* configuration attributes is controlled
                 by the ``return_unused_kwargs`` keyword parameter.
 
+        .. note::
+
+            Passing :obj:`use_auth_token=True` is required when you want to use a private model.
+
         Returns:
             :class:`PretrainedConfig`: The configuration object instantiated from this pretrained model.
 
@@ -373,6 +380,7 @@ class PretrainedConfig(object):
         force_download = kwargs.pop("force_download", False)
         resume_download = kwargs.pop("resume_download", False)
         proxies = kwargs.pop("proxies", None)
+        use_auth_token = kwargs.pop("use_auth_token", None)
         local_files_only = kwargs.pop("local_files_only", False)
         revision = kwargs.pop("revision", None)
 
@@ -395,6 +403,7 @@ class PretrainedConfig(object):
                 proxies=proxies,
                 resume_download=resume_download,
                 local_files_only=local_files_only,
+                use_auth_token=use_auth_token,
             )
             # Load config dict
             config_dict = cls._dict_from_json_file(resolved_config_file)

src/transformers/convert_slow_tokenizer.py

@@ -74,18 +74,6 @@ class BertConverter(Converter):
         vocab = self.original_tokenizer.vocab
         tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))
 
-        # # Let the tokenizer know about special tokens if they are part of the vocab
-        # if tokenizer.token_to_id(str(self.original_tokenizer.unk_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.unk_token)])
-        # if tokenizer.token_to_id(str(self.original_tokenizer.sep_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.sep_token)])
-        # if tokenizer.token_to_id(str(self.original_tokenizer.cls_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.cls_token)])
-        # if tokenizer.token_to_id(str(self.original_tokenizer.pad_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.pad_token)])
-        # if tokenizer.token_to_id(str(self.original_tokenizer.mask_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.mask_token)])
-
         tokenize_chinese_chars = False
         strip_accents = False
         do_lower_case = False
@@ -125,18 +113,6 @@ class FunnelConverter(Converter):
         vocab = self.original_tokenizer.vocab
         tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))
 
-        # # Let the tokenizer know about special tokens if they are part of the vocab
-        # if tokenizer.token_to_id(str(self.original_tokenizer.unk_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.unk_token)])
-        # if tokenizer.token_to_id(str(self.original_tokenizer.sep_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.sep_token)])
-        # if tokenizer.token_to_id(str(self.original_tokenizer.cls_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.cls_token)])
-        # if tokenizer.token_to_id(str(self.original_tokenizer.pad_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.pad_token)])
-        # if tokenizer.token_to_id(str(self.original_tokenizer.mask_token)) is not None:
-        #     tokenizer.add_special_tokens([str(self.original_tokenizer.mask_token)])
-
         tokenize_chinese_chars = False
         strip_accents = False
         do_lower_case = False
@@ -171,6 +147,45 @@ class FunnelConverter(Converter):
         return tokenizer
 
 
+class MPNetConverter(Converter):
+    def converted(self) -> Tokenizer:
+        vocab = self.original_tokenizer.vocab
+        tokenizer = Tokenizer(WordPiece(vocab, unk_token=str(self.original_tokenizer.unk_token)))
+
+        tokenize_chinese_chars = False
+        strip_accents = False
+        do_lower_case = False
+        if hasattr(self.original_tokenizer, "basic_tokenizer"):
+            tokenize_chinese_chars = self.original_tokenizer.basic_tokenizer.tokenize_chinese_chars
+            strip_accents = self.original_tokenizer.basic_tokenizer.strip_accents
+            do_lower_case = self.original_tokenizer.basic_tokenizer.do_lower_case
+
+        tokenizer.normalizer = normalizers.BertNormalizer(
+            clean_text=True,
+            handle_chinese_chars=tokenize_chinese_chars,
+            strip_accents=strip_accents,
+            lowercase=do_lower_case,
+        )
+        tokenizer.pre_tokenizer = pre_tokenizers.BertPreTokenizer()
+
+        cls = str(self.original_tokenizer.cls_token)
+        sep = str(self.original_tokenizer.sep_token)
+        cls_token_id = self.original_tokenizer.cls_token_id
+        sep_token_id = self.original_tokenizer.sep_token_id
+
+        tokenizer.post_processor = processors.TemplateProcessing(
+            single=f"{cls}:0 $A:0 {sep}:0",
+            pair=f"{cls}:0 $A:0 {sep}:0 {sep}:0 $B:1 {sep}:1",  # MPNet uses two [SEP] tokens
+            special_tokens=[
+                (cls, cls_token_id),
+                (sep, sep_token_id),
+            ],
+        )
+        tokenizer.decoder = decoders.WordPiece(prefix="##")
+
+        return tokenizer
+
+
 class OpenAIGPTConverter(Converter):
     def converted(self) -> Tokenizer:
         vocab = self.original_tokenizer.encoder
@@ -602,6 +617,7 @@ SLOW_TO_FAST_CONVERTERS = {
     "LongformerTokenizer": RobertaConverter,
     "LxmertTokenizer": BertConverter,
     "MBartTokenizer": MBartConverter,
+    "MPNetTokenizer": MPNetConverter,
     "MobileBertTokenizer": BertConverter,
     "OpenAIGPTTokenizer": OpenAIGPTConverter,
     "PegasusTokenizer": PegasusConverter,

src/transformers/dependency_versions_table.py

@@ -10,7 +10,7 @@ deps = {
     "fastapi": "fastapi",
     "filelock": "filelock",
     "flake8": "flake8>=3.8.3",
-    "flax": "flax==0.2.2",
+    "flax": "flax>=0.2.2",
     "fugashi": "fugashi>=1.0",
     "ipadic": "ipadic>=1.0.0,<2.0",
     "isort": "isort>=5.5.4",
@@ -40,8 +40,8 @@ deps = {
     "sphinx-rtd-theme": "sphinx-rtd-theme==0.4.3",
     "sphinx": "sphinx==3.2.1",
     "starlette": "starlette",
-    "tensorflow-cpu": "tensorflow-cpu>=2.0",
-    "tensorflow": "tensorflow>=2.0",
+    "tensorflow-cpu": "tensorflow-cpu>=2.3",
+    "tensorflow": "tensorflow>=2.3",
     "timeout-decorator": "timeout-decorator",
     "tokenizers": "tokenizers==0.9.4",
     "torch": "torch>=1.0",

src/transformers/file_utils.py

@@ -16,6 +16,7 @@ Utilities for working with the local dataset cache. Parts of this file is adapte
 https://github.com/allenai/allennlp.
 """
 
+import copy
 import fnmatch
 import io
 import json
@@ -42,6 +43,7 @@ import requests
 from filelock import FileLock
 
 from . import __version__
+from .hf_api import HfFolder
 from .utils import logging
 
 
@@ -216,6 +218,29 @@ except ImportError:
     _tokenizers_available = False
 
 
+try:
+    import pandas  # noqa: F401
+
+    _pandas_available = True
+
+except ImportError:
+    _pandas_available = False
+
+
+try:
+    import torch_scatter
+
+    # Check we're not importing a "torch_scatter" directory somewhere
+    _scatter_available = hasattr(torch_scatter, "__version__") and hasattr(torch_scatter, "scatter")
+    if _scatter_available:
+        logger.debug(f"Succesfully imported torch-scatter version {torch_scatter.__version__}")
+    else:
+        logger.debug("Imported a torch_scatter object but this doesn't seem to be the torch-scatter library.")
+
+except ImportError:
+    _scatter_available = False
+
+
 old_default_cache_path = os.path.join(torch_cache_home, "transformers")
 # New default cache, shared with the Datasets library
 hf_cache_home = os.path.expanduser(
@@ -247,6 +272,7 @@ TRANSFORMERS_CACHE = os.getenv("TRANSFORMERS_CACHE", PYTORCH_TRANSFORMERS_CACHE)
 WEIGHTS_NAME = "pytorch_model.bin"
 TF2_WEIGHTS_NAME = "tf_model.h5"
 TF_WEIGHTS_NAME = "model.ckpt"
+FLAX_WEIGHTS_NAME = "flax_model.msgpack"
 CONFIG_NAME = "config.json"
 MODEL_CARD_NAME = "modelcard.json"
 
@@ -325,6 +351,14 @@ def is_in_notebook():
     return _in_notebook
 
 
+def is_scatter_available():
+    return _scatter_available
+
+
+def is_pandas_available():
+    return _pandas_available
+
+
 def torch_only_method(fn):
     def wrapper(*args, **kwargs):
         if not _torch_available:
@@ -427,6 +461,20 @@ installation page: https://github.com/google/flax and follow the ones that match
 """
 
 
+# docstyle-ignore
+SCATTER_IMPORT_ERROR = """
+{0} requires the torch-scatter library but it was not found in your environment. You can install it with pip as
+explained here: https://github.com/rusty1s/pytorch_scatter.
+"""
+
+
+# docstyle-ignore
+PANDAS_IMPORT_ERROR = """
+{0} requires the pandas library but it was not found in your environment. You can install it with pip as
+explained here: https://pandas.pydata.org/pandas-docs/stable/getting_started/install.html.
+"""
+
+
 def requires_datasets(obj):
     name = obj.__name__ if hasattr(obj, "__name__") else obj.__class__.__name__
     if not is_datasets_available():
@@ -481,6 +529,18 @@ def requires_protobuf(obj):
         raise ImportError(PROTOBUF_IMPORT_ERROR.format(name))
 
 
+def requires_pandas(obj):
+    name = obj.__name__ if hasattr(obj, "__name__") else obj.__class__.__name__
+    if not is_pandas_available():
+        raise ImportError(PANDAS_IMPORT_ERROR.format(name))
+
+
+def requires_scatter(obj):
+    name = obj.__name__ if hasattr(obj, "__name__") else obj.__class__.__name__
+    if not is_scatter_available():
+        raise ImportError(SCATTER_IMPORT_ERROR.format(name))
+
+
 def add_start_docstrings(*docstr):
     def docstring_decorator(fn):
         fn.__doc__ = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")
@@ -979,6 +1039,7 @@ def cached_path(
     user_agent: Union[Dict, str, None] = None,
     extract_compressed_file=False,
     force_extract=False,
+    use_auth_token: Union[bool, str, None] = None,
     local_files_only=False,
 ) -> Optional[str]:
     """
@@ -991,6 +1052,8 @@ def cached_path(
         force_download: if True, re-download the file even if it's already cached in the cache dir.
         resume_download: if True, resume the download if incompletely received file is found.
         user_agent: Optional string or dict that will be appended to the user-agent on remote requests.
+        use_auth_token: Optional string or boolean to use as Bearer token for remote files. If True,
+            will get token from ~/.huggingface.
         extract_compressed_file: if True and the path point to a zip or tar file, extract the compressed
             file in a folder along the archive.
         force_extract: if True when extract_compressed_file is True and the archive was already extracted,
@@ -1018,6 +1081,7 @@ def cached_path(
             proxies=proxies,
             resume_download=resume_download,
             user_agent=user_agent,
+            use_auth_token=use_auth_token,
             local_files_only=local_files_only,
         )
     elif os.path.exists(url_or_filename):
@@ -1080,11 +1144,11 @@ def http_user_agent(user_agent: Union[Dict, str, None] = None) -> str:
     return ua
 
 
-def http_get(url: str, temp_file: BinaryIO, proxies=None, resume_size=0, user_agent: Union[Dict, str, None] = None):
+def http_get(url: str, temp_file: BinaryIO, proxies=None, resume_size=0, headers: Optional[Dict[str, str]] = None):
     """
     Donwload remote file. Do not gobble up errors.
     """
-    headers = {"user-agent": http_user_agent(user_agent)}
+    headers = copy.deepcopy(headers)
     if resume_size > 0:
         headers["Range"] = "bytes=%d-" % (resume_size,)
     r = requests.get(url, stream=True, proxies=proxies, headers=headers)
@@ -1114,6 +1178,7 @@ def get_from_cache(
     etag_timeout=10,
     resume_download=False,
     user_agent: Union[Dict, str, None] = None,
+    use_auth_token: Union[bool, str, None] = None,
     local_files_only=False,
 ) -> Optional[str]:
     """
@@ -1133,11 +1198,19 @@ def get_from_cache(
 
     os.makedirs(cache_dir, exist_ok=True)
 
+    headers = {"user-agent": http_user_agent(user_agent)}
+    if isinstance(use_auth_token, str):
+        headers["authorization"] = "Bearer {}".format(use_auth_token)
+    elif use_auth_token:
+        token = HfFolder.get_token()
+        if token is None:
+            raise EnvironmentError("You specified use_auth_token=True, but a huggingface token was not found.")
+        headers["authorization"] = "Bearer {}".format(token)
+
     url_to_download = url
     etag = None
     if not local_files_only:
         try:
-            headers = {"user-agent": http_user_agent(user_agent)}
             r = requests.head(url, headers=headers, allow_redirects=False, proxies=proxies, timeout=etag_timeout)
             r.raise_for_status()
             etag = r.headers.get("X-Linked-Etag") or r.headers.get("ETag")
@@ -1227,7 +1300,7 @@ def get_from_cache(
         with temp_file_manager() as temp_file:
             logger.info("%s not found in cache or force_download set to True, downloading to %s", url, temp_file.name)
 
-            http_get(url_to_download, temp_file, proxies=proxies, resume_size=resume_size, user_agent=user_agent)
+            http_get(url_to_download, temp_file, proxies=proxies, resume_size=resume_size, headers=headers)
 
         logger.info("storing %s in cache at %s", url, cache_path)
         os.replace(temp_file.name, cache_path)

src/transformers/generation_utils.py

@@ -156,7 +156,7 @@ class GenerationMixin:
         if is_encoder_decoder:
             assert encoder_outputs is not None
             encoder_outputs["last_hidden_state"] = encoder_outputs.last_hidden_state.index_select(
-                0, expanded_return_idx
+                0, expanded_return_idx.to(encoder_outputs.last_hidden_state.device)
             )
             model_kwargs["encoder_outputs"] = encoder_outputs
         return input_ids, model_kwargs
@@ -226,7 +226,7 @@ class GenerationMixin:
         For custom re-ordering of :obj:`past_key_values` or :obj:`mems`, the function should be implemented in
         subclasses of :class:`~transformers.PreTrainedModel`.
         """
-        return tuple(layer_past.index_select(1, beam_idx) for layer_past in past)
+        return tuple(layer_past.index_select(1, beam_idx.to(layer_past.device)) for layer_past in past)
 
     def _get_logits_warper(
         self, top_k: int = None, top_p: float = None, temperature: float = None, num_beams: int = None

src/transformers/hf_api.py

@@ -206,7 +206,7 @@ class HfApi:
 
     def model_list(self) -> List[ModelInfo]:
         """
-        Get the public list of all the models on huggingface, including the community models
+        Get the public list of all the models on huggingface.co
         """
         path = "{}/api/models".format(self.endpoint)
         r = requests.get(path)
@@ -214,7 +214,7 @@ class HfApi:
         d = r.json()
         return [ModelInfo(**x) for x in d]
 
-    def list_repos_objs(self, token: str, organization: Optional[str] = None) -> List[S3Obj]:
+    def list_repos_objs(self, token: str, organization: Optional[str] = None) -> List[RepoObj]:
         """
         HuggingFace git-based system, used for models.
 
@@ -228,7 +228,13 @@ class HfApi:
         return [RepoObj(**x) for x in d]
 
     def create_repo(
-        self, token: str, name: str, organization: Optional[str] = None, lfsmultipartthresh: Optional[int] = None
+        self,
+        token: str,
+        name: str,
+        organization: Optional[str] = None,
+        private: Optional[bool] = None,
+        exist_ok=False,
+        lfsmultipartthresh: Optional[int] = None,
     ) -> str:
         """
         HuggingFace git-based system, used for models.
@@ -236,10 +242,14 @@ class HfApi:
         Call HF API to create a whole repo.
 
         Params:
+            private: Whether the model repo should be private (requires a paid huggingface.co account)
+
+            exist_ok: Do not raise an error if repo already exists
+
             lfsmultipartthresh: Optional: internal param for testing purposes.
         """
         path = "{}/api/repos/create".format(self.endpoint)
-        json = {"name": name, "organization": organization}
+        json = {"name": name, "organization": organization, "private": private}
         if lfsmultipartthresh is not None:
             json["lfsmultipartthresh"] = lfsmultipartthresh
         r = requests.post(
@@ -247,6 +257,8 @@ class HfApi:
             headers={"authorization": "Bearer {}".format(token)},
             json=json,
         )
+        if exist_ok and r.status_code == 409:
+            return ""
         r.raise_for_status()
         d = r.json()
         return d["url"]

src/transformers/integrations.py

@@ -63,8 +63,16 @@ try:
     import ray  # noqa: F401
 
     _has_ray = True
+    try:
+        # Ray Tune has additional dependencies.
+        from ray import tune  # noqa: F401
+
+        _has_ray_tune = True
+    except (ImportError):
+        _has_ray_tune = False
 except (ImportError):
     _has_ray = False
+    _has_ray_tune = False
 
 try:
     from torch.utils.tensorboard import SummaryWriter  # noqa: F401
@@ -92,6 +100,13 @@ try:
 except ImportError:
     _has_mlflow = False
 
+try:
+    import fairscale  # noqa: F401
+
+    _has_fairscale = True
+except ImportError:
+    _has_fairscale = False
+
 # No transformer imports above this point
 
 from .file_utils import is_torch_tpu_available  # noqa: E402
@@ -120,6 +135,10 @@ def is_ray_available():
     return _has_ray
 
 
+def is_ray_tune_available():
+    return _has_ray_tune
+
+
 def is_azureml_available():
     return _has_azureml
 
@@ -128,11 +147,15 @@ def is_mlflow_available():
     return _has_mlflow
 
 
+def is_fairscale_available():
+    return _has_fairscale
+
+
 def hp_params(trial):
     if is_optuna_available():
         if isinstance(trial, optuna.Trial):
             return trial.params
-    if is_ray_available():
+    if is_ray_tune_available():
         if isinstance(trial, dict):
             return trial
 
@@ -142,7 +165,7 @@ def hp_params(trial):
 def default_hp_search_backend():
     if is_optuna_available():
         return "optuna"
-    elif is_ray_available():
+    elif is_ray_tune_available():
         return "ray"
 
 

src/transformers/modeling_flax_utils.py

@@ -15,64 +15,65 @@
 
 import os
 from abc import ABC, abstractmethod
+from functools import partial
 from pickle import UnpicklingError
-from typing import Dict
+from typing import Dict, Set, Tuple, Union
 
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from flax.serialization import to_bytes
-from flax.traverse_util import unflatten_dict
+from flax.core.frozen_dict import FrozenDict, freeze, unfreeze
+from flax.serialization import from_bytes, to_bytes
+from flax.traverse_util import flatten_dict, unflatten_dict
 from jax.random import PRNGKey
 
 from .configuration_utils import PretrainedConfig
-from .file_utils import WEIGHTS_NAME, cached_path, hf_bucket_url, is_remote_url
+from .file_utils import FLAX_WEIGHTS_NAME, WEIGHTS_NAME, cached_path, hf_bucket_url, is_remote_url
 from .utils import logging
 
 
 logger = logging.get_logger(__name__)
 
 
-@jax.jit
-def gelu(x):
-    r"""
-    Gaussian error linear unit activation function.
-
-    Computes the element-wise function:
-
-    .. math::
-      \mathrm{gelu}(x) = \frac{x}{2} \left(1 + \mathrm{tanh} \left(
-        \sqrt{\frac{2}{\pi}} \left(x + 0.044715 x^3 \right) \right) \right)
-
-    We explicitly use the approximation rather than the exact formulation for speed. For more information, see
-    `Gaussian Error Linear Units (GELUs) <https://arxiv.org/abs/1606.08415>`_, section 2.
-    """
-    return x * 0.5 * (1.0 + jax.lax.erf(x / jnp.sqrt(2.0)))
-
-
 ACT2FN = {
     "gelu": nn.gelu,
     "relu": nn.relu,
     "silu": nn.swish,
     "swish": nn.swish,
-    "gelu_new": gelu,
+    "gelu_new": partial(nn.gelu, approximate=True),
 }
 
 
 class FlaxPreTrainedModel(ABC):
+    r"""
+    Base class for all models.
+
+    :class:`~transformers.FlaxPreTrainedModel` takes care of storing the configuration of the models and handles
+    methods for loading, downloading and saving models.
+
+    Class attributes (overridden by derived classes):
+
+        - **config_class** (:class:`~transformers.PretrainedConfig`) -- A subclass of
+          :class:`~transformers.PretrainedConfig` to use as configuration class for this model architecture.
+        - **base_model_prefix** (:obj:`str`) -- A string indicating the attribute associated to the base model in
+          derived classes of the same architecture adding modules on top of the base model.
+    """
     config_class = None
-    pretrained_model_archive_map = {}
     base_model_prefix = ""
-    model_class = None
 
     def __init__(
-        self, config: PretrainedConfig, module: nn.Module, params: Dict, seed: int = 0, dtype: jnp.dtype = jnp.float32
+        self,
+        config: PretrainedConfig,
+        module: nn.Module,
+        input_shape: Tuple = (1, 1),
+        seed: int = 0,
+        dtype: jnp.dtype = jnp.float32,
     ):
         if config is None:
             raise ValueError("config cannot be None")
 
-        if params is None:
-            raise ValueError("state cannot be None")
+        if module is None:
+            raise ValueError("module cannot be None")
 
         # Those are private to be exposed as typed property on derived classes.
         self._config = config
@@ -80,9 +81,18 @@ class FlaxPreTrainedModel(ABC):
 
         # Those are public as their type is generic to every derived classes.
         self.key = PRNGKey(seed)
-        self.params = params
         self.dtype = dtype
 
+        # randomely initialized parameters
+        random_params = self.init(self.key, input_shape)
+
+        # save required_params as set
+        self._required_params = set(flatten_dict(unfreeze(random_params)).keys())
+        self.params = random_params
+
+    def init(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> Dict:
+        raise NotImplementedError(f"init method has to be implemented for {self}")
+
     @property
     def config(self) -> PretrainedConfig:
         return self._config
@@ -91,25 +101,132 @@ class FlaxPreTrainedModel(ABC):
     def module(self) -> nn.Module:
         return self._module
 
+    @property
+    def params(self) -> Union[Dict, FrozenDict]:
+        return self._params
+
+    @property
+    def required_params(self) -> Set:
+        return self._required_params
+
+    @params.setter
+    def params(self, params: Union[Dict, FrozenDict]):
+        if isinstance(params, FrozenDict):
+            params = unfreeze(params)
+        param_keys = set(flatten_dict(params).keys())
+        if len(self.required_params - param_keys) > 0:
+            raise ValueError(
+                "Some parameters are missing. Make sure that `params` include the following "
+                f"parameters {self.required_params - param_keys}"
+            )
+        self._params = freeze(params)
+
     @staticmethod
     @abstractmethod
     def convert_from_pytorch(pt_state: Dict, config: PretrainedConfig) -> Dict:
         raise NotImplementedError()
 
     @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path, dtype: jnp.dtype = jnp.float32, *model_args, **kwargs):
+    def from_pretrained(
+        cls,
+        pretrained_model_name_or_path: Union[str, os.PathLike],
+        dtype: jnp.dtype = jnp.float32,
+        *model_args,
+        **kwargs
+    ):
+
         r"""
-        Instantiate a pretrained Flax model from a pre-trained model configuration.
+        Instantiate a pretrained flax model from a pre-trained model configuration.
+
+        The warning `Weights from XXX not initialized from pretrained model` means that the weights of XXX do not come
+        pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
+        task.
+
+        The warning `Weights from XXX not used in YYY` means that the layer XXX is not used by YYY, therefore those
+        weights are discarded.
+
+        Parameters:
+            pretrained_model_name_or_path (:obj:`str` or :obj:`os.PathLike`):
+                Can be either:
+
+                    - A string, the `model id` of a pretrained model hosted inside a model repo on huggingface.co.
+                      Valid model ids can be located at the root-level, like ``bert-base-uncased``, or namespaced under
+                      a user or organization name, like ``dbmdz/bert-base-german-cased``.
+                    - A path to a `directory` containing model weights saved using
+                      :func:`~transformers.FlaxPreTrainedModel.save_pretrained`, e.g., ``./my_model_directory/``.
+                    - A path or url to a `pt index checkpoint file` (e.g, ``./tf_model/model.ckpt.index``). In this
+                      case, ``from_pt`` should be set to :obj:`True`.
+            model_args (sequence of positional arguments, `optional`):
+                All remaning positional arguments will be passed to the underlying model's ``__init__`` method.
+            config (:obj:`Union[PretrainedConfig, str, os.PathLike]`, `optional`):
+                Can be either:
+
+                    - an instance of a class derived from :class:`~transformers.PretrainedConfig`,
+                    - a string or path valid as input to :func:`~transformers.PretrainedConfig.from_pretrained`.
+
+                Configuration for the model to use instead of an automatically loaded configuation. Configuration can
+                be automatically loaded when:
+
+                    - The model is a model provided by the library (loaded with the `model id` string of a pretrained
+                      model).
+                    - The model was saved using :func:`~transformers.PreTrainedModel.save_pretrained` and is reloaded
+                      by supplying the save directory.
+                    - The model is loaded by supplying a local directory as ``pretrained_model_name_or_path`` and a
+                      configuration JSON file named `config.json` is found in the directory.
+            cache_dir (:obj:`Union[str, os.PathLike]`, `optional`):
+                Path to a directory in which a downloaded pretrained model configuration should be cached if the
+                standard cache should not be used.
+            from_pt (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Load the model weights from a PyTorch checkpoint save file (see docstring of
+                ``pretrained_model_name_or_path`` argument).
+            force_download (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
+                cached versions if they exist.
+            resume_download (:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to delete incompletely received files. Will attempt to resume the download if such a
+                file exists.
+            proxies (:obj:`Dict[str, str], `optional`):
+                A dictionary of proxy servers to use by protocol or endpoint, e.g., :obj:`{'http': 'foo.bar:3128',
+                'http://hostname': 'foo.bar:4012'}`. The proxies are used on each request.
+            local_files_only(:obj:`bool`, `optional`, defaults to :obj:`False`):
+                Whether or not to only look at local files (i.e., do not try to download the model).
+            revision(:obj:`str`, `optional`, defaults to :obj:`"main"`):
+                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
+                git-based system for storing models and other artifacts on huggingface.co, so ``revision`` can be any
+                identifier allowed by git.
+            kwargs (remaining dictionary of keyword arguments, `optional`):
+                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
+                :obj:`output_attentions=True`). Behaves differently depending on whether a ``config`` is provided or
+                automatically loaded:
+
+                    - If a configuration is provided with ``config``, ``**kwargs`` will be directly passed to the
+                      underlying model's ``__init__`` method (we assume all relevant updates to the configuration have
+                      already been done)
+                    - If a configuration is not provided, ``kwargs`` will be first passed to the configuration class
+                      initialization function (:func:`~transformers.PretrainedConfig.from_pretrained`). Each key of
+                      ``kwargs`` that corresponds to a configuration attribute will be used to override said attribute
+                      with the supplied ``kwargs`` value. Remaining keys that do not correspond to any configuration
+                      attribute will be passed to the underlying model's ``__init__`` function.
+
+        Examples::
+
+            >>> from transformers import BertConfig, FlaxBertModel
+            >>> # Download model and configuration from huggingface.co and cache.
+            >>> model = FlaxBertModel.from_pretrained('bert-base-cased')
+            >>> # Model was saved using `save_pretrained('./test/saved_model/')` (for example purposes, not runnable).
+            >>> model = FlaxBertModel.from_pretrained('./test/saved_model/')
+            >>> # Loading from a PyTorch checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable).
+            >>> config = BertConfig.from_json_file('./pt_model/config.json')
+            >>> model = FlaxBertModel.from_pretrained('./pt_model/pytorch_model.bin', from_pt=True, config=config)
         """
         config = kwargs.pop("config", None)
-        # state_dict = kwargs.pop("state_dict", None)
         cache_dir = kwargs.pop("cache_dir", None)
-        # from_tf = kwargs.pop("from_tf", False)
+        from_pt = kwargs.pop("from_pt", False)
         force_download = kwargs.pop("force_download", False)
         resume_download = kwargs.pop("resume_download", False)
         proxies = kwargs.pop("proxies", None)
-        # output_loading_info = kwargs.pop("output_loading_info", False)
         local_files_only = kwargs.pop("local_files_only", False)
+        use_auth_token = kwargs.pop("use_auth_token", None)
         revision = kwargs.pop("revision", None)
 
         # Load config if we don't provide a configuration
@@ -124,6 +241,7 @@ class FlaxPreTrainedModel(ABC):
                 resume_download=resume_download,
                 proxies=proxies,
                 local_files_only=local_files_only,
+                use_auth_token=use_auth_token,
                 revision=revision,
                 **kwargs,
             )
@@ -135,10 +253,28 @@ class FlaxPreTrainedModel(ABC):
 
         # Load model
         if pretrained_model_name_or_path is not None:
-            if os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
+            if os.path.isdir(pretrained_model_name_or_path):
+                if from_pt and os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
+                    # Load from a PyTorch checkpoint
+                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)):
+                    # Load from a Flax checkpoint
+                    archive_file = os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)
+                else:
+                    raise EnvironmentError(
+                        "Error no file named {} found in directory {} or `from_pt` set to False".format(
+                            [FLAX_WEIGHTS_NAME, WEIGHTS_NAME],
+                            pretrained_model_name_or_path,
+                        )
+                    )
+            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
                 archive_file = pretrained_model_name_or_path
             else:
-                archive_file = hf_bucket_url(pretrained_model_name_or_path, filename=WEIGHTS_NAME, revision=revision)
+                archive_file = hf_bucket_url(
+                    pretrained_model_name_or_path,
+                    filename=WEIGHTS_NAME if from_pt else FLAX_WEIGHTS_NAME,
+                    revision=revision,
+                )
 
             # redirect to the cache, if necessary
             try:
@@ -149,6 +285,7 @@ class FlaxPreTrainedModel(ABC):
                     proxies=proxies,
                     resume_download=resume_download,
                     local_files_only=local_files_only,
+                    use_auth_token=use_auth_token,
                 )
             except EnvironmentError as err:
                 logger.error(err)
@@ -169,31 +306,96 @@ class FlaxPreTrainedModel(ABC):
         # Instantiate model.
         with open(resolved_archive_file, "rb") as state_f:
             try:
-                from flax.serialization import from_bytes
-
-                state = from_bytes(cls.model_class, state_f)
-            except TypeError:
-                try:
+                if from_pt:
                     import torch
 
                     state = torch.load(state_f)
-                    state = {k: v.numpy() for k, v in state.items()}
-                    state = cls.convert_from_pytorch(state, config)
-                    state = unflatten_dict({tuple(k.split(".")[1:]): v for k, v in state.items()})
-                except UnpicklingError:
-                    raise EnvironmentError(
-                        f"Unable to convert model {archive_file} to Flax deserializable object. "
-                        "Supported format are PyTorch archive or Flax msgpack"
-                    )
 
-        return cls(config, state, *model_args, **model_kwargs)
+                    state = convert_state_dict_from_pt(cls, state, config)
+                else:
+                    state = from_bytes(cls, state_f.read())
+            except UnpicklingError:
+                raise EnvironmentError(
+                    f"Unable to convert pytorch model {archive_file} to Flax deserializable object. "
+                )
 
-    def save_pretrained(self, folder):
-        folder_abs = os.path.abspath(folder)
+        # init random models
+        model = cls(config, *model_args, **model_kwargs)
 
-        if not os.path.exists(folder_abs):
-            os.mkdir(folder_abs)
+        # if model is base model only use model_prefix key
+        if cls.base_model_prefix not in dict(model.params) and cls.base_model_prefix in state:
+            state = state[cls.base_model_prefix]
 
-        with open(os.path.join(folder_abs, f"{self._config.model_type}.flax", "wb")) as f:
+        # flatten dicts
+        state = flatten_dict(state)
+        random_state = flatten_dict(unfreeze(model.params))
+
+        missing_keys = model.required_params - set(state.keys())
+        unexpected_keys = set(state.keys()) - model.required_params
+
+        # add missing keys as random parameters
+        for missing_key in missing_keys:
+            state[missing_key] = random_state[missing_key]
+
+        if len(unexpected_keys) > 0:
+            logger.warning(
+                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when "
+                f"initializing {model.__class__.__name__}: {unexpected_keys}\n"
+                f"- This IS expected if you are initializing {model.__class__.__name__} from the checkpoint of a model trained on another task "
+                f"or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n"
+                f"- This IS NOT expected if you are initializing {model.__class__.__name__} from the checkpoint of a model that you expect "
+                f"to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
+            )
+        else:
+            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
+
+        if len(missing_keys) > 0:
+            logger.warning(
+                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at {pretrained_model_name_or_path} "
+                f"and are newly initialized: {missing_keys}\n"
+                f"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference."
+            )
+        else:
+            logger.info(
+                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at {pretrained_model_name_or_path}.\n"
+                f"If your task is similar to the task the model of the checkpoint was trained on, "
+                f"you can already use {model.__class__.__name__} for predictions without further training."
+            )
+
+        # set correct parameters
+        model.params = unflatten_dict(state)
+        return model
+
+    def save_pretrained(self, save_directory: Union[str, os.PathLike]):
+        """
+        Save a model and its configuration file to a directory, so that it can be re-loaded using the
+        `:func:`~transformers.FlaxPreTrainedModel.from_pretrained`` class method
+
+        Arguments:
+            save_directory (:obj:`str` or :obj:`os.PathLike`):
+                Directory to which to save. Will be created if it doesn't exist.
+        """
+        if os.path.isfile(save_directory):
+            logger.error("Provided path ({}) should be a directory, not a file".format(save_directory))
+            return
+        os.makedirs(save_directory, exist_ok=True)
+
+        # get abs dir
+        save_directory = os.path.abspath(save_directory)
+        # save config as well
+        self.config.save_pretrained(save_directory)
+
+        # save model
+        with open(os.path.join(save_directory, FLAX_WEIGHTS_NAME), "wb") as f:
             model_bytes = to_bytes(self.params)
             f.write(model_bytes)
+
+
+def convert_state_dict_from_pt(model_class: ABC, state: Dict, config: PretrainedConfig):
+    """
+    Converts a PyTorch parameter state dict to an equivalent Flax parameter state dict
+    """
+    state = {k: v.numpy() for k, v in state.items()}
+    state = model_class.convert_from_pytorch(state, config)
+    state = unflatten_dict({tuple(k.split(".")): v for k, v in state.items()})
+    return state

src/transformers/modeling_tf_utils.py

@@ -894,6 +894,9 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
                 Whether ot not to also return a dictionary containing missing keys, unexpected keys and error messages.
             local_files_only(:obj:`bool`, `optional`, defaults to :obj:`False`):
                 Whether or not to only look at local files (e.g., not try doanloading the model).
+            use_auth_token (:obj:`str` or `bool`, `optional`):
+                The token to use as HTTP bearer authorization for remote files. If :obj:`True`, will use the token
+                generated when running :obj:`transformers-cli login` (stored in :obj:`~/.huggingface`).
             revision(:obj:`str`, `optional`, defaults to :obj:`"main"`):
                 The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                 git-based system for storing models and other artifacts on huggingface.co, so ``revision`` can be any
@@ -916,6 +919,10 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
                       with the supplied ``kwargs`` value. Remaining keys that do not correspond to any configuration
                       attribute will be passed to the underlying model's ``__init__`` function.
 
+        .. note::
+
+            Passing :obj:`use_auth_token=True` is required when you want to use a private model.
+
         Examples::
 
             >>> from transformers import BertConfig, TFBertModel
@@ -939,6 +946,7 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
         proxies = kwargs.pop("proxies", None)
         output_loading_info = kwargs.pop("output_loading_info", False)
         local_files_only = kwargs.pop("local_files_only", False)
+        use_auth_token = kwargs.pop("use_auth_token", None)
         revision = kwargs.pop("revision", None)
         mirror = kwargs.pop("mirror", None)
 
@@ -954,6 +962,7 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
                 resume_download=resume_download,
                 proxies=proxies,
                 local_files_only=local_files_only,
+                use_auth_token=use_auth_token,
                 revision=revision,
                 **kwargs,
             )
@@ -996,6 +1005,7 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
                     proxies=proxies,
                     resume_download=resume_download,
                     local_files_only=local_files_only,
+                    use_auth_token=use_auth_token,
                 )
             except EnvironmentError as err:
                 logger.error(err)
@@ -1346,7 +1356,7 @@ def shape_list(tensor: tf.Tensor) -> List[int]:
     dynamic = tf.shape(tensor)
 
     if tensor.shape == tf.TensorShape(None):
-        return dynamic.as_list()
+        return dynamic
 
     static = tensor.shape.as_list()
 

src/transformers/modeling_utils.py

@@ -886,6 +886,9 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
                 Whether ot not to also return a dictionary containing missing keys, unexpected keys and error messages.
             local_files_only(:obj:`bool`, `optional`, defaults to :obj:`False`):
                 Whether or not to only look at local files (i.e., do not try to download the model).
+            use_auth_token (:obj:`str` or `bool`, `optional`):
+                The token to use as HTTP bearer authorization for remote files. If :obj:`True`, will use the token
+                generated when running :obj:`transformers-cli login` (stored in :obj:`~/.huggingface`).
             revision(:obj:`str`, `optional`, defaults to :obj:`"main"`):
                 The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
                 git-based system for storing models and other artifacts on huggingface.co, so ``revision`` can be any
@@ -908,6 +911,10 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
                       with the supplied ``kwargs`` value. Remaining keys that do not correspond to any configuration
                       attribute will be passed to the underlying model's ``__init__`` function.
 
+        .. note::
+
+            Passing :obj:`use_auth_token=True` is required when you want to use a private model.
+
         Examples::
 
             >>> from transformers import BertConfig, BertModel
@@ -931,6 +938,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
         proxies = kwargs.pop("proxies", None)
         output_loading_info = kwargs.pop("output_loading_info", False)
         local_files_only = kwargs.pop("local_files_only", False)
+        use_auth_token = kwargs.pop("use_auth_token", None)
         revision = kwargs.pop("revision", None)
         mirror = kwargs.pop("mirror", None)
 
@@ -946,6 +954,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
                 resume_download=resume_download,
                 proxies=proxies,
                 local_files_only=local_files_only,
+                use_auth_token=use_auth_token,
                 revision=revision,
                 **kwargs,
             )
@@ -998,6 +1007,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
                     proxies=proxies,
                     resume_download=resume_download,
                     local_files_only=local_files_only,
+                    use_auth_token=use_auth_token,
                 )
             except EnvironmentError as err:
                 logger.error(err)

src/transformers/models/auto/__init__.py

@@ -31,6 +31,7 @@ if is_torch_available():
         MODEL_FOR_QUESTION_ANSWERING_MAPPING,
         MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING,
         MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING,
+        MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING,
         MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING,
         MODEL_MAPPING,
         MODEL_WITH_LM_HEAD_MAPPING,
@@ -43,6 +44,7 @@ if is_torch_available():
         AutoModelForQuestionAnswering,
         AutoModelForSeq2SeqLM,
         AutoModelForSequenceClassification,
+        AutoModelForTableQuestionAnswering,
         AutoModelForTokenClassification,
         AutoModelWithLMHead,
     )

src/transformers/models/auto/configuration_auto.py

@@ -51,6 +51,7 @@ from ..retribert.configuration_retribert import RETRIBERT_PRETRAINED_CONFIG_ARCH
 from ..roberta.configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig
 from ..squeezebert.configuration_squeezebert import SQUEEZEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, SqueezeBertConfig
 from ..t5.configuration_t5 import T5_PRETRAINED_CONFIG_ARCHIVE_MAP, T5Config
+from ..tapas.configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
 from ..transfo_xl.configuration_transfo_xl import TRANSFO_XL_PRETRAINED_CONFIG_ARCHIVE_MAP, TransfoXLConfig
 from ..xlm.configuration_xlm import XLM_PRETRAINED_CONFIG_ARCHIVE_MAP, XLMConfig
 from ..xlm_prophetnet.configuration_xlm_prophetnet import (
@@ -95,6 +96,7 @@ ALL_PRETRAINED_CONFIG_ARCHIVE_MAP = dict(
         XLM_PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
         PROPHETNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
         MPNET_PRETRAINED_CONFIG_ARCHIVE_MAP,
+        TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP,
     ]
     for key, value, in pretrained_map.items()
 )
@@ -141,6 +143,7 @@ CONFIG_MAPPING = OrderedDict(
         ("dpr", DPRConfig),
         ("layoutlm", LayoutLMConfig),
         ("rag", RagConfig),
+        ("tapas", TapasConfig),
     ]
 )
 
@@ -185,6 +188,7 @@ MODEL_NAMES_MAPPING = OrderedDict(
         ("prophetnet", "ProphetNet"),
         ("mt5", "mT5"),
         ("mpnet", "MPNet"),
+        ("tapas", "TAPAS"),
     ]
 )
 

src/transformers/models/auto/modeling_auto.py

@@ -50,7 +50,7 @@ from ..bert.modeling_bert import (
     BertModel,
 )
 from ..bert_generation.modeling_bert_generation import BertGenerationDecoder, BertGenerationEncoder
-from ..blenderbot.modeling_blenderbot import BlenderbotForConditionalGeneration
+from ..blenderbot.modeling_blenderbot import BlenderbotForConditionalGeneration, BlenderbotModel
 from ..camembert.modeling_camembert import (
     CamembertForCausalLM,
     CamembertForMaskedLM,
@@ -111,7 +111,7 @@ from ..longformer.modeling_longformer import (
 )
 from ..lxmert.modeling_lxmert import LxmertForPreTraining, LxmertForQuestionAnswering, LxmertModel
 from ..marian.modeling_marian import MarianMTModel
-from ..mbart.modeling_mbart import MBartForConditionalGeneration
+from ..mbart.modeling_mbart import MBartForConditionalGeneration, MBartModel
 from ..mobilebert.modeling_mobilebert import (
     MobileBertForMaskedLM,
     MobileBertForMultipleChoice,
@@ -132,7 +132,7 @@ from ..mpnet.modeling_mpnet import (
 )
 from ..mt5.modeling_mt5 import MT5ForConditionalGeneration, MT5Model
 from ..openai.modeling_openai import OpenAIGPTForSequenceClassification, OpenAIGPTLMHeadModel, OpenAIGPTModel
-from ..pegasus.modeling_pegasus import PegasusForConditionalGeneration
+from ..pegasus.modeling_pegasus import PegasusForConditionalGeneration, PegasusModel
 from ..prophetnet.modeling_prophetnet import ProphetNetForCausalLM, ProphetNetForConditionalGeneration, ProphetNetModel
 from ..rag.modeling_rag import (  # noqa: F401 - need to import all RagModels to be in globals() function
     RagModel,
@@ -165,6 +165,12 @@ from ..squeezebert.modeling_squeezebert import (
     SqueezeBertModel,
 )
 from ..t5.modeling_t5 import T5ForConditionalGeneration, T5Model
+from ..tapas.modeling_tapas import (
+    TapasForMaskedLM,
+    TapasForQuestionAnswering,
+    TapasForSequenceClassification,
+    TapasModel,
+)
 from ..transfo_xl.modeling_transfo_xl import TransfoXLForSequenceClassification, TransfoXLLMHeadModel, TransfoXLModel
 from ..xlm.modeling_xlm import (
     XLMForMultipleChoice,
@@ -230,6 +236,7 @@ from .configuration_auto import (
     RobertaConfig,
     SqueezeBertConfig,
     T5Config,
+    TapasConfig,
     TransfoXLConfig,
     XLMConfig,
     XLMProphetNetConfig,
@@ -248,6 +255,10 @@ MODEL_MAPPING = OrderedDict(
         (RetriBertConfig, RetriBertModel),
         (MT5Config, MT5Model),
         (T5Config, T5Model),
+        (PegasusConfig, PegasusModel),
+        (MarianConfig, MarianMTModel),
+        (MBartConfig, MBartModel),
+        (BlenderbotConfig, BlenderbotModel),
         (DistilBertConfig, DistilBertModel),
         (AlbertConfig, AlbertModel),
         (CamembertConfig, CamembertModel),
@@ -277,6 +288,7 @@ MODEL_MAPPING = OrderedDict(
         (XLMProphetNetConfig, XLMProphetNetModel),
         (ProphetNetConfig, ProphetNetModel),
         (MPNetConfig, MPNetModel),
+        (TapasConfig, TapasModel),
     ]
 )
 
@@ -308,6 +320,7 @@ MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(
         (LxmertConfig, LxmertForPreTraining),
         (FunnelConfig, FunnelForPreTraining),
         (MPNetConfig, MPNetForMaskedLM),
+        (TapasConfig, TapasForMaskedLM),
     ]
 )
 
@@ -340,6 +353,7 @@ MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(
         (ReformerConfig, ReformerModelWithLMHead),
         (FunnelConfig, FunnelForMaskedLM),
         (MPNetConfig, MPNetForMaskedLM),
+        (TapasConfig, TapasForMaskedLM),
     ]
 )
 
@@ -386,6 +400,7 @@ MODEL_FOR_MASKED_LM_MAPPING = OrderedDict(
         (ReformerConfig, ReformerForMaskedLM),
         (FunnelConfig, FunnelForMaskedLM),
         (MPNetConfig, MPNetForMaskedLM),
+        (TapasConfig, TapasForMaskedLM),
     ]
 )
 
@@ -431,6 +446,7 @@ MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = OrderedDict(
         (CTRLConfig, CTRLForSequenceClassification),
         (TransfoXLConfig, TransfoXLForSequenceClassification),
         (MPNetConfig, MPNetForSequenceClassification),
+        (TapasConfig, TapasForSequenceClassification),
     ]
 )
 
@@ -458,6 +474,13 @@ MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
     ]
 )
 
+MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING = OrderedDict(
+    [
+        # Model for Table Question Answering mapping
+        (TapasConfig, TapasForQuestionAnswering),
+    ]
+)
+
 MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = OrderedDict(
     [
         # Model for Token Classification mapping
@@ -1371,6 +1394,106 @@ class AutoModelForQuestionAnswering:
         )
 
 
+class AutoModelForTableQuestionAnswering:
+    r"""
+    This is a generic model class that will be instantiated as one of the model classes of the library---with a table
+    question answering head---when created with the when created with the
+    :meth:`~transformers.AutoModeForTableQuestionAnswering.from_pretrained` class method or the
+    :meth:`~transformers.AutoModelForTableQuestionAnswering.from_config` class method.
+
+    This class cannot be instantiated directly using ``__init__()`` (throws an error).
+    """
+
+    def __init__(self):
+        raise EnvironmentError(
+            "AutoModelForQuestionAnswering is designed to be instantiated "
+            "using the `AutoModelForTableQuestionAnswering.from_pretrained(pretrained_model_name_or_path)` or "
+            "`AutoModelForTableQuestionAnswering.from_config(config)` methods."
+        )
+
+    @classmethod
+    @replace_list_option_in_docstrings(MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING, use_model_types=False)
+    def from_config(cls, config):
+        r"""
+        Instantiates one of the model classes of the library---with a table question answering head---from a
+        configuration.
+
+        Note:
+            Loading a model from its configuration file does **not** load the model weights. It only affects the
+            model's configuration. Use :meth:`~transformers.AutoModelForTableQuestionAnswering.from_pretrained` to load
+            the model weights.
+
+        Args:
+            config (:class:`~transformers.PretrainedConfig`):
+                The model class to instantiate is selected based on the configuration class:
+
+                List options
+
+        Examples::
+
+            >>> from transformers import AutoConfig, AutoModelForTableQuestionAnswering
+            >>> # Download configuration from huggingface.co and cache.
+            >>> config = AutoConfig.from_pretrained('google/tapas-base-finetuned-wtq')
+            >>> model = AutoModelForTableQuestionAnswering.from_config(config)
+        """
+        if type(config) in MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.keys():
+            return MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING[type(config)](config)
+
+        raise ValueError(
+            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
+            "Model type should be one of {}.".format(
+                config.__class__,
+                cls.__name__,
+                ", ".join(c.__name__ for c in MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.keys()),
+            )
+        )
+
+    @classmethod
+    @replace_list_option_in_docstrings(MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING)
+    @add_start_docstrings(
+        "Instantiate one of the model classes of the library---with a table question answering head---from a "
+        "pretrained model.",
+        AUTO_MODEL_PRETRAINED_DOCSTRING,
+    )
+    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+        r"""
+        Examples::
+
+            >>> from transformers import AutoConfig, AutoModelForTableQuestionAnswering
+
+            >>> # Download model and configuration from huggingface.co and cache.
+            >>> model = AutoModelForTableQuestionAnswering.from_pretrained('google/tapas-base-finetuned-wtq')
+
+            >>> # Update configuration during loading
+            >>> model = AutoModelForTableQuestionAnswering.from_pretrained('google/tapas-base-finetuned-wtq', output_attentions=True)
+            >>> model.config.output_attentions
+            True
+
+            >>> # Loading from a TF checkpoint file instead of a PyTorch model (slower)
+            >>> config = AutoConfig.from_json_file('./tf_model/tapas_tf_checkpoint.json')
+            >>> model = AutoModelForQuestionAnswering.from_pretrained('./tf_model/tapas_tf_checkpoint.ckpt.index', from_tf=True, config=config)
+        """
+        config = kwargs.pop("config", None)
+        if not isinstance(config, PretrainedConfig):
+            config, kwargs = AutoConfig.from_pretrained(
+                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
+            )
+
+        if type(config) in MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.keys():
+            return MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
+
+        raise ValueError(
+            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
+            "Model type should be one of {}.".format(
+                config.__class__,
+                cls.__name__,
+                ", ".join(c.__name__ for c in MODEL_FOR_TABLE_QUESTION_ANSWERING_MAPPING.keys()),
+            )
+        )
+
+
 class AutoModelForTokenClassification:
     r"""
     This is a generic model class that will be instantiated as one of the model classes of the library---with a token

src/transformers/models/auto/modeling_flax_auto.py

@@ -27,15 +27,6 @@ from .configuration_auto import AutoConfig, BertConfig, RobertaConfig
 logger = logging.get_logger(__name__)
 
 
-ALL_PRETRAINED_MODEL_ARCHIVE_MAP = dict(
-    (key, value)
-    for pretrained_map in [
-        FlaxBertModel.pretrained_model_archive_map,
-        FlaxRobertaModel.pretrained_model_archive_map,
-    ]
-    for key, value, in pretrained_map.items()
-)
-
 FLAX_MODEL_MAPPING = OrderedDict(
     [
         (RobertaConfig, FlaxRobertaModel),
@@ -114,10 +105,9 @@ class FlaxAutoModel(object):
                   organization name, like ``dbmdz/bert-base-german-cased``.
                 - a path to a `directory` containing model weights saved using
                   :func:`~transformers.FlaxPreTrainedModel.save_pretrained`, e.g.: ``./my_model_directory/``.
-                - a path or url to a `tensorflow index checkpoint file` (e.g. `./tf_model/model.ckpt.index`). In this
-                  case, ``from_tf`` should be set to True and a configuration object should be provided as ``config``
-                  argument. This loading path is slower than converting the TensorFlow checkpoint in a PyTorch model
-                  using the provided conversion scripts and loading the PyTorch model afterwards.
+                - a path or url to a `pytorch index checkpoint file` (e.g. `./pt_model/pytorch_model.bin`). In this
+                  case, ``from_pt`` should be set to True and a configuration object should be provided as ``config``
+                  argument.
 
             model_args: (`optional`) Sequence of positional arguments:
                 All remaining positional arguments will be passed to the underlying model's ``__init__`` method
@@ -133,13 +123,6 @@ class FlaxAutoModel(object):
                 - the model is loaded by supplying a local directory as ``pretrained_model_name_or_path`` and a
                   configuration JSON file named `config.json` is found in the directory.
 
-            state_dict: (`optional`) dict:
-                an optional state dictionary for the model to use instead of a state dictionary loaded from saved
-                weights file. This option can be used if you want to create a model from a pretrained configuration but
-                load your own weights. In this case though, you should check if using
-                :func:`~transformers.FlaxPreTrainedModel.save_pretrained` and
-                :func:`~transformers.FlaxPreTrainedModel.from_pretrained` is not a simpler option.
-
             cache_dir: (`optional`) string:
                 Path to a directory in which a downloaded pre-trained model configuration should be cached if the
                 standard cache should not be used.

src/transformers/models/auto/modeling_tf_auto.py

@@ -53,7 +53,7 @@ from ..camembert.modeling_tf_camembert import (
     TFCamembertForTokenClassification,
     TFCamembertModel,
 )
-from ..ctrl.modeling_tf_ctrl import TFCTRLLMHeadModel, TFCTRLModel
+from ..ctrl.modeling_tf_ctrl import TFCTRLForSequenceClassification, TFCTRLLMHeadModel, TFCTRLModel
 from ..distilbert.modeling_tf_distilbert import (
     TFDistilBertForMaskedLM,
     TFDistilBertForMultipleChoice,
@@ -120,7 +120,7 @@ from ..mpnet.modeling_tf_mpnet import (
     TFMPNetModel,
 )
 from ..mt5.modeling_tf_mt5 import TFMT5ForConditionalGeneration, TFMT5Model
-from ..openai.modeling_tf_openai import TFOpenAIGPTLMHeadModel, TFOpenAIGPTModel
+from ..openai.modeling_tf_openai import TFOpenAIGPTForSequenceClassification, TFOpenAIGPTLMHeadModel, TFOpenAIGPTModel
 from ..pegasus.modeling_tf_pegasus import TFPegasusForConditionalGeneration
 from ..roberta.modeling_tf_roberta import (
     TFRobertaForMaskedLM,
@@ -131,7 +131,11 @@ from ..roberta.modeling_tf_roberta import (
     TFRobertaModel,
 )
 from ..t5.modeling_tf_t5 import TFT5ForConditionalGeneration, TFT5Model
-from ..transfo_xl.modeling_tf_transfo_xl import TFTransfoXLLMHeadModel, TFTransfoXLModel
+from ..transfo_xl.modeling_tf_transfo_xl import (
+    TFTransfoXLForSequenceClassification,
+    TFTransfoXLLMHeadModel,
+    TFTransfoXLModel,
+)
 from ..xlm.modeling_tf_xlm import (
     TFXLMForMultipleChoice,
     TFXLMForQuestionAnsweringSimple,
@@ -341,6 +345,9 @@ TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = OrderedDict(
         (FunnelConfig, TFFunnelForSequenceClassification),
         (GPT2Config, TFGPT2ForSequenceClassification),
         (MPNetConfig, TFMPNetForSequenceClassification),
+        (OpenAIGPTConfig, TFOpenAIGPTForSequenceClassification),
+        (TransfoXLConfig, TFTransfoXLForSequenceClassification),
+        (CTRLConfig, TFCTRLForSequenceClassification),
     ]
 )
 

src/transformers/models/auto/tokenization_auto.py

@@ -47,6 +47,7 @@ from ..rag.tokenization_rag import RagTokenizer
 from ..retribert.tokenization_retribert import RetriBertTokenizer
 from ..roberta.tokenization_roberta import RobertaTokenizer
 from ..squeezebert.tokenization_squeezebert import SqueezeBertTokenizer
+from ..tapas.tokenization_tapas import TapasTokenizer
 from ..transfo_xl.tokenization_transfo_xl import TransfoXLTokenizer
 from ..xlm.tokenization_xlm import XLMTokenizer
 from .configuration_auto import (
@@ -84,6 +85,7 @@ from .configuration_auto import (
     RobertaConfig,
     SqueezeBertConfig,
     T5Config,
+    TapasConfig,
     TransfoXLConfig,
     XLMConfig,
     XLMProphetNetConfig,
@@ -223,6 +225,7 @@ TOKENIZER_MAPPING = OrderedDict(
         (XLMProphetNetConfig, (XLMProphetNetTokenizer, None)),
         (ProphetNetConfig, (ProphetNetTokenizer, None)),
         (MPNetConfig, (MPNetTokenizer, MPNetTokenizerFast)),
+        (TapasConfig, (TapasTokenizer, None)),
     ]
 )
 

src/transformers/models/bart/__init__.py

@@ -36,4 +36,4 @@ if is_torch_available():
     )
 
 if is_tf_available():
-    from .modeling_tf_bart import TFBartForConditionalGeneration, TFBartModel
+    from .modeling_tf_bart import TFBartForConditionalGeneration, TFBartModel, TFBartPretrainedModel

src/transformers/models/bart/modeling_bart.py

@@ -67,9 +67,16 @@ def shift_tokens_right(input_ids: torch.Tensor, pad_token_id: int):
     Shift input ids one token to the right, and wrap the last non pad token (usually <eos>).
     """
     prev_output_tokens = input_ids.clone()
-    index_of_eos = (input_ids.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
-    prev_output_tokens[:, 0] = input_ids.gather(1, index_of_eos).squeeze()
-    prev_output_tokens[:, 1:] = input_ids[:, :-1]
+
+    assert pad_token_id is not None, "self.model.config.pad_token_id has to be defined."
+    # replace possible -100 values in labels by `pad_token_id`
+    prev_output_tokens.masked_fill_(prev_output_tokens == -100, pad_token_id)
+
+    index_of_eos = (prev_output_tokens.ne(pad_token_id).sum(dim=1) - 1).unsqueeze(-1)
+    decoder_start_tokens = prev_output_tokens.gather(1, index_of_eos).squeeze()
+    prev_output_tokens[:, 1:] = prev_output_tokens[:, :-1].clone()
+    prev_output_tokens[:, 0] = decoder_start_tokens
+
     return prev_output_tokens
 
 
@@ -88,9 +95,7 @@ def _make_causal_mask(input_ids_shape: torch.Size, dtype: torch.dtype, past_key_
     return mask[None, None, :, :].expand(bsz, 1, tgt_len, tgt_len + past_key_values_length)
 
 
-def _expand_mask(
-    mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None, past_key_values_length: int = 0
-):
+def _expand_mask(mask: torch.Tensor, dtype: torch.dtype, tgt_len: Optional[int] = None):
     """
     Expands attention_mask from `[bsz, seq_len]` to `[bsz, 1, tgt_seq_len, src_seq_len]`.
     """
@@ -99,16 +104,6 @@ def _expand_mask(
 
     expanded_mask = mask[:, None, None, :].expand(bsz, 1, tgt_len, src_len).to(dtype)
 
-    if past_key_values_length > 0:
-        # concat fully attendend attention_mask to the beginning if `past_key_values` are used
-        expanded_mask = torch.cat(
-            [
-                torch.ones(bsz, 1, tgt_len, past_key_values_length, device=expanded_mask.device, dtype=dtype),
-                expanded_mask,
-            ],
-            dim=-1,
-        )
-
     inverted_mask = 1.0 - expanded_mask
 
     return inverted_mask.masked_fill(inverted_mask.bool(), torch.finfo(dtype).min)
@@ -215,10 +210,10 @@ class BartAttention(nn.Module):
 
     def forward(
         self,
-        hidden_states,
+        hidden_states: torch.Tensor,
         key_value_states: Optional[torch.Tensor] = None,
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
-        attn_mask: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
         output_attentions: bool = False,
     ) -> Tuple[torch.Tensor, Optional[torch.Tensor], Optional[Tuple[torch.Tensor]]]:
         """Input shape: Batch x Time x Channel"""
@@ -274,14 +269,14 @@ class BartAttention(nn.Module):
             src_len,
         ), f"Attention weights should be of size {(bsz * self.num_heads, tgt_len, src_len)}, but is {attn_weights.size()}"
 
-        if attn_mask is not None:
-            assert attn_mask.size() == (
+        if attention_mask is not None:
+            assert attention_mask.size() == (
                 bsz,
                 1,
                 tgt_len,
                 src_len,
-            ), f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attn_mask.size()}"
-            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attn_mask
+            ), f"Attention mask should be of size {(bsz, 1, tgt_len, src_len)}, but is {attention_mask.size()}"
+            attn_weights = attn_weights.view(bsz, self.num_heads, tgt_len, src_len) + attention_mask
             attn_weights = attn_weights.view(bsz * self.num_heads, tgt_len, src_len)
 
         attn_weights = F.softmax(attn_weights, dim=-1)
@@ -335,23 +330,19 @@ class BartEncoderLayer(nn.Module):
         self.fc2 = nn.Linear(config.encoder_ffn_dim, self.embed_dim)
         self.final_layer_norm = BartLayerNorm(self.embed_dim)
 
-    def forward(
-        self, hidden_states: torch.Tensor, encoder_padding_mask: torch.Tensor, output_attentions: bool = False
-    ):
+    def forward(self, hidden_states: torch.Tensor, attention_mask: torch.Tensor, output_attentions: bool = False):
         """
         Args:
             hidden_states (:obj:`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)`
-            encoder_padding_mask (:obj:`torch.FloatTensor`): attention mask of size
+            attention_mask (:obj:`torch.FloatTensor`): attention mask of size
                 `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
-
-        Returns:
-            encoded output of shape `(seq_len, batch, embed_dim)`
+            output_attentions (:obj:`bool`): Whether the base model outputs attentions. This requires the attentions tensor to be reshaped in this function.
         """
         residual = hidden_states
         if self.normalize_before:
             hidden_states = self.self_attn_layer_norm(hidden_states)
         hidden_states, attn_weights, _ = self.self_attn(
-            hidden_states=hidden_states, attn_mask=encoder_padding_mask, output_attentions=output_attentions
+            hidden_states=hidden_states, attention_mask=attention_mask, output_attentions=output_attentions
         )
         hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
         hidden_states = residual + hidden_states
@@ -405,24 +396,35 @@ class BartDecoderLayer(nn.Module):
     def forward(
         self,
         hidden_states: torch.Tensor,
-        encoder_hidden_states: torch.Tensor,
-        encoder_attn_mask: Optional[torch.Tensor] = None,
+        attention_mask: Optional[torch.Tensor] = None,
+        encoder_hidden_states: Optional[torch.Tensor] = None,
+        encoder_attention_mask: Optional[torch.Tensor] = None,
         past_key_value: Optional[Tuple[torch.Tensor]] = None,
-        attn_mask: Optional[torch.Tensor] = None,
         output_attentions: Optional[torch.Tensor] = False,
     ):
+        """
+        Args:
+            hidden_states (:obj:`torch.FloatTensor`): input to the layer of shape `(seq_len, batch, embed_dim)`
+            attention_mask (:obj:`torch.FloatTensor`): attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            encoder_hidden_states (:obj:`torch.FloatTensor`): cross attention input to the layer of shape `(seq_len, batch, embed_dim)`
+            encoder_attention_mask (:obj:`torch.FloatTensor`): encoder attention mask of size
+                `(batch, 1, tgt_len, src_len)` where padding elements are indicated by very large negative values.
+            past_key_value (:obj:`Tuple(torch.FloatTensor)`): cached past key and value projection states
+            output_attentions (:obj:`bool`): Whether the base model outputs attentions. This requires the attentions tensor to be reshaped in this function.
+        """
         residual = hidden_states
         if self.normalize_before:
             hidden_states = self.self_attn_layer_norm(hidden_states)
-        # Self Attention
 
+        # Self Attention
         # decoder uni-directional self-attention cached key/values tuple is at positions 1,2
         self_attn_past_key_value = past_key_value[:2] if past_key_value is not None else None
         # add present self-attn cache to positions 1,2 of present_key_value tuple
         hidden_states, self_attn_weights, present_key_value = self.self_attn(
             hidden_states=hidden_states,
             past_key_value=self_attn_past_key_value,
-            attn_mask=attn_mask,
+            attention_mask=attention_mask,
             output_attentions=output_attentions,
         )
         hidden_states = F.dropout(hidden_states, p=self.dropout, training=self.training)
@@ -443,7 +445,7 @@ class BartDecoderLayer(nn.Module):
             hidden_states, cross_attn_weights, cross_attn_present_key_value = self.encoder_attn(
                 hidden_states=hidden_states,
                 key_value_states=encoder_hidden_states,
-                attn_mask=encoder_attn_mask,
+                attention_mask=encoder_attention_mask,
                 past_key_value=cross_attn_past_key_value,
                 output_attentions=output_attentions,
             )
@@ -905,9 +907,9 @@ class BartDecoder(BartPretrainedModel):
 
         # create causal mask
         # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-        attn_mask = None
+        combined_attention_mask = None
         if input_shape[-1] > 1:
-            attn_mask = _make_causal_mask(
+            combined_attention_mask = _make_causal_mask(
                 input_shape, inputs_embeds.dtype, past_key_values_length=past_key_values_length
             ).to(self.device)
 
@@ -927,11 +929,21 @@ class BartDecoder(BartPretrainedModel):
             attention_mask = input_ids.ne(self.config.pad_token_id).to(torch.long)
             # never mask leading token, even if it is pad
             attention_mask[:, 0] = attention_mask[:, 1]
+            if past_key_values_length > 0:
+                attention_mask = torch.cat(
+                    [
+                        torch.ones(
+                            (input_shape[0], past_key_values_length), dtype=torch.long, device=input_ids.device
+                        ),
+                        attention_mask,
+                    ],
+                    dim=-1,
+                )
 
-        if attention_mask is not None and attn_mask is not None:
+        if attention_mask is not None and combined_attention_mask is not None:
             # [bsz, seq_len] -> [bsz, 1, tgt_seq_len, src_seq_len]
-            attn_mask = attn_mask + _expand_mask(
-                attention_mask, inputs_embeds.dtype, past_key_values_length=past_key_values_length
+            combined_attention_mask = combined_attention_mask + _expand_mask(
+                attention_mask, inputs_embeds.dtype, tgt_len=input_shape[-1]
             )
 
         # expand encoder attention mask
@@ -968,9 +980,9 @@ class BartDecoder(BartPretrainedModel):
 
             hidden_states, layer_self_attn, present_key_value, layer_cross_attn = decoder_layer(
                 hidden_states,
-                encoder_hidden_states,
-                encoder_attn_mask=encoder_attention_mask,
-                attn_mask=attn_mask,
+                attention_mask=combined_attention_mask,
+                encoder_hidden_states=encoder_hidden_states,
+                encoder_attention_mask=encoder_attention_mask,
                 past_key_value=past_key_value,
                 output_attentions=output_attentions,
             )

src/transformers/models/bert/modeling_flax_bert.py

@@ -20,10 +20,11 @@ import numpy as np
 import flax.linen as nn
 import jax
 import jax.numpy as jnp
+from flax.core.frozen_dict import FrozenDict
 from jax.random import PRNGKey
 
 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_flax_utils import FlaxPreTrainedModel, gelu
+from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel
 from ...utils import logging
 from .configuration_bert import BertConfig
 
@@ -205,7 +206,7 @@ class FlaxBertAttention(nn.Module):
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_states, attention_mask, deterministic: bool = True):
         # Attention mask comes in as attention_mask.shape == (*batch_sizes, kv_length)
         # FLAX expects: attention_mask.shape == (*batch_sizes, 1, 1, kv_length) such that it is broadcastable
         # with attn_weights.shape == (*batch_sizes, num_heads, q_length, kv_length)
@@ -219,27 +220,28 @@ class FlaxBertAttention(nn.Module):
             bias_init=jax.nn.initializers.zeros,
             name="self",
             dtype=self.dtype,
-        )(hidden_state, attention_mask)
+        )(hidden_states, attention_mask)
 
-        layer_norm = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_state)
+        layer_norm = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_states)
         return layer_norm
 
 
 class FlaxBertIntermediate(nn.Module):
     output_size: int
+    hidden_act: str = "gelu"
     kernel_init_scale: float = 0.2
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state):
-        # TODO: Add ACT2FN reference to change activation function
-        dense = nn.Dense(
+    def __call__(self, hidden_states):
+        hidden_states = nn.Dense(
             features=self.output_size,
             kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
             name="dense",
             dtype=self.dtype,
-        )(hidden_state)
-        return gelu(dense)
+        )(hidden_states)
+        hidden_states = ACT2FN[self.hidden_act](hidden_states)
+        return hidden_states
 
 
 class FlaxBertOutput(nn.Module):
@@ -249,27 +251,28 @@ class FlaxBertOutput(nn.Module):
 
     @nn.compact
     def __call__(self, intermediate_output, attention_output, deterministic: bool = True):
-        hidden_state = nn.Dense(
+        hidden_states = nn.Dense(
             attention_output.shape[-1],
             kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
             name="dense",
             dtype=self.dtype,
         )(intermediate_output)
-        hidden_state = nn.Dropout(rate=self.dropout_rate)(hidden_state, deterministic=deterministic)
-        hidden_state = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_state + attention_output)
-        return hidden_state
+        hidden_states = nn.Dropout(rate=self.dropout_rate)(hidden_states, deterministic=deterministic)
+        hidden_states = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_states + attention_output)
+        return hidden_states
 
 
 class FlaxBertLayer(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     dropout_rate: float = 0.0
     kernel_init_scale: float = 0.2
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_states, attention_mask, deterministic: bool = True):
         attention = FlaxBertAttention(
             self.num_heads,
             self.head_size,
@@ -277,9 +280,13 @@ class FlaxBertLayer(nn.Module):
             dropout_rate=self.dropout_rate,
             name="attention",
             dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
+        )(hidden_states, attention_mask, deterministic=deterministic)
         intermediate = FlaxBertIntermediate(
-            self.intermediate_size, kernel_init_scale=self.kernel_init_scale, name="intermediate", dtype=self.dtype
+            self.intermediate_size,
+            kernel_init_scale=self.kernel_init_scale,
+            hidden_act=self.hidden_act,
+            name="intermediate",
+            dtype=self.dtype,
         )(attention)
         output = FlaxBertOutput(
             kernel_init_scale=self.kernel_init_scale, dropout_rate=self.dropout_rate, name="output", dtype=self.dtype
@@ -297,6 +304,7 @@ class FlaxBertLayerCollection(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     dropout_rate: float = 0.0
     kernel_init_scale: float = 0.2
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
@@ -316,6 +324,7 @@ class FlaxBertLayerCollection(nn.Module):
                 self.intermediate_size,
                 kernel_init_scale=self.kernel_init_scale,
                 dropout_rate=self.dropout_rate,
+                hidden_act=self.hidden_act,
                 name=f"{i}",
                 dtype=self.dtype,
             )
@@ -328,22 +337,24 @@ class FlaxBertEncoder(nn.Module):
     num_heads: int
     head_size: int
     intermediate_size: int
+    hidden_act: str = "gelu"
     dropout_rate: float = 0.0
     kernel_init_scale: float = 0.2
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_states, attention_mask, deterministic: bool = True):
         layer = FlaxBertLayerCollection(
             self.num_layers,
             self.num_heads,
             self.head_size,
             self.intermediate_size,
+            hidden_act=self.hidden_act,
             kernel_init_scale=self.kernel_init_scale,
             dropout_rate=self.dropout_rate,
             name="layer",
             dtype=self.dtype,
-        )(hidden_state, attention_mask, deterministic=deterministic)
+        )(hidden_states, attention_mask, deterministic=deterministic)
         return layer
 
 
@@ -352,10 +363,10 @@ class FlaxBertPooler(nn.Module):
     dtype: jnp.dtype = jnp.float32  # the dtype of the computation
 
     @nn.compact
-    def __call__(self, hidden_state):
-        cls_token = hidden_state[:, 0]
+    def __call__(self, hidden_states):
+        cls_token = hidden_states[:, 0]
         out = nn.Dense(
-            hidden_state.shape[-1],
+            hidden_states.shape[-1],
             kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
             name="dense",
             dtype=self.dtype,
@@ -363,62 +374,20 @@ class FlaxBertPooler(nn.Module):
         return nn.tanh(out)
 
 
-class FlaxBertModule(nn.Module):
-    vocab_size: int
-    hidden_size: int
-    type_vocab_size: int
-    max_length: int
-    num_encoder_layers: int
-    num_heads: int
-    head_size: int
-    intermediate_size: int
-    dropout_rate: float = 0.0
-    kernel_init_scale: float = 0.2
-    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
-
-    @nn.compact
-    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic: bool = True):
-
-        # Embedding
-        embeddings = FlaxBertEmbeddings(
-            self.vocab_size,
-            self.hidden_size,
-            self.type_vocab_size,
-            self.max_length,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="embeddings",
-            dtype=self.dtype,
-        )(input_ids, token_type_ids, position_ids, attention_mask, deterministic=deterministic)
-
-        # N stacked encoding layers
-        encoder = FlaxBertEncoder(
-            self.num_encoder_layers,
-            self.num_heads,
-            self.head_size,
-            self.intermediate_size,
-            kernel_init_scale=self.kernel_init_scale,
-            dropout_rate=self.dropout_rate,
-            name="encoder",
-            dtype=self.dtype,
-        )(embeddings, attention_mask, deterministic=deterministic)
-
-        pooled = FlaxBertPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
-        return encoder, pooled
-
-
 class FlaxBertPredictionHeadTransform(nn.Module):
+    hidden_act: str = "gelu"
     dtype: jnp.dtype = jnp.float32
 
     @nn.compact
     def __call__(self, hidden_states):
         hidden_states = nn.Dense(hidden_states.shape[-1], name="dense", dtype=self.dtype)(hidden_states)
-        hidden_states = nn.elu(hidden_states)  # TODO: ACT2FN[config.hidden_act]
-        return FlaxBertLayerNorm(name="LayerNorm", dtype=self.dtype)(hidden_states)
+        hidden_states = ACT2FN[self.hidden_act](hidden_states)
+        return FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_states)
 
 
 class FlaxBertLMPredictionHead(nn.Module):
     vocab_size: int
+    hidden_act: str = "gelu"
     dtype: jnp.dtype = jnp.float32
 
     @nn.compact
@@ -428,64 +397,57 @@ class FlaxBertLMPredictionHead(nn.Module):
         #   Need a link between the two variables so that the bias is correctly
         #   resized with `resize_token_embeddings`
 
-        hidden_states = FlaxBertPredictionHeadTransform(name="transform", dtype=self.dtype)(hidden_states)
+        hidden_states = FlaxBertPredictionHeadTransform(
+            name="transform", hidden_act=self.hidden_act, dtype=self.dtype
+        )(hidden_states)
         hidden_states = nn.Dense(self.vocab_size, name="decoder", dtype=self.dtype)(hidden_states)
         return hidden_states
 
 
 class FlaxBertOnlyMLMHead(nn.Module):
     vocab_size: int
-    hidden_size: int
-    intermediate_size: int
-    head_size: int
-    num_heads: int
-    num_encoder_layers: int
-    type_vocab_size: int
-    max_length: int
-    dropout_rate: float = 0.0
+    hidden_act: str = "gelu"
     dtype: jnp.dtype = jnp.float32
 
     @nn.compact
-    def __call__(
-        self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, deterministic: bool = True
-    ):
-        # Model
-        encoder, pooled = FlaxBertModule(
-            vocab_size=self.vocab_size,
-            type_vocab_size=self.type_vocab_size,
-            hidden_size=self.hidden_size,
-            intermediate_size=self.intermediate_size,
-            head_size=self.hidden_size,
-            num_heads=self.num_heads,
-            num_encoder_layers=self.num_encoder_layers,
-            max_length=self.max_length,
-            dropout_rate=self.dropout_rate,
-            dtype=self.dtype,
-        )(input_ids, attention_mask, token_type_ids, position_ids, deterministic=deterministic)
-
-        # Compute the prediction scores
-        encoder = nn.Dropout(rate=self.dropout_rate)(encoder, deterministic=deterministic)
-        logits = FlaxBertLMPredictionHead(vocab_size=self.vocab_size, name="predictions", dtype=self.dtype)(encoder)
-
-        return logits, pooled
+    def __call__(self, hidden_states):
+        hidden_states = FlaxBertLMPredictionHead(
+            vocab_size=self.vocab_size, hidden_act=self.hidden_act, name="predictions", dtype=self.dtype
+        )(hidden_states)
+        return hidden_states
 
 
-@add_start_docstrings(
-    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
-    BERT_START_DOCSTRING,
-)
-class FlaxBertModel(FlaxPreTrainedModel):
+class FlaxBertPreTrainedModel(FlaxPreTrainedModel):
     """
-    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
-    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
-    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
-    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
+    models.
     """
 
-    model_class = FlaxBertModule
     config_class = BertConfig
     base_model_prefix = "bert"
 
+    def _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):
+        if token_type_ids is None:
+            token_type_ids = jnp.ones_like(input_ids)
+
+        if position_ids is None:
+            position_ids = jnp.arange(jnp.atleast_2d(input_ids).shape[-1])
+
+        if attention_mask is None:
+            attention_mask = jnp.ones_like(input_ids)
+
+        return input_ids, attention_mask, token_type_ids, position_ids
+
+    def init(self, rng: jax.random.PRNGKey, input_shape: Tuple) -> FrozenDict:
+        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
+            jnp.zeros(input_shape, dtype="i4"), None, None, None
+        )
+
+        params_rng, dropout_rng = jax.random.split(rng)
+        rngs = {"params": params_rng, "dropout": dropout_rng}
+
+        return self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
+
     @staticmethod
     def convert_from_pytorch(pt_state: Dict, config: BertConfig) -> Dict:
         jax_state = dict(pt_state)
@@ -501,6 +463,11 @@ class FlaxBertModel(FlaxPreTrainedModel):
                 key = key.replace("weight", "kernel")
                 jax_state[key] = tensor
 
+            if "decoder.weight" in key:
+                del jax_state[key]
+                key = key.replace("weight", "kernel")
+                jax_state[key] = tensor.T
+
             # SelfAttention needs also to replace "weight" by "kernel"
             if {"query", "key", "value"} & key_parts:
 
@@ -526,7 +493,7 @@ class FlaxBertModel(FlaxPreTrainedModel):
                 jax_state[key] = tensor
 
             # There are some transposed parameters w.r.t their PyTorch counterpart
-            if "intermediate.dense.kernel" in key or "output.dense.kernel" in key:
+            if "intermediate.dense.kernel" in key or "output.dense.kernel" in key or "transform.dense.kernel" in key:
                 jax_state[key] = tensor.T
 
             # Self Attention output projection needs to be transposed
@@ -539,6 +506,11 @@ class FlaxBertModel(FlaxPreTrainedModel):
             if "pooler.dense.kernel" in key:
                 jax_state[key] = tensor.T
 
+            # Hack to correctly load some pytorch models
+            if "predictions.bias" in key:
+                del jax_state[key]
+                jax_state[".".join(key.split(".")[:2]) + ".decoder.bias"] = tensor
+
             # Handle LayerNorm conversion
             if "LayerNorm" in key:
                 del jax_state[key]
@@ -555,7 +527,22 @@ class FlaxBertModel(FlaxPreTrainedModel):
 
         return jax_state
 
-    def __init__(self, config: BertConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32):
+
+@add_start_docstrings(
+    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
+    BERT_START_DOCSTRING,
+)
+class FlaxBertModel(FlaxBertPreTrainedModel):
+    """
+    The model can behave as an encoder (with only self-attention) as well as a decoder, in which case a layer of
+    cross-attention is added between the self-attention layers, following the architecture described in `Attention is
+    all you need <https://arxiv.org/abs/1706.03762>`__ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit,
+    Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
+    """
+
+    def __init__(
+        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
+    ):
         module = FlaxBertModule(
             vocab_size=config.vocab_size,
             hidden_size=config.hidden_size,
@@ -566,10 +553,12 @@ class FlaxBertModel(FlaxPreTrainedModel):
             head_size=config.hidden_size,
             intermediate_size=config.intermediate_size,
             dropout_rate=config.hidden_dropout_prob,
+            hidden_act=config.hidden_act,
             dtype=dtype,
+            **kwargs,
         )
 
-        super().__init__(config, module, state, seed)
+        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
 
     @add_start_docstrings_to_model_forward(BERT_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
     def __call__(
@@ -601,34 +590,62 @@ class FlaxBertModel(FlaxPreTrainedModel):
             rngs=rngs,
         )
 
-    def _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):
-        if token_type_ids is None:
-            token_type_ids = jnp.ones_like(input_ids)
 
-        if position_ids is None:
-            position_ids = jnp.arange(jnp.atleast_2d(input_ids).shape[-1])
+class FlaxBertModule(nn.Module):
+    vocab_size: int
+    hidden_size: int
+    type_vocab_size: int
+    max_length: int
+    num_encoder_layers: int
+    num_heads: int
+    head_size: int
+    intermediate_size: int
+    hidden_act: str = "gelu"
+    dropout_rate: float = 0.0
+    kernel_init_scale: float = 0.2
+    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
+    add_pooling_layer: bool = True
 
-        if attention_mask is None:
-            attention_mask = jnp.ones_like(input_ids)
+    @nn.compact
+    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic: bool = True):
 
-        return input_ids, attention_mask, token_type_ids, position_ids
+        # Embedding
+        embeddings = FlaxBertEmbeddings(
+            self.vocab_size,
+            self.hidden_size,
+            self.type_vocab_size,
+            self.max_length,
+            kernel_init_scale=self.kernel_init_scale,
+            dropout_rate=self.dropout_rate,
+            name="embeddings",
+            dtype=self.dtype,
+        )(input_ids, token_type_ids, position_ids, attention_mask, deterministic=deterministic)
 
-    def init(self, rng: jax.random.PRNGKey, input_shape: Tuple):
-        input_ids, attention_mask, token_type_ids, position_ids = self._check_inputs(
-            jnp.zeros(input_shape, dtype="i4"), None, None, None
-        )
+        # N stacked encoding layers
+        encoder = FlaxBertEncoder(
+            self.num_encoder_layers,
+            self.num_heads,
+            self.head_size,
+            self.intermediate_size,
+            kernel_init_scale=self.kernel_init_scale,
+            dropout_rate=self.dropout_rate,
+            hidden_act=self.hidden_act,
+            name="encoder",
+            dtype=self.dtype,
+        )(embeddings, attention_mask, deterministic=deterministic)
 
-        params_rng, dropout_rng = jax.random.split(rng)
-        rngs = {"params": params_rng, "dropout": dropout_rng}
+        if not self.add_pooling_layer:
+            return encoder
 
-        self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
+        pooled = FlaxBertPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
+        return encoder, pooled
 
 
-class FlaxBertForMaskedLM(FlaxBertModel):
-    def __init__(self, config: BertConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs):
-        super().__init__(config, state, seed, dtype)
-
-        self._module = FlaxBertOnlyMLMHead(
+class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
+    def __init__(
+        self, config: BertConfig, input_shape: Tuple = (1, 1), seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
+    ):
+        module = FlaxBertForMaskedLMModule(
             vocab_size=config.vocab_size,
             type_vocab_size=config.type_vocab_size,
             hidden_size=config.hidden_size,
@@ -636,10 +653,13 @@ class FlaxBertForMaskedLM(FlaxBertModel):
             head_size=config.hidden_size,
             num_heads=config.num_attention_heads,
             num_encoder_layers=config.num_hidden_layers,
-            max_length=config.max_length,
+            max_length=config.max_position_embeddings,
+            hidden_act=config.hidden_act,
             **kwargs,
         )
 
+        super().__init__(config, module, input_shape=input_shape, seed=seed, dtype=dtype)
+
     def __call__(
         self,
         input_ids,
@@ -659,7 +679,7 @@ class FlaxBertForMaskedLM(FlaxBertModel):
         if dropout_rng is not None:
             rngs["dropout"] = dropout_rng
 
-        pooled, logits = self.module.apply(
+        return self.module.apply(
             {"params": params or self.params},
             jnp.array(input_ids, dtype="i4"),
             jnp.array(attention_mask, dtype="i4"),
@@ -669,4 +689,45 @@ class FlaxBertForMaskedLM(FlaxBertModel):
             rngs=rngs,
         )
 
-        return logits, pooled
+
+class FlaxBertForMaskedLMModule(nn.Module):
+    vocab_size: int
+    hidden_size: int
+    intermediate_size: int
+    head_size: int
+    num_heads: int
+    num_encoder_layers: int
+    type_vocab_size: int
+    max_length: int
+    hidden_act: str
+    dropout_rate: float = 0.0
+    dtype: jnp.dtype = jnp.float32
+
+    @nn.compact
+    def __call__(
+        self, input_ids, attention_mask=None, token_type_ids=None, position_ids=None, deterministic: bool = True
+    ):
+        # Model
+        encoder = FlaxBertModule(
+            vocab_size=self.vocab_size,
+            type_vocab_size=self.type_vocab_size,
+            hidden_size=self.hidden_size,
+            intermediate_size=self.intermediate_size,
+            head_size=self.hidden_size,
+            num_heads=self.num_heads,
+            num_encoder_layers=self.num_encoder_layers,
+            max_length=self.max_length,
+            dropout_rate=self.dropout_rate,
+            hidden_act=self.hidden_act,
+            dtype=self.dtype,
+            add_pooling_layer=False,
+            name="bert",
+        )(input_ids, attention_mask, token_type_ids, position_ids, deterministic=deterministic)
+
+        # Compute the prediction scores
+        encoder = nn.Dropout(rate=self.dropout_rate)(encoder, deterministic=deterministic)
+        logits = FlaxBertOnlyMLMHead(
+            vocab_size=self.vocab_size, hidden_act=self.hidden_act, name="cls", dtype=self.dtype
+        )(encoder)
+
+        return (logits,)

src/transformers/models/bert/modeling_tf_bert.py

@@ -246,52 +246,52 @@ class TFBertSelfAttention(tf.keras.layers.Layer):
 
         if config.hidden_size % config.num_attention_heads != 0:
             raise ValueError(
-                "The hidden size (%d) is not a multiple of the number of attention "
-                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number "
+                f"of attention heads ({config.num_attention_heads})"
             )
 
         self.num_attention_heads = config.num_attention_heads
-        assert config.hidden_size % config.num_attention_heads == 0
         self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
-        self.all_head_size = self.num_attention_heads * self.attention_head_size
-        self.query = tf.keras.layers.Dense(
-            self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="query"
+        self.query = tf.keras.layers.experimental.EinsumDense(
+            equation="abc,cde->abde",
+            output_shape=(None, config.num_attention_heads, self.attention_head_size),
+            bias_axes="de",
+            kernel_initializer=get_initializer(initializer_range=config.initializer_range),
+            name="query",
         )
-        self.key = tf.keras.layers.Dense(
-            self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="key"
+        self.key = tf.keras.layers.experimental.EinsumDense(
+            equation="abc,cde->abde",
+            output_shape=(None, config.num_attention_heads, self.attention_head_size),
+            bias_axes="de",
+            kernel_initializer=get_initializer(initializer_range=config.initializer_range),
+            name="key",
         )
-        self.value = tf.keras.layers.Dense(
-            self.all_head_size, kernel_initializer=get_initializer(config.initializer_range), name="value"
+        self.value = tf.keras.layers.experimental.EinsumDense(
+            equation="abc,cde->abde",
+            output_shape=(None, config.num_attention_heads, self.attention_head_size),
+            bias_axes="de",
+            kernel_initializer=get_initializer(initializer_range=config.initializer_range),
+            name="value",
         )
-        self.dropout = tf.keras.layers.Dropout(config.attention_probs_dropout_prob)
+        self.dropout = tf.keras.layers.Dropout(rate=config.attention_probs_dropout_prob)
 
-    def transpose_for_scores(self, x, batch_size):
-        x = tf.reshape(x, (batch_size, -1, self.num_attention_heads, self.attention_head_size))
+    def call(self, hidden_states, attention_mask=None, head_mask=None, output_attentions=False, training=False):
+        query_layer = self.query(inputs=hidden_states)
+        key_layer = self.key(inputs=hidden_states)
+        value_layer = self.value(inputs=hidden_states)
 
-        return tf.transpose(x, perm=[0, 2, 1, 3])
-
-    def call(self, hidden_states, attention_mask, head_mask, output_attentions, training=False):
-        batch_size = shape_list(hidden_states)[0]
-        mixed_query_layer = self.query(hidden_states)
-        mixed_key_layer = self.key(hidden_states)
-        mixed_value_layer = self.value(hidden_states)
-        query_layer = self.transpose_for_scores(mixed_query_layer, batch_size)
-        key_layer = self.transpose_for_scores(mixed_key_layer, batch_size)
-        value_layer = self.transpose_for_scores(mixed_value_layer, batch_size)
-
-        # Take the dot product between "query" and "key" to get the raw attention scores.
-        attention_scores = tf.matmul(
-            query_layer, key_layer, transpose_b=True
-        )  # (batch size, num_heads, seq_len_q, seq_len_k)
-        dk = tf.cast(shape_list(key_layer)[-1], attention_scores.dtype)  # scale attention_scores
-        attention_scores = attention_scores / tf.math.sqrt(dk)
+        # Take the dot product between "query" and "key" to get the raw
+        # attention scores.
+        dk = tf.cast(x=self.attention_head_size, dtype=query_layer.dtype)
+        query_layer = tf.multiply(x=query_layer, y=tf.math.rsqrt(x=dk))
+        attention_scores = tf.einsum("aecd,abcd->acbe", key_layer, query_layer)
 
         if attention_mask is not None:
             # Apply the attention mask is (precomputed for all layers in TFBertModel call() function)
             attention_scores = attention_scores + attention_mask
 
         # Normalize the attention scores to probabilities.
-        attention_probs = tf.nn.softmax(attention_scores, axis=-1)
+        attention_probs = tf.nn.softmax(logits=attention_scores, axis=-1)
 
         # This is actually dropping out entire tokens to attend to, which might
         # seem a bit unusual, but is taken from the original Transformer paper.
@@ -299,14 +299,10 @@ class TFBertSelfAttention(tf.keras.layers.Layer):
 
         # Mask heads if we want to
         if head_mask is not None:
-            attention_probs = attention_probs * head_mask
+            attention_scores = attention_scores * head_mask
 
-        context_layer = tf.matmul(attention_probs, value_layer)
-        context_layer = tf.transpose(context_layer, perm=[0, 2, 1, 3])
-        context_layer = tf.reshape(
-            context_layer, (batch_size, -1, self.all_head_size)
-        )  # (batch_size, seq_len_q, all_head_size)
-        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+        attention_output = tf.einsum("acbe,aecd->abcd", attention_probs, value_layer)
+        outputs = (attention_output, attention_probs) if output_attentions else (attention_output,)
 
         return outputs
 
@@ -315,16 +311,29 @@ class TFBertSelfOutput(tf.keras.layers.Layer):
     def __init__(self, config, **kwargs):
         super().__init__(**kwargs)
 
-        self.dense = tf.keras.layers.Dense(
-            config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
+        if config.hidden_size % config.num_attention_heads != 0:
+            raise ValueError(
+                f"The hidden size ({config.hidden_size}) is not a multiple of the number "
+                f"of attention heads ({config.num_attention_heads})"
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+        self.dense = tf.keras.layers.experimental.EinsumDense(
+            equation="abcd,cde->abe",
+            output_shape=(None, self.all_head_size),
+            bias_axes="e",
+            kernel_initializer=get_initializer(initializer_range=config.initializer_range),
+            name="dense",
         )
         self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
-        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
+        self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
 
     def call(self, hidden_states, input_tensor, training=False):
-        hidden_states = self.dense(hidden_states)
-        hidden_states = self.dropout(hidden_states, training=training)
-        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        hidden_states = self.dense(inputs=hidden_states)
+        hidden_states = self.dropout(inputs=hidden_states, training=training)
+        hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
 
         return hidden_states
 
@@ -353,18 +362,22 @@ class TFBertIntermediate(tf.keras.layers.Layer):
     def __init__(self, config, **kwargs):
         super().__init__(**kwargs)
 
-        self.dense = tf.keras.layers.Dense(
-            config.intermediate_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
+        self.dense = tf.keras.layers.experimental.EinsumDense(
+            equation="abc,cd->abd",
+            output_shape=(None, config.intermediate_size),
+            bias_axes="d",
+            kernel_initializer=get_initializer(initializer_range=config.initializer_range),
+            name="dense",
         )
 
         if isinstance(config.hidden_act, str):
-            self.intermediate_act_fn = get_tf_activation(config.hidden_act)
+            self.intermediate_act_fn = get_tf_activation(activation_string=config.hidden_act)
         else:
             self.intermediate_act_fn = config.hidden_act
 
     def call(self, hidden_states):
-        hidden_states = self.dense(hidden_states)
-        hidden_states = self.intermediate_act_fn(hidden_states)
+        hidden_states = self.dense(inputs=hidden_states)
+        hidden_states = self.intermediate_act_fn(inputs=hidden_states)
 
         return hidden_states
 
@@ -373,16 +386,20 @@ class TFBertOutput(tf.keras.layers.Layer):
     def __init__(self, config, **kwargs):
         super().__init__(**kwargs)
 
-        self.dense = tf.keras.layers.Dense(
-            config.hidden_size, kernel_initializer=get_initializer(config.initializer_range), name="dense"
+        self.dense = tf.keras.layers.experimental.EinsumDense(
+            equation="abc,cd->abd",
+            bias_axes="d",
+            output_shape=(None, config.hidden_size),
+            kernel_initializer=get_initializer(config.initializer_range),
+            name="dense",
         )
         self.LayerNorm = tf.keras.layers.LayerNormalization(epsilon=config.layer_norm_eps, name="LayerNorm")
-        self.dropout = tf.keras.layers.Dropout(config.hidden_dropout_prob)
+        self.dropout = tf.keras.layers.Dropout(rate=config.hidden_dropout_prob)
 
     def call(self, hidden_states, input_tensor, training=False):
-        hidden_states = self.dense(hidden_states)
-        hidden_states = self.dropout(hidden_states, training=training)
-        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        hidden_states = self.dense(inputs=hidden_states)
+        hidden_states = self.dropout(inputs=hidden_states, training=training)
+        hidden_states = self.LayerNorm(inputs=hidden_states + input_tensor)
 
         return hidden_states
 

src/transformers/models/blenderbot/__init__.py

@@ -22,7 +22,11 @@ from .tokenization_blenderbot import BlenderbotSmallTokenizer, BlenderbotTokeniz
 
 
 if is_torch_available():
-    from .modeling_blenderbot import BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST, BlenderbotForConditionalGeneration
+    from .modeling_blenderbot import (
+        BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST,
+        BlenderbotForConditionalGeneration,
+        BlenderbotModel,
+    )
 
 if is_tf_available():
     from .modeling_tf_blenderbot import TFBlenderbotForConditionalGeneration

src/transformers/models/blenderbot/modeling_blenderbot.py

@@ -19,7 +19,7 @@
 import torch
 
 from ...file_utils import add_start_docstrings
-from ..bart.modeling_bart import BartForConditionalGeneration
+from ..bart.modeling_bart import BartForConditionalGeneration, BartModel
 from .configuration_blenderbot import BlenderbotConfig
 
 
@@ -39,7 +39,20 @@ BLENDERBOT_PRETRAINED_MODEL_ARCHIVE_LIST = ["facebook/blenderbot-3B", "facebook/
 
 
 @add_start_docstrings(
-    "The BART Model with a language modeling head. Can be used for summarization.", BLENDER_START_DOCSTRING
+    "The bare BlenderBot Model transformer outputting raw hidden-states without any specific head on top.",
+    BLENDER_START_DOCSTRING,
+)
+class BlenderbotModel(BartModel):
+    r"""
+    This class overrides :class:`~transformers.BartModel`. Please check the superclass for the appropriate
+    documentation alongside usage examples.
+    """
+
+    config_class = BlenderbotConfig
+
+
+@add_start_docstrings(
+    "The BlenderBot Model with a language modeling head. Can be used for summarization.", BLENDER_START_DOCSTRING
 )
 class BlenderbotForConditionalGeneration(BartForConditionalGeneration):
     """

src/transformers/models/ctrl/__init__.py

@@ -33,6 +33,7 @@ if is_torch_available():
 if is_tf_available():
     from .modeling_tf_ctrl import (
         TF_CTRL_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TFCTRLForSequenceClassification,
         TFCTRLLMHeadModel,
         TFCTRLModel,
         TFCTRLPreTrainedModel,

src/transformers/models/ctrl/modeling_tf_ctrl.py

@@ -19,11 +19,13 @@ import numpy as np
 import tensorflow as tf
 
 from ...file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_tf_outputs import TFBaseModelOutputWithPast, TFCausalLMOutputWithPast
+from ...modeling_tf_outputs import TFBaseModelOutputWithPast, TFCausalLMOutputWithPast, TFSequenceClassifierOutput
 from ...modeling_tf_utils import (
     TFCausalLanguageModelingLoss,
     TFPreTrainedModel,
+    TFSequenceClassificationLoss,
     TFSharedEmbeddings,
+    get_initializer,
     input_processing,
     keras_serializable,
     shape_list,
@@ -726,3 +728,160 @@ class TFCTRLLMHeadModel(TFCTRLPreTrainedModel, TFCausalLanguageModelingLoss):
             hidden_states=transformer_outputs.hidden_states,
             attentions=transformer_outputs.attentions,
         )
+
+
+@add_start_docstrings(
+    """
+    The CTRL Model transformer with a sequence classification head on top (linear layer).
+
+    :class:`~transformers.TFCTRLForSequenceClassification` uses the last token in order to do the classification, as
+    other causal models (e.g. GPT-1, GPT-2) do.
+
+    Since it does classification on the last token, it requires to know the position of the last token. If a
+    :obj:`pad_token_id` is defined in the configuration, it finds the last token that is not a padding token in each
+    row. If no :obj:`pad_token_id` is defined, it simply takes the last value in each row of the batch. Since it cannot
+    guess the padding tokens when :obj:`inputs_embeds` are passed instead of :obj:`input_ids`, it does the same (take
+    the last value in each row of the batch).
+    """,
+    CTRL_START_DOCSTRING,
+)
+class TFCTRLForSequenceClassification(TFCTRLPreTrainedModel, TFSequenceClassificationLoss):
+    def __init__(self, config, *inputs, **kwargs):
+        super().__init__(config, *inputs, **kwargs)
+        self.num_labels = config.num_labels
+        self.classifier = tf.keras.layers.Dense(
+            config.num_labels,
+            kernel_initializer=get_initializer(config.initializer_range),
+            name="classifier",
+            use_bias=False,
+        )
+        self.transformer = TFCTRLMainLayer(config, name="transformer")
+
+    def get_output_embeddings(self):
+        return self.transformer.w
+
+    @add_start_docstrings_to_model_forward(CTRL_INPUTS_DOCSTRING)
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="ctrl",
+        output_type=TFSequenceClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def call(
+        self,
+        input_ids=None,
+        past=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        use_cache=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+        labels=None,
+        training=False,
+        **kwargs,
+    ):
+        r"""
+        labels (:obj:`tf.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Labels for computing the cross entropy classification loss. Indices should be in ``[0, ...,
+            config.vocab_size - 1]``.
+        """
+        inputs = input_processing(
+            func=self.call,
+            config=self.config,
+            input_ids=input_ids,
+            past=past,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            use_cache=use_cache,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+            labels=labels,
+            training=training,
+            kwargs_call=kwargs,
+        )
+
+        transformer_outputs = self.transformer(
+            input_ids=inputs["input_ids"],
+            past=inputs["past"],
+            attention_mask=inputs["attention_mask"],
+            token_type_ids=inputs["token_type_ids"],
+            position_ids=inputs["position_ids"],
+            head_mask=inputs["head_mask"],
+            inputs_embeds=inputs["inputs_embeds"],
+            use_cache=inputs["use_cache"],
+            output_attentions=inputs["output_attentions"],
+            output_hidden_states=inputs["output_hidden_states"],
+            return_dict=inputs["return_dict"],
+            training=inputs["training"],
+        )
+
+        hidden_states = transformer_outputs[0]
+        logits = self.classifier(hidden_states)
+        logits_shape = shape_list(logits)
+        in_logits = None
+        if self.config.pad_token_id is None:
+            sequence_lengths = -1
+        else:
+            if inputs["input_ids"] is not None:
+                sequence_lengths = (
+                    tf.reduce_sum(
+                        tf.cast(tf.math.not_equal(inputs["input_ids"], self.config.pad_token_id), tf.int32),
+                        -1,
+                        keepdims=False,
+                    )
+                    - 1
+                )
+
+                def get_seq_element(sequence_position, input_batch):
+                    return tf.strided_slice(
+                        input_batch, [sequence_position, 0], [sequence_position + 1, input_batch.shape[-1]], [1, 1]
+                    )
+
+                result = tf.map_fn(
+                    fn=lambda t: get_seq_element(t[0], t[1]), elems=[sequence_lengths, logits], dtype="float"
+                )
+                in_logits = tf.reshape(result, [logits_shape[0], logits_shape[-1]])
+            else:
+                sequence_lengths = -1
+                logger.warning(
+                    f"{self.__class__.__name__} will not detect padding tokens in `inputs_embeds`. Results may be "
+                    f"unexpected if using padding tokens in conjunction with `inputs_embeds.`"
+                )
+        loss = None
+
+        if inputs["labels"] is not None:
+            if input_ids is not None:
+                batch_size, sequence_length = shape_list(inputs["input_ids"])[:2]
+            else:
+                batch_size, sequence_length = shape_list(inputs["inputs_embeds"])[:2]
+            assert (
+                self.config.pad_token_id is not None or batch_size == 1
+            ), "Cannot handle batch sizes > 1 if no padding token is defined."
+
+            if not tf.is_tensor(sequence_lengths):
+                in_logits = logits[0:batch_size, sequence_lengths]
+
+            loss = self.compute_loss(
+                tf.reshape(inputs["labels"], [-1, 1]), tf.reshape(in_logits, [-1, self.num_labels])
+            )
+
+        pooled_logits = in_logits if in_logits is not None else logits
+
+        if not inputs["return_dict"]:
+            output = (pooled_logits,) + transformer_outputs[1:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TFSequenceClassifierOutput(
+            loss=loss,
+            logits=pooled_logits,
+            hidden_states=transformer_outputs.hidden_states,
+            attentions=transformer_outputs.attentions,
+        )

src/transformers/models/distilbert/modeling_distilbert.py

@@ -627,7 +627,7 @@ class DistilBertForSequenceClassification(DistilBertPreTrainedModel):
         pooled_output = self.pre_classifier(pooled_output)  # (bs, dim)
         pooled_output = nn.ReLU()(pooled_output)  # (bs, dim)
         pooled_output = self.dropout(pooled_output)  # (bs, dim)
-        logits = self.classifier(pooled_output)  # (bs, dim)
+        logits = self.classifier(pooled_output)  # (bs, num_labels)
 
         loss = None
         if labels is not None: