Truncate when overflowing

Fix the masked LM model
Test the Masked LM model
2020-12-10 17:13:02 -05:00 · 2020-12-10 15:25:33 -05:00 · 2020-12-10 15:18:02 -05:00 · 2020-12-10 14:55:17 -05:00 · 2020-12-09 16:31:08 -05:00 · 2020-12-09 21:18:21 +01:00
55 changed files with 11101 additions and 1174 deletions
@@ -11,40 +11,21 @@ commands:
      - run:
          name: docs-only changes skip check
          command: |
-            if test -n "$CIRCLE_PR_NUMBER"
+            # pipeline.git.base_revision is not always defined, so only proceed if all external vars are defined
+            if test -n "<< pipeline.git.base_revision >>" && test -n "<< pipeline.git.revision >>" && test -n "$(git diff --name-only << pipeline.git.base_revision >>...<< pipeline.git.revision >>)"
            then
-                echo $CIRCLE_PR_NUMBER
-                resp=$(curl -Ls https://api.github.com/repos/huggingface/transformers/pulls/${CIRCLE_PR_NUMBER})
-                user=$(jq -r .user.login \<<< $resp)   # PR creator username
-                head_ref=$(jq -r .head.ref \<<< $resp) # PR user's branch name
-                echo head_ref=$head_ref, user=$user
-            fi
-            
-            if test -n "$user" && test -n "$head_ref"
-            then
-                git clone https://github.com/$user/transformers user-clone
-                cd user-clone
-                git checkout $head_ref
-                fork_point_sha=$(git merge-base --fork-point master)
-                cd -
-            fi
-
-            if test -n "$fork_point_sha" && test -n "$(git diff --name-only $fork_point_sha)"
-            then
-                git --no-pager diff --name-only $fork_point_sha
-                if git diff --name-only $fork_point_sha | egrep -qv '\.(md|rst)$'
+                if git diff --name-only << pipeline.git.base_revision >>...<< pipeline.git.revision >> | egrep -qv '\.(md|rst)$'
                then
                    echo "Non-docs were modified in this PR, proceeding normally"
                else
                    echo "Only docs were modified in this PR, quitting this job"
-                    # enable skipping once we get this sorted out
+                    # disable skipping for now, as circleCI's base_revision is inconsistent leading to invalid ranges
                    # circleci step halt
                fi
            else
-                echo "Not enough data to perform a skipping check - continuing the job"
+                echo "Can't perform skipping check w/o base_revision defined, continuing the job"
            fi

-
 # TPU REFERENCES
 references:
    checkout_ml_testing: &checkout_ml_testing
@@ -121,6 +102,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:
@@ -148,6 +130,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:
@@ -229,6 +212,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:
@@ -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)** 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.
@@ -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>`__ 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
@@ -32,9 +32,6 @@ generation.
 .. autoclass:: transformers.LogitsProcessorList
    :members: __call__

-.. autoclass:: transformers.LogitsWarper
-    :members: __call__
-
 .. autoclass:: transformers.MinLengthLogitsProcessor
    :members: __call__

@@ -70,10 +67,3 @@ BeamSearch

 .. autoclass:: transformers.BeamSearchScorer
    :members: process, finalize
-
-Utilities
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autofunction:: transformers.top_k_top_p_filtering
-
-.. autofunction:: transformers.tf_top_k_top_p_filtering
@@ -22,8 +22,6 @@ Utilities

 .. autoclass:: transformers.EvalPrediction

-.. autoclass:: transformers.EvaluationStrategy
-
 .. autofunction:: transformers.set_seed

 .. autofunction:: transformers.torch_distributed_zero_first
@@ -34,15 +32,8 @@ Callbacks internals

 .. autoclass:: transformers.trainer_callback.CallbackHandler

-
 Distributed Evaluation
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.trainer_pt_utils.DistributedTensorGatherer
    :members:
-
-
-Distributed Evaluation
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.HfArgumentParser
@@ -74,10 +74,6 @@ Learning Rate Schedules (Pytorch)
    :target: /imgs/warmup_linear_schedule.png
    :alt:

-
-.. autofunction:: transformers.get_polynomial_decay_schedule_with_warmup
-
-
 Warmup (TensorFlow)
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

@@ -73,9 +73,8 @@ FillMaskPipeline
 NerPipeline
 =======================================================================================================================

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

 QuestionAnsweringPipeline
 =======================================================================================================================
@@ -119,13 +118,6 @@ TokenClassificationPipeline
    :special-members: __call__
    :members:

-TranslationPipeline
-=======================================================================================================================
-
-.. autoclass:: transformers.TranslationPipeline
-    :special-members: __call__
-    :members:
-
 ZeroShotClassificationPipeline
 =======================================================================================================================

@@ -60,13 +60,6 @@ AlbertTokenizer
        create_token_type_ids_from_sequences, save_vocabulary


-AlbertTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.AlbertTokenizerFast
-    :members:
-
-
 Albert specific outputs
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -175,10 +175,3 @@ TFAutoModelForQuestionAnswering

 .. autoclass:: transformers.TFAutoModelForQuestionAnswering
    :members:
-
-
-FlaxAutoModel
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FlaxAutoModel
-    :members:
@@ -98,12 +98,6 @@ BartTokenizer
    :members:


-BartTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.BartTokenizerFast
-    :members:
-

 BartModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -51,9 +51,3 @@ BarthezTokenizer
 .. autoclass:: transformers.BarthezTokenizer
    :members:

-
-BarthezTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.BarthezTokenizerFast
-    :members:
@@ -207,10 +207,3 @@ FlaxBertModel

 .. autoclass:: transformers.FlaxBertModel
    :members: __call__
-
-
-FlaxBertForMaskedLM
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FlaxBertForMaskedLM
-    :members: __call__
@@ -54,13 +54,6 @@ CamembertTokenizer
        create_token_type_ids_from_sequences, save_vocabulary


-CamembertTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.CamembertTokenizerFast
-    :members:
-
-
 CamembertModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -57,13 +57,6 @@ LayoutLMTokenizer
    :members:


-LayoutLMTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.LayoutLMTokenizerFast
-    :members:
-
-
 LayoutLMModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -90,13 +90,6 @@ MBartTokenizer
    :members: build_inputs_with_special_tokens, prepare_seq2seq_batch


-MBartTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.MBartTokenizerFast
-    :members:
-
-
 MBartForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -1,15 +1,3 @@
-.. 
-    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.
-
 MPNet
 -----------------------------------------------------------------------------------------------------------------------

@@ -37,22 +37,6 @@ MT5Config
    :members:


-MT5Tokenizer
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.MT5Tokenizer
-
-See :class:`~transformers.T5Tokenizer` for all details.
-
-
-MT5TokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.MT5TokenizerFast
-
-See :class:`~transformers.T5TokenizerFast` for all details.
-
-
 MT5Model
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -112,13 +112,6 @@ warning: ``add_tokens`` does not work at the moment.
    :members: __call__, prepare_seq2seq_batch


-PegasusTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.PegasusTokenizerFast
-    :members:
-
-
 PegasusForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -163,13 +163,6 @@ ReformerTokenizer
    :members: save_vocabulary


-ReformerTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.ReformerTokenizerFast
-    :members:
-
-
 ReformerModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -107,13 +107,6 @@ T5Tokenizer
        create_token_type_ids_from_sequences, prepare_seq2seq_batch, save_vocabulary


-T5TokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.T5TokenizerFast
-    :members:
-
-
 T5Model
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -0,0 +1,395 @@
+TAPAS
+-----------------------------------------------------------------------------------------------------------------------
+
+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 (Sequential Question Answering by
+Microsoft), WTQ (Wiki Table Questions by Stanford University) and 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, 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://arxiv.org/abs/2010.00571>`__ by Julian Martin Eisenschlos, Syrine Krichene and
+Thomas Müller.
+
+The original code can be found `here <https://github.com/google-research/tapas>`__. This page also includes links to
+the datasets (SQA, WTQ, WikiSQL and TabFact).
+
+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 latest 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 a certain version when calling the `.from_pretrained`
+  method as explained in the model cards. 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). Metrics 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 training example one by one to the model, such that the `prev_labels` token type ids can
+  be overwritten by the predicted `label_ids` 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 datasets** | **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 is as
+easy as:
+
+.. code-block::
+
+        >>> from transformers import TapasForQuestionAnswering
+
+        >>> # for example, the base sized model 
+        >>> model = TapasForQuestionAnswering.from_pretrained('google/tapas-base-uncased')
+
+
+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-uncased', 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 :class:`~transformers.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``, ``label_ids``                         |
+------------------------------------+----------------------------------------------------------------------------------------------+
+| Weak supervision for aggregation   | ``input_ids``, ``attention_mask``, ``token_type_ids``, ``label_ids``, ``numeric_values``,    |
+|                                    | ``numeric_values_scale``, ``float_answer``                                                   |
+------------------------------------+----------------------------------------------------------------------------------------------+
+| Strong supervision for aggregation | ``input ids``, ``attention mask``, ``token type ids``, ``label ids``, ``aggregation_labels`` |
+------------------------------------+----------------------------------------------------------------------------------------------+
+
+:class:`~transformers.TapasTokenizer` creates the ``label_ids``, ``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-uncased'
+        >>> 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)], [(1, 0)], [(0, 2), (1, 2), (2, 2)]]
+        >>> answer_text = [["Brad Pitt"], ["69"], ["209"]]
+        >>> table = pd.Dataframe(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([[ ... ]]), label_ids: 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)
+        ...         encoding = self.tokenizer(table=table, 
+        ...                                 queries=item.question, 
+        ...                                 answer_coordinates=item.answer_coordinates, 
+        ...                                 answer_text=item.answer_text,
+        ...                                 padding="max_length",
+        ...                                 return_tensors="pt"
+        ...         )
+        ...         # we add the float_answer which is also required (weak supervision for aggregation)
+        ...         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.
+
+=================================================== STEP 4: Train (fine-tune) TapasForQuestionAnswering
+=======================================================================================================================
+
+You can then fine-tune :class:`~transformers.TapasForQuestionAnswering` using native PyTorch as follows:
+
+.. code-block::
+
+        >>> from transformers import TapasForQuestionAnswering
+
+        >>> model = TapasForQuestionAnswering.from_pretrained("google/tapas-base-uncased")
+
+        >>> for epoch in range(2):  # loop over the dataset multiple times
+        ...    for idx, batch in enumerate(train_dataloader):
+        ...         # get the inputs; 
+        ...         input_ids, attention_mask, token_type_ids, label_ids, numeric_values, numeric_values_scale, float_answer = batch
+
+        ...         # 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, 
+        ...                        label_ids=label_ids, 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-uncased-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(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, 
+        ...         output.logits, 
+        ...         outputs.logits_aggregation
+        ...)
+
+        >>> # 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(df.iat[coordinates[0]])
+        ...   else:
+        ...     # multiple cells
+        ...     cell_values = []
+        ...     for coordinate in coordinates:
+        ...        cell_values.append(df.iat[coordinate])
+        ...     answers.append(", ".join(cell_values))
+
+        >>> display(df)
+        >>> 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)    
+        When was Brad Pitt born?
+        Predicted answer: 18 december 1963
+        Which actor appeared in the least number of movies?
+        Predicted answer: Leonardo Di Caprio
+        What is the average number of movies?
+        Predicted answer: AVERAGE > 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 ``label_ids`` of the previous table-question
+pair.
+
+
+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:
+
+
+TapasForMaskedLM
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasForMaskedLM
+    :members:
+
+
+TapasForSequenceClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasForSequenceClassification
+    :members: forward
+
+
+TapasForQuestionAnswering
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TapasForQuestionAnswering
+    :members:
@@ -105,11 +105,3 @@ TFTransfoXLLMHeadModel

 .. autoclass:: transformers.TFTransfoXLLMHeadModel
    :members: call
-
-
-Internal Layers
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.AdaptiveEmbedding
-
-.. autoclass:: transformers.TFAdaptiveEmbedding
@@ -62,13 +62,6 @@ XLMRobertaTokenizer
        create_token_type_ids_from_sequences, save_vocabulary


-XLMRobertaTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.XLMRobertaTokenizerFast
-    :members:
-
-
 XLMRobertaModel
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -62,13 +62,6 @@ XLNetTokenizer
        create_token_type_ids_from_sequences, save_vocabulary


-XLNetTokenizerFast
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.XLNetTokenizerFast
-    :members:
-
-
 XLNet specific outputs
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -134,15 +134,6 @@ 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"
-        },
-    )
-    test_split_percentage: Optional[int] = field(
-        default=5, metadata={"help": "The percentage of the train set used as test set in case there's no test split"}
-    )
    max_seq_length: Optional[int] = field(
        default=None,
        metadata={
@@ -388,7 +379,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)

-        logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)[0]
+        pooled, logits = model(**batch, params=params, dropout_rng=dropout_rng, train=True)
        loss, weight_sum = cross_entropy(logits, targets, token_mask)
        return loss / weight_sum

@@ -410,7 +401,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)

    return compute_metrics(logits, targets, token_mask)

@@ -422,7 +413,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 = np.split(samples_idx, sections_split)
+    batch_idx = jnp.split(samples_idx, sections_split)
    return batch_idx


@@ -472,9 +463,7 @@ if __name__ == "__main__":

    # 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). If you choose this option, and the dataset
-    # does not have validation or test splits, this script will automatically partition the train set to create the
-    # missing splits. The arguments `validation_split_percentage` and `test_split_percentage` set the ratios.
+    # (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).
@@ -484,25 +473,6 @@ 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)
-        offset = 0
-        if "validation" not in datasets.keys():
-            datasets["validation"] = load_dataset(
-                data_args.dataset_name,
-                data_args.dataset_config_name,
-                split=f"train[{offset}%:{offset + data_args.validation_split_percentage}%]",
-            )
-            offset += data_args.validation_split_percentage
-        if "test" not in datasets.keys():
-            datasets["test"] = load_dataset(
-                data_args.dataset_name,
-                data_args.dataset_config_name,
-                split=f"train[{offset}%:{offset + data_args.test_split_percentage}%]",
-            )
-            offset += data_args.test_split_percentage
-        if offset >= 0:
-            datasets["train"] = load_dataset(
-                data_args.dataset_name, data_args.dataset_config_name, split=f"train[{offset}%:100%]"
-            )
    else:
        data_files = {}
        if data_args.train_file is not None:
@@ -521,21 +491,13 @@ if __name__ == "__main__":
    # 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())
-
-    if model_args.model_name_or_path:
-        model = FlaxBertForMaskedLM.from_pretrained(model_args.model_name_or_path, dtype=jnp.float32, dropout_rate=0.1)
-    elif model_args.config_name:
+    if model_args.config_name:
        config = AutoConfig.from_pretrained(model_args.config_name, cache_dir=model_args.cache_dir)
-        model = FlaxBertForMaskedLM(config)
-        model.init(jax.random.PRNGKey(training_args.seed), (training_args.train_batch_size, model.config.max_length))
+    elif model_args.model_name_or_path:
+        config = AutoConfig.from_pretrained(model_args.model_name_or_path, cache_dir=model_args.cache_dir)
    else:
        config = CONFIG_MAPPING[model_args.model_type]()
        logger.warning("You are instantiating a new config instance from scratch.")
-        model = FlaxBertForMaskedLM(config)
-        model.init(jax.random.PRNGKey(training_args.seed), (training_args.train_batch_size, model.config.max_length))

    if model_args.tokenizer_name:
        tokenizer = AutoTokenizer.from_pretrained(
@@ -558,15 +520,14 @@ if __name__ == "__main__":
    else:
        column_names = datasets["validation"].column_names
    text_column_name = "text" if "text" in column_names else column_names[0]
-    columns_to_remove = [column_name for column_name in column_names if column_name != text_column_name]

    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()]
+        examples["text"] = [line for line in examples["text"] if len(line) > 0 and not line.isspace()]
        return tokenizer(
-            examples,
+            examples["text"],
            return_special_tokens_mask=True,
            padding=padding,
            truncation=True,
@@ -575,10 +536,9 @@ 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=column_names,
+        remove_columns=[text_column_name],
        load_from_cache_file=not data_args.overwrite_cache,
    )

@@ -588,10 +548,6 @@ if __name__ == "__main__":

    # Data collator
    # This one will take care of randomly masking the tokens.
-    if training_args.tokenizers_parallelism:
-        os.environ["TOKENIZERS_PARALLELISM"] = "true"
-    else:
-        os.environ["TOKENIZERS_PARALLELISM"] = "false"
    data_collator = FlaxDataCollatorForLanguageModeling(tokenizer=tokenizer, mlm_probability=data_args.mlm_probability)

    # Initialize our training
@@ -610,8 +566,6 @@ if __name__ == "__main__":
    ).create(model.params)

    # Create learning rate scheduler
-    if training_args.warmup_steps == 0:
-        training_args.warmup_steps = 1
    lr_scheduler_fn = create_learning_rate_scheduler(
        base_learning_rate=training_args.learning_rate, warmup_steps=training_args.warmup_steps
    )
@@ -281,7 +281,7 @@ def main():
    # Add the chinese references if provided
    if data_args.train_ref_file is not None:
        tokenized_datasets["train"] = add_chinese_references(tokenized_datasets["train"], data_args.train_ref_file)
-    if data_args.validation_ref_file is not None:
+    if data_args.valid_ref_file is not None:
        tokenized_datasets["validation"] = add_chinese_references(
            tokenized_datasets["validation"], data_args.validation_ref_file
        )
@@ -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}
+}
+```
@@ -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,
@@ -605,6 +606,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,
@@ -216,6 +216,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,7 +270,6 @@ 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"

@@ -326,6 +348,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:
@@ -428,6 +458,13 @@ 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.
+"""
+
+
 def requires_datasets(obj):
    name = obj.__name__ if hasattr(obj, "__name__") else obj.__class__.__name__
    if not is_datasets_available():
@@ -482,6 +519,12 @@ def requires_protobuf(obj):
        raise ImportError(PROTOBUF_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 "")
@@ -16,17 +16,17 @@
 import os
 from abc import ABC, abstractmethod
 from pickle import UnpicklingError
-from typing import Dict, Optional, Union
+from typing import Dict

 import flax.linen as nn
 import jax
 import jax.numpy as jnp
-from flax.serialization import from_bytes, to_bytes
+from flax.serialization import to_bytes
 from flax.traverse_util import unflatten_dict
 from jax.random import PRNGKey

 from .configuration_utils import PretrainedConfig
-from .file_utils import FLAX_WEIGHTS_NAME, WEIGHTS_NAME, cached_path, hf_bucket_url, is_remote_url
+from .file_utils import WEIGHTS_NAME, cached_path, hf_bucket_url, is_remote_url
 from .utils import logging


@@ -34,8 +34,20 @@ logger = logging.get_logger(__name__)


@jax.jit
-def gelu_new(x):
-    return nn.gelu(x, approximate=True)
+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 = {
@@ -43,7 +55,7 @@ ACT2FN = {
    "relu": nn.relu,
    "silu": nn.swish,
    "swish": nn.swish,
-    "gelu_new": gelu_new,
+    "gelu_new": gelu,
 }


@@ -51,20 +63,16 @@ class FlaxPreTrainedModel(ABC):
    config_class = None
    pretrained_model_archive_map = {}
    base_model_prefix = ""
+    model_class = None

    def __init__(
-        self,
-        config: PretrainedConfig,
-        module: nn.Module,
-        params: Optional[Dict] = None,
-        seed: int = 0,
-        dtype: jnp.dtype = jnp.float32,
+        self, config: PretrainedConfig, module: nn.Module, params: Dict, seed: int = 0, dtype: jnp.dtype = jnp.float32
    ):
        if config is None:
            raise ValueError("config cannot be None")

-        if module is None:
-            raise ValueError("module cannot be None")
+        if params is None:
+            raise ValueError("state cannot be None")

        # Those are private to be exposed as typed property on derived classes.
        self._config = config
@@ -96,7 +104,7 @@ class FlaxPreTrainedModel(ABC):
        config = kwargs.pop("config", None)
        # state_dict = kwargs.pop("state_dict", None)
        cache_dir = kwargs.pop("cache_dir", None)
-        from_pt = kwargs.pop("from_pt", False)
+        # from_tf = kwargs.pop("from_tf", False)
        force_download = kwargs.pop("force_download", False)
        resume_download = kwargs.pop("resume_download", False)
        proxies = kwargs.pop("proxies", None)
@@ -127,28 +135,10 @@ class FlaxPreTrainedModel(ABC):

        # Load model
        if pretrained_model_name_or_path is not None:
-            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):
+            if 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 if from_pt else FLAX_WEIGHTS_NAME,
-                    revision=revision,
-                )
+                archive_file = hf_bucket_url(pretrained_model_name_or_path, filename=WEIGHTS_NAME, revision=revision)

            # redirect to the cache, if necessary
            try:
@@ -179,50 +169,31 @@ class FlaxPreTrainedModel(ABC):
        # Instantiate model.
        with open(resolved_archive_file, "rb") as state_f:
            try:
-                if from_pt:
+                from flax.serialization import from_bytes
+
+                state = from_bytes(cls.model_class, state_f)
+            except TypeError:
+                try:
                    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"
+                    )

-                    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. "
-                )
        return cls(config, state, *model_args, **model_kwargs)

-    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
+    def save_pretrained(self, folder):
+        folder_abs = os.path.abspath(folder)

-        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)
+        if not os.path.exists(folder_abs):
+            os.mkdir(folder_abs)

-        # 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:
+        with open(os.path.join(folder_abs, f"{self._config.model_type}.flax", "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
@@ -309,9 +309,7 @@ def booleans_processing(config, **kwargs):

 def input_processing(func, config, input_ids, **kwargs):
    """
-    Process the input of each TensorFlow model including the booleans. In case of a list of symbolic inputs, each input
-    has to be named accordingly to the parameters name, i.e. `input_ids = tf.keras.Input(shape=(128,), dtype='int32',
-    name="input_ids")` otherwise the order of the tensors will not be guaranteed during the training.
+    Process the input of each TensorFlow model including the booleans.

    Args:
        func (:obj:`callable`):
@@ -367,7 +365,9 @@ def input_processing(func, config, input_ids, **kwargs):
                if tensor_name in parameter_names:
                    output[tensor_name] = input
                else:
-                    output[parameter_names[i]] = input
+                    raise ValueError(
+                        f"The tensor named {input.name} does not belong to the authorized list of names {parameter_names}."
+                    )
            elif isinstance(input, allowed_types) or input is None:
                output[parameter_names[i]] = input
            else:
@@ -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"),
    ]
 )

@@ -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,
@@ -277,6 +284,7 @@ MODEL_MAPPING = OrderedDict(
        (XLMProphetNetConfig, XLMProphetNetModel),
        (ProphetNetConfig, ProphetNetModel),
        (MPNetConfig, MPNetModel),
+        (TapasConfig, TapasModel),
    ]
 )

@@ -308,6 +316,7 @@ MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(
        (LxmertConfig, LxmertForPreTraining),
        (FunnelConfig, FunnelForPreTraining),
        (MPNetConfig, MPNetForMaskedLM),
+        (TapasConfig, TapasForMaskedLM),
    ]
 )

@@ -340,6 +349,7 @@ MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(
        (ReformerConfig, ReformerModelWithLMHead),
        (FunnelConfig, FunnelForMaskedLM),
        (MPNetConfig, MPNetForMaskedLM),
+        (TapasConfig, TapasForMaskedLM),
    ]
 )

@@ -386,6 +396,7 @@ MODEL_FOR_MASKED_LM_MAPPING = OrderedDict(
        (ReformerConfig, ReformerForMaskedLM),
        (FunnelConfig, FunnelForMaskedLM),
        (MPNetConfig, MPNetForMaskedLM),
+        (TapasConfig, TapasForMaskedLM),
    ]
 )

@@ -431,6 +442,7 @@ MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING = OrderedDict(
        (CTRLConfig, CTRLForSequenceClassification),
        (TransfoXLConfig, TransfoXLForSequenceClassification),
        (MPNetConfig, MPNetForSequenceClassification),
+        (TapasConfig, TapasForSequenceClassification),
    ]
 )

@@ -455,6 +467,7 @@ MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
        (FunnelConfig, FunnelForQuestionAnswering),
        (LxmertConfig, LxmertForQuestionAnswering),
        (MPNetConfig, MPNetForQuestionAnswering),
+        (TapasConfig, TapasForQuestionAnswering),
    ]
 )

@@ -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)),
    ]
 )

@@ -235,9 +235,9 @@ class TFEncoderLayer(tf.keras.layers.Layer):
        )
        self.normalize_before = config.normalize_before
        self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
-        self.dropout = tf.keras.layers.Dropout(config.dropout)
+        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
-        self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout)
+        self.activation_dropout = config.activation_dropout
        self.fc1 = tf.keras.layers.Dense(config.encoder_ffn_dim, name="fc1")
        self.fc2 = tf.keras.layers.Dense(self.embed_dim, name="fc2")
        self.final_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="final_layer_norm")
@@ -261,7 +261,7 @@ class TFEncoderLayer(tf.keras.layers.Layer):
        assert shape_list(x) == shape_list(
            residual
        ), f"Self attn modified the shape of query {shape_list(residual)} to {shape_list(x)}"
-        x = self.dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
        if not self.normalize_before:
            x = self.self_attn_layer_norm(x)
@@ -270,9 +270,9 @@ class TFEncoderLayer(tf.keras.layers.Layer):
        if self.normalize_before:
            x = self.final_layer_norm(x)
        x = self.activation_fn(self.fc1(x))
-        x = self.activation_dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.self.activation_dropout if training else 0)
        x = self.fc2(x)
-        x = self.dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
        if not self.normalize_before:
            x = self.final_layer_norm(x)
@@ -293,7 +293,7 @@ class TFBartEncoder(tf.keras.layers.Layer):
    def __init__(self, config: BartConfig, embed_tokens: TFSharedEmbeddings, **kwargs):
        super().__init__(**kwargs)

-        self.dropout = tf.keras.layers.Dropout(config.dropout)
+        self.dropout = config.dropout
        self.layerdrop = config.encoder_layerdrop
        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions
@@ -370,7 +370,7 @@ class TFBartEncoder(tf.keras.layers.Layer):
        embed_pos = self.embed_positions(input_ids)
        x = inputs_embeds + embed_pos
        x = self.layernorm_embedding(x)
-        x = self.dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)

        # B x T x C -> T x B x C
        x = tf.transpose(x, perm=[1, 0, 2])
@@ -413,9 +413,9 @@ class TFDecoderLayer(tf.keras.layers.Layer):
            dropout=config.attention_dropout,
            name="self_attn",
        )
-        self.dropout = tf.keras.layers.Dropout(config.dropout)
+        self.dropout = config.dropout
        self.activation_fn = ACT2FN[config.activation_function]
-        self.activation_dropout = tf.keras.layers.Dropout(config.activation_dropout)
+        self.activation_dropout = config.activation_dropout
        self.normalize_before = config.normalize_before

        self.self_attn_layer_norm = tf.keras.layers.LayerNormalization(epsilon=1e-5, name="self_attn_layer_norm")
@@ -467,7 +467,7 @@ class TFDecoderLayer(tf.keras.layers.Layer):
            attn_mask=causal_mask,
            key_padding_mask=decoder_padding_mask,
        )
-        x = self.dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
        if not self.normalize_before:
            x = self.self_attn_layer_norm(x)
@@ -481,7 +481,7 @@ class TFDecoderLayer(tf.keras.layers.Layer):
            key_padding_mask=encoder_attn_mask,
            layer_state=layer_state,  # mutates layer state
        )
-        x = self.dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
        if not self.normalize_before:
            x = self.encoder_attn_layer_norm(x)
@@ -490,9 +490,9 @@ class TFDecoderLayer(tf.keras.layers.Layer):
        if self.normalize_before:
            x = self.final_layer_norm(x)
        x = self.activation_fn(self.fc1(x))
-        x = self.activation_dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.activation_dropout if training else 0)
        x = self.fc2(x)
-        x = self.dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)
        x = residual + x
        if not self.normalize_before:
            x = self.final_layer_norm(x)
@@ -545,7 +545,7 @@ class TFBartDecoder(tf.keras.layers.Layer):
            else None
        )

-        self.dropout = tf.keras.layers.Dropout(config.dropout)
+        self.dropout = config.dropout
        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions
        self.use_cache = config.use_cache
@@ -588,7 +588,7 @@ class TFBartDecoder(tf.keras.layers.Layer):
            x = self.layernorm_embedding(x) + positions
        else:
            x = self.layernorm_embedding(x + positions)
-        x = self.dropout(x, training=training)
+        x = tf.nn.dropout(x, rate=self.dropout if training else 0)

        # Convert to Bart output format: (BS, seq_len, model_dim) ->  (seq_len, BS, model_dim)
        x = tf.transpose(x, perm=(1, 0, 2))
@@ -674,7 +674,7 @@ class TFAttention(tf.keras.layers.Layer):
        self.embed_dim = embed_dim

        self.num_heads = num_heads
-        self.dropout = tf.keras.layers.Dropout(dropout)
+        self.dropout = dropout
        self.head_dim = embed_dim // num_heads
        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
        self.scaling = self.head_dim ** -0.5
@@ -772,7 +772,7 @@ class TFAttention(tf.keras.layers.Layer):
            attn_weights = tf.reshape(attn_weights, (bsz * self.num_heads, tgt_len, src_len))

        attn_weights = tf.nn.softmax(attn_weights, axis=-1)
-        attn_probs = self.dropout(attn_weights, training=training)
+        attn_probs = tf.nn.dropout(attn_weights, rate=self.dropout if training else 0.0)

        attn_output = tf.matmul(attn_probs, v)  # shape: (bsz * self.num_heads, tgt_len, self.head_dim)
        attn_output = tf.transpose(attn_output, perm=(1, 0, 2))
@@ -23,7 +23,7 @@ import jax.numpy as jnp
 from jax.random import PRNGKey

 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel
+from ...modeling_flax_utils import FlaxPreTrainedModel, gelu
 from ...utils import logging
 from .configuration_bert import BertConfig

@@ -205,7 +205,7 @@ class FlaxBertAttention(nn.Module):
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation

    @nn.compact
-    def __call__(self, hidden_states, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, 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,28 +219,27 @@ class FlaxBertAttention(nn.Module):
            bias_init=jax.nn.initializers.zeros,
            name="self",
            dtype=self.dtype,
-        )(hidden_states, attention_mask)
+        )(hidden_state, attention_mask)

-        layer_norm = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_states)
+        layer_norm = FlaxBertLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_state)
        return layer_norm


 class FlaxBertIntermediate(nn.Module):
    output_size: int
-    hidden_act: str
    kernel_init_scale: float = 0.2
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation

    @nn.compact
-    def __call__(self, hidden_states):
-        hidden_states = nn.Dense(
+    def __call__(self, hidden_state):
+        # TODO: Add ACT2FN reference to change activation function
+        dense = nn.Dense(
            features=self.output_size,
            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
            name="dense",
            dtype=self.dtype,
-        )(hidden_states)
-        hidden_states = ACT2FN[self.hidden_act](hidden_states)
-        return hidden_states
+        )(hidden_state)
+        return gelu(dense)


 class FlaxBertOutput(nn.Module):
@@ -250,28 +249,27 @@ class FlaxBertOutput(nn.Module):

    @nn.compact
    def __call__(self, intermediate_output, attention_output, deterministic: bool = True):
-        hidden_states = nn.Dense(
+        hidden_state = 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_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
+        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


 class FlaxBertLayer(nn.Module):
    num_heads: int
    head_size: int
    intermediate_size: int
-    hidden_act: str
    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_states, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
        attention = FlaxBertAttention(
            self.num_heads,
            self.head_size,
@@ -279,13 +277,9 @@ class FlaxBertLayer(nn.Module):
            dropout_rate=self.dropout_rate,
            name="attention",
            dtype=self.dtype,
-        )(hidden_states, attention_mask, deterministic=deterministic)
+        )(hidden_state, attention_mask, deterministic=deterministic)
        intermediate = FlaxBertIntermediate(
-            self.intermediate_size,
-            kernel_init_scale=self.kernel_init_scale,
-            hidden_act=self.hidden_act,
-            name="intermediate",
-            dtype=self.dtype,
+            self.intermediate_size, kernel_init_scale=self.kernel_init_scale, 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
@@ -303,7 +297,6 @@ class FlaxBertLayerCollection(nn.Module):
    num_heads: int
    head_size: int
    intermediate_size: int
-    hidden_act: str
    dropout_rate: float = 0.0
    kernel_init_scale: float = 0.2
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
@@ -323,7 +316,6 @@ 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,
            )
@@ -336,24 +328,22 @@ class FlaxBertEncoder(nn.Module):
    num_heads: int
    head_size: int
    intermediate_size: int
-    hidden_act: str
    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_states, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, 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_states, attention_mask, deterministic=deterministic)
+        )(hidden_state, attention_mask, deterministic=deterministic)
        return layer


@@ -362,10 +352,10 @@ class FlaxBertPooler(nn.Module):
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation

    @nn.compact
-    def __call__(self, hidden_states):
-        cls_token = hidden_states[:, 0]
+    def __call__(self, hidden_state):
+        cls_token = hidden_state[:, 0]
        out = nn.Dense(
-            hidden_states.shape[-1],
+            hidden_state.shape[-1],
            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
            name="dense",
            dtype=self.dtype,
@@ -373,20 +363,62 @@ 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
    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 = ACT2FN[self.hidden_act](hidden_states)
-        return FlaxBertLayerNorm(name="layer_norm", 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)


 class FlaxBertLMPredictionHead(nn.Module):
    vocab_size: int
-    hidden_act: str
    dtype: jnp.dtype = jnp.float32

    @nn.compact
@@ -396,57 +428,64 @@ 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", hidden_act=self.hidden_act, dtype=self.dtype
-        )(hidden_states)
+        hidden_states = FlaxBertPredictionHeadTransform(name="transform", 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_act: str
+    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
    dtype: jnp.dtype = jnp.float32

    @nn.compact
-    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
+    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


-class FlaxBertPreTrainedModel(FlaxPreTrainedModel):
+@add_start_docstrings(
+    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
+    BERT_START_DOCSTRING,
+)
+class FlaxBertModel(FlaxPreTrainedModel):
    """
-    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
-    models.
+    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 = 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):
-        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}
-
-        self.params = 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)
@@ -462,11 +501,6 @@ class FlaxBertPreTrainedModel(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:

@@ -492,7 +526,7 @@ class FlaxBertPreTrainedModel(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 or "transform.dense.kernel" in key:
+            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
@@ -521,20 +555,7 @@ class FlaxBertPreTrainedModel(FlaxPreTrainedModel):

        return jax_state

-
-@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, state: dict = None, seed: int = 0, dtype: jnp.dtype = jnp.float32):
+    def __init__(self, config: BertConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32):
        module = FlaxBertModule(
            vocab_size=config.vocab_size,
            hidden_size=config.hidden_size,
@@ -545,7 +566,6 @@ class FlaxBertModel(FlaxBertPreTrainedModel):
            head_size=config.hidden_size,
            intermediate_size=config.intermediate_size,
            dropout_rate=config.hidden_dropout_prob,
-            hidden_act=config.hidden_act,
            dtype=dtype,
        )

@@ -581,62 +601,34 @@ class FlaxBertModel(FlaxBertPreTrainedModel):
            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)

-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 position_ids is None:
+            position_ids = jnp.arange(jnp.atleast_2d(input_ids).shape[-1])

-    @nn.compact
-    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic: bool = True):
+        if attention_mask is None:
+            attention_mask = jnp.ones_like(input_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)
+        return input_ids, attention_mask, token_type_ids, position_ids

-        # 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)
+    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
+        )

-        if not self.add_pooling_layer:
-            return encoder
+        params_rng, dropout_rng = jax.random.split(rng)
+        rngs = {"params": params_rng, "dropout": dropout_rng}

-        pooled = FlaxBertPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
-        return encoder, pooled
+        self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]


-class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
-    def __init__(
-        self, config: BertConfig, state: dict = None, seed: int = 0, dtype: jnp.dtype = jnp.float32, **kwargs
-    ):
-        module = FlaxBertForMaskedLMModule(
+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(
            vocab_size=config.vocab_size,
            type_vocab_size=config.type_vocab_size,
            hidden_size=config.hidden_size,
@@ -645,12 +637,9 @@ class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
            num_heads=config.num_attention_heads,
            num_encoder_layers=config.num_hidden_layers,
            max_length=config.max_length,
-            hidden_act=config.hidden_act,
            **kwargs,
        )

-        super().__init__(config, module, state, seed)
-
    def __call__(
        self,
        input_ids,
@@ -670,7 +659,7 @@ class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
        if dropout_rng is not None:
            rngs["dropout"] = dropout_rng

-        logits = self.module.apply(
+        pooled, logits = self.module.apply(
            {"params": params or self.params},
            jnp.array(input_ids, dtype="i4"),
            jnp.array(attention_mask, dtype="i4"),
@@ -680,47 +669,4 @@ class FlaxBertForMaskedLM(FlaxBertPreTrainedModel):
            rngs=rngs,
        )

-        return (logits,)
-
-
-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
+        return logits, pooled
@@ -562,7 +562,7 @@ LAYOUTLM_START_DOCSTRING = r"""

 LAYOUTLM_INPUTS_DOCSTRING = r"""
    Args:
-        input_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`):
+        input_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`):
            Indices of input sequence tokens in the vocabulary.

            Indices can be obtained using :class:`transformers.LayoutLMTokenizer`. See
@@ -570,20 +570,22 @@ LAYOUTLM_INPUTS_DOCSTRING = r"""
            details.

            `What are input IDs? <../glossary.html#input-ids>`__
-        bbox (:obj:`torch.LongTensor` of shape :obj:`({0}, 4)`, `optional`):
+        bbox (:obj:`torch.LongTensor` of shape :obj:`{0}`, `optional`):
            Bounding Boxes of each input sequence tokens. Selected in the range ``[0, config.max_2d_position_embeddings
            - 1]``.
-        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
+
+            `What are bboxes? <../glossary.html#position-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 MASKED tokens.

            `What are attention masks? <../glossary.html#attention-mask>`__
-        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
+        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`):
+        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]``.

@@ -641,7 +643,7 @@ class LayoutLMModel(LayoutLMPreTrainedModel):
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

-    @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="layoutlm-base-uncased",
@@ -782,7 +784,7 @@ class LayoutLMForMaskedLM(LayoutLMPreTrainedModel):
    def set_output_embeddings(self, new_embeddings):
        self.cls.predictions.decoder = new_embeddings

-    @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="layoutlm-base-uncased",
@@ -870,7 +872,7 @@ class LayoutLMForTokenClassification(LayoutLMPreTrainedModel):
    def get_input_embeddings(self):
        return self.layoutlm.embeddings.word_embeddings

-    @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("batch_size, sequence_length"))
+    @add_start_docstrings_to_model_forward(LAYOUTLM_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="layoutlm-base-uncased",
@@ -2,20 +2,6 @@
 # There's no way to ignore "F401 '...' imported but unused" warnings in this
 # module, but to preserve other warnings. So, don't check this module at all.

-# 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.
-
 from ...file_utils import is_flax_available, is_tf_available, is_tokenizers_available, is_torch_available
 from .configuration_mpnet import MPNET_PRETRAINED_CONFIG_ARCHIVE_MAP, MPNetConfig
 from .tokenization_mpnet import MPNetTokenizer
@@ -22,7 +22,7 @@ import jax.numpy as jnp
 from jax.random import PRNGKey

 from ...file_utils import add_start_docstrings, add_start_docstrings_to_model_forward
-from ...modeling_flax_utils import ACT2FN, FlaxPreTrainedModel
+from ...modeling_flax_utils import FlaxPreTrainedModel, gelu
 from ...utils import logging
 from .configuration_roberta import RobertaConfig

@@ -33,23 +33,6 @@ _CONFIG_FOR_DOC = "RobertaConfig"
 _TOKENIZER_FOR_DOC = "RobertaTokenizer"


-def create_position_ids_from_input_ids(input_ids, padding_idx):
-    """
-    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
-    are ignored. This is modified from fairseq's `utils.make_positions`.
-
-    Args:
-        input_ids: jnp.ndarray
-        padding_idx: int
-
-    Returns: jnp.ndarray
-    """
-    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
-    mask = (input_ids != padding_idx).astype("i4")
-    incremental_indices = jnp.cumsum(mask, axis=1).astype("i4") * mask
-    return incremental_indices.astype("i4") + padding_idx
-
-
 ROBERTA_START_DOCSTRING = r"""

    This model inherits from :class:`~transformers.FlaxPreTrainedModel`. Check the superclass documentation for the
@@ -225,7 +208,7 @@ class FlaxRobertaAttention(nn.Module):
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation

    @nn.compact
-    def __call__(self, hidden_states, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, 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)
@@ -239,29 +222,28 @@ class FlaxRobertaAttention(nn.Module):
            bias_init=jax.nn.initializers.zeros,
            name="self",
            dtype=self.dtype,
-        )(hidden_states, attention_mask)
+        )(hidden_state, attention_mask)

-        layer_norm = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_states)
+        layer_norm = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(self_att + hidden_state)
        return layer_norm


 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertIntermediate with Bert->Roberta
 class FlaxRobertaIntermediate(nn.Module):
    output_size: int
-    hidden_act: str
    kernel_init_scale: float = 0.2
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation

    @nn.compact
-    def __call__(self, hidden_states):
-        hidden_states = nn.Dense(
+    def __call__(self, hidden_state):
+        # TODO: Add ACT2FN reference to change activation function
+        dense = nn.Dense(
            features=self.output_size,
            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
            name="dense",
            dtype=self.dtype,
-        )(hidden_states)
-        hidden_states = ACT2FN[self.hidden_act](hidden_states)
-        return hidden_states
+        )(hidden_state)
+        return gelu(dense)


 # Copied from transformers.models.bert.modeling_flax_bert.FlaxBertOutput with Bert->Roberta
@@ -272,28 +254,27 @@ class FlaxRobertaOutput(nn.Module):

    @nn.compact
    def __call__(self, intermediate_output, attention_output, deterministic: bool = True):
-        hidden_states = nn.Dense(
+        hidden_state = 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_states = nn.Dropout(rate=self.dropout_rate)(hidden_states, deterministic=deterministic)
-        hidden_states = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_states + attention_output)
-        return hidden_states
+        hidden_state = nn.Dropout(rate=self.dropout_rate)(hidden_state, deterministic=deterministic)
+        hidden_state = FlaxRobertaLayerNorm(name="layer_norm", dtype=self.dtype)(hidden_state + attention_output)
+        return hidden_state


 class FlaxRobertaLayer(nn.Module):
    num_heads: int
    head_size: int
    intermediate_size: int
-    hidden_act: str
    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_states, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
        attention = FlaxRobertaAttention(
            self.num_heads,
            self.head_size,
@@ -301,11 +282,10 @@ class FlaxRobertaLayer(nn.Module):
            dropout_rate=self.dropout_rate,
            name="attention",
            dtype=self.dtype,
-        )(hidden_states, attention_mask, deterministic=deterministic)
+        )(hidden_state, attention_mask, deterministic=deterministic)
        intermediate = FlaxRobertaIntermediate(
            self.intermediate_size,
            kernel_init_scale=self.kernel_init_scale,
-            hidden_act=self.hidden_act,
            name="intermediate",
            dtype=self.dtype,
        )(attention)
@@ -326,7 +306,6 @@ class FlaxRobertaLayerCollection(nn.Module):
    num_heads: int
    head_size: int
    intermediate_size: int
-    hidden_act: str
    dropout_rate: float = 0.0
    kernel_init_scale: float = 0.2
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation
@@ -346,7 +325,6 @@ class FlaxRobertaLayerCollection(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,
            )
@@ -360,24 +338,22 @@ class FlaxRobertaEncoder(nn.Module):
    num_heads: int
    head_size: int
    intermediate_size: int
-    hidden_act: str
    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_states, attention_mask, deterministic: bool = True):
+    def __call__(self, hidden_state, attention_mask, deterministic: bool = True):
        layer = FlaxRobertaLayerCollection(
            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_states, attention_mask, deterministic=deterministic)
+        )(hidden_state, attention_mask, deterministic=deterministic)
        return layer


@@ -387,10 +363,10 @@ class FlaxRobertaPooler(nn.Module):
    dtype: jnp.dtype = jnp.float32  # the dtype of the computation

    @nn.compact
-    def __call__(self, hidden_states):
-        cls_token = hidden_states[:, 0]
+    def __call__(self, hidden_state):
+        cls_token = hidden_state[:, 0]
        out = nn.Dense(
-            hidden_states.shape[-1],
+            hidden_state.shape[-1],
            kernel_init=jax.nn.initializers.normal(self.kernel_init_scale, self.dtype),
            name="dense",
            dtype=self.dtype,
@@ -398,12 +374,64 @@ class FlaxRobertaPooler(nn.Module):
        return nn.tanh(out)


-class FlaxRobertaPreTrainedModel(FlaxPreTrainedModel):
+# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertModule with Bert->Roberta
+class FlaxRobertaModule(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 = FlaxRobertaEmbeddings(
+            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 = FlaxRobertaEncoder(
+            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 = FlaxRobertaPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
+        return encoder, pooled
+
+
+@add_start_docstrings(
+    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
+    ROBERTA_START_DOCSTRING,
+)
+class FlaxRobertaModel(FlaxPreTrainedModel):
    """
-    An abstract class to handle weights initialization and a simple interface for downloading and loading pretrained
-    models.
+    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`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
+    Kaiser and Illia Polosukhin.
    """

+    model_class = FlaxRobertaModule
    config_class = RobertaConfig
    base_model_prefix = "roberta"

@@ -476,42 +504,7 @@ class FlaxRobertaPreTrainedModel(FlaxPreTrainedModel):

        return jax_state

-    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
-        )
-
-        params_rng, dropout_rng = jax.random.split(rng)
-        rngs = {"params": params_rng, "dropout": dropout_rng}
-
-        self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]
-
-    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 = create_position_ids_from_input_ids(input_ids, self.config.pad_token_id)
-
-        if attention_mask is None:
-            attention_mask = jnp.ones_like(input_ids)
-
-        return input_ids, attention_mask, token_type_ids, position_ids
-
-
-@add_start_docstrings(
-    "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
-    ROBERTA_START_DOCSTRING,
-)
-class FlaxRobertaModel(FlaxRobertaPreTrainedModel):
-    """
-    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`_ by Ashish Vaswani, Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz
-    Kaiser and Illia Polosukhin.
-    """
-
-    def __init__(self, config: RobertaConfig, state: dict = None, seed: int = 0, dtype: jnp.dtype = jnp.float32):
+    def __init__(self, config: RobertaConfig, state: dict, seed: int = 0, dtype: jnp.dtype = jnp.float32):
        module = FlaxRobertaModule(
            vocab_size=config.vocab_size,
            hidden_size=config.hidden_size,
@@ -520,7 +513,6 @@ class FlaxRobertaModel(FlaxRobertaPreTrainedModel):
            num_encoder_layers=config.num_hidden_layers,
            num_heads=config.num_attention_heads,
            head_size=config.hidden_size,
-            hidden_act=config.hidden_act,
            intermediate_size=config.intermediate_size,
            dropout_rate=config.hidden_dropout_prob,
            dtype=dtype,
@@ -558,53 +550,42 @@ class FlaxRobertaModel(FlaxRobertaPreTrainedModel):
            rngs=rngs,
        )

+    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
+        )

-# Copied from transformers.models.bert.modeling_flax_bert.FlaxBertModule with Bert->Roberta
-class FlaxRobertaModule(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
+        params_rng, dropout_rng = jax.random.split(rng)
+        rngs = {"params": params_rng, "dropout": dropout_rng}

-    @nn.compact
-    def __call__(self, input_ids, attention_mask, token_type_ids, position_ids, deterministic: bool = True):
+        self.params = self.module.init(rngs, input_ids, attention_mask, token_type_ids, position_ids)["params"]

-        # Embedding
-        embeddings = FlaxRobertaEmbeddings(
-            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 _check_inputs(self, input_ids, attention_mask, token_type_ids, position_ids):

-        # N stacked encoding layers
-        encoder = FlaxRobertaEncoder(
-            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)
+        if token_type_ids is None:
+            token_type_ids = jnp.ones_like(input_ids)

-        if not self.add_pooling_layer:
-            return encoder
+        if position_ids is None:
+            position_ids = create_position_ids_from_input_ids(input_ids, self.config.pad_token_id)

-        pooled = FlaxRobertaPooler(kernel_init_scale=self.kernel_init_scale, name="pooler", dtype=self.dtype)(encoder)
-        return encoder, pooled
+        if attention_mask is None:
+            attention_mask = jnp.ones_like(input_ids)
+
+        return input_ids, attention_mask, token_type_ids, position_ids
+
+
+def create_position_ids_from_input_ids(input_ids, padding_idx):
+    """
+    Replace non-padding symbols with their position numbers. Position numbers begin at padding_idx+1. Padding symbols
+    are ignored. This is modified from fairseq's `utils.make_positions`.
+
+    Args:
+        input_ids: jnp.ndarray
+        padding_idx: int
+
+    Returns: jnp.ndarray
+    """
+    # The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+    mask = (input_ids != padding_idx).astype("i4")
+    incremental_indices = jnp.cumsum(mask, axis=1).astype("i4") * mask
+    return incremental_indices.astype("i4") + padding_idx
@@ -0,0 +1,17 @@
+# flake8: noqa
+# There's no way to ignore "F401 '...' imported but unused" warnings in this
+# module, but to preserve other warnings. So, don't check this module at all.
+
+from ...file_utils import is_torch_available
+from .configuration_tapas import TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP, TapasConfig
+from .tokenization_tapas import TapasTokenizer
+
+
+if is_torch_available():
+    from .modeling_tapas import (
+        TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST,
+        TapasForMaskedLM,
+        TapasForQuestionAnswering,
+        TapasForSequenceClassification,
+        TapasModel,
+    )
@@ -0,0 +1,211 @@
+# coding=utf-8
+# Copyright 2020 Google Research 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.
+""" TAPAS configuration. Adds additional hyperparameters to the configuration of BERT."""
+
+
+from ...configuration_utils import PretrainedConfig
+
+
+TAPAS_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "nielsr/tapas-base-finetuned-sqa": "https://huggingface.co/nielsr/tapas-base-finetuned-sqa/resolve/main/config.json",
+    "nielsr/tapas-base-finetuned-wtq": "https://huggingface.co/nielsr/tapas-base-finetuned-wtq/resolve/main/config.json",
+    "nielsr/tapas-base-finetuned-wikisql-supervised": "https://huggingface.co/nielsr/tapas-base-finetuned-wikisql-supervised/resolve/main/config.json",
+    "nielsr/tapas-base-finetuned-tabfact": "https://huggingface.co/nielsr/tapas-base-finetuned-tabfact/resolve/main/config.json",
+}
+
+
+class TapasConfig(PretrainedConfig):
+    r"""
+    This is the configuration class to store the configuration of a :class:`~transformers.TapasModel`. It is used to
+    instantiate a TAPAS model according to the specified arguments, defining the model architecture. Instantiating a
+    configuration with the defaults will yield a similar configuration to that of the TAPAS `tapas-base-finetuned-sqa`
+    architecture. Configuration objects inherit from :class:`~transformers.PreTrainedConfig` and can be used to control
+    the model outputs. Read the documentation from :class:`~transformers.PretrainedConfig` for more information.
+
+    Hyperparameters additional to BERT are taken from run_task_main.py and hparam_utils.py of the original
+    implementation. Original implementation available at https://github.com/google-research/tapas/tree/master.
+
+    Args:
+        vocab_size (:obj:`int`, `optional`, defaults to 30522):
+            Vocabulary size of the TAPAS model. Defines the number of different tokens that can be represented by the
+            :obj:`inputs_ids` passed when calling :class:`~transformers.TapasModel`.
+        hidden_size (:obj:`int`, `optional`, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_hidden_layers (:obj:`int`, `optional`, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (:obj:`int`, `optional`, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (:obj:`int`, `optional`, defaults to 3072):
+            Dimensionality of the "intermediate" (often named feed-forward) layer in the Transformer encoder.
+        hidden_act (:obj:`str` or :obj:`Callable`, `optional`, defaults to :obj:`"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler. If string,
+            :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
+            The dropout probability for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
+            The dropout ratio for the attention probabilities.
+        max_position_embeddings (:obj:`int`, `optional`, defaults to 1024):
+            The maximum sequence length that this model might ever be used with. Typically set this to something large
+            just in case (e.g., 512 or 1024 or 2048).
+        type_vocab_sizes (:obj:`List[int]`, `optional`, defaults to [3, 256, 256, 2, 256, 256, 10]):
+            The vocabulary sizes of the :obj:`token_type_ids` passed when calling :class:`~transformers.TapasModel`.
+        initializer_range (:obj:`float`, `optional`, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_eps (:obj:`float`, `optional`, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to use gradient checkpointing to save memory at the expense of a slower backward pass.
+        positive_label_weight (:obj:`float`, `optional`, defaults to 10.0):
+            Weight for positive labels.
+        num_aggregation_labels (:obj:`int`, `optional`, defaults to 0):
+            The number of aggregation operators to predict.
+        aggregation_loss_weight (:obj:`float`, `optional`, defaults to 1.0):
+            Importance weight for the aggregation loss.
+        use_answer_as_supervision (:obj:`bool`, `optional`):
+            Whether to use the answer as the only supervision for aggregation examples.
+        answer_loss_importance (:obj:`float`, `optional`, defaults to 1.0):
+            Importance weight for the regression loss.
+        use_normalized_answer_loss (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to normalize the answer loss by the maximum of the predicted and expected value.
+        huber_loss_delta: (:obj:`float`, `optional`):
+            Delta parameter used to calculate the regression loss.
+        temperature: (:obj:`float`, `optional`, defaults to 1.0):
+            Value used to control (OR change) the skewness of cell logits probabilities.
+        aggregation_temperature: (:obj:`float`, `optional`, defaults to 1.0):
+            Scales aggregation logits to control the skewness of probabilities.
+        use_gumbel_for_cells: (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to apply Gumbel-Softmax to cell selection.
+        use_gumbel_for_aggregation: (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to apply Gumbel-Softmax to aggregation selection.
+        average_approximation_function: (:obj:`string`, `optional`, defaults to :obj:`"ratio"`):
+            Method to calculate the expected average of cells in the weak supervision case. One of :obj:`"ratio"`,
+            :obj:`"first_order"` or :obj:`"second_order"`.
+        cell_selection_preference: (:obj:`float`, `optional`):
+            Preference for cell selection in ambiguous cases. Only applicable in case of weak supervision for
+            aggregation (WTQ, WikiSQL). If the total mass of the aggregation probabilities (excluding the "NONE"
+            operator) is higher than this hyperparameter, then aggregation is predicted for an example.
+        answer_loss_cutoff: (:obj:`float`, `optional`):
+            Ignore examples with answer loss larger than cutoff.
+        max_num_rows: (:obj:`int`, `optional`, defaults to 64):
+            Maximum number of rows.
+        max_num_columns: (:obj:`int`, `optional`, defaults to 32):
+            Maximum number of columns.
+        average_logits_per_cell: (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to average logits per cell.
+        select_one_column: (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether to constrain the model to only select cells from a single column.
+        allow_empty_column_selection: (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to allow not to select any column.
+        init_cell_selection_weights_to_zero: (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to initialize cell selection weights to 0 so that the initial probabilities are 50%.
+        reset_position_index_per_cell: (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Whether to restart position indexes at every cell (i.e. use relative position embeddings).
+        disable_per_token_loss: (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to disable any (strong or weak) supervision on cells.
+
+    Example::
+
+        >>> from transformers import TapasModel, TapasConfig
+        >>> # Initializing a Tapas configuration
+        >>> configuration = TapasConfig()
+        >>> # Initializing a model from the configuration
+        >>> model = TapasModel(configuration)
+        >>> # Accessing the model configuration
+        >>> configuration = model.config
+    """
+
+    model_type = "tapas"
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=1024,
+        type_vocab_sizes=[3, 256, 256, 2, 256, 256, 10],
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        pad_token_id=0,
+        gradient_checkpointing=False,
+        positive_label_weight=10.0,
+        num_aggregation_labels=0,
+        aggregation_loss_weight=1.0,
+        use_answer_as_supervision=None,
+        answer_loss_importance=1.0,
+        use_normalized_answer_loss=False,
+        huber_loss_delta=None,
+        temperature=1.0,
+        aggregation_temperature=1.0,
+        use_gumbel_for_cells=False,
+        use_gumbel_for_aggregation=False,
+        average_approximation_function="ratio",
+        cell_selection_preference=None,
+        answer_loss_cutoff=None,
+        max_num_rows=64,
+        max_num_columns=32,
+        average_logits_per_cell=False,
+        select_one_column=True,
+        allow_empty_column_selection=False,
+        init_cell_selection_weights_to_zero=False,
+        reset_position_index_per_cell=True,
+        disable_per_token_loss=False,
+        **kwargs
+    ):
+
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+        # BERT hyperparameters (with updated max_position_embeddings and type_vocab_sizes)
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_sizes = type_vocab_sizes
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.gradient_checkpointing = gradient_checkpointing
+
+        # Fine-tuning task hyperparameters
+        self.positive_label_weight = positive_label_weight
+        self.num_aggregation_labels = num_aggregation_labels
+        self.aggregation_loss_weight = aggregation_loss_weight
+        self.use_answer_as_supervision = use_answer_as_supervision
+        self.answer_loss_importance = answer_loss_importance
+        self.use_normalized_answer_loss = use_normalized_answer_loss
+        self.huber_loss_delta = huber_loss_delta
+        self.temperature = temperature
+        self.aggregation_temperature = aggregation_temperature
+        self.use_gumbel_for_cells = use_gumbel_for_cells
+        self.use_gumbel_for_aggregation = use_gumbel_for_aggregation
+        self.average_approximation_function = average_approximation_function
+        self.cell_selection_preference = cell_selection_preference
+        self.answer_loss_cutoff = answer_loss_cutoff
+        self.max_num_rows = max_num_rows
+        self.max_num_columns = max_num_columns
+        self.average_logits_per_cell = average_logits_per_cell
+        self.select_one_column = select_one_column
+        self.allow_empty_column_selection = allow_empty_column_selection
+        self.init_cell_selection_weights_to_zero = init_cell_selection_weights_to_zero
+        self.reset_position_index_per_cell = reset_position_index_per_cell
+        self.disable_per_token_loss = disable_per_token_loss
@@ -0,0 +1,116 @@
+# coding=utf-8
+# Copyright 2018 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.
+"""Convert TAPAS checkpoint."""
+
+
+import argparse
+
+import torch
+
+from transformers import TapasForQuestionAnswering  # noqa F401
+from transformers import TapasForSequenceClassification  # noqa F401
+from transformers import TapasConfig, TapasModel, load_tf_weights_in_tapas
+from transformers.utils import logging
+
+
+logging.set_verbosity_info()
+
+
+def convert_tf_checkpoint_to_pytorch(tf_checkpoint_path, tapas_config_file, pytorch_dump_path):
+    # Initialise PyTorch model. Defaults to TapasForQuestionAnswering with default SQA config.
+    # Uncomment another config and/or model to change this. If you want to convert a checkpoint
+    # that has absolute position embeddings, make sure to set reset_position_index_per_cell of
+    # TapasConfig to False.
+
+    # WTQ config
+    # config = TapasConfig(
+    #         # run_task_main.py hparams
+    #         num_aggregation_labels = 4,
+    #         use_answer_as_supervision = True,
+    #         # hparam_utils.py hparams
+    #         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,
+    # )
+
+    # WikiSQL config
+    # config = TapasConfig(
+    #         # run_task_main.py hparams
+    #         num_aggregation_labels = 4,
+    #         use_answer_as_supervision = True,
+    #         # hparam_utils.py hparams
+    #         answer_loss_cutoff = 0.185567,
+    #         cell_selection_preference = 0.611754,
+    #         huber_loss_delta = 1265.74,
+    #         init_cell_selection_weights_to_zero = False,
+    #         select_one_column = False,
+    #         allow_empty_column_selection = False,
+    #         temperature = 0.107515,
+    # )
+
+    # WikiSQL-supervised config
+    # config = TapasConfig(
+    #         # run_task_main.py hparams
+    #         num_aggregation_labels = 4,
+    #         use_answer_as_supervision = False,
+    #         # hparam_utils.py hparams
+    #         answer_loss_cutoff = 36.4519,
+    #         cell_selection_preference = 0.903421,
+    #         huber_loss_delta = 222.088,
+    #         init_cell_selection_weights_to_zero = True,
+    #         select_one_column = True,
+    #         allow_empty_column_selection = True,
+    #         temperature = 0.763141,
+    # )
+
+    # SQA config
+    config = TapasConfig()
+
+    print("Building PyTorch model from configuration: {}".format(str(config)))
+    model = TapasModel(config)
+    # model = TapasForQuestionAnswering(config)
+    # model = TapasForSequenceClassification(config)
+
+    # Load weights from tf checkpoint
+    load_tf_weights_in_tapas(model, config, tf_checkpoint_path)
+
+    # Save pytorch-model
+    print("Save PyTorch model to {}".format(pytorch_dump_path))
+    torch.save(model.state_dict(), pytorch_dump_path)
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    # Required parameters
+    parser.add_argument(
+        "--tf_checkpoint_path", default=None, type=str, required=True, help="Path to the TensorFlow checkpoint path."
+    )
+    parser.add_argument(
+        "--tapas_config_file",
+        default=None,
+        type=str,
+        required=True,
+        help="The config json file corresponding to the pre-trained TAPAS model. \n"
+        "This specifies the model architecture.",
+    )
+    parser.add_argument(
+        "--pytorch_dump_path", default=None, type=str, required=True, help="Path to the output PyTorch model."
+    )
+    args = parser.parse_args()
+    convert_tf_checkpoint_to_pytorch(args.tf_checkpoint_path, args.tapas_config_file, args.pytorch_dump_path)
@@ -28,6 +28,8 @@ from .file_utils import (
    _datasets_available,
    _faiss_available,
    _flax_available,
+    _pandas_available,
+    _scatter_available,
    _sentencepiece_available,
    _tf_available,
    _tokenizers_available,
@@ -221,6 +223,30 @@ def require_tokenizers(test_case):
        return test_case


+def require_pandas(test_case):
+    """
+    Decorator marking a test that requires pandas.
+
+    These tests are skipped when pandas isn't installed.
+
+    """
+    if not _pandas_available:
+        return unittest.skip("test requires pandas")(test_case)
+    else:
+        return test_case
+
+
+def require_scatter(test_case):
+    """
+    Decorator marking a test that requires PyTorch Scatter. These tests are skipped when PyTorch Scatter isn't
+    installed.
+    """
+    if not _scatter_available:
+        return unittest.skip("test requires PyTorch Scatter")(test_case)
+    else:
+        return test_case
+
+
 def require_torch_multi_gpu(test_case):
    """
    Decorator marking a test that requires a multi-GPU setup (in PyTorch).
@@ -50,7 +50,6 @@ if is_tf_available():

 if is_torch_available():
    import torch
-
 if is_flax_available():
    import jax.numpy as jnp

@@ -23,12 +23,14 @@ from .file_utils import cached_property, is_torch_available, is_torch_tpu_availa
 from .trainer_utils import EvaluationStrategy
 from .utils import logging

+
 if is_torch_available():
    import torch

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

+
 logger = logging.get_logger(__name__)


@@ -82,8 +84,6 @@ class TrainingArguments:
                * :obj:`"steps"`: Evaluation is done (and logged) every :obj:`eval_steps`.
                * :obj:`"epoch"`: Evaluation is done at the end of each epoch.

-        initial_evaluation (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether to perform evaluation before doing any training, for example for fine-tuning.
        prediction_loss_only (:obj:`bool`, `optional`, defaults to `False`):
            When performing evaluation and predictions, only returns the loss.
        per_device_train_batch_size (:obj:`int`, `optional`, defaults to 8):
@@ -238,8 +238,6 @@ class TrainingArguments:
        default="no",
        metadata={"help": "Run evaluation during training at each logging step."},
    )
-    initial_evaluation: bool = field(default=False, metadata={
-        "help": "Whether to perform evaluation before doing any training, for example for fine-tuning."})
    prediction_loss_only: bool = field(
        default=False,
        metadata={"help": "When performing evaluation and predictions, only returns the loss."},
@@ -256,14 +254,14 @@ class TrainingArguments:
        default=None,
        metadata={
            "help": "Deprecated, the use of `--per_device_train_batch_size` is preferred. "
-                    "Batch size per GPU/TPU core/CPU for training."
+            "Batch size per GPU/TPU core/CPU for training."
        },
    )
    per_gpu_eval_batch_size: Optional[int] = field(
        default=None,
        metadata={
            "help": "Deprecated, the use of `--per_device_eval_batch_size` is preferred."
-                    "Batch size per GPU/TPU core/CPU for evaluation."
+            "Batch size per GPU/TPU core/CPU for evaluation."
        },
    )

@@ -376,12 +374,6 @@ class TrainingArguments:
            "help": "When resuming training, whether or not to skip the first epochs and batches to get to the same training data."
        },
    )
-    tokenizers_parallelism: bool = field(
-        default=True,
-        metadata={
-            "help": "Sets the TOKENIZERS_PARALLELISM variable to True or False to prevent warnings"
-        },
-    )

    def __post_init__(self):
        if self.disable_tqdm is None:
@@ -1867,6 +1867,45 @@ def load_tf_weights_in_t5(*args, **kwargs):
    requires_pytorch(load_tf_weights_in_t5)


+TAPAS_PRETRAINED_MODEL_ARCHIVE_LIST = None
+
+
+class TapasForMaskedLM:
+    def __init__(self, *args, **kwargs):
+        requires_pytorch(self)
+
+    @classmethod
+    def from_pretrained(self, *args, **kwargs):
+        requires_pytorch(self)
+
+
+class TapasForQuestionAnswering:
+    def __init__(self, *args, **kwargs):
+        requires_pytorch(self)
+
+    @classmethod
+    def from_pretrained(self, *args, **kwargs):
+        requires_pytorch(self)
+
+
+class TapasForSequenceClassification:
+    def __init__(self, *args, **kwargs):
+        requires_pytorch(self)
+
+    @classmethod
+    def from_pretrained(self, *args, **kwargs):
+        requires_pytorch(self)
+
+
+class TapasModel:
+    def __init__(self, *args, **kwargs):
+        requires_pytorch(self)
+
+    @classmethod
+    def from_pretrained(self, *args, **kwargs):
+        requires_pytorch(self)
+
+
 TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST = None


@@ -14,117 +14,70 @@

 import unittest

-from transformers import BertConfig, is_flax_available
-from transformers.testing_utils import require_flax, slow
+from numpy import ndarray

-from .test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask
+from transformers import BertTokenizerFast, TensorType, is_flax_available, is_torch_available
+from transformers.testing_utils import require_flax, require_torch


 if is_flax_available():
-    from transformers.models.bert.modeling_flax_bert import FlaxBertForMaskedLM, FlaxBertModel
+    import os

+    os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.12"  # assumed parallelism: 8

-class FlaxBertModelTester(unittest.TestCase):
-    def __init__(
-        self,
-        parent,
-        batch_size=13,
-        seq_length=7,
-        is_training=True,
-        use_attention_mask=True,
-        use_token_type_ids=True,
-        use_labels=True,
-        vocab_size=99,
-        hidden_size=32,
-        num_hidden_layers=5,
-        num_attention_heads=4,
-        intermediate_size=37,
-        hidden_act="gelu",
-        hidden_dropout_prob=0.1,
-        attention_probs_dropout_prob=0.1,
-        max_position_embeddings=512,
-        type_vocab_size=16,
-        type_sequence_label_size=2,
-        initializer_range=0.02,
-    ):
-        self.parent = parent
-        self.batch_size = batch_size
-        self.seq_length = seq_length
-        self.is_training = is_training
-        self.use_attention_mask = use_attention_mask
-        self.use_token_type_ids = use_token_type_ids
-        self.use_labels = use_labels
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.intermediate_size = intermediate_size
-        self.hidden_act = hidden_act
-        self.hidden_dropout_prob = hidden_dropout_prob
-        self.attention_probs_dropout_prob = attention_probs_dropout_prob
-        self.max_position_embeddings = max_position_embeddings
-        self.type_vocab_size = type_vocab_size
-        self.type_sequence_label_size = type_sequence_label_size
-        self.initializer_range = initializer_range
+    import jax
+    from transformers.models.bert.modeling_flax_bert import FlaxBertModel

-    def prepare_config_and_inputs(self):
-        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
+if is_torch_available():
+    import torch

-        attention_mask = None
-        if self.use_attention_mask:
-            attention_mask = random_attention_mask([self.batch_size, self.seq_length])
-
-        token_type_ids = None
-        if self.use_token_type_ids:
-            token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
-
-        config = BertConfig(
-            vocab_size=self.vocab_size,
-            hidden_size=self.hidden_size,
-            num_hidden_layers=self.num_hidden_layers,
-            num_attention_heads=self.num_attention_heads,
-            intermediate_size=self.intermediate_size,
-            hidden_act=self.hidden_act,
-            hidden_dropout_prob=self.hidden_dropout_prob,
-            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
-            max_position_embeddings=self.max_position_embeddings,
-            type_vocab_size=self.type_vocab_size,
-            is_decoder=False,
-            initializer_range=self.initializer_range,
-        )
-
-        return config, input_ids, token_type_ids, attention_mask
-
-    def prepare_config_and_inputs_for_common(self):
-        config_and_inputs = self.prepare_config_and_inputs()
-        config, input_ids, token_type_ids, attention_mask = config_and_inputs
-        inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
-        return config, inputs_dict
+    from transformers.models.bert.modeling_bert import BertModel


@require_flax
-class FlaxBertModelTest(FlaxModelTesterMixin, unittest.TestCase):
+@require_torch
+class FlaxBertModelTest(unittest.TestCase):
+    def assert_almost_equals(self, a: ndarray, b: ndarray, tol: float):
+        diff = (a - b).sum()
+        self.assertLessEqual(diff, tol, f"Difference between torch and flax is {diff} (>= {tol})")

-    all_model_classes = (FlaxBertModel, FlaxBertForMaskedLM) if is_flax_available() else ()
+    def test_from_pytorch(self):
+        with torch.no_grad():
+            with self.subTest("bert-base-cased"):
+                tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")
+                fx_model = FlaxBertModel.from_pretrained("bert-base-cased")
+                pt_model = BertModel.from_pretrained("bert-base-cased")

-    def setUp(self):
-        self.model_tester = FlaxBertModelTester(self)
+                # Check for simple input
+                pt_inputs = tokenizer.encode_plus("This is a simple input", return_tensors=TensorType.PYTORCH)
+                fx_inputs = tokenizer.encode_plus("This is a simple input", return_tensors=TensorType.JAX)
+                pt_outputs = pt_model(**pt_inputs).to_tuple()
+                fx_outputs = fx_model(**fx_inputs)

-    def test_naming_convention(self):
-        for model_class in self.all_model_classes:
-            model_class_name = model_class.__name__
-            module_class_name = (
-                model_class_name[:-5] + "Module" if model_class_name[-5:] == "Model" else model_class_name + "Module"
-            )
-            bert_modeling_flax_module = __import__(
-                "transformers.models.bert.modeling_flax_bert", fromlist=[module_class_name]
-            )
-            module_cls = getattr(bert_modeling_flax_module, module_class_name)
+                self.assertEqual(len(fx_outputs), len(pt_outputs), "Output lengths differ between Flax and PyTorch")

-            self.assertIsNotNone(module_cls)
+                for fx_output, pt_output in zip(fx_outputs, pt_outputs):
+                    self.assert_almost_equals(fx_output, pt_output.numpy(), 5e-3)

-    @slow
-    def test_model_from_pretrained(self):
-        for model_class_name in self.all_model_classes:
-            model = model_class_name.from_pretrained("bert-base-cased")
-            self.assertIsNotNone(model)
+    def test_multiple_sequences(self):
+        tokenizer = BertTokenizerFast.from_pretrained("bert-base-cased")
+        model = FlaxBertModel.from_pretrained("bert-base-cased")
+
+        sequences = ["this is an example sentence", "this is another", "and a third one"]
+        encodings = tokenizer(sequences, return_tensors=TensorType.JAX, padding=True, truncation=True)
+
+        @jax.jit
+        def model_jitted(input_ids, attention_mask=None, token_type_ids=None):
+            return model(input_ids, attention_mask, token_type_ids)
+
+        with self.subTest("JIT Disabled"):
+            with jax.disable_jit():
+                tokens, pooled = model_jitted(**encodings)
+                self.assertEqual(tokens.shape, (3, 7, 768))
+                self.assertEqual(pooled.shape, (3, 768))
+
+        with self.subTest("JIT Enabled"):
+            jitted_tokens, jitted_pooled = model_jitted(**encodings)
+
+            self.assertEqual(jitted_tokens.shape, (3, 7, 768))
+            self.assertEqual(jitted_pooled.shape, (3, 768))
@@ -1,144 +0,0 @@
-# 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.
-
-import random
-import tempfile
-
-import numpy as np
-
-import transformers
-from transformers import is_flax_available, is_torch_available
-from transformers.testing_utils import require_flax, require_torch
-
-
-if is_flax_available():
-    import os
-
-    import jax
-    import jax.numpy as jnp
-    from transformers.modeling_flax_utils import convert_state_dict_from_pt
-
-    os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.12"  # assumed parallelism: 8
-
-if is_torch_available():
-    import torch
-
-
-def ids_tensor(shape, vocab_size, rng=None):
-    """Creates a random int32 tensor of the shape within the vocab size."""
-    if rng is None:
-        rng = random.Random()
-
-    total_dims = 1
-    for dim in shape:
-        total_dims *= dim
-
-    values = []
-    for _ in range(total_dims):
-        values.append(rng.randint(0, vocab_size - 1))
-
-    output = np.array(values, dtype=jnp.int32).reshape(shape)
-
-    return output
-
-
-def random_attention_mask(shape, rng=None):
-    attn_mask = ids_tensor(shape, vocab_size=2, rng=rng)
-    # make sure that at least one token is attended to for each batch
-    attn_mask[:, -1] = 1
-    return attn_mask
-
-
-@require_flax
-class FlaxModelTesterMixin:
-    model_tester = None
-    all_model_classes = ()
-
-    def assert_almost_equals(self, a: np.ndarray, b: np.ndarray, tol: float):
-        diff = np.abs((a - b)).max()
-        self.assertLessEqual(diff, tol, f"Difference between torch and flax is {diff} (>= {tol}).")
-
-    @require_torch
-    def test_equivalence_flax_pytorch(self):
-        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
-        for model_class in self.all_model_classes:
-            with self.subTest(model_class.__name__):
-                pt_model_class_name = model_class.__name__[4:]  # Skip the "Flax" at the beginning
-                pt_model_class = getattr(transformers, pt_model_class_name)
-                pt_model = pt_model_class(config).eval()
-
-                fx_state = convert_state_dict_from_pt(model_class, pt_model.state_dict(), config)
-                fx_model = model_class(config, fx_state, dtype=jnp.float32)
-
-                pt_inputs = {k: torch.tensor(v.tolist()) for k, v in inputs_dict.items()}
-
-                with torch.no_grad():
-                    pt_outputs = pt_model(**pt_inputs).to_tuple()
-                fx_outputs = fx_model(**inputs_dict)
-                self.assertEqual(len(fx_outputs), len(pt_outputs), "Output lengths differ between Flax and PyTorch")
-                for fx_output, pt_output in zip(fx_outputs, pt_outputs):
-                    self.assert_almost_equals(fx_output, pt_output.numpy(), 1e-3)
-
-                with tempfile.TemporaryDirectory() as tmpdirname:
-                    pt_model.save_pretrained(tmpdirname)
-                    fx_model_loaded = model_class.from_pretrained(tmpdirname, from_pt=True)
-
-                fx_outputs_loaded = fx_model_loaded(**inputs_dict)
-                self.assertEqual(
-                    len(fx_outputs_loaded), len(pt_outputs), "Output lengths differ between Flax and PyTorch"
-                )
-                for fx_output_loaded, pt_output in zip(fx_outputs_loaded, pt_outputs):
-                    self.assert_almost_equals(fx_output_loaded, pt_output.numpy(), 5e-3)
-
-    def test_from_pretrained_save_pretrained(self):
-        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
-        for model_class in self.all_model_classes:
-            with self.subTest(model_class.__name__):
-                model = model_class(config)
-                model.init(model.key, inputs_dict["input_ids"].shape)
-
-                outputs = model(**inputs_dict)
-
-                with tempfile.TemporaryDirectory() as tmpdirname:
-                    model.save_pretrained(tmpdirname)
-                    model_loaded = model_class.from_pretrained(tmpdirname)
-
-                outputs_loaded = model_loaded(**inputs_dict)
-                for output_loaded, output in zip(outputs_loaded, outputs):
-                    self.assert_almost_equals(output_loaded, output, 5e-3)
-
-    def test_jit_compilation(self):
-        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
-
-        for model_class in self.all_model_classes:
-            with self.subTest(model_class.__name__):
-                model = model_class(config)
-                model.init(model.key, inputs_dict["input_ids"].shape)
-
-                @jax.jit
-                def model_jitted(input_ids, attention_mask=None, token_type_ids=None):
-                    return model(input_ids, attention_mask, token_type_ids)
-
-                with self.subTest("JIT Disabled"):
-                    with jax.disable_jit():
-                        outputs = model_jitted(**inputs_dict)
-
-                with self.subTest("JIT Enabled"):
-                    jitted_outputs = model_jitted(**inputs_dict)
-
-                self.assertEqual(len(outputs), len(jitted_outputs))
-                for jitted_output, output in zip(jitted_outputs, outputs):
-                    self.assertEqual(jitted_output.shape, output.shape)
@@ -14,117 +14,70 @@

 import unittest

-from transformers import RobertaConfig, is_flax_available
-from transformers.testing_utils import require_flax, slow
+from numpy import ndarray

-from .test_modeling_flax_common import FlaxModelTesterMixin, ids_tensor, random_attention_mask
+from transformers import RobertaTokenizerFast, TensorType, is_flax_available, is_torch_available
+from transformers.testing_utils import require_flax, require_torch


 if is_flax_available():
+    import os
+
+    os.environ["XLA_PYTHON_CLIENT_MEM_FRACTION"] = "0.12"  # assumed parallelism: 8
+
+    import jax
    from transformers.models.roberta.modeling_flax_roberta import FlaxRobertaModel

+if is_torch_available():
+    import torch

-class FlaxRobertaModelTester(unittest.TestCase):
-    def __init__(
-        self,
-        parent,
-        batch_size=13,
-        seq_length=7,
-        is_training=True,
-        use_attention_mask=True,
-        use_token_type_ids=True,
-        use_labels=True,
-        vocab_size=99,
-        hidden_size=32,
-        num_hidden_layers=5,
-        num_attention_heads=4,
-        intermediate_size=37,
-        hidden_act="gelu",
-        hidden_dropout_prob=0.1,
-        attention_probs_dropout_prob=0.1,
-        max_position_embeddings=512,
-        type_vocab_size=16,
-        type_sequence_label_size=2,
-        initializer_range=0.02,
-    ):
-        self.parent = parent
-        self.batch_size = batch_size
-        self.seq_length = seq_length
-        self.is_training = is_training
-        self.use_attention_mask = use_attention_mask
-        self.use_token_type_ids = use_token_type_ids
-        self.use_labels = use_labels
-        self.vocab_size = vocab_size
-        self.hidden_size = hidden_size
-        self.num_hidden_layers = num_hidden_layers
-        self.num_attention_heads = num_attention_heads
-        self.intermediate_size = intermediate_size
-        self.hidden_act = hidden_act
-        self.hidden_dropout_prob = hidden_dropout_prob
-        self.attention_probs_dropout_prob = attention_probs_dropout_prob
-        self.max_position_embeddings = max_position_embeddings
-        self.type_vocab_size = type_vocab_size
-        self.type_sequence_label_size = type_sequence_label_size
-        self.initializer_range = initializer_range
-
-    def prepare_config_and_inputs(self):
-        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
-
-        attention_mask = None
-        if self.use_attention_mask:
-            attention_mask = random_attention_mask([self.batch_size, self.seq_length])
-
-        token_type_ids = None
-        if self.use_token_type_ids:
-            token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
-
-        config = RobertaConfig(
-            vocab_size=self.vocab_size,
-            hidden_size=self.hidden_size,
-            num_hidden_layers=self.num_hidden_layers,
-            num_attention_heads=self.num_attention_heads,
-            intermediate_size=self.intermediate_size,
-            hidden_act=self.hidden_act,
-            hidden_dropout_prob=self.hidden_dropout_prob,
-            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
-            max_position_embeddings=self.max_position_embeddings,
-            type_vocab_size=self.type_vocab_size,
-            is_decoder=False,
-            initializer_range=self.initializer_range,
-        )
-
-        return config, input_ids, token_type_ids, attention_mask
-
-    def prepare_config_and_inputs_for_common(self):
-        config_and_inputs = self.prepare_config_and_inputs()
-        config, input_ids, token_type_ids, attention_mask = config_and_inputs
-        inputs_dict = {"input_ids": input_ids, "token_type_ids": token_type_ids, "attention_mask": attention_mask}
-        return config, inputs_dict
+    from transformers.models.roberta.modeling_roberta import RobertaModel


@require_flax
-class FlaxRobertaModelTest(FlaxModelTesterMixin, unittest.TestCase):
+@require_torch
+class FlaxRobertaModelTest(unittest.TestCase):
+    def assert_almost_equals(self, a: ndarray, b: ndarray, tol: float):
+        diff = (a - b).sum()
+        self.assertLessEqual(diff, tol, f"Difference between torch and flax is {diff} (>= {tol})")

-    all_model_classes = (FlaxRobertaModel,) if is_flax_available() else ()
+    def test_from_pytorch(self):
+        with torch.no_grad():
+            with self.subTest("roberta-base"):
+                tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")
+                fx_model = FlaxRobertaModel.from_pretrained("roberta-base")
+                pt_model = RobertaModel.from_pretrained("roberta-base")

-    def setUp(self):
-        self.model_tester = FlaxRobertaModelTester(self)
+                # Check for simple input
+                pt_inputs = tokenizer.encode_plus("This is a simple input", return_tensors=TensorType.PYTORCH)
+                fx_inputs = tokenizer.encode_plus("This is a simple input", return_tensors=TensorType.JAX)
+                pt_outputs = pt_model(**pt_inputs)
+                fx_outputs = fx_model(**fx_inputs)

-    def test_naming_convention(self):
-        for model_class in self.all_model_classes:
-            model_class_name = model_class.__name__
-            module_class_name = (
-                model_class_name[:-5] + "Module" if model_class_name[-5:] == "Model" else model_class_name + "Module"
-            )
-            bert_modeling_flax_module = __import__(
-                "transformers.models.roberta.modeling_flax_roberta", fromlist=[module_class_name]
-            )
-            module_cls = getattr(bert_modeling_flax_module, module_class_name)
+                self.assertEqual(len(fx_outputs), len(pt_outputs), "Output lengths differ between Flax and PyTorch")

-            self.assertIsNotNone(module_cls)
+                for fx_output, pt_output in zip(fx_outputs, pt_outputs.to_tuple()):
+                    self.assert_almost_equals(fx_output, pt_output.numpy(), 5e-3)

-    @slow
-    def test_model_from_pretrained(self):
-        for model_class_name in self.all_model_classes:
-            model = model_class_name.from_pretrained("roberta-base")
-            self.assertIsNotNone(model)
+    def test_multiple_sequences(self):
+        tokenizer = RobertaTokenizerFast.from_pretrained("roberta-base")
+        model = FlaxRobertaModel.from_pretrained("roberta-base")
+
+        sequences = ["this is an example sentence", "this is another", "and a third one"]
+        encodings = tokenizer(sequences, return_tensors=TensorType.JAX, padding=True, truncation=True)
+
+        @jax.jit
+        def model_jitted(input_ids, attention_mask=None, token_type_ids=None):
+            return model(input_ids, attention_mask, token_type_ids)
+
+        with self.subTest("JIT Disabled"):
+            with jax.disable_jit():
+                tokens, pooled = model_jitted(**encodings)
+                self.assertEqual(tokens.shape, (3, 7, 768))
+                self.assertEqual(pooled.shape, (3, 768))
+
+        with self.subTest("JIT Enabled"):
+            jitted_tokens, jitted_pooled = model_jitted(**encodings)
+
+            self.assertEqual(jitted_tokens.shape, (3, 7, 768))
+            self.assertEqual(jitted_pooled.shape, (3, 768))
@@ -584,7 +584,7 @@ class TokenizerTesterMixin:

                # We want to have sequence 0 and sequence 1 are tagged
                # respectively with 0 and 1 token_ids
-                # (regardeless of weither the model use token type ids)
+                # (regardless of whether the model use token type ids)
                # We use this assumption in the QA pipeline among other place
                output = tokenizer(seq_0, return_token_type_ids=True)
                self.assertIn(0, output["token_type_ids"])
@@ -600,7 +600,7 @@ class TokenizerTesterMixin:

                # We want to have sequence 0 and sequence 1 are tagged
                # respectively with 0 and 1 token_ids
-                # (regardeless of weither the model use token type ids)
+                # (regardless of whether the model use token type ids)
                # We use this assumption in the QA pipeline among other place
                output = tokenizer(seq_0)
                self.assertIn(0, output.sequence_ids())
@@ -17,7 +17,6 @@ import importlib
 import inspect
 import os
 import re
-from pathlib import Path


 # All paths are set with the intent you should run this script from the root of the repo with the command
@@ -58,6 +57,28 @@ TEST_FILES_WITH_NO_COMMON_TESTS = [
    "test_modeling_xlm_roberta.py",
 ]

+# Update this list for models that are not documented with a comment explaining the reason it should not be.
+# Being in this list is an exception and should **not** be the rule.
+IGNORE_NON_DOCUMENTED = [
+    "BartDecoder",  # Building part of bigger (documented) model.
+    "BartEncoder",  # Building part of bigger (documented) model.
+    "DPREncoder",  # Building part of bigger (documented) model.
+    "DPRSpanPredictor",  # Building part of bigger (documented) model.
+    "T5Stack",  # Building part of bigger (documented) model.
+    "TFDPREncoder",  # Building part of bigger (documented) model.
+    "TFDPRSpanPredictor",  # Building part of bigger (documented) model.
+]
+
+# Update this dict with any special correspondance model name (used in modeling_xxx.py) to doc file.
+MODEL_NAME_TO_DOC_FILE = {
+    "openai": "gpt.rst",
+    "transfo_xl": "transformerxl.rst",
+    "xlm_prophetnet": "xlmprophetnet.rst",
+    "xlm_roberta": "xlmroberta.rst",
+    "bert_generation": "bertgeneration.rst",
+    "marian": "marian.rst",
+}
+
 # Update this list for models that are not in any of the auto MODEL_XXX_MAPPING. Being in this list is an exception and
 # should **not** be the rule.
 IGNORE_NON_AUTO_CONFIGURED = [
@@ -171,6 +192,22 @@ def get_model_test_files():
    return test_files


+# If some doc source files should be ignored when checking models are all documented, they should be added in the
+# nested list _ignore_modules of this function.
+def get_model_doc_files():
+    """ Get the model doc files."""
+    _ignore_modules = [
+        "auto",
+        "dialogpt",
+        "retribert",
+    ]
+    doc_files = []
+    for filename in os.listdir(PATH_TO_DOC):
+        if os.path.isfile(f"{PATH_TO_DOC}/{filename}") and not os.path.splitext(filename)[0] in _ignore_modules:
+            doc_files.append(filename)
+    return doc_files
+
+
 # This is a bit hacky but I didn't find a way to import the test_file as a module and read inside the tester class
 # for the all_model_classes variable.
 def find_tested_models(test_file):
@@ -232,6 +269,58 @@ def check_all_models_are_tested():
        raise Exception(f"There were {len(failures)} failures:\n" + "\n".join(failures))


+def find_documented_classes(doc_file):
+    """ Parse the content of doc_file to detect which classes it documents"""
+    with open(os.path.join(PATH_TO_DOC, doc_file), "r", encoding="utf-8", newline="\n") as f:
+        content = f.read()
+    return re.findall(r"autoclass:: transformers.(\S+)\s+", content)
+
+
+def check_models_are_documented(module, doc_file):
+    """ Check models defined in module are documented in doc_file."""
+    defined_models = get_models(module)
+    documented_classes = find_documented_classes(doc_file)
+    failures = []
+    for model_name, _ in defined_models:
+        if model_name not in documented_classes and model_name not in IGNORE_NON_DOCUMENTED:
+            failures.append(
+                f"{model_name} is defined in {module.__name__} but is not documented in "
+                + f"{os.path.join(PATH_TO_DOC, doc_file)}. Add it to that file."
+                + "If this model should not be documented, add its name to `IGNORE_NON_DOCUMENTED`"
+                + "in the file `utils/check_repo.py`."
+            )
+    return failures
+
+
+def _get_model_name(module):
+    """ Get the model name for the module defining it."""
+    module_name = module.__name__.split(".")[-1]
+    splits = module_name.split("_")
+    splits = splits[(2 if splits[1] in ["flax", "tf"] else 1) :]
+    return "_".join(splits)
+
+
+def check_all_models_are_documented():
+    """ Check all models are properly documented."""
+    modules = get_model_modules()
+    doc_files = get_model_doc_files()
+    failures = []
+    for module in modules:
+        model_name = _get_model_name(module)
+        doc_file = MODEL_NAME_TO_DOC_FILE.get(model_name, f"{model_name}.rst")
+        if doc_file not in doc_files:
+            failures.append(
+                f"{module.__name__} does not have its corresponding doc file {doc_file}. "
+                + f"If the doc file exists but isn't named {doc_file}, update `MODEL_NAME_TO_DOC_FILE` "
+                + "in the file `utils/check_repo.py`."
+            )
+        new_failures = check_models_are_documented(module, doc_file)
+        if new_failures is not None:
+            failures += new_failures
+    if len(failures) > 0:
+        raise Exception(f"There were {len(failures)} failures:\n" + "\n".join(failures))
+
+
 def get_all_auto_configured_models():
    """ Return the list of all models in at least one auto class."""
    result = set()  # To avoid duplicates we concatenate all model classes in a set.
@@ -307,154 +396,13 @@ def check_all_decorator_order():
        )


-def find_all_documented_objects():
-    """ Parse the content of all doc files to detect which classes and functions it documents"""
-    documented_obj = []
-    for doc_file in Path(PATH_TO_DOC).glob("**/*.rst"):
-        with open(doc_file) as f:
-            content = f.read()
-        raw_doc_objs = re.findall(r"(?:autoclass|autofunction):: transformers.(\S+)\s+", content)
-        documented_obj += [obj.split(".")[-1] for obj in raw_doc_objs]
-    return documented_obj
-
-
-# One good reason for not being documented is to be deprecated. Put in this list deprecated objects.
-DEPRECATED_OBJECTS = [
-    "AutoModelWithLMHead",
-    "GlueDataset",
-    "GlueDataTrainingArguments",
-    "LineByLineTextDataset",
-    "LineByLineWithRefDataset",
-    "LineByLineWithSOPTextDataset",
-    "PretrainedBartModel",
-    "PretrainedFSMTModel",
-    "SingleSentenceClassificationProcessor",
-    "SquadDataTrainingArguments",
-    "SquadDataset",
-    "SquadExample",
-    "SquadFeatures",
-    "SquadV1Processor",
-    "SquadV2Processor",
-    "TFAutoModelWithLMHead",
-    "TextDataset",
-    "TextDatasetForNextSentencePrediction",
-    "glue_compute_metrics",
-    "glue_convert_examples_to_features",
-    "glue_output_modes",
-    "glue_processors",
-    "glue_tasks_num_labels",
-    "squad_convert_examples_to_features",
-    "xnli_compute_metrics",
-    "xnli_output_modes",
-    "xnli_processors",
-    "xnli_tasks_num_labels",
-]
-
-# Exceptionally, some objects should not be documented after all rules passed.
-# ONLY PUT SOMETHING IN THIS LIST AS A LAST RESORT!
-UNDOCUMENTED_OBJECTS = [
-    "AddedToken",  # This is a tokenizers class.
-    "BasicTokenizer",  # Internal, should never have been in the main init.
-    "DPRPretrainedReader",  # Like an Encoder.
-    "ModelCard",  # Internal type.
-    "SqueezeBertModule",  # Internal building block (should have been called SqueezeBertLayer)
-    "TFDPRPretrainedReader",  # Like an Encoder.
-    "TransfoXLCorpus",  # Internal type.
-    "WordpieceTokenizer",  # Internal, should never have been in the main init.
-    "absl",  # External module
-    "add_end_docstrings",  # Internal, should never have been in the main init.
-    "add_start_docstrings",  # Internal, should never have been in the main init.
-    "cached_path",  # Internal used for downloading models.
-    "convert_tf_weight_name_to_pt_weight_name",  # Internal used to convert model weights
-    "logger",  # Internal logger
-    "logging",  # External module
-]
-
-# This list should be empty. Objects in it should get their own doc page.
-SHOULD_HAVE_THEIR_OWN_PAGE = [
-    # bert-japanese
-    "BertJapaneseTokenizer",
-    "CharacterTokenizer",
-    "MecabTokenizer",
-    # Bertweet
-    "BertweetTokenizer",
-    # Herbert
-    "HerbertTokenizer",
-    "HerbertTokenizerFast",
-    # Phoebus
-    "PhobertTokenizer",
-    # Benchmarks
-    "PyTorchBenchmark",
-    "PyTorchBenchmarkArguments",
-    "TensorFlowBenchmark",
-    "TensorFlowBenchmarkArguments",
-]
-
-
-def ignore_undocumented(name):
-    """Rules to determine if `name` should be undocumented."""
-    # NOT DOCUMENTED ON PURPOSE.
-    # Magic attributes are not documented.
-    if name.startswith("__"):
-        return True
-    # Constants uppercase are not documented.
-    if name.isupper():
-        return True
-    # PreTrainedModels / Encoders / Decoders / Layers / Embeddings / Attention are not documented.
-    if (
-        name.endswith("PreTrainedModel")
-        or name.endswith("Decoder")
-        or name.endswith("Encoder")
-        or name.endswith("Layer")
-        or name.endswith("Embeddings")
-        or name.endswith("Attention")
-    ):
-        return True
-    # Submodules are not documented.
-    if os.path.isdir(os.path.join(PATH_TO_TRANSFORMERS, name)) or os.path.isfile(
-        os.path.join(PATH_TO_TRANSFORMERS, f"{name}.py")
-    ):
-        return True
-    # All load functions are not documented.
-    if name.startswith("load_tf") or name.startswith("load_pytorch"):
-        return True
-    # is_xxx_available functions are not documented.
-    if name.startswith("is_") and name.endswith("_available"):
-        return True
-    # Deprecated objects are not documented.
-    if name in DEPRECATED_OBJECTS or name in UNDOCUMENTED_OBJECTS:
-        return True
-    # MMBT model does not really work.
-    if name.startswith("MMBT"):
-        return True
-
-    # NOT DOCUMENTED BUT NOT ON PURPOSE, SHOULD BE FIXED!
-    # All data collators should be documented
-    if name.startswith("DataCollator") or name.endswith("data_collator"):
-        return True
-    if name in SHOULD_HAVE_THEIR_OWN_PAGE:
-        return True
-    return False
-
-
-def check_all_objects_are_documented():
-    """ Check all models are properly documented."""
-    documented_objs = find_all_documented_objects()
-    undocumented_objs = [c for c in dir(transformers) if c not in documented_objs and not ignore_undocumented(c)]
-    if len(undocumented_objs) > 0:
-        raise Exception(
-            "The following objects are in the public init so should be documented:\n - "
-            + "\n - ".join(undocumented_objs)
-        )
-
-
 def check_repo_quality():
    """ Check all models are properly tested and documented."""
    print("Checking all models are properly tested.")
    check_all_decorator_order()
    check_all_models_are_tested()
-    print("Checking all objects are properly documented.")
-    check_all_objects_are_documented
+    print("Checking all models are properly documented.")
+    check_all_models_are_documented()
    print("Checking all models are in at least one auto class.")
    check_all_models_are_auto_configured()
Author	SHA1	Message	Date
LysandreJik	93843345dd	Truncate when overflowing	2020-12-10 17:13:02 -05:00
LysandreJik	f872d4c5ed	Fix the masked LM model	2020-12-10 15:25:33 -05:00
LysandreJik	a1ec0a6a07	Test the Masked LM model	2020-12-10 15:18:02 -05:00
LysandreJik	29d7e4a1ae	Address CI needs	2020-12-10 14:55:17 -05:00
LysandreJik	d240f471a0	Quality	2020-12-09 16:31:08 -05:00
Rogge Niels	4ceab1bc88	Make style & make quality	2020-12-09 21:18:21 +01:00
Lysandre	1f20a93df2	Disable torchscript test	2020-12-09 21:18:21 +01:00
Lysandre	212f4b0d59	Fix truncation when no max length is specified	2020-12-09 21:18:21 +01:00
Rogge Niels	71890e5082	Added more model names	2020-12-09 21:18:21 +01:00
Rogge Niels	abb1e4cb3b	Update PRETRAINED_VOCAB_FILES_MAP and PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES	2020-12-09 21:18:21 +01:00
Rogge Niels	0b8a882226	Comment out positional embedding sizes	2020-12-09 21:18:21 +01:00
Rogge Niels	9e7f856ad8	Change position embeddings sizes to 512 instead of 1024	2020-12-09 21:18:21 +01:00
LysandreJik	de8308d41f	Fix remaining tests	2020-12-09 21:18:21 +01:00
Rogge Niels	9fa0881bc2	Incorporate suggestions by Google authors	2020-12-09 21:18:21 +01:00
Rogge Niels	4cb0339945	Add test for convert_logits_to_predictions method of TapasTokenizer	2020-12-09 21:18:21 +01:00
Rogge Niels	0a40f4300e	Remove print statement	2020-12-09 21:18:21 +01:00
Rogge Niels	0072132a11	All model integration tests pass	2020-12-09 21:18:21 +01:00
Rogge Niels	15788934bc	Fix multiple bugs with integration tests	2020-12-09 21:18:21 +01:00
Rogge Niels	40afab66e6	Remove model cards as they are already on the hub	2020-12-09 21:18:21 +01:00
Rogge Niels	3da59dd1a6	First draft model cards of base sized models	2020-12-09 21:18:21 +01:00
Rogge Niels	ebf850205b	Use Python 3 inheritance rather than Python 2	2020-12-09 21:18:21 +01:00
Rogge Niels	53eade1f89	Improve testing by adding torch_device to inputs and adding dependency on scatter	2020-12-09 21:18:21 +01:00
Rogge Niels	e1c74ea4ac	Fix multiple bugs including soft dependency and new structure of the library	2020-12-09 21:18:21 +01:00
Rogge Niels	a149897eef	First commit: adding all files from tapas_v3	2020-12-09 21:18:15 +01:00