boom boom

.circleci/config.yml

@@ -77,14 +77,13 @@ jobs:
                       - v0.3-torch_and_tf-{{ checksum "setup.py" }}
                       - v0.3-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,tf-cpu,torch,testing]
             - run: pip install codecov pytest-cov
             - save_cache:
                 key: v0.3-{{ checksum "setup.py" }}
                 paths:
                     - '~/.cache/pip'
-            - run: python -m pytest -n 8 --dist=loadfile -rA -s ./tests/ --cov  | tee output.txt
+            - run: python -m pytest -n 8 --dist=loadfile -s ./tests/ --cov  | tee output.txt
             - run: codecov
             - store_artifacts:
                   path: ~/transformers/output.txt
@@ -104,13 +103,12 @@ jobs:
                       - v0.3-torch-{{ checksum "setup.py" }}
                       - v0.3-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,torch,testing]
             - save_cache:
                   key: v0.3-torch-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: python -m pytest -n 8 --dist=loadfile -rA -s ./tests/ | tee output.txt
+            - run: python -m pytest -n 8 --dist=loadfile -s ./tests/ | tee output.txt
             - store_artifacts:
                   path: ~/transformers/output.txt
                   destination: test_output.txt
@@ -129,13 +127,12 @@ jobs:
                       - v0.3-tf-{{ checksum "setup.py" }}
                       - v0.3-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install git+https://github.com/huggingface/datasets
             - run: pip install .[sklearn,tf-cpu,testing]
             - save_cache:
                   key: v0.3-tf-{{ checksum "setup.py" }}
                   paths:
                       - '~/.cache/pip'
-            - run: python -m pytest -n 8 --dist=loadfile -rA -s ./tests/ | tee output.txt
+            - run: python -m pytest -n 8 --dist=loadfile -s ./tests/ | tee output.txt
             - store_artifacts:
                path: ~/transformers/output.txt
                destination: test_output.txt
@@ -238,7 +235,8 @@ jobs:
                       - v0.3-code_quality-{{ checksum "setup.py" }}
                       - v0.3-{{ checksum "setup.py" }}
             - run: pip install --upgrade pip
-            - run: pip install isort
+            # we need a version of isort with https://github.com/timothycrosley/isort/pull/1000
+            - run: pip install git+git://github.com/timothycrosley/isort.git@e63ae06ec7d70b06df9e528357650281a3d3ec22#egg=isort
             - run: pip install .[tf,torch,quality]
             - save_cache:
                   key: v0.3-code_quality-{{ checksum "setup.py" }}

.circleci/deploy.sh

@@ -47,5 +47,4 @@ deploy_doc "e7cfc1a" v2.9.0
 deploy_doc "7cb203f" v2.9.1
 deploy_doc "10d7239" v2.10.0 
 deploy_doc "b42586e" v2.11.0
-deploy_doc "7fb8bdf" v3.0.2
-deploy_doc "4b3ee9c" # v3.1.0 Latest stable release
+deploy_doc "7fb8bdf" #v3.0.2 Latest stable release

.github/PULL_REQUEST_TEMPLATE.md

@@ -1,2 +0,0 @@
-<!-- This line specifies which issue to close after the pull request is merged. -->
-Fixes #{issue number}

.github/workflows/self-push.yml

@@ -45,8 +45,7 @@ jobs:
         source .env/bin/activate
         pip install --upgrade pip
         pip install torch!=1.6.0
-        pip install .[sklearn,testing,onnxruntime]
-        pip install git+https://github.com/huggingface/datasets
+        pip install .[sklearn,testing]
 
     - name: Are GPUs recognized by our DL frameworks
       run: |

.github/workflows/self-scheduled.yml

@@ -42,8 +42,7 @@ jobs:
         source .env/bin/activate
         pip install --upgrade pip
         pip install torch!=1.6.0
-        pip install .[sklearn,testing,onnxruntime]
-        pip install git+https://github.com/huggingface/datasets
+        pip install .[sklearn,testing]
 
     - name: Are GPUs recognized by our DL frameworks
       run: |
@@ -59,7 +58,6 @@ jobs:
       run: |
         source .env/bin/activate
         python -m pytest -n 1 --dist=loadfile -s ./tests/
-
     - name: Run examples tests on GPU
       env:
         TF_FORCE_GPU_ALLOW_GROWTH: "true"

.gitignore

@@ -11,7 +11,6 @@ __pycache__/
 # tests and logs
 tests/fixtures
 logs/
-lightning_logs/
 
 # Distribution / packaging
 .Python
@@ -140,7 +139,6 @@ runs
 /wandb
 /examples/runs
 /examples/**/*.args
-/examples/rag/sweep
 
 # data
 /data

CONTRIBUTING.md

@@ -134,18 +134,12 @@ Follow these steps to start contributing:
    it with `pip uninstall transformers` before reinstalling it in editable
    mode with the `-e` flag.)
 
-   To run the full test suite, you might need the additional dependency on `datasets` which requires a separate source
-   install:
+   Right now, we need an unreleased version of `isort` to avoid a
+   [bug](https://github.com/timothycrosley/isort/pull/1000):
 
    ```bash
-   $ git clone https://github.com/huggingface/datasets
-   $ cd datasets
-   $ pip install -e .
+   $ pip install -U git+git://github.com/timothycrosley/isort.git@e63ae06ec7d70b06df9e528357650281a3d3ec22#egg=isort
    ```
-
-   If you have already cloned that repo, you might need to `git pull` to get the most recent changes in the `datasets`
-   library.
-
 5. Develop the features on your branch.
 
    As you work on the features, you should make sure that the test suite
@@ -177,16 +171,6 @@ Follow these steps to start contributing:
    $ make quality
    ```
 
-   If you're modifying documents under `docs/source`, make sure to validate that
-   they can still be built. This check also runs in CI. To run a local check
-   make sure you have installed the documentation builder requirements, by
-   running `pip install .[tf,torch,docs]` once from the root of this repository
-   and then run:
-
-   ```bash
-   $ make docs
-   ```
-
    Once you're happy with your changes, add changed files using `git add` and
    make a commit with `git commit` to record your changes locally:
 

Makefile

@@ -1,10 +1,10 @@
-.PHONY: quality style test test-examples docs
+.PHONY: quality style test test-examples
 
 # Check that source code meets quality standards
 
 quality:
 	black --check --line-length 119 --target-version py35 examples templates tests src utils
-	isort --check-only examples templates tests src utils
+	isort --check-only --recursive examples templates tests src utils
 	flake8 examples templates tests src utils
 	python utils/check_repo.py
 
@@ -12,7 +12,7 @@ quality:
 
 style:
 	black --line-length 119 --target-version py35 examples templates tests src utils
-	isort examples templates tests src utils
+	isort --recursive examples templates tests src utils
 
 # Run tests for the library
 
@@ -23,8 +23,3 @@ test:
 
 test-examples:
 	python -m pytest -n auto --dist=loadfile -s -v ./examples/
-
-# Check that docs can build
-
-docs:
-	cd docs && make html SPHINXOPTS="-W"

README.md

@@ -172,10 +172,8 @@ for Open-Domain Question Answering](https://arxiv.org/abs/2004.04906) by Vladimi
 Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih.
 23. **[Pegasus](https://github.com/google-research/pegasus)** (from Google) released with the paper [PEGASUS: Pre-training with Extracted Gap-sentences for Abstractive Summarization](https://arxiv.org/abs/1912.08777)> by Jingqing Zhang, Yao Zhao, Mohammad Saleh and Peter J. Liu.
 24. **[MBart](https://github.com/pytorch/fairseq/tree/master/examples/mbart)** (from Facebook) released with the paper  [Multilingual Denoising Pre-training for Neural Machine Translation](https://arxiv.org/abs/2001.08210) by Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov, Marjan Ghazvininejad, Mike Lewis, Luke Zettlemoyer.  
-25. **[LXMERT](https://github.com/airsplay/lxmert)** (from UNC Chapel Hill) released with the paper [LXMERT: Learning Cross-Modality Encoder Representations from Transformers for Open-Domain Question Answering](https://arxiv.org/abs/1908.07490) by Hao Tan and Mohit Bansal.
-26. **[Funnel Transformer](https://github.com/laiguokun/Funnel-Transformer)** (from CMU/Google Brain) released with the paper [Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing](https://arxiv.org/abs/2006.03236) by Zihang Dai, Guokun Lai, Yiming Yang, Quoc V. Le.
-27. **[Other community models](https://huggingface.co/models)**, contributed by the [community](https://huggingface.co/users).
-28. Want to contribute a new model? We have added a **detailed guide and templates** to guide you in the process of adding a new model. You can find them in the [`templates`](./templates) folder of the repository. Be sure to check the [contributing guidelines](./CONTRIBUTING.md) and contact the maintainers or open an issue to collect feedbacks before starting your PR.
+25. **[Other community models](https://huggingface.co/models)**, contributed by the [community](https://huggingface.co/users).
+26. Want to contribute a new model? We have added a **detailed guide and templates** to guide you in the process of adding a new model. You can find them in the [`templates`](./templates) folder of the repository. Be sure to check the [contributing guidelines](./CONTRIBUTING.md) and contact the maintainers or open an issue to collect feedbacks before starting your PR.
 
 These implementations have been tested on several datasets (see the example scripts) and should match the performances of the original implementations (e.g. ~93 F1 on SQuAD for BERT Whole-Word-Masking, ~88 F1 on RocStories for OpenAI GPT, ~18.3 perplexity on WikiText 103 for Transformer-XL, ~0.916 Pearson R coefficient on STS-B for XLNet). You can find more details on the performances in the Examples section of the [documentation](https://huggingface.co/transformers/examples.html).
 
@@ -627,7 +625,7 @@ Breaking change in the `from_pretrained()` method:
 
 1. Models are now set in evaluation mode by default when instantiated with the `from_pretrained()` method. To train them, don't forget to set them back in training mode (`model.train()`) to activate the dropout modules.
 
-2. The additional `*input` and `**kwargs` arguments supplied to the `from_pretrained()` method used to be directly passed to the underlying model's class `__init__()` method. They are now used to update the model configuration attribute instead, which can break derived model classes built based on the previous `BertForSequenceClassification` examples. We are working on a way to mitigate this breaking change in [#866](https://github.com/huggingface/transformers/pull/866) by forwarding the model's `__init__()` method (i) the provided positional arguments and (ii) the keyword arguments which do not match any configuration class attributes.
+2. The additional `*input` and `**kwargs` arguments supplied to the `from_pretrained()` method used to be directly passed to the underlying model's class `__init__()` method. They are now used to update the model configuration attribute instead, which can break derived model classes built based on the previous `BertForSequenceClassification` examples. We are working on a way to mitigate this breaking change in [#866](https://github.com/huggingface/transformers/pull/866) by forwarding the the model's `__init__()` method (i) the provided positional arguments and (ii) the keyword arguments which do not match any configuration class attributes.
 
 Also, while not a breaking change, the serialization methods have been standardized and you probably should switch to the new method `save_pretrained(save_directory)` if you were using any other serialization method before.
 
@@ -700,11 +698,11 @@ for batch in train_data:
 
 ## Citation
 
-We now have a [paper](https://arxiv.org/abs/1910.03771) you can cite for the 🤗 Transformers library:
+We now have a paper you can cite for the 🤗 Transformers library:
 ```bibtex
 @article{Wolf2019HuggingFacesTS,
   title={HuggingFace's Transformers: State-of-the-art Natural Language Processing},
-  author={Thomas Wolf and Lysandre Debut and Victor Sanh and Julien Chaumond and Clement Delangue and Anthony Moi and Pierric Cistac and Tim Rault and Rémi Louf and Morgan Funtowicz and Joe Davison and Sam Shleifer and Patrick von Platen and Clara Ma and Yacine Jernite and Julien Plu and Canwen Xu and Teven Le Scao and Sylvain Gugger and Mariama Drame and Quentin Lhoest and Alexander M. Rush},
+  author={Thomas Wolf and Lysandre Debut and Victor Sanh and Julien Chaumond and Clement Delangue and Anthony Moi and Pierric Cistac and Tim Rault and R'emi Louf and Morgan Funtowicz and Jamie Brew},
   journal={ArXiv},
   year={2019},
   volume={abs/1910.03771}

docs/source/_static/js/custom.js

@@ -1,11 +1,10 @@
 // These two things need to be updated at each release for the version selector.
 // Last stable version
-const stableVersion = "v3.1.0"
+const stableVersion = "v3.0.2"
 // Dictionary doc folder to label
 const versionMapping = {
     "master": "master",
-    "": "v3.1.0 (stable)",
-    "v3.0.2": "v3.0.0/v3.0.1/v3.0.2 (stable)",
+    "": "v3.0.0/v3.0.1/v3.0.2 (stable)",
     "v2.11.0": "v2.11.0",
     "v2.10.0": "v2.10.0",
     "v2.9.1": "v2.9.0/v2.9.1",

docs/source/conf.py

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

docs/source/glossary.rst

@@ -128,7 +128,7 @@ The encoded versions have different lengths:
     >>> len(encoded_sequence_a), len(encoded_sequence_b)
     (8, 19)
 
-Therefore, we can't put them together in the same tensor as-is. The first sequence needs to be padded up to the length
+Therefore, we can't be put then together in a same tensor as-is. The first sequence needs to be padded up to the length
 of the second one, or the second one needs to be truncated down to the length of the first one.
 
 In the first case, the list of IDs will be extended by the padding indices. We can pass a list to the tokenizer and ask

docs/source/index.rst

@@ -128,16 +128,7 @@ conversion utilities for the following models:
     <https://arxiv.org/abs/1912.08777>`_ by Jingqing Zhang, Yao Zhao, Mohammad Saleh and Peter J. Liu.
 24. `MBart <https://github.com/pytorch/fairseq/tree/master/examples/mbart>`_ (from Facebook) released with the paper  `Multilingual Denoising Pre-training for Neural Machine Translation <https://arxiv.org/abs/2001.08210>`_ by Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov,
     Marjan Ghazvininejad, Mike Lewis, Luke Zettlemoyer.  
-25. `LXMERT <https://github.com/airsplay/lxmert>`_ (from UNC Chapel Hill) released with the paper `LXMERT: Learning
-    Cross-Modality Encoder Representations from Transformers for Open-Domain Question
-    Answering <https://arxiv.org/abs/1908.07490>`_ by Hao Tan and Mohit Bansal.
-26. `Funnel Transformer <https://github.com/laiguokun/Funnel-Transformer>`_ (from CMU/Google Brain) released with the paper
-    `Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing
-    <https://arxiv.org/abs/2006.03236>`_ by Zihang Dai, Guokun Lai, Yiming Yang, Quoc V. Le.
-27. `Bert For Sequence Generation <https://tfhub.dev/s?module-type=text-generation&subtype=module,placeholder>`_ (from Google) released with the paper
-    `Leveraging Pre-trained Checkpoints for Sequence Generation Tasks
-    <https://arxiv.org/abs/1907.12461>`_ by Sascha Rothe, Shashi Narayan, Aliaksei Severyn.
-28. `Other community models <https://huggingface.co/models>`_, contributed by the `community
+25. `Other community models <https://huggingface.co/models>`_, contributed by the `community
     <https://huggingface.co/users>`_.
 
 .. toctree::
@@ -194,7 +185,6 @@ conversion utilities for the following models:
     main_classes/trainer
     main_classes/optimizer_schedules
     main_classes/processors
-    main_classes/logging
     model_doc/auto
     model_doc/encoderdecoder
     model_doc/bert
@@ -222,9 +212,6 @@ conversion utilities for the following models:
     model_doc/dpr
     model_doc/pegasus
     model_doc/mbart
-    model_doc/funnel
-    model_doc/lxmert
-    model_doc/bertgeneration
     internal/modeling_utils
     internal/tokenization_utils
     internal/pipelines_utils

docs/source/main_classes/configuration.rst

@@ -1,13 +1,12 @@
 Configuration
 ----------------------------------------------------
 
-The base class :class:`~transformers.PretrainedConfig` implements the common methods for loading/saving a configuration
-either from a local file or directory, or from a pretrained model configuration provided by the library (downloaded
-from HuggingFace's AWS S3 repository).
+The base class ``PretrainedConfig`` implements the common methods for loading/saving a configuration either from a
+local file or directory, or from a pretrained model configuration provided by the library (downloaded from
+HuggingFace's AWS S3 repository).
 
-
-PretrainedConfig
-~~~~~~~~~~~~~~~~
+``PretrainedConfig``
+~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: transformers.PretrainedConfig
     :members:

docs/source/main_classes/logging.rst

@@ -1,52 +0,0 @@
-Logging
--------
-
-🤗 Transformers has a centralized logging system, so that you can setup the verbosity of the library easily.
-
-Currently the default verbosity of the library is ``WARNING``.
-
-To change the level of verbosity, just use one of the direct setters. For instance, here is how to change the verbosity to the INFO level.
-
-.. code-block:: python
-
-    import transformers
-    transformers.logging.set_verbosity_info()
-
-You can also use the environment variable ``TRANSFORMERS_VERBOSITY`` to override the default verbosity. You can set it to one of the following: ``debug``, ``info``, ``warning``, ``error``, ``critical``. For example:
-
-.. code-block:: bash
-               
-    TRANSFORMERS_VERBOSITY=error ./myprogram.py
-
-All the methods of this logging module are documented below, the main ones are
-:func:`transformers.logging.get_verbosity` to get the current level of verbosity in the logger and
-:func:`transformers.logging.set_verbosity` to set the verbosity to the level of your choice. In order (from the least
-verbose to the most verbose), those levels (with their corresponding int values in parenthesis) are:
-
-- :obj:`transformers.logging.CRITICAL` or :obj:`transformers.logging.FATAL` (int value, 50): only report the most
-  critical errors.
-- :obj:`transformers.logging.ERROR` (int value, 40): only report errors.
-- :obj:`transformers.logging.WARNING` or :obj:`transformers.logging.WARN` (int value, 30): only reports error and
-  warnings. This the default level used by the library.
-- :obj:`transformers.logging.INFO` (int value, 20): reports error, warnings and basic information.
-- :obj:`transformers.logging.DEBUG` (int value, 10): report all information.
-
-Base setters
-~~~~~~~~~~~~
-
-.. autofunction:: transformers.logging.set_verbosity_error
-
-.. autofunction:: transformers.logging.set_verbosity_warning
-
-.. autofunction:: transformers.logging.set_verbosity_info
-
-.. autofunction:: transformers.logging.set_verbosity_debug
-
-Other functions
-~~~~~~~~~~~~~~~
-
-.. autofunction:: transformers.logging.get_verbosity
-
-.. autofunction:: transformers.logging.set_verbosity
-
-.. autofunction:: transformers.logging.get_logger

docs/source/main_classes/optimizer_schedules.rst

@@ -13,11 +13,6 @@ The ``.optimization`` module provides:
 .. autoclass:: transformers.AdamW
     :members:
 
-``AdaFactor`` (PyTorch)
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.Adafactor
-
 ``AdamWeightDecay`` (TensorFlow)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/main_classes/pipelines.rst

@@ -21,7 +21,6 @@ There are two categories of pipeline abstractions to be aware about:
     - :class:`~transformers.TokenClassificationPipeline`
     - :class:`~transformers.TranslationPipeline`
     - :class:`~transformers.ZeroShotClassificationPipeline`
-    - :class:`~transformers.Text2TextGenerationPipeline`
 
 The pipeline abstraction
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -92,13 +91,6 @@ TextGenerationPipeline
     :special-members: __call__
     :members:
 
-Text2TextGenerationPipeline
-==========================================
-
-.. autoclass:: transformers.Text2TextGenerationPipeline
-    :special-members: __call__
-    :members:
-
 TokenClassificationPipeline
 ==========================================
 
@@ -113,6 +105,7 @@ ZeroShotClassificationPipeline
     :special-members: __call__
     :members:
 
+
 Parent class: :obj:`Pipeline`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 

docs/source/main_classes/trainer.rst

@@ -21,25 +21,12 @@ previous features. To inject custom behavior you can subclass them and override
 - **setup_wandb** -- Setups wandb (see `here <https://docs.wandb.com/huggingface>`__ for more information).
 - **create_optimizer_and_scheduler** -- Setups the optimizer and learning rate scheduler if they were not passed at
   init.
-- **compute_loss** - Computes the loss on a batch of training inputs.
 - **training_step** -- Performs a training step.
 - **prediction_step** -- Performs an evaluation/test step.
 - **run_model** (TensorFlow only) -- Basic pass through the model.
 - **evaluate** -- Runs an evaluation loop and returns metrics.
 - **predict** -- Returns predictions (with metrics if labels are available) on a test set.
 
-Here is an example of how to customize :class:`~transformers.Trainer` using a custom loss function:
-
-.. code-block:: python
-
-    from transformers import Trainer
-    class MyTrainer(Trainer):
-        def compute_loss(self, model, inputs):
-            labels = inputs.pop("labels")
-            outputs = models(**inputs)
-            logits = outputs[0]
-            return my_custom_loss(logits, labels)
-
 
 ``Trainer`` 
 ~~~~~~~~~~~

docs/source/model_doc/bart.rst

@@ -32,7 +32,7 @@ BartForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: transformers.BartForConditionalGeneration
-    :members: forward
+    :members: generate, forward
 
 
 BartConfig

docs/source/model_doc/bertgeneration.rst

@@ -1,82 +0,0 @@
-BertGeneration
-----------------------------------------------------
-
-Overview
-~~~~~~~~~~~~~~~~~~~~~
-
-The BertGeneration model is a BERT model that can be leveraged for sequence-to-sequence tasks using :class:`~transformers.EncoderDecoderModel` as proposed in `Leveraging Pre-trained Checkpoints for Sequence Generation Tasks <https://arxiv.org/abs/1907.12461>`__ by Sascha Rothe, Shashi Narayan, Aliaksei Severyn.
-
-The abstract from the paper is the following:
-
-*Unsupervised pre-training of large neural models has recently revolutionized Natural Language Processing. By warm-starting from the publicly released checkpoints, NLP practitioners have pushed the state-of-the-art on multiple benchmarks while saving significant amounts of compute time. So far the focus has been mainly on the Natural Language Understanding tasks. In this paper, we demonstrate the efficacy of pre-trained checkpoints for Sequence Generation. We developed a Transformer-based sequence-to-sequence model that is compatible with publicly available pre-trained BERT, GPT-2 and RoBERTa checkpoints and conducted an extensive empirical study on the utility of initializing our model, both encoder and decoder, with these checkpoints. Our models result in new state-of-the-art results on Machine Translation, Text Summarization, Sentence Splitting, and Sentence Fusion.*
-
-Usage:
-
-- The model can be used in combination with the :class:`~transformers.EncoderDecoderModel` to leverage two bert pretrained bert checkpoints for subsequent fine-tuning.
-
-::
-  
-  # leverage checkpoints for Bert2Bert model...
-  encoder = BertGenerationEncoder.from_pretrained("bert-large-uncased", bos_token_id=101, eos_token_id=102)  # use BERT's cls token as BOS token and sep token as EOS token
-  decoder = BertGenerationDecoder.from_pretrained("bert-large-uncased", add_cross_attention=True, is_decoder=True, bos_token_id=101, eos_token_id=102)  # add cross attention layers and use BERT's cls token as BOS token and sep token as EOS token
-  bert2bert = EncoderDecoderModel(encoder=encoder, decoder=decoder)
-  
-  # create tokenizer...
-  tokenizer = BertTokenizer.from_pretrained("bert-large-uncased")
-
-  input_ids = tokenizer('This is a long article to summarize', add_special_tokens=False, return_tensors="pt").input_ids
-  labels = tokenizer('This is a short summary', return_tensors="pt").input_ids
-
-  # train...
-  loss = bert2bert(input_ids=input_ids, decoder_input_ids=labels, labels=labels, return_dict=True).loss
-  loss.backward()
-
-
-- Pretrained :class:`~transformers.EncoderDecoderModel` are also directly available in the model hub, *e.g.*:
-
-
-::
-
-  # instantiate sentence fusion model
-  sentence_fuser = EncoderDecoderModel.from_pretrained("google/roberta2roberta_L-24_discofuse")
-  tokenizer = AutoTokenizer.from_pretrained("google/roberta2roberta_L-24_discofuse")
-
-  input_ids = tokenizer('This is the first sentence. This is the second sentence.', add_special_tokens=False, return_tensors="pt").input_ids
-
-  outputs = sentence_fuser.generate(input_ids)
-
-  print(tokenizer.decode(outputs[0]))
-
-
-Tips:
-
-- :class:`~transformers.BertGenerationEncoder` and :class:`~transformers.BertGenerationDecoder`  should be used in combination with :class:`~transformers.EncoderDecoder`.
-- For summarization, sentence splitting, sentence fusion and translation, no special tokens are required for the input. Therefore, no EOS token should be added to the end of the input.
-
-The original code can be found `here <https://tfhub.dev/s?module-type=text-generation&subtype=module,placeholder>`__.
-
-BertGenerationConfig
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.BertGenerationConfig
-    :members:
-
-
-BertGenerationTokenizer
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.BertGenerationTokenizer
-    :members: 
-
-BertGenerationEncoder
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.BertGenerationEncoder
-    :members:
-
-
-BertGenerationDecoder
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.BertGenerationDecoder
-    :members:

docs/source/model_doc/encoderdecoder.rst

@@ -1,13 +1,12 @@
 Encoder Decoder Models
 ------------------------
 
-The :class:`~transformers.EncoderDecoderModel` can be used to initialize a sequence-to-sequence model with any pre-trained autoencoding model as the encoder and any pre-trained autoregressive model as the decoder.
+This class can wrap an encoder model, such as ``BertModel`` and a decoder modeling with a language modeling head, such as ``BertForMaskedLM`` into a encoder-decoder model.
 
-The effectiveness of initializing sequence-to-sequence models with pre-trained checkpoints for sequence generation tasks was shown in `Leveraging Pre-trained Checkpoints for Sequence Generation Tasks <https://arxiv.org/abs/1907.12461>`__ by Sascha Rothe, Shashi Narayan, Aliaksei Severyn.
+The ``EncoderDecoderModel`` class allows to instantiate a encoder decoder model using the ``from_encoder_decoder_pretrain`` class method taking a pretrained encoder and pretrained decoder model as an input. 
+The ``EncoderDecoderModel`` is saved using the standard ``save_pretrained()`` method and can also again be loaded using the standard ``from_pretrained()`` method. 
 
-After such an :class:`~transformers.EncoderDecoderModel` has been trained / fine-tuned, it can be saved / loaded just like any other models (see Examples for more information).
-
-An application of this architecture could be to leverage two pre-trained :obj:`transformers.BertModel` models as the encoder and decoder for a summarization model as was shown in: `Text Summarization with Pretrained Encoders <https://arxiv.org/abs/1910.13461>`_ by Yang Liu and Mirella Lapata. 
+An application of this architecture could be *summarization* using two pretrained Bert models as is shown in the paper: `Text Summarization with Pretrained Encoders <https://arxiv.org/abs/1910.13461>`_ by Yang Liu and Mirella Lapata. 
 
 
 ``EncoderDecoderConfig``

docs/source/model_doc/funnel.rst

@@ -1,185 +0,0 @@
-Funnel Transformer
-------------------
-
-Overview
-~~~~~~~~~~~~~~~~~~~~~
-
-The Funnel Transformer model was proposed in the paper
-`Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing
-<https://arxiv.org/abs/2006.03236>`__.
-It is a bidirectional transformer model, like BERT, but with a pooling operation after each block of layers, a bit
-like in traditional convolutional neural networks (CNN) in computer vision.
-
-The abstract from the paper is the following:
-
-*With the success of language pretraining, it is highly desirable to develop more efficient architectures of good
-scalability that can exploit the abundant unlabeled data at a lower cost. To improve the efficiency, we examine the
-much-overlooked redundancy in maintaining a full-length token-level presentation, especially for tasks that only
-require a single-vector presentation of the sequence. With this intuition, we propose Funnel-Transformer which
-gradually compresses the sequence of hidden states to a shorter one and hence reduces the computation cost. More
-importantly, by re-investing the saved FLOPs from length reduction in constructing a deeper or wider model, we further
-improve the model capacity. In addition, to perform token-level predictions as required by common pretraining
-objectives, Funnel-Transformer is able to recover a deep representation for each token from the reduced hidden sequence
-via a decoder. Empirically, with comparable or fewer FLOPs, Funnel-Transformer outperforms the standard Transformer on
-a wide variety of sequence-level prediction tasks, including text classification, language understanding, and reading
-comprehension.*
-
-Tips:
-
-- Since Funnel Transformer uses pooling, the sequence length of the hidden states changes after each block of layers.
-  The base model therefore has a final sequence length that is a quarter of the original one. This model can be used
-  directly for tasks that just require a sentence summary (like sequence classification or multiple choice). For other
-  tasks, the full model is used; this full model has a decoder that upsamples the final hidden states to the same
-  sequence length as the input.
-- The Funnel Transformer checkpoints are all available with a full version and a base version. The first ones should
-  be used for :class:`~transformers.FunnelModel`, :class:`~transformers.FunnelForPreTraining`,
-  :class:`~transformers.FunnelForMaskedLM`, :class:`~transformers.FunnelForTokenClassification` and
-  class:`~transformers.FunnelForQuestionAnswering`. The second ones should be used for
-  :class:`~transformers.FunnelBaseModel`, :class:`~transformers.FunnelForSequenceClassification` and
-  :class:`~transformers.FunnelForMultipleChoice`.
-
-The original code can be found `here <https://github.com/laiguokun/Funnel-Transformer>`_.
-
-
-FunnelConfig
-~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelConfig
-    :members:
-
-
-FunnelTokenizer
-~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelTokenizer
-    :members: build_inputs_with_special_tokens, get_special_tokens_mask,
-        create_token_type_ids_from_sequences, save_vocabulary
-
-
-FunnelTokenizerFast
-~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelTokenizerFast
-    :members:
-
-
-Funnel specific outputs
-~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.modeling_funnel.FunnelForPreTrainingOutput
-    :members:
-
-.. autoclass:: transformers.modeling_tf_funnel.TFFunnelForPreTrainingOutput
-    :members:
-
-
-FunnelBaseModel
-~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelBaseModel
-    :members:
-
-
-FunnelModel
-~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelModel
-    :members:
-
-
-FunnelModelForPreTraining
-~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelForPreTraining
-    :members:
-
-
-FunnelForMaskedLM
-~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelForMaskedLM
-    :members:
-
-
-FunnelForSequenceClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelForSequenceClassification
-    :members:
-
-
-FunnelForMultipleChoice
-~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelForMultipleChoice
-    :members:
-
-
-FunnelForTokenClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelForTokenClassification
-    :members:
-
-
-FunnelForQuestionAnswering
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FunnelForQuestionAnswering
-    :members:
-
-
-TFFunnelBaseModel
-~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFFunnelBaseModel
-    :members:
-
-
-TFFunnelModel
-~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFFunnelModel
-    :members:
-
-
-TFFunnelModelForPreTraining
-~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFFunnelForPreTraining
-    :members:
-
-
-TFFunnelForMaskedLM
-~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFFunnelForMaskedLM
-    :members:
-
-
-TFFunnelForSequenceClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFFunnelForSequenceClassification
-    :members:
-
-
-TFFunnelForMultipleChoice
-~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFFunnelForMultipleChoice
-    :members:
-
-
-TFFunnelForTokenClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFFunnelForTokenClassification
-    :members:
-
-
-TFFunnelForQuestionAnswering
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFFunnelForQuestionAnswering
-    :members:

docs/source/model_doc/lxmert.rst

@@ -1,109 +0,0 @@
-LXMERT
-----------------------------------------------------
-
-Overview
-~~~~~~~~~~~~~~~~~~~~~
-
-The LXMERT model was proposed in `LXMERT: Learning Cross-Modality Encoder Representations from Transformers <https://arxiv.org/abs/1908.07490>`__
-by Hao Tan & Mohit Bansal. It is a series of bidirectional transformer encoders (one for the vision modality, one for the language modality, and then one to fuse both modalities)
-pre-trained using a combination of masked language modeling, visual-language text alignment, ROI-feature regression, masked visual-attribute modeling, masked visual-object modeling, and visual-question answering objectives.
-The pretraining consists of multiple multi-modal datasets: MSCOCO, Visual-Genome + Visual-Genome Question Answering, VQA 2.0, and GQA.
-
-The abstract from the paper is the following:
-
-*Vision-and-language reasoning requires an understanding of visual concepts, language semantics, and, most importantly, the alignment and relationships between these two
-modalities. We thus propose the LXMERT
-(Learning Cross-Modality Encoder Representations from Transformers) framework to learn
-these vision-and-language connections. In
-LXMERT, we build a large-scale Transformer
-model that consists of three encoders: an object relationship encoder, a language encoder,
-and a cross-modality encoder. Next, to endow our model with the capability of connecting vision and language semantics, we
-pre-train the model with large amounts of
-image-and-sentence pairs, via five diverse representative pre-training tasks: masked language modeling, masked object prediction
-(feature regression and label classification),
-cross-modality matching, and image question answering. These tasks help in learning both intra-modality and cross-modality relationships. After fine-tuning from our pretrained parameters, our model achieves the
-state-of-the-art results on two visual question answering datasets (i.e., VQA and GQA).
-We also show the generalizability of our pretrained cross-modality model by adapting it to
-a challenging visual-reasoning task, NLVR
-,
-and improve the previous best result by 22%
-absolute (54% to 76%). Lastly, we demonstrate detailed ablation studies to prove that
-both our novel model components and pretraining strategies significantly contribute to
-our strong results; and also present several
-attention visualizations for the different encoders*
-
-Tips:
-
-- Bounding boxes are not necessary to be used in the visual feature embeddings, any kind of visual-spacial features will work.
-- Both the language hidden states and the visual hidden states that LXMERT outputs are passed through the cross-modality layer, so they
-  contain information from both modalities. To access a modality that only attends to itself, select the vision/language hidden states from the first input in the tuple.
-- The bi-directional cross-modality encoder attention only returns attention values when the language modality is used as the input and the vision modality is used as the context vector. Further,
-  while the cross-modality encoder contains self-attention for each respective modality and cross-attention, only the cross attention is returned and both self attention outputs are disregarded.
-
-The code can be found `here <https://github.com/airsplay/lxmert>`__
-
-
-LxmertConfig
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.LxmertConfig
-    :members:
-
-
-LxmertTokenizer
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.LxmertTokenizer
-    :members: build_inputs_with_special_tokens, get_special_tokens_mask,
-        create_token_type_ids_from_sequences, save_vocabulary
-
-
-Lxmert specific outputs
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.modeling_lxmert.LxmertModelOutput
-    :members:
-
-.. autoclass:: transformers.modeling_lxmert.LxmertForPreTrainingOutput
-    :members:
-
-.. autoclass:: transformers.modeling_lxmert.LxmertForQuestionAnsweringOutput
-    :members:
-
-.. autoclass:: transformers.modeling_tf_lxmert.TFLxmertModelOutput
-    :members:
-
-.. autoclass:: transformers.modeling_tf_lxmert.TFLxmertForPreTrainingOutput
-    :members:
-
-
-LxmertModel
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.LxmertModel
-    :members:
-
-LxmertForPreTraining
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.LxmertForPreTraining
-    :members:
-
-LxmertForQuestionAnswering
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.LxmertForQuestionAnswering
-    :members:
-
-
-TFLxmertModel
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFLxmertModel
-    :members:
-
-TFLxmertForPreTraining
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFLxmertForPreTraining
-    :members:

docs/source/model_doc/pegasus.rst

@@ -48,7 +48,6 @@ Usage Example
 .. code-block:: python
 
     from transformers import PegasusForConditionalGeneration, PegasusTokenizer
-    import torch
     src_text = [
         """ PG&E stated it scheduled the blackouts in response to forecasts for high winds amid dry conditions. The aim is to reduce the risk of wildfires. Nearly 800 thousand customers were scheduled to be affected by the shutoffs which were expected to last through at least midday tomorrow."""
     ]
@@ -60,7 +59,7 @@ Usage Example
     batch = tokenizer.prepare_seq2seq_batch(src_text, truncation=True, padding='longest').to(torch_device)
     translated = model.generate(**batch)
     tgt_text = tokenizer.batch_decode(translated, skip_special_tokens=True)
-    assert tgt_text[0] == "California's largest electricity provider has turned off power to hundreds of thousands of customers."
+    assert tgt_text[0] == "California's largest electricity provider has turned off power to tens of thousands of customers."
 
 PegasusForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

docs/source/model_doc/xlmroberta.rst

@@ -56,13 +56,6 @@ XLMRobertaModel
     :members:
 
 
-XLMRobertaForCausalLM
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.XLMRobertaForCausalLM
-    :members:
-
-
 XLMRobertaForMaskedLM
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
@@ -137,4 +130,4 @@ TFXLMRobertaForQuestionAnswering
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
 
 .. autoclass:: transformers.TFXLMRobertaForQuestionAnswering
-    :members:
+    :members:

docs/source/model_sharing.rst

@@ -136,13 +136,6 @@ Then log in using the same credentials as on huggingface.co. To upload your mode
 
 This will upload the folder containing the weights, tokenizer and configuration we prepared in the previous section.
 
-By default you will be prompted to confirm that you want these files to be uploaded. If you are uploading multiple models and need to script that process, you can add `-y` to bypass the prompt. For example:
-
-::
-
-    transformers-cli upload -y path/to/awesome-name-you-picked/
-
-
 If you want to upload a single file (a new version of your model, or the other framework checkpoint you want to add),
 just type:
 

docs/source/model_summary.rst

@@ -14,7 +14,6 @@ Each one of the models in the library falls into one of the following categories
   * :ref:`autoencoding-models`
   * :ref:`seq-to-seq-models`
   * :ref:`multimodal-models`
-  * :ref:`retrieval-based-models`
 
 Autoregressive models are pretrained on the classic language modeling task: guess the next token having read all the
 previous ones. They correspond to the decoder of the original transformer model, and a mask is used on top of the full
@@ -416,38 +415,6 @@ traditional GAN setting) then the ELECTRA model is trained for a few steps.
 The library provides a version of the model for masked language modeling, token classification and sentence
 classification.
 
-Funnel Transformer
-----------------------------------------------
-
-.. raw:: html
-
-   <a href="https://huggingface.co/models?filter=funnel">
-       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-funnel-blueviolet">
-   </a>
-   <a href="model_doc/funnel.html">
-       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-funnel-blueviolet">
-   </a>
-
-`Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing
-<https://arxiv.org/abs/2006.03236>`_, Zihang Dai et al.
-
-Funnel Transformer is a transformer model using pooling, a bit like a ResNet model: layers are grouped in blocks, and
-at the beginning of each block (except the first one), the hidden states are pooled among the sequence dimension. This
-way, their length is divided by 2, which speeds up the computation of the next hidden states. All pretrained models
-have three blocks, which means the final hidden state has a sequence length that is one fourth of the original sequence
-length.
-
-For tasks such as classification, this is not a problem, but for tasks like masked language modeling or token
-classification, we need a hidden state with the same sequence length as the original input. In those cases, the final
-hidden states are upsampled to the input sequence length and go through two additional layers. That's why there are two
-versions of each checkpoint. The version suffixed with "-base" contains only the three blocks, while the version
-without that suffix contains the three blocks and the upsampling head with its additional layers.
-
-The pretrained models available use the same pretraining objective as ELECTRA.
-
-The library provides a version of the model for masked language modeling, token classification, sentence
-classification, multiple choice classification and question answering.
-
 .. _longformer:
 
 Longformer
@@ -638,40 +605,6 @@ The pretrained model only works for classification.
     More information in this :doc:`model documentation </model_doc/mmbt.html>`.
     TODO: write this page
 
-.. _retrieval-based-models:
-
-Retrieval-based models
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-Some models use documents retrieval during (pre)training and inference for open-domain question answering, for example.
-
-
-DPR
-----------------------------------------------
-
-.. raw:: html
-
-   <a href="https://huggingface.co/models?filter=dpr">
-       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-dpr-blueviolet">
-   </a>
-   <a href="model_doc/ctrl.dpr">
-       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-dpr-blueviolet">
-   </a>
-
-`Dense Passage Retrieval for Open-Domain Question Answering <https://arxiv.org/abs/2004.04906>`_,
-Vladimir Karpukhin et al.
-
-Dense Passage Retrieval (DPR) - is a set of tools and models for state-of-the-art open-domain question-answering research.
-
-
-DPR consists in three models:
-
-  * Question encoder: encode questions as vectors
-  * Context encoder: encode contexts as vectors
-  * Reader: extract the answer of the questions inside retrieved contexts, along with a relevance score (high if the inferred span actually answers the question).
-
-DPR's pipeline (not implemented yet) uses a retrieval step to find the top k contexts given a certain question, and then it calls the reader with the question and the retrieved documents to get the answer.
-
 More technical aspects
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

docs/source/pretrained_models.rst

@@ -5,406 +5,362 @@ Here is the full list of the currently provided pretrained models together with
 
 For a list that includes community-uploaded models, refer to `https://huggingface.co/models <https://huggingface.co/models>`__.
 
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Architecture       | Shortcut name                                              | Details of the model                                                                                                                  |
-+====================+============================================================+=======================================================================================================================================+
-| BERT               | ``bert-base-uncased``                                      | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on lower-cased English text.                                                                                                |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-large-uncased``                                     | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
-|                    |                                                            | | Trained on lower-cased English text.                                                                                                |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-base-cased``                                        | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on cased English text.                                                                                                      |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-large-cased``                                       | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
-|                    |                                                            | | Trained on cased English text.                                                                                                      |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-base-multilingual-uncased``                         | | (Original, not recommended) 12-layer, 768-hidden, 12-heads, 110M parameters.                                                        |
-|                    |                                                            | | Trained on lower-cased text in the top 102 languages with the largest Wikipedias                                                    |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/bert/blob/master/multilingual.md>`__).                                              |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-base-multilingual-cased``                           | | (New, **recommended**) 12-layer, 768-hidden, 12-heads, 110M parameters.                                                             |
-|                    |                                                            | | Trained on cased text in the top 104 languages with the largest Wikipedias                                                          |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/bert/blob/master/multilingual.md>`__).                                              |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-base-chinese``                                      | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on cased Chinese Simplified and Traditional text.                                                                           |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-base-german-cased``                                 | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on cased German text by Deepset.ai                                                                                          |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on deepset.ai website <https://deepset.ai/german-bert>`__).                                                             |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-large-uncased-whole-word-masking``                  | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
-|                    |                                                            | | Trained on lower-cased English text using Whole-Word-Masking                                                                        |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/bert/#bert>`__).                                                                    |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-large-cased-whole-word-masking``                    | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
-|                    |                                                            | | Trained on cased English text using Whole-Word-Masking                                                                              |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/bert/#bert>`__).                                                                    |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-large-uncased-whole-word-masking-finetuned-squad``  | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
-|                    |                                                            | | The ``bert-large-uncased-whole-word-masking`` model fine-tuned on SQuAD                                                             |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see details of fine-tuning in the `example section <https://github.com/huggingface/transformers/tree/master/examples>`__).           |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-large-cased-whole-word-masking-finetuned-squad``    | | 24-layer, 1024-hidden, 16-heads, 340M parameters                                                                                    |
-|                    |                                                            | | The ``bert-large-cased-whole-word-masking`` model fine-tuned on SQuAD                                                               |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details of fine-tuning in the example section <https://huggingface.co/transformers/examples.html>`__)                           |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-base-cased-finetuned-mrpc``                         | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | The ``bert-base-cased`` model fine-tuned on MRPC                                                                                    |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details of fine-tuning in the example section <https://huggingface.co/transformers/examples.html>`__)                           |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-base-german-dbmdz-cased``                           | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on cased German text by DBMDZ                                                                                               |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on dbmdz repository <https://github.com/dbmdz/german-bert>`__).                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``bert-base-german-dbmdz-uncased``                         | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on uncased German text by DBMDZ                                                                                             |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on dbmdz repository <https://github.com/dbmdz/german-bert>`__).                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``cl-tohoku/bert-base-japanese``                           | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on Japanese text. Text is tokenized with MeCab and WordPiece and this requires some extra dependencies,                     |
-|                    |                                                            | | `fugashi <https://github.com/polm/fugashi>`__ which is a wrapper around `MeCab <https://taku910.github.io/mecab/>`__.               |
-|                    |                                                            | | Use ``pip install transformers["ja"]`` (or ``pip install -e .["ja"]`` if you install from source) to install them.                  |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on cl-tohoku repository <https://github.com/cl-tohoku/bert-japanese>`__).                                               |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``cl-tohoku/bert-base-japanese-whole-word-masking``        | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on Japanese text. Text is tokenized with MeCab and WordPiece and this requires some extra dependencies,                     |
-|                    |                                                            | | `fugashi <https://github.com/polm/fugashi>`__ which is a wrapper around `MeCab <https://taku910.github.io/mecab/>`__.               |
-|                    |                                                            | | Use ``pip install transformers["ja"]`` (or ``pip install -e .["ja"]`` if you install from source) to install them.                  |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on cl-tohoku repository <https://github.com/cl-tohoku/bert-japanese>`__).                                               |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``cl-tohoku/bert-base-japanese-char``                      | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on Japanese text. Text is tokenized into characters.                                                                        |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on cl-tohoku repository <https://github.com/cl-tohoku/bert-japanese>`__).                                               |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``cl-tohoku/bert-base-japanese-char-whole-word-masking``   | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on Japanese text using Whole-Word-Masking. Text is tokenized into characters.                                               |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on cl-tohoku repository <https://github.com/cl-tohoku/bert-japanese>`__).                                               |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``TurkuNLP/bert-base-finnish-cased-v1``                    | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on cased Finnish text.                                                                                                      |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on turkunlp.org <http://turkunlp.org/FinBERT/>`__).                                                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``TurkuNLP/bert-base-finnish-uncased-v1``                  | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on uncased Finnish text.                                                                                                    |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on turkunlp.org <http://turkunlp.org/FinBERT/>`__).                                                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``wietsedv/bert-base-dutch-cased``                         | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | Trained on cased Dutch text.                                                                                                        |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details on wietsedv repository <https://github.com/wietsedv/bertje/>`__).                                                       |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| GPT                | ``openai-gpt``                                             | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | OpenAI GPT English model                                                                                                            |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| GPT-2              | ``gpt2``                                                   | | 12-layer, 768-hidden, 12-heads, 117M parameters.                                                                                    |
-|                    |                                                            | | OpenAI GPT-2 English model                                                                                                          |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``gpt2-medium``                                            | | 24-layer, 1024-hidden, 16-heads, 345M parameters.                                                                                   |
-|                    |                                                            | | OpenAI's Medium-sized GPT-2 English model                                                                                           |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``gpt2-large``                                             | | 36-layer, 1280-hidden, 20-heads, 774M parameters.                                                                                   |
-|                    |                                                            | | OpenAI's Large-sized GPT-2 English model                                                                                            |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``gpt2-xl``                                                | | 48-layer, 1600-hidden, 25-heads, 1558M parameters.                                                                                  |
-|                    |                                                            | | OpenAI's XL-sized GPT-2 English model                                                                                               |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Transformer-XL     | ``transfo-xl-wt103``                                       | | 18-layer, 1024-hidden, 16-heads, 257M parameters.                                                                                   |
-|                    |                                                            | | English model trained on wikitext-103                                                                                               |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| XLNet              | ``xlnet-base-cased``                                       | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
-|                    |                                                            | | XLNet English model                                                                                                                 |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlnet-large-cased``                                      | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
-|                    |                                                            | | XLNet Large English model                                                                                                           |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| XLM                | ``xlm-mlm-en-2048``                                        | | 12-layer, 2048-hidden, 16-heads                                                                                                     |
-|                    |                                                            | | XLM English model                                                                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-mlm-ende-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
-|                    |                                                            | | XLM English-German model trained on the concatenation of English and German wikipedia                                               |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-mlm-enfr-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
-|                    |                                                            | | XLM English-French model trained on the concatenation of English and French wikipedia                                               |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-mlm-enro-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
-|                    |                                                            | | XLM English-Romanian Multi-language model                                                                                           |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-mlm-xnli15-1024``                                    | | 12-layer, 1024-hidden, 8-heads                                                                                                      |
-|                    |                                                            | | XLM Model pre-trained with MLM on the `15 XNLI languages <https://github.com/facebookresearch/XNLI>`__.                             |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-mlm-tlm-xnli15-1024``                                | | 12-layer, 1024-hidden, 8-heads                                                                                                      |
-|                    |                                                            | | XLM Model pre-trained with MLM + TLM on the `15 XNLI languages <https://github.com/facebookresearch/XNLI>`__.                       |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-clm-enfr-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
-|                    |                                                            | | XLM English-French model trained with CLM (Causal Language Modeling) on the concatenation of English and French wikipedia           |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-clm-ende-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
-|                    |                                                            | | XLM English-German model trained with CLM (Causal Language Modeling) on the concatenation of English and German wikipedia           |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-mlm-17-1280``                                        | | 16-layer, 1280-hidden, 16-heads                                                                                                     |
-|                    |                                                            | | XLM model trained with MLM (Masked Language Modeling) on 17 languages.                                                              |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-mlm-100-1280``                                       | | 16-layer, 1280-hidden, 16-heads                                                                                                     |
-|                    |                                                            | | XLM model trained with MLM (Masked Language Modeling) on 100 languages.                                                             |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| RoBERTa            | ``roberta-base``                                           | | 12-layer, 768-hidden, 12-heads, 125M parameters                                                                                     |
-|                    |                                                            | | RoBERTa using the BERT-base architecture                                                                                            |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/roberta>`__)                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``roberta-large``                                          | | 24-layer, 1024-hidden, 16-heads, 355M parameters                                                                                    |
-|                    |                                                            | | RoBERTa using the BERT-large architecture                                                                                           |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/roberta>`__)                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``roberta-large-mnli``                                     | | 24-layer, 1024-hidden, 16-heads, 355M parameters                                                                                    |
-|                    |                                                            | | ``roberta-large`` fine-tuned on `MNLI <http://www.nyu.edu/projects/bowman/multinli/>`__.                                            |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/roberta>`__)                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``distilroberta-base``                                     | | 6-layer, 768-hidden, 12-heads, 82M parameters                                                                                       |
-|                    |                                                            | | The DistilRoBERTa model distilled from the RoBERTa model `roberta-base` checkpoint.                                                 |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``roberta-base-openai-detector``                           | | 12-layer, 768-hidden, 12-heads, 125M parameters                                                                                     |
-|                    |                                                            | | ``roberta-base`` fine-tuned by OpenAI on the outputs of the 1.5B-parameter GPT-2 model.                                             |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/openai/gpt-2-output-dataset/tree/master/detector>`__)                                               |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``roberta-large-openai-detector``                          | | 24-layer, 1024-hidden, 16-heads, 355M parameters                                                                                    |
-|                    |                                                            | | ``roberta-large`` fine-tuned by OpenAI on the outputs of the 1.5B-parameter GPT-2 model.                                            |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/openai/gpt-2-output-dataset/tree/master/detector>`__)                                               |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| DistilBERT         | ``distilbert-base-uncased``                                | | 6-layer, 768-hidden, 12-heads, 66M parameters                                                                                       |
-|                    |                                                            | | The DistilBERT model distilled from the BERT model `bert-base-uncased` checkpoint                                                   |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``distilbert-base-uncased-distilled-squad``                | | 6-layer, 768-hidden, 12-heads, 66M parameters                                                                                       |
-|                    |                                                            | | The DistilBERT model distilled from the BERT model `bert-base-uncased` checkpoint, with an additional linear layer.                 |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``distilbert-base-cased``                                  | | 6-layer, 768-hidden, 12-heads, 65M parameters                                                                                       |
-|                    |                                                            | | The DistilBERT model distilled from the BERT model `bert-base-cased` checkpoint                                                     |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``distilbert-base-cased-distilled-squad``                  | | 6-layer, 768-hidden, 12-heads, 65M parameters                                                                                       |
-|                    |                                                            | | The DistilBERT model distilled from the BERT model `bert-base-cased` checkpoint, with an additional question answering layer.       |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``distilgpt2``                                             | | 6-layer, 768-hidden, 12-heads, 82M parameters                                                                                       |
-|                    |                                                            | | The DistilGPT2 model distilled from the GPT2 model `gpt2` checkpoint.                                                               |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``distilbert-base-german-cased``                           | | 6-layer, 768-hidden, 12-heads, 66M parameters                                                                                       |
-|                    |                                                            | | The German DistilBERT model distilled from the German DBMDZ BERT model `bert-base-german-dbmdz-cased` checkpoint.                   |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``distilbert-base-multilingual-cased``                     | | 6-layer, 768-hidden, 12-heads, 134M parameters                                                                                      |
-|                    |                                                            | | The multilingual DistilBERT model distilled from the Multilingual BERT model `bert-base-multilingual-cased` checkpoint.             |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| CTRL               | ``ctrl``                                                   | | 48-layer, 1280-hidden, 16-heads, 1.6B parameters                                                                                    |
-|                    |                                                            | | Salesforce's Large-sized CTRL English model                                                                                         |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| CamemBERT          | ``camembert-base``                                         | | 12-layer, 768-hidden, 12-heads, 110M parameters                                                                                     |
-|                    |                                                            | | CamemBERT using the BERT-base architecture                                                                                          |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/camembert>`__)                                                 |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| ALBERT             | ``albert-base-v1``                                         | | 12 repeating layers, 128 embedding, 768-hidden, 12-heads, 11M parameters                                                            |
-|                    |                                                            | | ALBERT base model                                                                                                                   |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``albert-large-v1``                                        | | 24 repeating layers, 128 embedding, 1024-hidden, 16-heads, 17M parameters                                                           |
-|                    |                                                            | | ALBERT large model                                                                                                                  |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``albert-xlarge-v1``                                       | | 24 repeating layers, 128 embedding, 2048-hidden, 16-heads, 58M parameters                                                           |
-|                    |                                                            | | ALBERT xlarge model                                                                                                                 |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``albert-xxlarge-v1``                                      | | 12 repeating layer, 128 embedding, 4096-hidden, 64-heads, 223M parameters                                                           |
-|                    |                                                            | | ALBERT xxlarge model                                                                                                                |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``albert-base-v2``                                         | | 12 repeating layers, 128 embedding, 768-hidden, 12-heads, 11M parameters                                                            |
-|                    |                                                            | | ALBERT base model with no dropout, additional training data and longer training                                                     |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``albert-large-v2``                                        | | 24 repeating layers, 128 embedding, 1024-hidden, 16-heads, 17M parameters                                                           |
-|                    |                                                            | | ALBERT large model with no dropout, additional training data and longer training                                                    |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``albert-xlarge-v2``                                       | | 24 repeating layers, 128 embedding, 2048-hidden, 16-heads, 58M parameters                                                           |
-|                    |                                                            | | ALBERT xlarge model with no dropout, additional training data and longer training                                                   |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``albert-xxlarge-v2``                                      | | 12 repeating layer, 128 embedding, 4096-hidden, 64-heads, 223M parameters                                                           |
-|                    |                                                            | | ALBERT xxlarge model with no dropout, additional training data and longer training                                                  |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| T5                 | ``t5-small``                                               | | ~60M parameters with 6-layers, 512-hidden-state, 2048 feed-forward hidden-state, 8-heads,                                           |
-|                    |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``t5-base``                                                | | ~220M parameters with 12-layers, 768-hidden-state, 3072 feed-forward hidden-state, 12-heads,                                        |
-|                    |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``t5-large``                                               | | ~770M parameters with 24-layers, 1024-hidden-state, 4096 feed-forward hidden-state, 16-heads,                                       |
-|                    |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``t5-3B``                                                  | | ~2.8B parameters with 24-layers, 1024-hidden-state, 16384 feed-forward hidden-state, 32-heads,                                      |
-|                    |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``t5-11B``                                                 | | ~11B parameters with 24-layers, 1024-hidden-state, 65536 feed-forward hidden-state, 128-heads,                                      |
-|                    |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| XLM-RoBERTa        | ``xlm-roberta-base``                                       | | ~125M parameters with 12-layers, 768-hidden-state, 3072 feed-forward hidden-state, 8-heads,                                         |
-|                    |                                                            | | Trained on on 2.5 TB of newly created clean CommonCrawl data in 100 languages                                                       |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``xlm-roberta-large``                                      | | ~355M parameters with 24-layers, 1027-hidden-state, 4096 feed-forward hidden-state, 16-heads,                                       |
-|                    |                                                            | | Trained on 2.5 TB of newly created clean CommonCrawl data in 100 languages                                                          |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| FlauBERT           | ``flaubert/flaubert_small_cased``                          | | 6-layer, 512-hidden, 8-heads, 54M parameters                                                                                        |
-|                    |                                                            | | FlauBERT small architecture                                                                                                         |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/getalp/Flaubert>`__)                                                                                |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``flaubert/flaubert_base_uncased``                         | | 12-layer, 768-hidden, 12-heads, 137M parameters                                                                                     |
-|                    |                                                            | | FlauBERT base architecture with uncased vocabulary                                                                                  |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/getalp/Flaubert>`__)                                                                                |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``flaubert/flaubert_base_cased``                           | | 12-layer, 768-hidden, 12-heads, 138M parameters                                                                                     |
-|                    |                                                            | | FlauBERT base architecture with cased vocabulary                                                                                    |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/getalp/Flaubert>`__)                                                                                |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``flaubert/flaubert_large_cased``                          | | 24-layer, 1024-hidden, 16-heads, 373M parameters                                                                                    |
-|                    |                                                            | | FlauBERT large architecture                                                                                                         |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/getalp/Flaubert>`__)                                                                                |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Bart               | ``facebook/bart-large``                                    | | 24-layer, 1024-hidden, 16-heads, 406M parameters                                                                                    |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/bart>`_)                                                       |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``facebook/bart-base``                                     | | 12-layer, 768-hidden, 16-heads, 139M parameters                                                                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``facebook/bart-large-mnli``                               | | Adds a 2 layer classification head with 1 million parameters                                                                        |
-|                    |                                                            | | bart-large base architecture with a classification head, finetuned on MNLI                                                          |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``facebook/bart-large-cnn``                                | | 12-layer, 1024-hidden, 16-heads, 406M parameters       (same as base)                                                               |
-|                    |                                                            | | bart-large base architecture finetuned on cnn summarization task                                                                    |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| DialoGPT           | ``DialoGPT-small``                                         | | 12-layer, 768-hidden, 12-heads, 124M parameters                                                                                     |
-|                    |                                                            | | Trained on English text: 147M conversation-like exchanges extracted from Reddit.                                                    |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``DialoGPT-medium``                                        | | 24-layer, 1024-hidden, 16-heads, 355M parameters                                                                                    |
-|                    |                                                            | | Trained on English text: 147M conversation-like exchanges extracted from Reddit.                                                    |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``DialoGPT-large``                                         | | 36-layer, 1280-hidden, 20-heads, 774M parameters                                                                                    |
-|                    |                                                            | | Trained on English text: 147M conversation-like exchanges extracted from Reddit.                                                    |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Reformer           | ``reformer-enwik8``                                        | | 12-layer, 1024-hidden, 8-heads, 149M parameters                                                                                     |
-|                    |                                                            | | Trained on English Wikipedia data - enwik8.                                                                                         |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``reformer-crime-and-punishment``                          | | 6-layer, 256-hidden, 2-heads, 3M parameters                                                                                         |
-|                    |                                                            | | Trained on English text: Crime and Punishment novel by Fyodor Dostoyevsky.                                                          |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| MarianMT           | ``Helsinki-NLP/opus-mt-{src}-{tgt}``                       | | 12-layer, 512-hidden, 8-heads, ~74M parameter Machine translation models. Parameter counts vary depending on vocab size.            |
-|                    |                                                            | | (see `model list <https://huggingface.co/Helsinki-NLP>`_)                                                                           |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Pegasus            | ``google/pegasus-{dataset}``                               | | 16-layer, 1024-hidden, 16-heads, ~568M parameter, 2.2 GB for summary. `model list <https://huggingface.co/models?search=pegasus>`__ |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Longformer         | ``allenai/longformer-base-4096``                           | | 12-layer, 768-hidden, 12-heads, ~149M parameters                                                                                    |
-|                    |                                                            | | Starting from RoBERTa-base checkpoint, trained on documents of max length 4,096                                                     |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``allenai/longformer-large-4096``                          | | 24-layer, 1024-hidden, 16-heads, ~435M parameters                                                                                   |
-|                    |                                                            | | Starting from RoBERTa-large checkpoint, trained on documents of max length 4,096                                                    |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| MBart              | ``facebook/mbart-large-cc25``                              | | 24-layer, 1024-hidden, 16-heads, 610M parameters                                                                                    |
-|                    |                                                            | | mBART (bart-large architecture) model trained on 25 languages' monolingual corpus                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``facebook/mbart-large-en-ro``                             | | 24-layer, 1024-hidden, 16-heads, 610M parameters                                                                                    |
-|                    |                                                            | | mbart-large-cc25 model finetuned on WMT english romanian translation.                                                               |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Lxmert             | ``lxmert-base-uncased``                                    | | 9-language layers, 9-relationship layers, and 12-cross-modality layers                                                              |
-|                    |                                                            | | 768-hidden, 12-heads (for each layer) ~ 228M parameters                                                                             |
-|                    |                                                            | | Starting from lxmert-base checkpoint, trained on over 9 million image-text couplets from COCO, VisualGenome, GQA, VQA               |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-| Funnel Transformer | ``funnel-transformer/small``                               | | 14 layers: 3 blocks of 4 layers then 2 layers decoder, 768-hidden, 12-heads, 130M parameters                                        |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/small-base``                          | | 12 layers: 3 blocks of 4 layers (no decoder), 768-hidden, 12-heads, 115M parameters                                                 |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/medium``                              | | 14 layers: 3 blocks 6, 3x2, 3x2 layers then 2 layers decoder, 768-hidden, 12-heads, 130M parameters                                 |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/medium-base``                         | | 12 layers: 3 blocks 6, 3x2, 3x2 layers(no decoder), 768-hidden, 12-heads, 115M parameters                                           |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/intermediate``                        | | 20 layers: 3 blocks of 6 layers then 2 layers decoder, 768-hidden, 12-heads, 177M parameters                                        |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/intermediate-base``                   | | 18 layers: 3 blocks of 6 layers (no decoder), 768-hidden, 12-heads, 161M parameters                                                 |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/large``                               | | 26 layers: 3 blocks of 8 layers then 2 layers decoder, 1024-hidden, 12-heads, 386M parameters                                       |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/large-base``                          | | 24 layers: 3 blocks of 8 layers (no decoder), 1024-hidden, 12-heads, 358M parameters                                                |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/xlarge``                              | | 32 layers: 3 blocks of 10 layers then 2 layers decoder, 1024-hidden, 12-heads, 468M parameters                                      |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-|                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
-|                    | ``funnel-transformer/xlarge-base``                         | | 30 layers: 3 blocks of 10 layers (no decoder), 1024-hidden, 12-heads, 440M parameters                                               |
-|                    |                                                            |                                                                                                                                       |
-|                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
-+--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| Architecture      | Shortcut name                                              | Details of the model                                                                                                                  |
++===================+============================================================+=======================================================================================================================================+
+| BERT              | ``bert-base-uncased``                                      | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on lower-cased English text.                                                                                                |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-large-uncased``                                     | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
+|                   |                                                            | | Trained on lower-cased English text.                                                                                                |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-base-cased``                                        | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on cased English text.                                                                                                      |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-large-cased``                                       | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
+|                   |                                                            | | Trained on cased English text.                                                                                                      |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-base-multilingual-uncased``                         | | (Original, not recommended) 12-layer, 768-hidden, 12-heads, 110M parameters.                                                        |
+|                   |                                                            | | Trained on lower-cased text in the top 102 languages with the largest Wikipedias                                                    |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/bert/blob/master/multilingual.md>`__).                                              |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-base-multilingual-cased``                           | | (New, **recommended**) 12-layer, 768-hidden, 12-heads, 110M parameters.                                                             |
+|                   |                                                            | | Trained on cased text in the top 104 languages with the largest Wikipedias                                                          |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/bert/blob/master/multilingual.md>`__).                                              |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-base-chinese``                                      | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on cased Chinese Simplified and Traditional text.                                                                           |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-base-german-cased``                                 | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on cased German text by Deepset.ai                                                                                          |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on deepset.ai website <https://deepset.ai/german-bert>`__).                                                             |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-large-uncased-whole-word-masking``                  | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
+|                   |                                                            | | Trained on lower-cased English text using Whole-Word-Masking                                                                        |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/bert/#bert>`__).                                                                    |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-large-cased-whole-word-masking``                    | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
+|                   |                                                            | | Trained on cased English text using Whole-Word-Masking                                                                              |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/bert/#bert>`__).                                                                    |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-large-uncased-whole-word-masking-finetuned-squad``  | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
+|                   |                                                            | | The ``bert-large-uncased-whole-word-masking`` model fine-tuned on SQuAD                                                             |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see details of fine-tuning in the `example section <https://github.com/huggingface/transformers/tree/master/examples>`__).           |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-large-cased-whole-word-masking-finetuned-squad``    | | 24-layer, 1024-hidden, 16-heads, 340M parameters                                                                                    |
+|                   |                                                            | | The ``bert-large-cased-whole-word-masking`` model fine-tuned on SQuAD                                                               |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details of fine-tuning in the example section <https://huggingface.co/transformers/examples.html>`__)                           |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-base-cased-finetuned-mrpc``                         | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | The ``bert-base-cased`` model fine-tuned on MRPC                                                                                    |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details of fine-tuning in the example section <https://huggingface.co/transformers/examples.html>`__)                           |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-base-german-dbmdz-cased``                           | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on cased German text by DBMDZ                                                                                               |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on dbmdz repository <https://github.com/dbmdz/german-bert>`__).                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``bert-base-german-dbmdz-uncased``                         | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on uncased German text by DBMDZ                                                                                             |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on dbmdz repository <https://github.com/dbmdz/german-bert>`__).                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``cl-tohoku/bert-base-japanese``                           | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on Japanese text. Text is tokenized with MeCab and WordPiece and this requires some extra dependencies,                     |
+|                   |                                                            | | `fugashi <https://github.com/polm/fugashi>`__ which is a wrapper around `MeCab <https://taku910.github.io/mecab/>`__.               |
+|                   |                                                            | | Use ``pip install transformers["ja"]`` (or ``pip install -e .["ja"]`` if you install from source) to install them.                  |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on cl-tohoku repository <https://github.com/cl-tohoku/bert-japanese>`__).                                               |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``cl-tohoku/bert-base-japanese-whole-word-masking``        | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on Japanese text. Text is tokenized with MeCab and WordPiece and this requires some extra dependencies,                     |
+|                   |                                                            | | `fugashi <https://github.com/polm/fugashi>`__ which is a wrapper around `MeCab <https://taku910.github.io/mecab/>`__.               |
+|                   |                                                            | | Use ``pip install transformers["ja"]`` (or ``pip install -e .["ja"]`` if you install from source) to install them.                  |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on cl-tohoku repository <https://github.com/cl-tohoku/bert-japanese>`__).                                               |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``cl-tohoku/bert-base-japanese-char``                      | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on Japanese text. Text is tokenized into characters.                                                                        |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on cl-tohoku repository <https://github.com/cl-tohoku/bert-japanese>`__).                                               |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``cl-tohoku/bert-base-japanese-char-whole-word-masking``   | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on Japanese text using Whole-Word-Masking. Text is tokenized into characters.                                               |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on cl-tohoku repository <https://github.com/cl-tohoku/bert-japanese>`__).                                               |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``TurkuNLP/bert-base-finnish-cased-v1``                    | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on cased Finnish text.                                                                                                      |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on turkunlp.org <http://turkunlp.org/FinBERT/>`__).                                                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``TurkuNLP/bert-base-finnish-uncased-v1``                  | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on uncased Finnish text.                                                                                                    |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on turkunlp.org <http://turkunlp.org/FinBERT/>`__).                                                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``wietsedv/bert-base-dutch-cased``                         | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | Trained on cased Dutch text.                                                                                                        |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details on wietsedv repository <https://github.com/wietsedv/bertje/>`__).                                                       |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| GPT               | ``openai-gpt``                                             | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | OpenAI GPT English model                                                                                                            |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| GPT-2             | ``gpt2``                                                   | | 12-layer, 768-hidden, 12-heads, 117M parameters.                                                                                    |
+|                   |                                                            | | OpenAI GPT-2 English model                                                                                                          |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``gpt2-medium``                                            | | 24-layer, 1024-hidden, 16-heads, 345M parameters.                                                                                   |
+|                   |                                                            | | OpenAI's Medium-sized GPT-2 English model                                                                                           |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``gpt2-large``                                             | | 36-layer, 1280-hidden, 20-heads, 774M parameters.                                                                                   |
+|                   |                                                            | | OpenAI's Large-sized GPT-2 English model                                                                                            |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``gpt2-xl``                                                | | 48-layer, 1600-hidden, 25-heads, 1558M parameters.                                                                                  |
+|                   |                                                            | | OpenAI's XL-sized GPT-2 English model                                                                                               |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| Transformer-XL    | ``transfo-xl-wt103``                                       | | 18-layer, 1024-hidden, 16-heads, 257M parameters.                                                                                   |
+|                   |                                                            | | English model trained on wikitext-103                                                                                               |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| XLNet             | ``xlnet-base-cased``                                       | | 12-layer, 768-hidden, 12-heads, 110M parameters.                                                                                    |
+|                   |                                                            | | XLNet English model                                                                                                                 |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlnet-large-cased``                                      | | 24-layer, 1024-hidden, 16-heads, 340M parameters.                                                                                   |
+|                   |                                                            | | XLNet Large English model                                                                                                           |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| XLM               | ``xlm-mlm-en-2048``                                        | | 12-layer, 2048-hidden, 16-heads                                                                                                     |
+|                   |                                                            | | XLM English model                                                                                                                   |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-mlm-ende-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
+|                   |                                                            | | XLM English-German model trained on the concatenation of English and German wikipedia                                               |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-mlm-enfr-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
+|                   |                                                            | | XLM English-French model trained on the concatenation of English and French wikipedia                                               |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-mlm-enro-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
+|                   |                                                            | | XLM English-Romanian Multi-language model                                                                                           |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-mlm-xnli15-1024``                                    | | 12-layer, 1024-hidden, 8-heads                                                                                                      |
+|                   |                                                            | | XLM Model pre-trained with MLM on the `15 XNLI languages <https://github.com/facebookresearch/XNLI>`__.                             |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-mlm-tlm-xnli15-1024``                                | | 12-layer, 1024-hidden, 8-heads                                                                                                      |
+|                   |                                                            | | XLM Model pre-trained with MLM + TLM on the `15 XNLI languages <https://github.com/facebookresearch/XNLI>`__.                       |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-clm-enfr-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
+|                   |                                                            | | XLM English-French model trained with CLM (Causal Language Modeling) on the concatenation of English and French wikipedia           |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-clm-ende-1024``                                      | | 6-layer, 1024-hidden, 8-heads                                                                                                       |
+|                   |                                                            | | XLM English-German model trained with CLM (Causal Language Modeling) on the concatenation of English and German wikipedia           |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-mlm-17-1280``                                        | | 16-layer, 1280-hidden, 16-heads                                                                                                     |
+|                   |                                                            | | XLM model trained with MLM (Masked Language Modeling) on 17 languages.                                                              |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-mlm-100-1280``                                       | | 16-layer, 1280-hidden, 16-heads                                                                                                     |
+|                   |                                                            | | XLM model trained with MLM (Masked Language Modeling) on 100 languages.                                                             |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| RoBERTa           | ``roberta-base``                                           | | 12-layer, 768-hidden, 12-heads, 125M parameters                                                                                     |
+|                   |                                                            | | RoBERTa using the BERT-base architecture                                                                                            |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/roberta>`__)                                                   |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``roberta-large``                                          | | 24-layer, 1024-hidden, 16-heads, 355M parameters                                                                                    |
+|                   |                                                            | | RoBERTa using the BERT-large architecture                                                                                           |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/roberta>`__)                                                   |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``roberta-large-mnli``                                     | | 24-layer, 1024-hidden, 16-heads, 355M parameters                                                                                    |
+|                   |                                                            | | ``roberta-large`` fine-tuned on `MNLI <http://www.nyu.edu/projects/bowman/multinli/>`__.                                            |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/roberta>`__)                                                   |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``distilroberta-base``                                     | | 6-layer, 768-hidden, 12-heads, 82M parameters                                                                                       |
+|                   |                                                            | | The DistilRoBERTa model distilled from the RoBERTa model `roberta-base` checkpoint.                                                 |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``roberta-base-openai-detector``                           | | 12-layer, 768-hidden, 12-heads, 125M parameters                                                                                     |
+|                   |                                                            | | ``roberta-base`` fine-tuned by OpenAI on the outputs of the 1.5B-parameter GPT-2 model.                                             |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/openai/gpt-2-output-dataset/tree/master/detector>`__)                                               |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``roberta-large-openai-detector``                          | | 24-layer, 1024-hidden, 16-heads, 355M parameters                                                                                    |
+|                   |                                                            | | ``roberta-large`` fine-tuned by OpenAI on the outputs of the 1.5B-parameter GPT-2 model.                                            |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/openai/gpt-2-output-dataset/tree/master/detector>`__)                                               |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| DistilBERT        | ``distilbert-base-uncased``                                | | 6-layer, 768-hidden, 12-heads, 66M parameters                                                                                       |
+|                   |                                                            | | The DistilBERT model distilled from the BERT model `bert-base-uncased` checkpoint                                                   |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``distilbert-base-uncased-distilled-squad``                | | 6-layer, 768-hidden, 12-heads, 66M parameters                                                                                       |
+|                   |                                                            | | The DistilBERT model distilled from the BERT model `bert-base-uncased` checkpoint, with an additional linear layer.                 |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``distilbert-base-cased``                                  | | 6-layer, 768-hidden, 12-heads, 65M parameters                                                                                       |
+|                   |                                                            | | The DistilBERT model distilled from the BERT model `bert-base-cased` checkpoint                                                     |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``distilbert-base-cased-distilled-squad``                  | | 6-layer, 768-hidden, 12-heads, 65M parameters                                                                                       |
+|                   |                                                            | | The DistilBERT model distilled from the BERT model `bert-base-cased` checkpoint, with an additional question answering layer.       |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``distilgpt2``                                             | | 6-layer, 768-hidden, 12-heads, 82M parameters                                                                                       |
+|                   |                                                            | | The DistilGPT2 model distilled from the GPT2 model `gpt2` checkpoint.                                                               |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``distilbert-base-german-cased``                           | | 6-layer, 768-hidden, 12-heads, 66M parameters                                                                                       |
+|                   |                                                            | | The German DistilBERT model distilled from the German DBMDZ BERT model `bert-base-german-dbmdz-cased` checkpoint.                   |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``distilbert-base-multilingual-cased``                     | | 6-layer, 768-hidden, 12-heads, 134M parameters                                                                                      |
+|                   |                                                            | | The multilingual DistilBERT model distilled from the Multilingual BERT model `bert-base-multilingual-cased` checkpoint.             |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/huggingface/transformers/tree/master/examples/distillation>`__)                                     |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| CTRL              | ``ctrl``                                                   | | 48-layer, 1280-hidden, 16-heads, 1.6B parameters                                                                                    |
+|                   |                                                            | | Salesforce's Large-sized CTRL English model                                                                                         |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| CamemBERT         | ``camembert-base``                                         | | 12-layer, 768-hidden, 12-heads, 110M parameters                                                                                     |
+|                   |                                                            | | CamemBERT using the BERT-base architecture                                                                                          |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/camembert>`__)                                                 |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| ALBERT            | ``albert-base-v1``                                         | | 12 repeating layers, 128 embedding, 768-hidden, 12-heads, 11M parameters                                                            |
+|                   |                                                            | | ALBERT base model                                                                                                                   |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``albert-large-v1``                                        | | 24 repeating layers, 128 embedding, 1024-hidden, 16-heads, 17M parameters                                                           |
+|                   |                                                            | | ALBERT large model                                                                                                                  |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``albert-xlarge-v1``                                       | | 24 repeating layers, 128 embedding, 2048-hidden, 16-heads, 58M parameters                                                           |
+|                   |                                                            | | ALBERT xlarge model                                                                                                                 |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``albert-xxlarge-v1``                                      | | 12 repeating layer, 128 embedding, 4096-hidden, 64-heads, 223M parameters                                                           |
+|                   |                                                            | | ALBERT xxlarge model                                                                                                                |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``albert-base-v2``                                         | | 12 repeating layers, 128 embedding, 768-hidden, 12-heads, 11M parameters                                                            |
+|                   |                                                            | | ALBERT base model with no dropout, additional training data and longer training                                                     |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``albert-large-v2``                                        | | 24 repeating layers, 128 embedding, 1024-hidden, 16-heads, 17M parameters                                                           |
+|                   |                                                            | | ALBERT large model with no dropout, additional training data and longer training                                                    |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``albert-xlarge-v2``                                       | | 24 repeating layers, 128 embedding, 2048-hidden, 16-heads, 58M parameters                                                           |
+|                   |                                                            | | ALBERT xlarge model with no dropout, additional training data and longer training                                                   |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``albert-xxlarge-v2``                                      | | 12 repeating layer, 128 embedding, 4096-hidden, 64-heads, 223M parameters                                                           |
+|                   |                                                            | | ALBERT xxlarge model with no dropout, additional training data and longer training                                                  |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/google-research/ALBERT>`__)                                                                         |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| T5                | ``t5-small``                                               | | ~60M parameters with 6-layers, 512-hidden-state, 2048 feed-forward hidden-state, 8-heads,                                           |
+|                   |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``t5-base``                                                | | ~220M parameters with 12-layers, 768-hidden-state, 3072 feed-forward hidden-state, 12-heads,                                        |
+|                   |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``t5-large``                                               | | ~770M parameters with 24-layers, 1024-hidden-state, 4096 feed-forward hidden-state, 16-heads,                                       |
+|                   |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``t5-3B``                                                  | | ~2.8B parameters with 24-layers, 1024-hidden-state, 16384 feed-forward hidden-state, 32-heads,                                      |
+|                   |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``t5-11B``                                                 | | ~11B parameters with 24-layers, 1024-hidden-state, 65536 feed-forward hidden-state, 128-heads,                                      |
+|                   |                                                            | | Trained on English text: the Colossal Clean Crawled Corpus (C4)                                                                     |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| XLM-RoBERTa       | ``xlm-roberta-base``                                       | | ~125M parameters with 12-layers, 768-hidden-state, 3072 feed-forward hidden-state, 8-heads,                                         |
+|                   |                                                            | | Trained on on 2.5 TB of newly created clean CommonCrawl data in 100 languages                                                       |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``xlm-roberta-large``                                      | | ~355M parameters with 24-layers, 1027-hidden-state, 4096 feed-forward hidden-state, 16-heads,                                       |
+|                   |                                                            | | Trained on 2.5 TB of newly created clean CommonCrawl data in 100 languages                                                          |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| FlauBERT          | ``flaubert/flaubert_small_cased``                          | | 6-layer, 512-hidden, 8-heads, 54M parameters                                                                                        |
+|                   |                                                            | | FlauBERT small architecture                                                                                                         |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/getalp/Flaubert>`__)                                                                                |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``flaubert/flaubert_base_uncased``                         | | 12-layer, 768-hidden, 12-heads, 137M parameters                                                                                     |
+|                   |                                                            | | FlauBERT base architecture with uncased vocabulary                                                                                  |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/getalp/Flaubert>`__)                                                                                |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``flaubert/flaubert_base_cased``                           | | 12-layer, 768-hidden, 12-heads, 138M parameters                                                                                     |
+|                   |                                                            | | FlauBERT base architecture with cased vocabulary                                                                                    |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/getalp/Flaubert>`__)                                                                                |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``flaubert/flaubert_large_cased``                          | | 24-layer, 1024-hidden, 16-heads, 373M parameters                                                                                    |
+|                   |                                                            | | FlauBERT large architecture                                                                                                         |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/getalp/Flaubert>`__)                                                                                |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| Bart              | ``facebook/bart-large``                                    | | 24-layer, 1024-hidden, 16-heads, 406M parameters                                                                                    |
+|                   |                                                            |                                                                                                                                       |
+|                   |                                                            | (see `details <https://github.com/pytorch/fairseq/tree/master/examples/bart>`_)                                                       |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``facebook/bart-base``                                     | | 12-layer, 768-hidden, 16-heads, 139M parameters                                                                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``facebook/bart-large-mnli``                               | | Adds a 2 layer classification head with 1 million parameters                                                                        |
+|                   |                                                            | | bart-large base architecture with a classification head, finetuned on MNLI                                                          |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``facebook/bart-large-cnn``                                | | 12-layer, 1024-hidden, 16-heads, 406M parameters       (same as base)                                                               |
+|                   |                                                            | | bart-large base architecture finetuned on cnn summarization task                                                                    |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| DialoGPT          | ``DialoGPT-small``                                         | | 12-layer, 768-hidden, 12-heads, 124M parameters                                                                                     |
+|                   |                                                            | | Trained on English text: 147M conversation-like exchanges extracted from Reddit.                                                    |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``DialoGPT-medium``                                        | | 24-layer, 1024-hidden, 16-heads, 355M parameters                                                                                    |
+|                   |                                                            | | Trained on English text: 147M conversation-like exchanges extracted from Reddit.                                                    |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``DialoGPT-large``                                         | | 36-layer, 1280-hidden, 20-heads, 774M parameters                                                                                    |
+|                   |                                                            | | Trained on English text: 147M conversation-like exchanges extracted from Reddit.                                                    |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| Reformer          | ``reformer-enwik8``                                        | | 12-layer, 1024-hidden, 8-heads, 149M parameters                                                                                     |
+|                   |                                                            | | Trained on English Wikipedia data - enwik8.                                                                                         |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``reformer-crime-and-punishment``                          | | 6-layer, 256-hidden, 2-heads, 3M parameters                                                                                         |
+|                   |                                                            | | Trained on English text: Crime and Punishment novel by Fyodor Dostoyevsky.                                                          |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| MarianMT          | ``Helsinki-NLP/opus-mt-{src}-{tgt}``                       | | 12-layer, 512-hidden, 8-heads, ~74M parameter Machine translation models. Parameter counts vary depending on vocab size.            |
+|                   |                                                            | | (see `model list <https://huggingface.co/Helsinki-NLP>`_)                                                                           |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| Pegasus           | ``google/pegasus-{dataset}``                               | | 16-layer, 1024-hidden, 16-heads, ~568M parameter, 2.2 GB for summary. `model list <https://huggingface.co/models?search=pegasus>`__ |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| Longformer        | ``allenai/longformer-base-4096``                           | | 12-layer, 768-hidden, 12-heads, ~149M parameters                                                                                    |
+|                   |                                                            | | Starting from RoBERTa-base checkpoint, trained on documents of max length 4,096                                                     |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``allenai/longformer-large-4096``                          | | 24-layer, 1024-hidden, 16-heads, ~435M parameters                                                                                   |
+|                   |                                                            | | Starting from RoBERTa-large checkpoint, trained on documents of max length 4,096                                                    |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| MBart             | ``facebook/mbart-large-cc25``                              | | 24-layer, 1024-hidden, 16-heads, 610M parameters                                                                                    |
+|                   |                                                            | | mBART (bart-large architecture) model trained on 25 languages' monolingual corpus                                                   |
+|                   +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                   | ``facebook/mbart-large-en-ro``                             | | 24-layer, 1024-hidden, 16-heads, 610M parameters                                                                                    |
+|                   |                                                            | | mbart-large-cc25 model finetuned on WMT english romanian translation.                                                               |
++-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+

docs/source/quicktour.rst

@@ -214,7 +214,7 @@ Using the model
 
 Once your input has been preprocessed by the tokenizer, you can send it directly to the model. As we mentioned, it will
 contain all the relevant information the model needs. If you're using a TensorFlow model, you can pass the
-dictionary keys directly to tensors, for a PyTorch model, you need to unpack the dictionary by adding :obj:`**`.
+dictionary keys directly to tensor, for a PyTorch model, you need to unpack the dictionary by adding :obj:`**`.
 
 .. code-block::
 

docs/source/serialization.rst

@@ -130,12 +130,13 @@ Pytorch's two modules `JIT and TRACE <https://pytorch.org/docs/stable/jit.html>`
 their model to be re-used in other programs, such as efficiency-oriented C++ programs.
 
 We have provided an interface that allows the export of 🤗 Transformers models to TorchScript so that they can
-be reused in a different environment than a Pytorch-based python program. Here we explain how to export and use our models using TorchScript.
+be reused in a different environment than a Pytorch-based python program. Here we explain how to use our models so that
+they can be exported, and what to be mindful of when using these models with TorchScript.
 
-Exporting a model requires two things:
+Exporting a model needs two things:
 
-* a forward pass with dummy inputs.
-* model instantiation with the ``torchscript`` flag.
+* dummy inputs to execute a model forward pass.
+* the model needs to be instantiated with the ``torchscript`` flag.
 
 These necessities imply several things developers should be careful about. These are detailed below.
 
@@ -146,8 +147,8 @@ Implications
 TorchScript flag and tied weights
 ------------------------------------------------
 This flag is necessary because most of the language models in this repository have tied weights between their
-``Embedding`` layer and their ``Decoding`` layer. TorchScript does not allow the export of models that have tied weights, therefore
-it is necessary to untie and clone the weights beforehand.
+``Embedding`` layer and their ``Decoding`` layer. TorchScript does not allow the export of models that have tied weights,
+it is therefore necessary to untie the weights beforehand.
 
 This implies that models instantiated with the ``torchscript`` flag have their ``Embedding`` layer and ``Decoding`` layer
 separate, which means that they should not be trained down the line. Training would de-synchronize the two layers,
@@ -180,7 +181,7 @@ when exporting varying sequence-length models.
 Using TorchScript in Python
 -------------------------------------------------
 
-Below is an example, showing how to save, load models as well as how to use the trace for inference.
+Below are examples of using the Python to save, load models as well as how to use the trace for inference.
 
 Saving a model
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -236,10 +237,10 @@ We are re-using the previously initialised ``dummy_input``.
 
 .. code-block:: python
 
-    loaded_model = torch.jit.load("traced_bert.pt")
+    loaded_model = torch.jit.load("traced_model.pt")
     loaded_model.eval()
 
-    all_encoder_layers, pooled_output = loaded_model(*dummy_input)
+    all_encoder_layers, pooled_output = loaded_model(dummy_input)
 
 Using a traced model for inference
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

examples/adversarial/utils_hans.py

@@ -20,8 +20,8 @@ from dataclasses import dataclass
 from typing import List, Optional, Union
 
 import tqdm
-
 from filelock import FileLock
+
 from transformers import (
     BartTokenizer,
     BartTokenizerFast,
@@ -112,10 +112,7 @@ if is_torch_available():
             cached_features_file = os.path.join(
                 data_dir,
                 "cached_{}_{}_{}_{}".format(
-                    "dev" if evaluate else "train",
-                    tokenizer.__class__.__name__,
-                    str(max_seq_length),
-                    task,
+                    "dev" if evaluate else "train", tokenizer.__class__.__name__, str(max_seq_length), task,
                 ),
             )
             label_list = processor.get_labels()
@@ -281,10 +278,7 @@ class HansProcessor(DataProcessor):
 
 
 def hans_convert_examples_to_features(
-    examples: List[InputExample],
-    label_list: List[str],
-    max_length: int,
-    tokenizer: PreTrainedTokenizer,
+    examples: List[InputExample], label_list: List[str], max_length: int, tokenizer: PreTrainedTokenizer,
 ):
     """
     Loads a data file into a list of ``InputFeatures``

examples/benchmarking/plot_csv_file.py

@@ -20,9 +20,7 @@ class PlotArguments:
     Arguments pertaining to which model/config/tokenizer we are going to fine-tune, or train from scratch.
     """
 
-    csv_file: str = field(
-        metadata={"help": "The csv file to plot."},
-    )
+    csv_file: str = field(metadata={"help": "The csv file to plot."},)
     plot_along_batch: bool = field(
         default=False,
         metadata={"help": "Whether to plot along batch size or sequence lengh. Defaults to sequence length."},
@@ -32,8 +30,7 @@ class PlotArguments:
         metadata={"help": "Whether the csv file has time results or memory results. Defaults to memory results."},
     )
     no_log_scale: bool = field(
-        default=False,
-        metadata={"help": "Disable logarithmic scale when plotting"},
+        default=False, metadata={"help": "Disable logarithmic scale when plotting"},
     )
     is_train: bool = field(
         default=False,
@@ -42,8 +39,7 @@ class PlotArguments:
         },
     )
     figure_png_file: Optional[str] = field(
-        default=None,
-        metadata={"help": "Filename under which the plot will be saved. If unused no plot is saved."},
+        default=None, metadata={"help": "Filename under which the plot will be saved. If unused no plot is saved."},
     )
     short_model_names: Optional[List[str]] = list_field(
         default=None, metadata={"help": "List of model names that are used instead of the ones in the csv file."}

examples/bert-loses-patience/pabee/modeling_pabee_albert.py

@@ -101,30 +101,30 @@ class AlbertModelWithPabee(AlbertModel):
         regression=False,
     ):
         r"""
-        Return:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
-            last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
-                Sequence of hidden-states at the output of the last layer of the model.
-            pooler_output (:obj:`torch.FloatTensor`: of shape :obj:`(batch_size, hidden_size)`):
-                Last layer hidden-state of the first token of the sequence (classification token)
-                further processed by a Linear layer and a Tanh activation function. The Linear
-                layer weights are trained from the next sentence prediction (classification)
-                objective during pre-training.
+    Return:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
+        last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        pooler_output (:obj:`torch.FloatTensor`: of shape :obj:`(batch_size, hidden_size)`):
+            Last layer hidden-state of the first token of the sequence (classification token)
+            further processed by a Linear layer and a Tanh activation function. The Linear
+            layer weights are trained from the next sentence prediction (classification)
+            objective during pre-training.
 
-                This output is usually *not* a good summary
-                of the semantic content of the input, you're often better with averaging or pooling
-                the sequence of hidden-states for the whole input sequence.
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+            This output is usually *not* a good summary
+            of the semantic content of the input, you're often better with averaging or pooling
+            the sequence of hidden-states for the whole input sequence.
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
         """
 
         if input_ids is not None and inputs_embeds is not None:
@@ -157,10 +157,7 @@ class AlbertModelWithPabee(AlbertModel):
             res = []
             for i in range(self.config.num_hidden_layers):
                 encoder_outputs = self.encoder.adaptive_forward(
-                    encoder_outputs,
-                    current_layer=i,
-                    attention_mask=extended_attention_mask,
-                    head_mask=head_mask,
+                    encoder_outputs, current_layer=i, attention_mask=extended_attention_mask, head_mask=head_mask,
                 )
 
                 pooled_output = self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]))
@@ -177,10 +174,7 @@ class AlbertModelWithPabee(AlbertModel):
             for i in range(self.config.num_hidden_layers):
                 calculated_layer_num += 1
                 encoder_outputs = self.encoder.adaptive_forward(
-                    encoder_outputs,
-                    current_layer=i,
-                    attention_mask=extended_attention_mask,
-                    head_mask=head_mask,
+                    encoder_outputs, current_layer=i, attention_mask=extended_attention_mask, head_mask=head_mask,
                 )
 
                 pooled_output = self.pooler_activation(self.pooler(encoder_outputs[0][:, 0]))
@@ -242,42 +236,42 @@ class AlbertForSequenceClassificationWithPabee(AlbertPreTrainedModel):
         labels=None,
     ):
         r"""
-            labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
-                Labels for computing the sequence classification/regression loss.
-                Indices should be in ``[0, ..., config.num_labels - 1]``.
-                If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
-                If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the sequence classification/regression loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
+            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).
 
-        Returns:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
-            loss: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
-                Classification (or regression if config.num_labels==1) loss.
-            logits ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
-                Classification (or regression if config.num_labels==1) scores (before SoftMax).
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+    Returns:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
+        loss: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+            Classification (or regression if config.num_labels==1) loss.
+        logits ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
+            Classification (or regression if config.num_labels==1) scores (before SoftMax).
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
 
-            Examples::
+        Examples::
 
-                from transformers import AlbertTokenizer
-                from pabee import AlbertForSequenceClassificationWithPabee
-                import torch
+            from transformers import AlbertTokenizer
+            from pabee import AlbertForSequenceClassificationWithPabee
+            import torch
 
-                tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
-                model = AlbertForSequenceClassificationWithPabee.from_pretrained('albert-base-v2')
-                input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
-                labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
-                outputs = model(input_ids, labels=labels)
-                loss, logits = outputs[:2]
+            tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
+            model = AlbertForSequenceClassificationWithPabee.from_pretrained('albert-base-v2')
+            input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
+            labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
+            outputs = model(input_ids, labels=labels)
+            loss, logits = outputs[:2]
 
         """
 

examples/bert-loses-patience/pabee/modeling_pabee_bert.py

@@ -108,30 +108,30 @@ class BertModelWithPabee(BertModel):
         regression=False,
     ):
         r"""
-        Return:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
-            last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
-                Sequence of hidden-states at the output of the last layer of the model.
-            pooler_output (:obj:`torch.FloatTensor`: of shape :obj:`(batch_size, hidden_size)`):
-                Last layer hidden-state of the first token of the sequence (classification token)
-                further processed by a Linear layer and a Tanh activation function. The Linear
-                layer weights are trained from the next sentence prediction (classification)
-                objective during pre-training.
+    Return:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        pooler_output (:obj:`torch.FloatTensor`: of shape :obj:`(batch_size, hidden_size)`):
+            Last layer hidden-state of the first token of the sequence (classification token)
+            further processed by a Linear layer and a Tanh activation function. The Linear
+            layer weights are trained from the next sentence prediction (classification)
+            objective during pre-training.
 
-                This output is usually *not* a good summary
-                of the semantic content of the input, you're often better with averaging or pooling
-                the sequence of hidden-states for the whole input sequence.
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+            This output is usually *not* a good summary
+            of the semantic content of the input, you're often better with averaging or pooling
+            the sequence of hidden-states for the whole input sequence.
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
         """
 
         if input_ids is not None and inputs_embeds is not None:
@@ -266,44 +266,44 @@ class BertForSequenceClassificationWithPabee(BertPreTrainedModel):
         labels=None,
     ):
         r"""
-            labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
-                Labels for computing the sequence classification/regression loss.
-                Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
-                If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
-                If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the sequence classification/regression loss.
+            Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
+            If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
 
-        Returns:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
-            loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
-                Classification (or regression if config.num_labels==1) loss.
-            logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
-                Classification (or regression if config.num_labels==1) scores (before SoftMax).
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+    Returns:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
+            Classification (or regression if config.num_labels==1) loss.
+        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
+            Classification (or regression if config.num_labels==1) scores (before SoftMax).
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
 
-        Examples::
+    Examples::
 
-            from transformers import BertTokenizer, BertForSequenceClassification
-            from pabee import BertForSequenceClassificationWithPabee
-            import torch
+        from transformers import BertTokenizer, BertForSequenceClassification
+        from pabee import BertForSequenceClassificationWithPabee
+        import torch
 
-            tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
-            model = BertForSequenceClassificationWithPabee.from_pretrained('bert-base-uncased')
+        tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+        model = BertForSequenceClassificationWithPabee.from_pretrained('bert-base-uncased')
 
-            input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
-            labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
-            outputs = model(input_ids, labels=labels)
+        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+        labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
+        outputs = model(input_ids, labels=labels)
 
-            loss, logits = outputs[:2]
+        loss, logits = outputs[:2]
 
         """
 

examples/bert-loses-patience/run_glue_with_pabee.py

@@ -120,10 +120,7 @@ def train(args, train_dataset, model, tokenizer):
     # Distributed training (should be after apex fp16 initialization)
     if args.local_rank != -1:
         model = torch.nn.parallel.DistributedDataParallel(
-            model,
-            device_ids=[args.local_rank],
-            output_device=args.local_rank,
-            find_unused_parameters=True,
+            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True,
         )
 
     # Train!
@@ -154,17 +151,13 @@ def train(args, train_dataset, model, tokenizer):
         logger.info("  Continuing training from epoch %d", epochs_trained)
         logger.info("  Continuing training from global step %d", global_step)
         logger.info(
-            "  Will skip the first %d steps in the first epoch",
-            steps_trained_in_current_epoch,
+            "  Will skip the first %d steps in the first epoch", steps_trained_in_current_epoch,
         )
 
     tr_loss, logging_loss = 0.0, 0.0
     model.zero_grad()
     train_iterator = trange(
-        epochs_trained,
-        int(args.num_train_epochs),
-        desc="Epoch",
-        disable=args.local_rank not in [-1, 0],
+        epochs_trained, int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0],
     )
     set_seed(args)  # Added here for reproductibility
     for _ in train_iterator:
@@ -379,11 +372,7 @@ def load_and_cache_examples(args, task, tokenizer, evaluate=False):
             processor.get_dev_examples(args.data_dir) if evaluate else processor.get_train_examples(args.data_dir)
         )
         features = convert_examples_to_features(
-            examples,
-            tokenizer,
-            label_list=label_list,
-            max_length=args.max_seq_length,
-            output_mode=output_mode,
+            examples, tokenizer, label_list=label_list, max_length=args.max_seq_length, output_mode=output_mode,
         )
         if args.local_rank in [-1, 0]:
             logger.info("Saving features into cached file %s", cached_features_file)
@@ -445,24 +434,15 @@ def main():
         help="The output directory where the model predictions and checkpoints will be written.",
     )
     parser.add_argument(
-        "--patience",
-        default="0",
-        type=str,
-        required=False,
+        "--patience", default="0", type=str, required=False,
     )
     parser.add_argument(
-        "--regression_threshold",
-        default=0,
-        type=float,
-        required=False,
+        "--regression_threshold", default=0, type=float, required=False,
     )
 
     # Other parameters
     parser.add_argument(
-        "--config_name",
-        default="",
-        type=str,
-        help="Pretrained config name or path if not the same as model_name",
+        "--config_name", default="", type=str, help="Pretrained config name or path if not the same as model_name",
     )
     parser.add_argument(
         "--tokenizer_name",
@@ -486,27 +466,17 @@ def main():
     parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
     parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the dev set.")
     parser.add_argument(
-        "--evaluate_during_training",
-        action="store_true",
-        help="Run evaluation during training at each logging step.",
+        "--evaluate_during_training", action="store_true", help="Run evaluation during training at each logging step.",
     )
     parser.add_argument(
-        "--do_lower_case",
-        action="store_true",
-        help="Set this flag if you are using an uncased model.",
+        "--do_lower_case", action="store_true", help="Set this flag if you are using an uncased model.",
     )
 
     parser.add_argument(
-        "--per_gpu_train_batch_size",
-        default=8,
-        type=int,
-        help="Batch size per GPU/CPU for training.",
+        "--per_gpu_train_batch_size", default=8, type=int, help="Batch size per GPU/CPU for training.",
     )
     parser.add_argument(
-        "--per_gpu_eval_batch_size",
-        default=1,
-        type=int,
-        help="Batch size per GPU/CPU for evaluation.",
+        "--per_gpu_eval_batch_size", default=1, type=int, help="Batch size per GPU/CPU for evaluation.",
     )
     parser.add_argument(
         "--gradient_accumulation_steps",
@@ -515,19 +485,13 @@ def main():
         help="Number of updates steps to accumulate before performing a backward/update pass.",
     )
     parser.add_argument(
-        "--learning_rate",
-        default=5e-5,
-        type=float,
-        help="The initial learning rate for Adam.",
+        "--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.",
     )
     parser.add_argument("--weight_decay", default=0.0, type=float, help="Weight decay if we apply some.")
     parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
     parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
     parser.add_argument(
-        "--num_train_epochs",
-        default=3.0,
-        type=float,
-        help="Total number of training epochs to perform.",
+        "--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform.",
     )
     parser.add_argument(
         "--max_steps",
@@ -539,10 +503,7 @@ def main():
 
     parser.add_argument("--logging_steps", type=int, default=500, help="Log every X updates steps.")
     parser.add_argument(
-        "--save_steps",
-        type=int,
-        default=500,
-        help="Save checkpoint every X updates steps.",
+        "--save_steps", type=int, default=500, help="Save checkpoint every X updates steps.",
     )
     parser.add_argument(
         "--eval_all_checkpoints",
@@ -551,14 +512,10 @@ def main():
     )
     parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
     parser.add_argument(
-        "--overwrite_output_dir",
-        action="store_true",
-        help="Overwrite the content of the output directory",
+        "--overwrite_output_dir", action="store_true", help="Overwrite the content of the output directory",
     )
     parser.add_argument(
-        "--overwrite_cache",
-        action="store_true",
-        help="Overwrite the cached training and evaluation sets",
+        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets",
     )
     parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
 
@@ -575,10 +532,7 @@ def main():
         "See details at https://nvidia.github.io/apex/amp.html",
     )
     parser.add_argument(
-        "--local_rank",
-        type=int,
-        default=-1,
-        help="For distributed training: local_rank",
+        "--local_rank", type=int, default=-1, help="For distributed training: local_rank",
     )
     parser.add_argument("--server_ip", type=str, default="", help="For distant debugging.")
     parser.add_argument("--server_port", type=str, default="", help="For distant debugging.")
@@ -680,8 +634,7 @@ def main():
     print("Output Layers Parameters:", output_layers_param_num)
     single_output_layer_param_num = sum(param.numel() for param in model.classifiers[0].parameters())
     print(
-        "Added Output Layers Parameters:",
-        output_layers_param_num - single_output_layer_param_num,
+        "Added Output Layers Parameters:", output_layers_param_num - single_output_layer_param_num,
     )
 
     logger.info("Training/evaluation parameters %s", args)

examples/bertology/run_bertology.py

@@ -66,9 +66,9 @@ def print_2d_tensor(tensor):
 def compute_heads_importance(
     args, model, eval_dataloader, compute_entropy=True, compute_importance=True, head_mask=None, actually_pruned=False
 ):
-    """This method shows how to compute:
-    - head attention entropy
-    - head importance scores according to http://arxiv.org/abs/1905.10650
+    """ This method shows how to compute:
+        - head attention entropy
+        - head importance scores according to http://arxiv.org/abs/1905.10650
     """
     # Prepare our tensors
     n_layers, n_heads = model.config.num_hidden_layers, model.config.num_attention_heads
@@ -150,8 +150,8 @@ def compute_heads_importance(
 
 
 def mask_heads(args, model, eval_dataloader):
-    """This method shows how to mask head (set some heads to zero), to test the effect on the network,
-    based on the head importance scores, as described in Michel et al. (http://arxiv.org/abs/1905.10650)
+    """ This method shows how to mask head (set some heads to zero), to test the effect on the network,
+        based on the head importance scores, as described in Michel et al. (http://arxiv.org/abs/1905.10650)
     """
     _, head_importance, preds, labels = compute_heads_importance(args, model, eval_dataloader, compute_entropy=False)
     preds = np.argmax(preds, axis=1) if args.output_mode == "classification" else np.squeeze(preds)
@@ -201,8 +201,8 @@ def mask_heads(args, model, eval_dataloader):
 
 
 def prune_heads(args, model, eval_dataloader, head_mask):
-    """This method shows how to prune head (remove heads weights) based on
-    the head importance scores as described in Michel et al. (http://arxiv.org/abs/1905.10650)
+    """ This method shows how to prune head (remove heads weights) based on
+        the head importance scores as described in Michel et al. (http://arxiv.org/abs/1905.10650)
     """
     # Try pruning and test time speedup
     # Pruning is like masking but we actually remove the masked weights
@@ -395,8 +395,7 @@ def main():
         cache_dir=args.cache_dir,
     )
     tokenizer = AutoTokenizer.from_pretrained(
-        args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
-        cache_dir=args.cache_dir,
+        args.tokenizer_name if args.tokenizer_name else args.model_name_or_path, cache_dir=args.cache_dir,
     )
     model = AutoModelForSequenceClassification.from_pretrained(
         args.model_name_or_path,

examples/contrib/mm-imdb/utils_mmimdb.py

@@ -138,9 +138,6 @@ def get_image_transforms():
             transforms.Resize(256),
             transforms.CenterCrop(224),
             transforms.ToTensor(),
-            transforms.Normalize(
-                mean=[0.46777044, 0.44531429, 0.40661017],
-                std=[0.12221994, 0.12145835, 0.14380469],
-            ),
+            transforms.Normalize(mean=[0.46777044, 0.44531429, 0.40661017], std=[0.12221994, 0.12145835, 0.14380469],),
         ]
     )

examples/contrib/run_camembert.py

@@ -30,11 +30,7 @@ def fill_mask(masked_input, model, tokenizer, topk=5):
             )
         else:
             topk_filled_outputs.append(
-                (
-                    masked_input.replace(masked_token, predicted_token),
-                    values[index].item(),
-                    predicted_token,
-                )
+                (masked_input.replace(masked_token, predicted_token), values[index].item(), predicted_token,)
             )
     return topk_filled_outputs
 

examples/contrib/run_openai_gpt.py

@@ -71,10 +71,10 @@ def load_rocstories_dataset(dataset_path):
 
 
 def pre_process_datasets(encoded_datasets, input_len, cap_length, start_token, delimiter_token, clf_token):
-    """Pre-process datasets containing lists of tuples(story, 1st continuation, 2nd continuation, label)
+    """ Pre-process datasets containing lists of tuples(story, 1st continuation, 2nd continuation, label)
 
-    To Transformer inputs of shape (n_batch, n_alternative, length) comprising for each batch, continuation:
-    input_ids[batch, alternative, :] = [start_token] + story[:cap_length] + [delimiter_token] + cont1[:cap_length] + [clf_token]
+        To Transformer inputs of shape (n_batch, n_alternative, length) comprising for each batch, continuation:
+        input_ids[batch, alternative, :] = [start_token] + story[:cap_length] + [delimiter_token] + cont1[:cap_length] + [clf_token]
     """
     tensor_datasets = []
     for dataset in encoded_datasets:
@@ -83,10 +83,7 @@ def pre_process_datasets(encoded_datasets, input_len, cap_length, start_token, d
         mc_token_ids = np.zeros((n_batch, 2), dtype=np.int64)
         lm_labels = np.full((n_batch, 2, input_len), fill_value=-100, dtype=np.int64)
         mc_labels = np.zeros((n_batch,), dtype=np.int64)
-        for (
-            i,
-            (story, cont1, cont2, mc_label),
-        ) in enumerate(dataset):
+        for i, (story, cont1, cont2, mc_label), in enumerate(dataset):
             with_cont1 = [start_token] + story[:cap_length] + [delimiter_token] + cont1[:cap_length] + [clf_token]
             with_cont2 = [start_token] + story[:cap_length] + [delimiter_token] + cont2[:cap_length] + [clf_token]
             input_ids[i, 0, : len(with_cont1)] = with_cont1

examples/contrib/run_swag.py

@@ -629,9 +629,7 @@ def main():
         torch.distributed.barrier()  # Make sure only the first process in distributed training will download model & vocab
 
     config = AutoConfig.from_pretrained(args.config_name if args.config_name else args.model_name_or_path)
-    tokenizer = AutoTokenizer.from_pretrained(
-        args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,
-    )
+    tokenizer = AutoTokenizer.from_pretrained(args.tokenizer_name if args.tokenizer_name else args.model_name_or_path,)
     model = AutoModelForMultipleChoice.from_pretrained(
         args.model_name_or_path, from_tf=bool(".ckpt" in args.model_name_or_path), config=config
     )

examples/deebert/run_glue_deebert.py

@@ -358,11 +358,7 @@ def load_and_cache_examples(args, task, tokenizer, evaluate=False):
             processor.get_dev_examples(args.data_dir) if evaluate else processor.get_train_examples(args.data_dir)
         )
         features = convert_examples_to_features(
-            examples,
-            tokenizer,
-            label_list=label_list,
-            max_length=args.max_seq_length,
-            output_mode=output_mode,
+            examples, tokenizer, label_list=label_list, max_length=args.max_seq_length, output_mode=output_mode,
         )
         if args.local_rank in [-1, 0]:
             logger.info("Saving features into cached file %s", cached_features_file)

examples/deebert/src/modeling_highway_bert.py

@@ -14,7 +14,8 @@ from transformers.modeling_bert import (
 
 
 def entropy(x):
-    """Calculate entropy of a pre-softmax logit Tensor"""
+    """ Calculate entropy of a pre-softmax logit Tensor
+    """
     exp_x = torch.exp(x)
     A = torch.sum(exp_x, dim=1)  # sum of exp(x_i)
     B = torch.sum(x * exp_x, dim=1)  # sum of x_i * exp(x_i)
@@ -103,8 +104,7 @@ class DeeBertEncoder(nn.Module):
 
 
 @add_start_docstrings(
-    "The Bert Model transformer with early exiting (DeeBERT). ",
-    BERT_START_DOCSTRING,
+    "The Bert Model transformer with early exiting (DeeBERT). ", BERT_START_DOCSTRING,
 )
 class DeeBertModel(BertPreTrainedModel):
     def __init__(self, config):
@@ -127,9 +127,9 @@ class DeeBertModel(BertPreTrainedModel):
         self.embeddings.word_embeddings = value
 
     def _prune_heads(self, heads_to_prune):
-        """Prunes heads of the model.
-        heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
-        See base class PreTrainedModel
+        """ Prunes heads of the model.
+            heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
+            See base class PreTrainedModel
         """
         for layer, heads in heads_to_prune.items():
             self.encoder.layer[layer].attention.prune_heads(heads)
@@ -147,33 +147,33 @@ class DeeBertModel(BertPreTrainedModel):
         encoder_attention_mask=None,
     ):
         r"""
-        Return:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
-            last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
-                Sequence of hidden-states at the output of the last layer of the model.
-            pooler_output (:obj:`torch.FloatTensor`: of shape :obj:`(batch_size, hidden_size)`):
-                Last layer hidden-state of the first token of the sequence (classification token)
-                further processed by a Linear layer and a Tanh activation function. The Linear
-                layer weights are trained from the next sentence prediction (classification)
-                objective during pre-training.
+    Return:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
+            Sequence of hidden-states at the output of the last layer of the model.
+        pooler_output (:obj:`torch.FloatTensor`: of shape :obj:`(batch_size, hidden_size)`):
+            Last layer hidden-state of the first token of the sequence (classification token)
+            further processed by a Linear layer and a Tanh activation function. The Linear
+            layer weights are trained from the next sentence prediction (classification)
+            objective during pre-training.
 
-                This output is usually *not* a good summary
-                of the semantic content of the input, you're often better with averaging or pooling
-                the sequence of hidden-states for the whole input sequence.
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+            This output is usually *not* a good summary
+            of the semantic content of the input, you're often better with averaging or pooling
+            the sequence of hidden-states for the whole input sequence.
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
-            highway_exits (:obj:`tuple(tuple(torch.Tensor))`:
-                Tuple of each early exit's results (total length: number of layers)
-                Each tuple is again, a tuple of length 2 - the first entry is logits and the second entry is hidden states.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        highway_exits (:obj:`tuple(tuple(torch.Tensor))`:
+            Tuple of each early exit's results (total length: number of layers)
+            Each tuple is again, a tuple of length 2 - the first entry is logits and the second entry is hidden states.
         """
         if input_ids is not None and inputs_embeds is not None:
             raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
@@ -302,32 +302,32 @@ class DeeBertForSequenceClassification(BertPreTrainedModel):
         train_highway=False,
     ):
         r"""
-            labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
-                Labels for computing the sequence classification/regression loss.
-                Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
-                If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
-                If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the sequence classification/regression loss.
+            Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
+            If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
 
-        Returns:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
-            loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
-                Classification (or regression if config.num_labels==1) loss.
-            logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
-                Classification (or regression if config.num_labels==1) scores (before SoftMax).
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+    Returns:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
+            Classification (or regression if config.num_labels==1) loss.
+        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
+            Classification (or regression if config.num_labels==1) scores (before SoftMax).
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
-            highway_exits (:obj:`tuple(tuple(torch.Tensor))`:
-                Tuple of each early exit's results (total length: number of layers)
-                Each tuple is again, a tuple of length 2 - the first entry is logits and the second entry is hidden states.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        highway_exits (:obj:`tuple(tuple(torch.Tensor))`:
+            Tuple of each early exit's results (total length: number of layers)
+            Each tuple is again, a tuple of length 2 - the first entry is logits and the second entry is hidden states.
         """
 
         exit_layer = self.num_layers

examples/deebert/src/modeling_highway_roberta.py

@@ -11,8 +11,7 @@ from .modeling_highway_bert import BertPreTrainedModel, DeeBertModel, HighwayExc
 
 
 @add_start_docstrings(
-    "The RoBERTa Model transformer with early exiting (DeeRoBERTa). ",
-    ROBERTA_START_DOCSTRING,
+    "The RoBERTa Model transformer with early exiting (DeeRoBERTa). ", ROBERTA_START_DOCSTRING,
 )
 class DeeRobertaModel(DeeBertModel):
 
@@ -59,32 +58,32 @@ class DeeRobertaForSequenceClassification(BertPreTrainedModel):
         train_highway=False,
     ):
         r"""
-            labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
-                Labels for computing the sequence classification/regression loss.
-                Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
-                If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
-                If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the sequence classification/regression loss.
+            Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
+            If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
 
-        Returns:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
-            loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
-                Classification (or regression if config.num_labels==1) loss.
-            logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
-                Classification (or regression if config.num_labels==1) scores (before SoftMax).
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+    Returns:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
+            Classification (or regression if config.num_labels==1) loss.
+        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
+            Classification (or regression if config.num_labels==1) scores (before SoftMax).
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_attentions=True`` is passed or when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
-            highway_exits (:obj:`tuple(tuple(torch.Tensor))`:
-                Tuple of each early exit's results (total length: number of layers)
-                Each tuple is again, a tuple of length 2 - the first entry is logits and the second entry is hidden states.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
+        highway_exits (:obj:`tuple(tuple(torch.Tensor))`:
+            Tuple of each early exit's results (total length: number of layers)
+            Each tuple is again, a tuple of length 2 - the first entry is logits and the second entry is hidden states.
         """
 
         exit_layer = self.num_layers

examples/distillation/lm_seqs_dataset.py

@@ -101,7 +101,7 @@ class LmSeqsDataset(Dataset):
 
     def remove_empty_sequences(self):
         """
-        Too short sequences are simply removed. This could be tuned.
+        Too short sequences are simply removed. This could be tunedd.
         """
         init_size = len(self)
         indices = self.lengths > 11

examples/distillation/run_squad_w_distillation.py

@@ -228,20 +228,14 @@ def train(args, train_dataset, model, tokenizer, teacher=None):
                 assert end_logits_tea.size() == end_logits_stu.size()
 
                 loss_fct = nn.KLDivLoss(reduction="batchmean")
-                loss_start = (
-                    loss_fct(
-                        F.log_softmax(start_logits_stu / args.temperature, dim=-1),
-                        F.softmax(start_logits_tea / args.temperature, dim=-1),
-                    )
-                    * (args.temperature ** 2)
-                )
-                loss_end = (
-                    loss_fct(
-                        F.log_softmax(end_logits_stu / args.temperature, dim=-1),
-                        F.softmax(end_logits_tea / args.temperature, dim=-1),
-                    )
-                    * (args.temperature ** 2)
-                )
+                loss_start = loss_fct(
+                    F.log_softmax(start_logits_stu / args.temperature, dim=-1),
+                    F.softmax(start_logits_tea / args.temperature, dim=-1),
+                ) * (args.temperature ** 2)
+                loss_end = loss_fct(
+                    F.log_softmax(end_logits_stu / args.temperature, dim=-1),
+                    F.softmax(end_logits_tea / args.temperature, dim=-1),
+                ) * (args.temperature ** 2)
                 loss_ce = (loss_start + loss_end) / 2.0
 
                 loss = args.alpha_ce * loss_ce + args.alpha_squad * loss

examples/distillation/utils.py

@@ -118,8 +118,7 @@ def init_gpu_params(params):
     if params.multi_gpu:
         logger.info("Initializing PyTorch distributed")
         torch.distributed.init_process_group(
-            init_method="env://",
-            backend="nccl",
+            init_method="env://", backend="nccl",
         )
 
 

examples/language-modeling/run_language_modeling.py

@@ -125,22 +125,13 @@ class DataTrainingArguments:
     )
 
 
-def get_dataset(
-    args: DataTrainingArguments,
-    tokenizer: PreTrainedTokenizer,
-    evaluate: bool = False,
-    cache_dir: Optional[str] = None,
-):
+def get_dataset(args: DataTrainingArguments, tokenizer: PreTrainedTokenizer, evaluate=False):
     file_path = args.eval_data_file if evaluate else args.train_data_file
     if args.line_by_line:
         return LineByLineTextDataset(tokenizer=tokenizer, file_path=file_path, block_size=args.block_size)
     else:
         return TextDataset(
-            tokenizer=tokenizer,
-            file_path=file_path,
-            block_size=args.block_size,
-            overwrite_cache=args.overwrite_cache,
-            cache_dir=cache_dir,
+            tokenizer=tokenizer, file_path=file_path, block_size=args.block_size, overwrite_cache=args.overwrite_cache
         )
 
 
@@ -238,19 +229,11 @@ def main():
 
     # Get datasets
 
-    train_dataset = (
-        get_dataset(data_args, tokenizer=tokenizer, cache_dir=model_args.cache_dir) if training_args.do_train else None
-    )
-    eval_dataset = (
-        get_dataset(data_args, tokenizer=tokenizer, evaluate=True, cache_dir=model_args.cache_dir)
-        if training_args.do_eval
-        else None
-    )
+    train_dataset = get_dataset(data_args, tokenizer=tokenizer) if training_args.do_train else None
+    eval_dataset = get_dataset(data_args, tokenizer=tokenizer, evaluate=True) if training_args.do_eval else None
     if config.model_type == "xlnet":
         data_collator = DataCollatorForPermutationLanguageModeling(
-            tokenizer=tokenizer,
-            plm_probability=data_args.plm_probability,
-            max_span_length=data_args.max_span_length,
+            tokenizer=tokenizer, plm_probability=data_args.plm_probability, max_span_length=data_args.max_span_length,
         )
     else:
         data_collator = DataCollatorForLanguageModeling(

examples/lightning_base.py

@@ -22,7 +22,6 @@ from transformers import (
     PreTrainedTokenizer,
 )
 from transformers.optimization import (
-    Adafactor,
     get_cosine_schedule_with_warmup,
     get_cosine_with_hard_restarts_schedule_with_warmup,
     get_linear_schedule_with_warmup,
@@ -138,15 +137,7 @@ class BaseTransformer(pl.LightningModule):
                 "weight_decay": 0.0,
             },
         ]
-        if self.hparams.adafactor:
-            optimizer = Adafactor(
-                optimizer_grouped_parameters, lr=self.hparams.learning_rate, scale_parameter=False, relative_step=False
-            )
-
-        else:
-            optimizer = AdamW(
-                optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon
-            )
+        optimizer = AdamW(optimizer_grouped_parameters, lr=self.hparams.learning_rate, eps=self.hparams.adam_epsilon)
         self.opt = optimizer
 
         scheduler = self.get_lr_scheduler()
@@ -178,10 +169,10 @@ class BaseTransformer(pl.LightningModule):
         return self.train_loader
 
     def val_dataloader(self):
-        return self.get_dataloader("dev", self.hparams.eval_batch_size, shuffle=False)
+        return self.get_dataloader("dev", self.hparams.eval_batch_size)
 
     def test_dataloader(self):
-        return self.get_dataloader("test", self.hparams.eval_batch_size, shuffle=False)
+        return self.get_dataloader("test", self.hparams.eval_batch_size)
 
     def _feature_file(self, mode):
         return os.path.join(
@@ -235,14 +226,10 @@ class BaseTransformer(pl.LightningModule):
             help="Decoder layer dropout probability (Optional). Goes into model.config",
         )
         parser.add_argument(
-            "--dropout",
-            type=float,
-            help="Dropout probability (Optional). Goes into model.config",
+            "--dropout", type=float, help="Dropout probability (Optional). Goes into model.config",
         )
         parser.add_argument(
-            "--attention_dropout",
-            type=float,
-            help="Attention dropout probability (Optional). Goes into model.config",
+            "--attention_dropout", type=float, help="Attention dropout probability (Optional). Goes into model.config",
         )
         parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
         parser.add_argument(
@@ -260,7 +247,6 @@ class BaseTransformer(pl.LightningModule):
         parser.add_argument("--num_train_epochs", dest="max_epochs", default=3, type=int)
         parser.add_argument("--train_batch_size", default=32, type=int)
         parser.add_argument("--eval_batch_size", default=32, type=int)
-        parser.add_argument("--adafactor", action="store_true")
 
 
 class LoggingCallback(pl.Callback):
@@ -366,8 +352,6 @@ def generic_train(
     if args.gpus > 1:
         train_params["distributed_backend"] = "ddp"
 
-    train_params["accumulate_grad_batches"] = args.accumulate_grad_batches
-
     trainer = pl.Trainer.from_argparse_args(
         args,
         weights_summary=None,

examples/longform-qa/README.md

@@ -1,5 +1,5 @@
 # Long Form Question Answering
 
-This folder contains the code for the Long Form Question answering [demo](http://35.226.96.115:8080/) as well as methods to train and use a fully end-to-end Long Form Question Answering system using the [🤗transformers](https://github.com/huggingface/transformers) and [🤗datasets](https://github.com/huggingface/datasets) libraries.
+This folder contains the code for the Long Form Question answering [demo](http://35.226.96.115:8080/) as well as methods to train and use a fully end-to-end Long Form Question Answering system using the [🤗transformers](https://github.com/huggingface/transformers) and [🤗nlp](https://github.com/huggingface/nlp) libraries.
 
 You can use these methods to train your own system by following along the associate [notebook](https://github.com/huggingface/notebooks/blob/master/longform-qa/Long_Form_Question_Answering_with_ELI5_and_Wikipedia.ipynb) or [blog post](https://yjernite.github.io/lfqa.html).

examples/longform-qa/eli5_app.py

@@ -1,5 +1,5 @@
-import datasets
 import faiss
+import nlp
 import numpy as np
 import streamlit as st
 import torch
@@ -45,7 +45,7 @@ def load_models():
 def load_indexes():
     if LOAD_DENSE_INDEX:
         faiss_res = faiss.StandardGpuResources()
-        wiki40b_passages = datasets.load_dataset(path="wiki_snippets", name="wiki40b_en_100_0")["train"]
+        wiki40b_passages = nlp.load_dataset(path="wiki_snippets", name="wiki40b_en_100_0")["train"]
         wiki40b_passage_reps = np.memmap(
             "wiki40b_passages_reps_32_l-8_h-768_b-512-512.dat",
             dtype="float32",
@@ -63,7 +63,7 @@ def load_indexes():
 
 @st.cache(allow_output_mutation=True)
 def load_train_data():
-    eli5 = datasets.load_dataset("eli5", name="LFQA_reddit")
+    eli5 = nlp.load_dataset("eli5", name="LFQA_reddit")
     eli5_train = eli5["train_eli5"]
     eli5_train_q_reps = np.memmap(
         "eli5_questions_reps.dat", dtype="float32", mode="r", shape=(eli5_train.num_rows, 128)
@@ -95,10 +95,7 @@ def make_support(question, source="wiki40b", method="dense", n_results=10):
             )
         else:
             support_doc, hit_lst = query_es_index(
-                question,
-                es_client,
-                index_name="english_wiki40b_snippets_100w",
-                n_results=n_results,
+                question, es_client, index_name="english_wiki40b_snippets_100w", n_results=n_results,
             )
     support_list = [
         (res["article_title"], res["section_title"].strip(), res["score"], res["passage_text"]) for res in hit_lst
@@ -157,8 +154,7 @@ header_full = """
     header_html,
 )
 st.sidebar.markdown(
-    header_full,
-    unsafe_allow_html=True,
+    header_full, unsafe_allow_html=True,
 )
 
 # Long Form QA with ELI5 and Wikipedia
@@ -177,17 +173,9 @@ action_list = [
 ]
 demo_options = st.sidebar.checkbox("Demo options")
 if demo_options:
-    action_st = st.sidebar.selectbox(
-        "",
-        action_list,
-        index=3,
-    )
+    action_st = st.sidebar.selectbox("", action_list, index=3,)
     action = action_list.index(action_st)
-    show_type = st.sidebar.selectbox(
-        "",
-        ["Show full text of passages", "Show passage section titles"],
-        index=0,
-    )
+    show_type = st.sidebar.selectbox("", ["Show full text of passages", "Show passage section titles"], index=0,)
     show_passages = show_type == "Show full text of passages"
 else:
     action = 3
@@ -262,9 +250,7 @@ questions_list = [
     "How does New Zealand have so many large bird predators?",
 ]
 question_s = st.selectbox(
-    "What would you like to ask? ---- select <MY QUESTION> to enter a new query",
-    questions_list,
-    index=1,
+    "What would you like to ask? ---- select <MY QUESTION> to enter a new query", questions_list, index=1,
 )
 if question_s == "<MY QUESTION>":
     question = st.text_input("Enter your question here:", "")

examples/longform-qa/eli5_utils.py

@@ -4,8 +4,8 @@ import os  # noqa: F401
 from random import choice, randint
 from time import time
 
-import datasets  # noqa: F401
 import faiss  # noqa: F401
+import nlp  # noqa: F401
 import numpy as np
 import pandas as pd
 import torch
@@ -48,11 +48,7 @@ def make_es_index_snippets(es_client, passages_dset, index_name="english_wiki_ki
             yield passage
 
     # create the ES index
-    for ok, action in streaming_bulk(
-        client=es_client,
-        index=index_name,
-        actions=passage_generator(),
-    ):
+    for ok, action in streaming_bulk(client=es_client, index=index_name, actions=passage_generator(),):
         progress.update(1)
         successes += ok
     print("Indexed %d documents" % (successes,))
@@ -141,11 +137,7 @@ class RetrievalQAEmbedder(torch.nn.Module):
 
             # define function for checkpointing
             def partial_encode(*inputs):
-                encoder_outputs = self.sent_encoder.encoder(
-                    inputs[0],
-                    attention_mask=inputs[1],
-                    head_mask=head_mask,
-                )
+                encoder_outputs = self.sent_encoder.encoder(inputs[0], attention_mask=inputs[1], head_mask=head_mask,)
                 sequence_output = encoder_outputs[0]
                 pooled_output = self.sent_encoder.pooler(sequence_output)
                 return pooled_output
@@ -242,11 +234,7 @@ def train_qa_retriever_epoch(model, dataset, tokenizer, optimizer, scheduler, ar
         if step % args.print_freq == 0 or step == 1:
             print(
                 "{:2d} {:5d} of {:5d} \t L: {:.3f} \t -- {:.3f}".format(
-                    e,
-                    step,
-                    len(dataset) // args.batch_size,
-                    loc_loss / loc_steps,
-                    time() - st_time,
+                    e, step, len(dataset) // args.batch_size, loc_loss / loc_steps, time() - st_time,
                 )
             )
             loc_loss = 0
@@ -285,11 +273,7 @@ def train_qa_retriever_joint_epoch(model, dataset_list, tokenizer, optimizer, sc
         if step % args.print_freq == 0:
             print(
                 "{:2d} {:5d} of {:5d} \t L: {:.3f} \t -- {:.3f}".format(
-                    e,
-                    step,
-                    len(dataset_list[0]) // args.batch_size,
-                    loc_loss / loc_steps,
-                    time() - st_time,
+                    e, step, len(dataset_list[0]) // args.batch_size, loc_loss / loc_steps, time() - st_time,
                 )
             )
             loc_loss = 0
@@ -370,8 +354,7 @@ class ELI5DatasetS2S(Dataset):
             self.document_cache[q_id] = self.document_cache.get(q_id, self.make_doc_function(example["title"]))
         document = self.document_cache[q_id]
         in_st = "question: {} context: {}".format(
-            question.lower().replace(" --t--", "").strip(),
-            document.lower().strip(),
+            question.lower().replace(" --t--", "").strip(), document.lower().strip(),
         )
         out_st = answer
         return (in_st, out_st)
@@ -444,11 +427,7 @@ def train_qa_s2s_epoch(model, dataset, tokenizer, optimizer, scheduler, args, e=
         if step % args.print_freq == 0 or step == 1:
             print(
                 "{:2d} {:5d} of {:5d} \t L: {:.3f} \t -- {:.3f}".format(
-                    e,
-                    step,
-                    len(dataset) // args.batch_size,
-                    loc_loss / loc_steps,
-                    time() - st_time,
+                    e, step, len(dataset) // args.batch_size, loc_loss / loc_steps, time() - st_time,
                 )
             )
             loc_loss = 0
@@ -477,18 +456,10 @@ def eval_qa_s2s_epoch(model, dataset, tokenizer, args):
             if step % args.print_freq == 0:
                 print(
                     "{:5d} of {:5d} \t L: {:.3f} \t -- {:.3f}".format(
-                        step,
-                        len(dataset) // args.batch_size,
-                        loc_loss / loc_steps,
-                        time() - st_time,
+                        step, len(dataset) // args.batch_size, loc_loss / loc_steps, time() - st_time,
                     )
                 )
-    print(
-        "Total \t L: {:.3f} \t -- {:.3f}".format(
-            loc_loss / loc_steps,
-            time() - st_time,
-        )
-    )
+    print("Total \t L: {:.3f} \t -- {:.3f}".format(loc_loss / loc_steps, time() - st_time,))
 
 
 def train_qa_s2s(qa_s2s_model, qa_s2s_tokenizer, s2s_train_dset, s2s_valid_dset, s2s_args):
@@ -535,12 +506,7 @@ def qa_s2s_generate(
     max_input_length=512,
     device="cuda:0",
 ):
-    model_inputs = make_qa_s2s_batch(
-        [(question_doc, "A")],
-        qa_s2s_tokenizer,
-        max_input_length,
-        device=device,
-    )
+    model_inputs = make_qa_s2s_batch([(question_doc, "A")], qa_s2s_tokenizer, max_input_length, device=device,)
     n_beams = num_answers if num_beams is None else max(num_beams, num_answers)
     generated_ids = qa_s2s_model.generate(
         input_ids=model_inputs["input_ids"],

examples/movement-pruning/emmental/modeling_bert_masked.py

@@ -37,7 +37,8 @@ logger = logging.getLogger(__name__)
 
 
 class BertEmbeddings(nn.Module):
-    """Construct the embeddings from word, position and token_type embeddings."""
+    """Construct the embeddings from word, position and token_type embeddings.
+    """
 
     def __init__(self, config):
         super().__init__()
@@ -384,8 +385,8 @@ class BertPooler(nn.Module):
 
 
 class MaskedBertPreTrainedModel(PreTrainedModel):
-    """An abstract class to handle weights initialization and
-    a simple interface for downloading and loading pretrained models.
+    """ An abstract class to handle weights initialization and
+        a simple interface for downloading and loading pretrained models.
     """
 
     config_class = MaskedBertConfig
@@ -426,35 +427,35 @@ MASKED_BERT_INPUTS_DOCSTRING = r"""
             :func:`transformers.PreTrainedTokenizer.__call__` for details.
 
             `What are input IDs? <../glossary.html#input-ids>`__
-        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
             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:`(batch_size, sequence_length)`, `optional`):
+        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
             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:`(batch_size, sequence_length)`, `optional`):
+        position_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
             Indices of positions of each input sequence tokens in the position embeddings.
             Selected in the range ``[0, config.max_position_embeddings - 1]``.
 
             `What are position IDs? <../glossary.html#position-ids>`_
-        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
+        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`):
             Mask to nullify selected heads of the self-attention modules.
             Mask values selected in ``[0, 1]``:
             :obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
-        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
             Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
             This is useful if you want more control over how to convert `input_ids` indices into associated vectors
             than the model's internal embedding lookup matrix.
-        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
             Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
             if the model is configured as a decoder.
-        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
             Mask to avoid performing attention on the padding token indices of the encoder input. This mask
             is used in the cross-attention if the model is configured as a decoder.
             Mask values selected in ``[0, 1]``:
@@ -491,9 +492,9 @@ class MaskedBertModel(MaskedBertPreTrainedModel):
         self.embeddings.word_embeddings = value
 
     def _prune_heads(self, heads_to_prune):
-        """Prunes heads of the model.
-        heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
-        See base class PreTrainedModel
+        """ Prunes heads of the model.
+            heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
+            See base class PreTrainedModel
         """
         for layer, heads in heads_to_prune.items():
             self.encoder.layer[layer].attention.prune_heads(heads)
@@ -684,31 +685,31 @@ class MaskedBertForSequenceClassification(MaskedBertPreTrainedModel):
         threshold=None,
     ):
         r"""
-            labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
-                Labels for computing the sequence classification/regression loss.
-                Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
-                If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
-                If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
-            threshold (:obj:`float`):
-                Threshold value (see :class:`~emmental.MaskedLinear`).
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the sequence classification/regression loss.
+            Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
+            If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+            If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+        threshold (:obj:`float`):
+            Threshold value (see :class:`~emmental.MaskedLinear`).
 
-        Returns:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~emmental.MaskedBertConfig`) and inputs:
-            loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
-                Classification (or regression if config.num_labels==1) loss.
-            logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
-                Classification (or regression if config.num_labels==1) scores (before SoftMax).
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+    Returns:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~emmental.MaskedBertConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
+            Classification (or regression if config.num_labels==1) loss.
+        logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
+            Classification (or regression if config.num_labels==1) scores (before SoftMax).
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
         """
 
         outputs = self.bert(
@@ -769,32 +770,32 @@ class MaskedBertForMultipleChoice(MaskedBertPreTrainedModel):
         threshold=None,
     ):
         r"""
-            labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
-                Labels for computing the multiple choice classification loss.
-                Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
-                of the input tensors. (see `input_ids` above)
-            threshold (:obj:`float`):
-                Threshold value (see :class:`~emmental.MaskedLinear`).
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the multiple choice classification loss.
+            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
+            of the input tensors. (see `input_ids` above)
+        threshold (:obj:`float`):
+            Threshold value (see :class:`~emmental.MaskedLinear`).
 
-        Returns:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~emmental.MaskedBertConfig`) and inputs:
-            loss (:obj:`torch.FloatTensor` of shape `(1,)`, `optional`, returned when :obj:`labels` is provided):
-                Classification loss.
-            classification_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices)`):
-                `num_choices` is the second dimension of the input tensors. (see `input_ids` above).
+    Returns:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~emmental.MaskedBertConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape `(1,)`, `optional`, returned when :obj:`labels` is provided):
+            Classification loss.
+        classification_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices)`):
+            `num_choices` is the second dimension of the input tensors. (see `input_ids` above).
 
-                Classification scores (before SoftMax).
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+            Classification scores (before SoftMax).
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
 
         """
         num_choices = input_ids.shape[1]
@@ -859,29 +860,29 @@ class MaskedBertForTokenClassification(MaskedBertPreTrainedModel):
         threshold=None,
     ):
         r"""
-            labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
-                Labels for computing the token classification loss.
-                Indices should be in ``[0, ..., config.num_labels - 1]``.
-            threshold (:obj:`float`):
-                Threshold value (see :class:`~emmental.MaskedLinear`).
+        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+            Labels for computing the token classification loss.
+            Indices should be in ``[0, ..., config.num_labels - 1]``.
+        threshold (:obj:`float`):
+            Threshold value (see :class:`~emmental.MaskedLinear`).
 
-        Returns:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~emmental.MaskedBertConfig`) and inputs:
-            loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided) :
-                Classification loss.
-            scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`)
-                Classification scores (before SoftMax).
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+    Returns:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~emmental.MaskedBertConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided) :
+            Classification loss.
+        scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`)
+            Classification scores (before SoftMax).
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
         """
 
         outputs = self.bert(
@@ -946,36 +947,36 @@ class MaskedBertForQuestionAnswering(MaskedBertPreTrainedModel):
         threshold=None,
     ):
         r"""
-            start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
-                Labels for position (index) of the start of the labelled span for computing the token classification loss.
-                Positions are clamped to the length of the sequence (`sequence_length`).
-                Position outside of the sequence are not taken into account for computing the loss.
-            end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`):
-                Labels for position (index) of the end of the labelled span for computing the token classification loss.
-                Positions are clamped to the length of the sequence (`sequence_length`).
-                Position outside of the sequence are not taken into account for computing the loss.
-            threshold (:obj:`float`):
-                Threshold value (see :class:`~emmental.MaskedLinear`).
+        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the start of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+            Labels for position (index) of the end of the labelled span for computing the token classification loss.
+            Positions are clamped to the length of the sequence (`sequence_length`).
+            Position outside of the sequence are not taken into account for computing the loss.
+        threshold (:obj:`float`):
+            Threshold value (see :class:`~emmental.MaskedLinear`).
 
-        Returns:
-            :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~emmental.MaskedBertConfig`) and inputs:
-            loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
-                Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
-            start_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
-                Span-start scores (before SoftMax).
-            end_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
-                Span-end scores (before SoftMax).
-            hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
-                of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+    Returns:
+        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~emmental.MaskedBertConfig`) and inputs:
+        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
+            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
+        start_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
+            Span-start scores (before SoftMax).
+        end_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
+            Span-end scores (before SoftMax).
+        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+            of shape :obj:`(batch_size, sequence_length, hidden_size)`.
 
-                Hidden-states of the model at the output of each layer plus the initial embedding outputs.
-            attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
-                Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
-                :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
 
-                Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
-                heads.
+            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+            heads.
         """
 
         outputs = self.bert(
@@ -995,10 +996,7 @@ class MaskedBertForQuestionAnswering(MaskedBertPreTrainedModel):
         start_logits = start_logits.squeeze(-1)
         end_logits = end_logits.squeeze(-1)
 
-        outputs = (
-            start_logits,
-            end_logits,
-        ) + outputs[2:]
+        outputs = (start_logits, end_logits,) + outputs[2:]
         if start_positions is not None and end_positions is not None:
             # If we are on multi-GPU, split add a dimension
             if len(start_positions.size()) > 1:

examples/movement-pruning/masked_run_glue.py

@@ -173,10 +173,7 @@ def train(args, train_dataset, model, tokenizer, teacher=None):
     # Distributed training (should be after apex fp16 initialization)
     if args.local_rank != -1:
         model = torch.nn.parallel.DistributedDataParallel(
-            model,
-            device_ids=[args.local_rank],
-            output_device=args.local_rank,
-            find_unused_parameters=True,
+            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True,
         )
 
     # Train!
@@ -220,10 +217,7 @@ def train(args, train_dataset, model, tokenizer, teacher=None):
     tr_loss, logging_loss = 0.0, 0.0
     model.zero_grad()
     train_iterator = trange(
-        epochs_trained,
-        int(args.num_train_epochs),
-        desc="Epoch",
-        disable=args.local_rank not in [-1, 0],
+        epochs_trained, int(args.num_train_epochs), desc="Epoch", disable=args.local_rank not in [-1, 0],
     )
     set_seed(args)  # Added here for reproductibility
     for _ in train_iterator:
@@ -286,14 +280,11 @@ def train(args, train_dataset, model, tokenizer, teacher=None):
                         attention_mask=inputs["attention_mask"],
                     )
 
-                loss_logits = (
-                    F.kl_div(
-                        input=F.log_softmax(logits_stu / args.temperature, dim=-1),
-                        target=F.softmax(logits_tea / args.temperature, dim=-1),
-                        reduction="batchmean",
-                    )
-                    * (args.temperature ** 2)
-                )
+                loss_logits = F.kl_div(
+                    input=F.log_softmax(logits_stu / args.temperature, dim=-1),
+                    target=F.softmax(logits_tea / args.temperature, dim=-1),
+                    reduction="batchmean",
+                ) * (args.temperature ** 2)
 
                 loss = args.alpha_distil * loss_logits + args.alpha_ce * loss
 
@@ -538,11 +529,7 @@ def load_and_cache_examples(args, task, tokenizer, evaluate=False):
             processor.get_dev_examples(args.data_dir) if evaluate else processor.get_train_examples(args.data_dir)
         )
         features = convert_examples_to_features(
-            examples,
-            tokenizer,
-            max_length=args.max_seq_length,
-            label_list=label_list,
-            output_mode=output_mode,
+            examples, tokenizer, max_length=args.max_seq_length, label_list=label_list, output_mode=output_mode,
         )
         if args.local_rank in [-1, 0]:
             logger.info("Saving features into cached file %s", cached_features_file)
@@ -605,10 +592,7 @@ def main():
     )
     # Other parameters
     parser.add_argument(
-        "--config_name",
-        default="",
-        type=str,
-        help="Pretrained config name or path if not the same as model_name",
+        "--config_name", default="", type=str, help="Pretrained config name or path if not the same as model_name",
     )
     parser.add_argument(
         "--tokenizer_name",
@@ -632,27 +616,17 @@ def main():
     parser.add_argument("--do_train", action="store_true", help="Whether to run training.")
     parser.add_argument("--do_eval", action="store_true", help="Whether to run eval on the dev set.")
     parser.add_argument(
-        "--evaluate_during_training",
-        action="store_true",
-        help="Run evaluation during training at each logging step.",
+        "--evaluate_during_training", action="store_true", help="Run evaluation during training at each logging step.",
     )
     parser.add_argument(
-        "--do_lower_case",
-        action="store_true",
-        help="Set this flag if you are using an uncased model.",
+        "--do_lower_case", action="store_true", help="Set this flag if you are using an uncased model.",
     )
 
     parser.add_argument(
-        "--per_gpu_train_batch_size",
-        default=8,
-        type=int,
-        help="Batch size per GPU/CPU for training.",
+        "--per_gpu_train_batch_size", default=8, type=int, help="Batch size per GPU/CPU for training.",
     )
     parser.add_argument(
-        "--per_gpu_eval_batch_size",
-        default=8,
-        type=int,
-        help="Batch size per GPU/CPU for evaluation.",
+        "--per_gpu_eval_batch_size", default=8, type=int, help="Batch size per GPU/CPU for evaluation.",
     )
     parser.add_argument("--learning_rate", default=5e-5, type=float, help="The initial learning rate for Adam.")
 
@@ -749,10 +723,7 @@ def main():
     parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
     parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
     parser.add_argument(
-        "--num_train_epochs",
-        default=3.0,
-        type=float,
-        help="Total number of training epochs to perform.",
+        "--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform.",
     )
     parser.add_argument(
         "--max_steps",
@@ -771,14 +742,10 @@ def main():
     )
     parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
     parser.add_argument(
-        "--overwrite_output_dir",
-        action="store_true",
-        help="Overwrite the content of the output directory",
+        "--overwrite_output_dir", action="store_true", help="Overwrite the content of the output directory",
     )
     parser.add_argument(
-        "--overwrite_cache",
-        action="store_true",
-        help="Overwrite the cached training and evaluation sets",
+        "--overwrite_cache", action="store_true", help="Overwrite the cached training and evaluation sets",
     )
     parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
 

examples/movement-pruning/masked_run_squad.py

@@ -181,10 +181,7 @@ def train(args, train_dataset, model, tokenizer, teacher=None):
     # Distributed training (should be after apex fp16 initialization)
     if args.local_rank != -1:
         model = torch.nn.parallel.DistributedDataParallel(
-            model,
-            device_ids=[args.local_rank],
-            output_device=args.local_rank,
-            find_unused_parameters=True,
+            model, device_ids=[args.local_rank], output_device=args.local_rank, find_unused_parameters=True,
         )
 
     # Train!
@@ -307,22 +304,16 @@ def train(args, train_dataset, model, tokenizer, teacher=None):
                         attention_mask=inputs["attention_mask"],
                     )
 
-                loss_start = (
-                    F.kl_div(
-                        input=F.log_softmax(start_logits_stu / args.temperature, dim=-1),
-                        target=F.softmax(start_logits_tea / args.temperature, dim=-1),
-                        reduction="batchmean",
-                    )
-                    * (args.temperature ** 2)
-                )
-                loss_end = (
-                    F.kl_div(
-                        input=F.log_softmax(end_logits_stu / args.temperature, dim=-1),
-                        target=F.softmax(end_logits_tea / args.temperature, dim=-1),
-                        reduction="batchmean",
-                    )
-                    * (args.temperature ** 2)
-                )
+                loss_start = F.kl_div(
+                    input=F.log_softmax(start_logits_stu / args.temperature, dim=-1),
+                    target=F.softmax(start_logits_tea / args.temperature, dim=-1),
+                    reduction="batchmean",
+                ) * (args.temperature ** 2)
+                loss_end = F.kl_div(
+                    input=F.log_softmax(end_logits_stu / args.temperature, dim=-1),
+                    target=F.softmax(end_logits_tea / args.temperature, dim=-1),
+                    reduction="batchmean",
+                ) * (args.temperature ** 2)
                 loss_logits = (loss_start + loss_end) / 2.0
 
                 loss = args.alpha_distil * loss_logits + args.alpha_ce * loss
@@ -868,10 +859,7 @@ def main():
     parser.add_argument("--adam_epsilon", default=1e-8, type=float, help="Epsilon for Adam optimizer.")
     parser.add_argument("--max_grad_norm", default=1.0, type=float, help="Max gradient norm.")
     parser.add_argument(
-        "--num_train_epochs",
-        default=3.0,
-        type=float,
-        help="Total number of training epochs to perform.",
+        "--num_train_epochs", default=3.0, type=float, help="Total number of training epochs to perform.",
     )
     parser.add_argument(
         "--max_steps",

examples/multiple-choice/utils_multiple_choice.py

@@ -26,8 +26,8 @@ from enum import Enum
 from typing import List, Optional
 
 import tqdm
-
 from filelock import FileLock
+
 from transformers import PreTrainedTokenizer, is_tf_available, is_torch_available
 
 
@@ -100,12 +100,7 @@ if is_torch_available():
 
             cached_features_file = os.path.join(
                 data_dir,
-                "cached_{}_{}_{}_{}".format(
-                    mode.value,
-                    tokenizer.__class__.__name__,
-                    str(max_seq_length),
-                    task,
-                ),
+                "cached_{}_{}_{}_{}".format(mode.value, tokenizer.__class__.__name__, str(max_seq_length), task,),
             )
 
             # Make sure only the first process in distributed training processes the dataset,
@@ -126,12 +121,7 @@ if is_torch_available():
                     else:
                         examples = processor.get_train_examples(data_dir)
                     logger.info("Training examples: %s", len(examples))
-                    self.features = convert_examples_to_features(
-                        examples,
-                        label_list,
-                        max_seq_length,
-                        tokenizer,
-                    )
+                    self.features = convert_examples_to_features(examples, label_list, max_seq_length, tokenizer,)
                     logger.info("Saving features into cached file %s", cached_features_file)
                     torch.save(self.features, cached_features_file)
 
@@ -174,12 +164,7 @@ if is_tf_available():
                 examples = processor.get_train_examples(data_dir)
             logger.info("Training examples: %s", len(examples))
 
-            self.features = convert_examples_to_features(
-                examples,
-                label_list,
-                max_seq_length,
-                tokenizer,
-            )
+            self.features = convert_examples_to_features(examples, label_list, max_seq_length, tokenizer,)
 
             def gen():
                 for (ex_index, ex) in tqdm.tqdm(enumerate(self.features), desc="convert examples to features"):
@@ -506,10 +491,7 @@ class ArcProcessor(DataProcessor):
 
 
 def convert_examples_to_features(
-    examples: List[InputExample],
-    label_list: List[str],
-    max_length: int,
-    tokenizer: PreTrainedTokenizer,
+    examples: List[InputExample], label_list: List[str], max_length: int, tokenizer: PreTrainedTokenizer,
 ) -> List[InputFeatures]:
     """
     Loads a data file into a list of `InputFeatures`

examples/question-answering/run_squad.py

@@ -187,7 +187,7 @@ def train(args, train_dataset, model, tokenizer):
                 "end_positions": batch[4],
             }
 
-            if args.model_type in ["xlm", "roberta", "distilbert", "camembert", "bart"]:
+            if args.model_type in ["xlm", "roberta", "distilbert", "camembert"]:
                 del inputs["token_type_ids"]
 
             if args.model_type in ["xlnet", "xlm"]:
@@ -300,7 +300,7 @@ def evaluate(args, model, tokenizer, prefix=""):
                 "token_type_ids": batch[2],
             }
 
-            if args.model_type in ["xlm", "roberta", "distilbert", "camembert", "bart"]:
+            if args.model_type in ["xlm", "roberta", "distilbert", "camembert"]:
                 del inputs["token_type_ids"]
 
             feature_indices = batch[3]

examples/question-answering/run_squad_trainer.py

@@ -137,12 +137,7 @@ def main():
     )
 
     # Initialize our Trainer
-    trainer = Trainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        eval_dataset=eval_dataset,
-    )
+    trainer = Trainer(model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset,)
 
     # Training
     if training_args.do_train:

examples/question-answering/run_tf_squad.py

@@ -231,12 +231,7 @@ def main():
     eval_dataset = eval_dataset.apply(tf.data.experimental.assert_cardinality(len(eval_examples)))
 
     # Initialize our Trainer
-    trainer = TFTrainer(
-        model=model,
-        args=training_args,
-        train_dataset=train_dataset,
-        eval_dataset=eval_dataset,
-    )
+    trainer = TFTrainer(model=model, args=training_args, train_dataset=train_dataset, eval_dataset=eval_dataset,)
 
     # Training
     if training_args.do_train:

examples/rag/README.md

@@ -1,85 +0,0 @@
-# Intro
-RAG (for Retrieval Augmented Generation) is a seq2seq model which encapsulates two core components: a question encoder and a generator. During a forward pass, we encode the input with the question encoder and pass it
-to the retriever to extract relevant context documents. The documents are then prepended to the input. Such contextualized input is passed to the generator. See [the paper](https://arxiv.org/pdf/2005.11401.pdf) for mored details.
-
-We implement two variants of the model, both presented in the paper - `RagSequenceForGeneration. and `RagTokenForGeneration`. In both cases we use `DPRQuestionEncoder` as the question encoder. As for the generator, two compatible architectures have been tested: `BartForConditionalGeneration`  and `T5ForConditionalGeneration`.
-
-Key files:
-- `modeling_rag.py`, `tokenization_rag.py`, `configuration_rag.py` the core model implementation
-- `retrieval_rag.py` - a distributed retriever built on top of the `torch.distributed` communication package. The retriever is an interface between the model and the faiss index of the encoded documents. During training, all workers initialize their own instance of the retriever, however, only the main worker loads the index into memory, which prevents OOMs on machines with multiple GPUs (we store the index in RAM). The index itself is based on the `nlp.Datasets`. We also implement a variant compatible with indices built using the original DPR implementation (https://github.com/facebookresearch/DPR)
-- `eval_rag.py` - an evaluation script which allows to perform the evaluation end to end (measures the exact match and F1 on the downstream task) as well as the evaluation of the retrieval component alone (measures precision@k).
-- `finetune.py` - a training script for finetuning RAG models.
-
-
-# Finetuning
-Our finetuning logic is based on scripts from [`examples/seq2seq`](https://github.com/huggingface/transformers/tree/master/examples/seq2seq).
-Follow instructions there regarding data preprocessing. A sample finetuning command:
-
-```
-python examples/rag/finetune.py \
-    --data_dir $DATA_DIR \
-    --output_dir $OUTPUT_DIR \
-    --model_name_or_path $MODEL_NAME_OR_PATH \
-    --model_type rag_sequence \
-    --fp16 \
-    --gpus 8
-```
-
-
-# Evaluation
-Apart for parameters specifying the model that's being evaluated and some extra parameters, the evaluation script expects paths to two files:
-- `evaluation_set` - a path file specifying the input dataset for evaluation, a single datapoint per line, e.g.
-```who is the owner of reading football club```
-- `gold_data_path` - a path to a file contaning ground truth answers for samples from the `evaluation_set`.
-
-We expect the following formats of the gold data file:
-
-- for e2e evaluation, we support two formats of gold files:
-    - `qa` - where a single line in the following format: input [tab] output_list, e.g.:
-    ```
-    who is the owner of reading football club	['Xiu Li Dai', 'Dai Yongge', 'Dai Xiuli', 'Yongge Dai']
-    ```
-    - `ans` - where a single line of the gold file contains the expected output string,
-    ```
-    Xiu Li Dai
-    ```
-
-- for retrieval evaluation, we expect a tab-separated list of Wikipedia page titles constituting positive contexts for a given query, e.g. given a question `who sings does he love me with reba`, a line with ground truth retrieval data could look as follows:
-```
-Does He Love You	Does He Love You	Red Sandy Spika dress of Reba McEntire	Greatest Hits Volume Two (Reba McEntire album)	Shoot for the Moon (album)
-```
-
-## Retrieval evaluation
-
-We demonstrate how to evaluate retrieval against DPR evaluation data. You can download respective files from links listed [here](https://github.com/facebookresearch/DPR/blob/master/data/download_data.py#L39-L45).
-
-1. Download and unzip the gold data file. We use the `biencoder-nq-dev` from https://dl.fbaipublicfiles.com/dpr/data/retriever/biencoder-nq-dev.json.gz.
-2. Parse the unziped file using the `parse_dpr_relevance_data.py`
-```
-python examples/rag/parse_dpr_relevance_data.py --src_path path/to/unziped/biencoder-nq-dev.json --evaluation_set path/to/output/biencoder-nq-dev.questions --gold_data_path path/to/output/biencoder-nq-dev.pages
-```
-3. Run evaluation:
-```
-python examples/rag/eval_rag.py \
-    --model_name_or_path $MODEL_NAME_OR_PATH \ # model name or path of the model we're evaluating
-    --model_type rag_sequence \ # RAG model type (rag_token or rag_sequence)
-    --evaluation_set path/to/output/biencoder-nq-dev.questions \ # an input dataset for evaluation
-    --gold_data_path path/to/output/biencoder-nq-dev.pages \ # a dataset containing ground truth answers for samples from the evaluation_set
-    --predictions_filename retrieval_preds.tsv  \ # name of file in which predictions will be stored
-    --eval_mode retrieval  \ # indicates whether we're performing retrieval evaluation or e2e evaluation
-    --recalculate # if predictions_filename already exists, and this option is set - we regenerate the answers, otherwise we reuse the predicsion file to calculate metrics.
-```
-
-
-## End-to-end evaluation
-```
-python examples/rag/eval_rag.py \
-    --model_name_or_path /private/home/piktus/rag_huggingface/data/repro-rag-sequence-63/ \
-    --model_type rag_sequence \
-    --evaluation_set path/to/test.source \
-    --gold_data_path path/to/gold_data \
-    --predictions_filename e2e_preds.txt \
-    --eval_mode e2e  \ # indicates whether we're performing retrieval evaluation or e2e evaluation (default)
-    --n_docs 5 \ # You can experiment with retrieving different number of documents at evaluation time
-    --print_predictions
-```

examples/rag/callbacks.py

@@ -1,30 +0,0 @@
-import logging
-import os
-
-from pytorch_lightning.callbacks import ModelCheckpoint
-
-
-logger = logging.getLogger(__name__)
-
-
-def get_checkpoint_callback(output_dir, metric):
-    """Saves the best model by validation ROUGE2 score."""
-    if metric == "rouge2":
-        exp = "{val_avg_rouge2:.4f}-{step_count}"
-    elif metric == "bleu":
-        exp = "{val_avg_bleu:.4f}-{step_count}"
-    elif metric == "em":
-        exp = "{val_avg_em:.4f}-{step_count}"
-    else:
-        raise NotImplementedError(
-            f"seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this function."
-        )
-
-    checkpoint_callback = ModelCheckpoint(
-        filepath=os.path.join(output_dir, exp),
-        monitor=f"val_{metric}",
-        mode="max",
-        save_top_k=10,
-        period=0,  # maybe save a checkpoint every time val is run, not just end of epoch.
-    )
-    return checkpoint_callback

examples/rag/eval_rag.py

@@ -1,317 +0,0 @@
-""" Evaluation script for RAG models."""
-
-import argparse
-import ast
-import logging
-import os
-import sys
-
-import pandas as pd
-import torch
-from tqdm import tqdm
-
-from transformers import (
-    BartForConditionalGeneration,
-    BartTokenizer,
-    RagRetriever,
-    RagSequenceForGeneration,
-    RagTokenForGeneration,
-)
-from transformers import logging as transformers_logging
-
-
-sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # isort:skip
-from examples.rag.utils import exact_match_score, f1_score  # noqa: E402 # isort:skip
-
-
-logger = logging.getLogger(__name__)
-logging.basicConfig(level=logging.INFO)
-
-transformers_logging.set_verbosity_info()
-
-
-def infer_model_type(model_name_or_path):
-    if "token" in model_name_or_path:
-        return "rag_token"
-    if "sequence" in model_name_or_path:
-        return "rag_sequence"
-    if "bart" in model_name_or_path:
-        return "bart"
-    return None
-
-
-def metric_max_over_ground_truths(metric_fn, prediction, ground_truths):
-    scores_for_ground_truths = []
-    for ground_truth in ground_truths:
-        score = metric_fn(prediction, ground_truth)
-        scores_for_ground_truths.append(score)
-    return max(scores_for_ground_truths)
-
-
-def get_scores(args, preds_path, gold_data_path):
-    hypos = [line.strip() for line in open(preds_path, "r").readlines()]
-    answers = []
-
-    if args.gold_data_mode == "qa":
-        data = pd.read_csv(gold_data_path, sep="\t", header=None)
-        for answer_list in data[1]:
-            ground_truths = ast.literal_eval(answer_list)
-            answers.append(ground_truths)
-    else:
-        references = [line.strip() for line in open(gold_data_path, "r").readlines()]
-        answers = [[reference] for reference in references]
-
-    f1 = em = total = 0
-    for prediction, ground_truths in zip(hypos, answers):
-        total += 1
-        em += metric_max_over_ground_truths(exact_match_score, prediction, ground_truths)
-        f1 += metric_max_over_ground_truths(f1_score, prediction, ground_truths)
-
-    em = 100.0 * em / total
-    f1 = 100.0 * f1 / total
-
-    logger.info("F1: {}".format(f1))
-    logger.info("EM: {}".format(em))
-
-
-def get_precision_at_k(args, preds_path, gold_data_path):
-    k = args.k
-    hypos = [line.strip() for line in open(preds_path, "r").readlines()]
-    references = [line.strip() for line in open(gold_data_path, "r").readlines()]
-
-    em = total = 0
-    for hypo, reference in zip(hypos, references):
-        hypo_provenance = set(hypo.split("\t")[:k])
-        ref_provenance = set(reference.split("\t")[1 : (k + 1)])
-        total += 1
-        em += len(hypo_provenance & ref_provenance) / k
-
-    em = 100.0 * em / total
-    logger.info("Precision@{}: {}".format(k, em))
-
-
-def evaluate_batch_retrieval(args, rag_model, tokenizer, retriever, questions):
-    def strip_title(title):
-        if title.startswith('"'):
-            title = title[1:]
-        if title.endswith('"'):
-            title = title[:-1]
-        return title
-
-    retriever_inputs = tokenizer.batch_encode_plus(
-        questions,
-        return_tensors="pt",
-        padding=True,
-        truncation=True,
-    )
-    retriever_input_embs = rag_model.model.question_encoder(retriever_inputs["input_ids"].to(args.device))[0]
-
-    _, all_docs = retriever.retrieve(retriever_input_embs.numpy(), rag_model.config.n_docs)
-
-    provenance_strings = []
-    for docs in all_docs:
-        provenance = [strip_title(title) for title in docs["title"]]
-        provenance_strings.append("\t".join(provenance))
-    return provenance_strings
-
-
-def evaluate_batch_e2e(args, rag_model, tokenizer, retriever, questions):
-    with torch.no_grad():
-        input_ids = tokenizer.batch_encode_plus(questions, return_tensors="pt", padding=True, truncation=True)[
-            "input_ids"
-        ].to(args.device)
-        outputs = rag_model.generate(
-            input_ids,
-            retriever=retriever,
-            num_beams=args.num_beams,
-            min_length=args.min_length,
-            max_length=args.max_length,
-            early_stopping=False,
-            num_return_sequences=1,
-            bad_words_ids=[[0, 0]],  # BART likes to repeat BOS tokens, dont allow it to generate more than one
-            clean_up_tokenization=True,
-            print_docs=args.print_docs,
-        )
-        answers = tokenizer.batch_decode(outputs, skip_special_tokens=True)
-
-        if args.print_predictions:
-            for q, a in zip(questions, answers):
-                logger.info("Q: {} - A: {}".format(q, a))
-
-        return answers
-
-
-def get_args():
-    parser = argparse.ArgumentParser()
-    parser.add_argument(
-        "--model_type",
-        choices=["rag_sequence", "rag_token", "bart"],
-        type=str,
-        help="RAG model type: rag_sequence, rag_token or bart, if none specified, the type is inferred from the model_name_or_path",
-    )
-    parser.add_argument(
-        "--retriever_type",
-        default=None,
-        choices=["hf_retriever", "legacy_retriever"],
-        type=str,
-        help="RAG model retriever type",
-    )
-    parser.add_argument(
-        "--index_path",
-        default=None,
-        type=str,
-        help="Path to the retrieval index",
-    )
-    parser.add_argument("--n_docs", default=5, type=int, help="Number of retrieved docs")
-    parser.add_argument(
-        "--model_name_or_path",
-        default=None,
-        type=str,
-        required=True,
-        help="Path to pretrained checkpoints or model identifier from huggingface.co/models",
-    )
-    parser.add_argument(
-        "--eval_mode",
-        choices=["e2e", "retrieval"],
-        default="e2e",
-        type=str,
-        help="Evaluation mode, e2e calculates exact match and F1 of the downstream task, retrieval calulates precision@k.",
-    )
-    parser.add_argument("--k", default=1, type=int, help="k for the precision@k calculation")
-    parser.add_argument(
-        "--evaluation_set",
-        default=None,
-        type=str,
-        required=True,
-        help="Path to a file containing evaluation samples",
-    )
-    parser.add_argument(
-        "--gold_data_path",
-        default=None,
-        type=str,
-        required=True,
-        help="Path to a tab-separated file with gold samples",
-    )
-    parser.add_argument(
-        "--gold_data_mode",
-        default="qa",
-        type=str,
-        choices=["qa", "ans"],
-        help="Format of the gold data file"
-        "qa - a single line in the following format: question [tab] answer_list"
-        "ans - a single line of the gold file contains the expected answer string",
-    )
-    parser.add_argument(
-        "--predictions_filename",
-        type=str,
-        default="predictions.txt",
-        help="Name of the predictions file, to be stored in the checkpoints directry",
-    )
-    parser.add_argument(
-        "--eval_all_checkpoints",
-        action="store_true",
-        help="Evaluate all checkpoints starting with the same prefix as model_name ending and ending with step number",
-    )
-    parser.add_argument(
-        "--eval_batch_size",
-        default=8,
-        type=int,
-        help="Batch size per GPU/CPU for evaluation.",
-    )
-    parser.add_argument(
-        "--recalculate",
-        help="Recalculate predictions even if the prediction file exists",
-        action="store_true",
-    )
-    parser.add_argument(
-        "--num_beams",
-        default=4,
-        type=int,
-        help="Number of beams to be used when generating answers",
-    )
-    parser.add_argument("--min_length", default=1, type=int, help="Min length of the generated answers")
-    parser.add_argument("--max_length", default=50, type=int, help="Max length of the generated answers")
-
-    parser.add_argument(
-        "--print_predictions",
-        action="store_true",
-        help="If True, prints predictions while evaluating.",
-    )
-    parser.add_argument(
-        "--print_docs",
-        action="store_true",
-        help="If True, prints docs retried while generating.",
-    )
-    args = parser.parse_args()
-    args.device = torch.device("cuda" if torch.cuda.is_available() else "cpu")
-    return args
-
-
-def main(args):
-    model_kwargs = {}
-    if args.model_type is None:
-        args.model_type = infer_model_type(args.model_name_or_path)
-        assert args.model_type is not None
-    if args.model_type.startswith("rag"):
-        model_class = RagTokenForGeneration if args.model_type == "rag_token" else RagSequenceForGeneration
-        model_kwargs["n_docs"] = args.n_docs
-        if args.retriever_type is not None:
-            model_kwargs["retriever_type"] = args.retriever_type
-        if args.index_path is not None:
-            model_kwargs["index_path"] = args.index_path
-    else:
-        model_class = BartForConditionalGeneration
-
-    checkpoints = (
-        [f.path for f in os.scandir(args.model_name_or_path) if f.is_dir()]
-        if args.eval_all_checkpoints
-        else [args.model_name_or_path]
-    )
-
-    logger.info("Evaluate the following checkpoints: %s", checkpoints)
-
-    score_fn = get_scores if args.eval_mode == "e2e" else get_precision_at_k
-    evaluate_batch_fn = evaluate_batch_e2e if args.eval_mode == "e2e" else evaluate_batch_retrieval
-
-    for checkpoint in checkpoints:
-        predictions_path = os.path.join(checkpoint, args.predictions_filename)
-        if os.path.exists(predictions_path) and (not args.recalculate):
-            logger.info("Calculating metrics based on an existing predictions file: {}".format(predictions_path))
-            score_fn(args, predictions_path, args.gold_data_path)
-            continue
-
-        logger.info("***** Running evaluation for {} *****".format(checkpoint))
-        logger.info("  Batch size = %d", args.eval_batch_size)
-        logger.info("  Predictions will be stored under {}".format(predictions_path))
-
-        model = model_class.from_pretrained(checkpoint, **model_kwargs)
-        model.to(args.device)
-        retriever = RagRetriever(model.config)  # TODO: add tokenizers
-        tokenizer = (
-            retriever.generator_tokenizer
-            if args.model_type != "bart" and args.eval_mode == "e2e"
-            else retriever.question_encoder_tokenizer
-            if args.model_type != "bart" and args.eval_mode == "retrieval"
-            else BartTokenizer.from_pretrained("facebook/bart-large")
-        )
-
-        with open(args.evaluation_set, "r") as eval_file, open(predictions_path, "w") as preds_file:
-            questions = []
-            for line in tqdm(eval_file):
-                questions.append(line.strip())
-                if len(questions) == args.eval_batch_size:
-                    answers = evaluate_batch_fn(args, model, tokenizer, retriever, questions)
-                    preds_file.write("\n".join(answers) + "\n")
-                    preds_file.flush()
-                    questions = []
-            if len(questions) > 0:
-                answers = evaluate_batch_fn(args, model, tokenizer, retriever, questions)
-                preds_file.write("\n".join(answers))
-                preds_file.flush()
-
-            score_fn(args, predictions_path, args.gold_data_path)
-
-
-if __name__ == "__main__":
-    args = get_args()
-    main(args)

examples/rag/finetune.py

@@ -1,476 +0,0 @@
-"""Finetuning script for RAG models. Adapted from examples.seq2seq.finetune.py"""
-
-import argparse
-import glob
-import logging
-import os
-import sys
-import time
-import warnings
-from collections import defaultdict
-from pathlib import Path
-from typing import Any, Dict, List, Tuple
-
-import numpy as np
-import pytorch_lightning as pl
-import torch
-import torch.distributed as dist
-from torch.utils.data import DataLoader
-
-from transformers import (
-    AutoConfig,
-    AutoTokenizer,
-    BartForConditionalGeneration,
-    RagConfig,
-    RagPyTorchDistributedRetriever,
-    RagSequenceForGeneration,
-    RagTokenForGeneration,
-    T5ForConditionalGeneration,
-    get_linear_schedule_with_warmup,
-)
-from transformers import logging as transformers_logging
-
-
-sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip
-
-from examples.lightning_base import BaseTransformer, add_generic_args, generic_train  # noqa: E402 # isort:skip
-from examples.rag.callbacks import get_checkpoint_callback  # noqa: E402 # isort:skip
-from examples.rag.utils import (  # noqa: E402 # isort:skip
-    Seq2SeqDataset,
-    calculate_exact_match,
-    is_rag_model,
-    set_extra_model_params,
-)
-from examples.seq2seq.callbacks import Seq2SeqLoggingCallback, get_early_stopping_callback  # noqa: E402 # isort:skip
-from examples.seq2seq.utils import (  # noqa: E402 # isort:skip
-    flatten_list,
-    get_git_info,
-    lmap,
-    pickle_save,
-    save_git_info,
-    save_json,
-)
-
-logging.basicConfig(level=logging.INFO)
-logger = logging.getLogger(__name__)
-
-transformers_logging.set_verbosity_info()
-
-
-class AttrDict(dict):
-    def __init__(self, *args, **kwargs):
-        super(AttrDict, self).__init__(*args, **kwargs)
-        self.__dict__ = self
-
-
-class GenerativeQAModule(BaseTransformer):
-    mode = "generative_qa"
-    loss_names = ["loss"]
-    metric_names = ["em"]
-    val_metric = "em"
-
-    def __init__(self, hparams, **kwargs):
-        # when loading from a pytorch lightning checkpoint, hparams are passed as dict
-        if isinstance(hparams, dict):
-            hparams = AttrDict(hparams)
-        if hparams.model_type == "rag_sequence":
-            self.model_class = RagSequenceForGeneration
-        elif hparams.model_type == "rag_token":
-            self.model_class = RagTokenForGeneration
-        elif hparams.model_type == "bart":
-            self.model_class = BartForConditionalGeneration
-        else:
-            self.model_class = T5ForConditionalGeneration
-        self.is_rag_model = is_rag_model(hparams.model_type)
-
-        config_class = RagConfig if self.is_rag_model else AutoConfig
-        config = config_class.from_pretrained(hparams.model_name_or_path)
-
-        # set extra_model_params for generator configs and load_model
-        extra_model_params = ("encoder_layerdrop", "decoder_layerdrop", "attention_dropout", "dropout")
-        if self.is_rag_model:
-            generator_config = AutoConfig.from_pretrained(
-                config.pretrained_generator_name_or_path, prefix=config.prefix
-            )
-            hparams, generator_config = set_extra_model_params(extra_model_params, hparams, generator_config)
-            model = self.model_class.from_pretrained(
-                hparams.model_name_or_path, config=config, generator_config=generator_config
-            )
-        else:
-            if args.prefix is not None:
-                setattr(config, "prefix", args.prefix)
-            hparams, config = set_extra_model_params(extra_model_params, hparams, config)
-            model = self.model_class.from_pretrained(hparams.model_name_or_path, config=config)
-            generator_config = config
-
-        tokenizer = (
-            AutoTokenizer.from_pretrained(config.pretrained_generator_tokenizer_name_or_path)
-            if self.is_rag_model
-            else AutoTokenizer.from_pretrained(hparams.model_name_or_path)
-        )
-
-        super().__init__(hparams, config=config, tokenizer=tokenizer, model=model)
-
-        self.retriever = (
-            RagPyTorchDistributedRetriever(self.model.config) if self.is_rag_model else None
-        )  # TODO add tokenizers
-
-        save_git_info(self.hparams.output_dir)
-        self.output_dir = Path(self.hparams.output_dir)
-        self.metrics_save_path = Path(self.output_dir) / "metrics.json"
-        self.hparams_save_path = Path(self.output_dir) / "hparams.pkl"
-        pickle_save(self.hparams, self.hparams_save_path)
-        self.step_count = 0
-        self.metrics = defaultdict(list)
-
-        self.dataset_kwargs: dict = dict(
-            data_dir=self.hparams.data_dir,
-            max_source_length=self.hparams.max_source_length,
-            prefix=generator_config.prefix or "",
-        )
-        n_observations_per_split = {
-            "train": self.hparams.n_train,
-            "val": self.hparams.n_val,
-            "test": self.hparams.n_test,
-        }
-        self.n_obs = {k: v if v >= 0 else None for k, v in n_observations_per_split.items()}
-
-        self.target_lens = {
-            "train": self.hparams.max_target_length,
-            "val": self.hparams.val_max_target_length,
-            "test": self.hparams.test_max_target_length,
-        }
-        assert self.target_lens["train"] <= self.target_lens["val"], f"target_lens: {self.target_lens}"
-        assert self.target_lens["train"] <= self.target_lens["test"], f"target_lens: {self.target_lens}"
-
-        self.hparams.git_sha = get_git_info()["repo_sha"]
-        self.num_workers = hparams.num_workers
-        self.distributed_port = self.hparams.distributed_port
-
-    def init_ddp_connection(self, global_rank: int, world_size: int, is_slurm_managing_tasks: bool = True):
-        logger.info("Custom init_ddp_connection.")
-        os.environ["MASTER_PORT"] = str(self.distributed_port)
-        super().init_ddp_connection(global_rank, world_size, is_slurm_managing_tasks)
-        if self.is_rag_model:
-            self.retriever.init_retrieval(self.distributed_port)
-
-    def forward(self, input_ids, **kwargs):
-        return self.model(input_ids, **kwargs)
-
-    def ids_to_clean_text(self, generated_ids: List[int]):
-        gen_text = self.tokenizer.batch_decode(
-            generated_ids, skip_special_tokens=True, clean_up_tokenization_spaces=True
-        )
-        return lmap(str.strip, gen_text)
-
-    def _step(self, batch: dict) -> Tuple:
-        source_ids, source_mask, target_ids = batch["input_ids"], batch["attention_mask"], batch["decoder_input_ids"]
-
-        if isinstance(self.model, T5ForConditionalGeneration):
-            decoder_input_ids = self.model._shift_right(target_ids)
-            lm_labels = target_ids
-        elif isinstance(self.model, BartForConditionalGeneration):
-            decoder_input_ids = target_ids[:, :-1].contiguous()
-            lm_labels = target_ids[:, 1:].clone()
-        else:
-            assert self.is_rag_model
-            generator = self.model.model.generator
-            if isinstance(generator, T5ForConditionalGeneration):
-                decoder_start_token_id = generator.config.decoder_start_token_id
-                decoder_input_ids = (
-                    torch.cat(
-                        [torch.Tensor([[decoder_start_token_id]] * target_ids.shape[0]).to(target_ids), target_ids],
-                        dim=1,
-                    )
-                    if target_ids.shape[0] < self.target_lens["train"]
-                    else generator._shift_right(target_ids)
-                )
-            elif isinstance(generator, BartForConditionalGeneration):
-                decoder_input_ids = target_ids
-            lm_labels = None
-
-        assert decoder_input_ids is not None
-
-        if lm_labels is not None:
-            outputs = self(
-                source_ids,
-                attention_mask=source_mask,
-                decoder_input_ids=decoder_input_ids,
-                use_cache=False,
-                labels=lm_labels,
-                return_dict=True,
-            )
-        else:  # RAG models
-            outputs = self(
-                source_ids,
-                retriever=self.retriever,
-                attention_mask=source_mask,
-                decoder_input_ids=decoder_input_ids,
-                use_cache=False,
-                return_loss=True,
-                reduce=True,
-                label_smoothing=self.hparams.label_smoothing,
-            )
-
-        loss = outputs["loss"]
-        return (loss,)
-
-    @property
-    def pad(self) -> int:
-        return self.tokenizer.pad_token_id
-
-    def training_step(self, batch, batch_idx) -> Dict:
-        loss_tensors = self._step(batch)
-
-        logs = {name: loss for name, loss in zip(self.loss_names, loss_tensors)}
-        # tokens per batch
-        logs["tpb"] = batch["input_ids"].ne(self.pad).sum() + batch["decoder_input_ids"].ne(self.pad).sum()
-
-        return {"loss": loss_tensors[0], "log": logs}
-
-    def validation_step(self, batch, batch_idx) -> Dict:
-        return self._generative_step(batch)
-
-    def validation_epoch_end(self, outputs, prefix="val") -> Dict:
-        self.step_count += 1
-        losses = {k: torch.stack([x[k] for x in outputs]).mean() for k in self.loss_names}
-        loss = losses["loss"]
-        gen_metrics = {
-            k: np.array([x[k] for x in outputs]).mean() for k in self.metric_names + ["gen_time", "gen_len"]
-        }
-        metrics_tensor: torch.FloatTensor = torch.tensor(gen_metrics[self.val_metric]).type_as(loss)
-        gen_metrics.update({k: v.item() for k, v in losses.items()})
-
-        # fix for https://github.com/PyTorchLightning/pytorch-lightning/issues/2424
-        if dist.is_initialized():
-            dist.all_reduce(metrics_tensor, op=dist.ReduceOp.SUM)
-            metrics_tensor = metrics_tensor / dist.get_world_size()
-            gen_metrics.update({self.val_metric: metrics_tensor.item()})
-
-        losses.update(gen_metrics)
-        metrics = {f"{prefix}_avg_{k}": x for k, x in losses.items()}
-        metrics["step_count"] = self.step_count
-        self.save_metrics(metrics, prefix)  # writes to self.metrics_save_path
-        preds = flatten_list([x["preds"] for x in outputs])
-        return {"log": metrics, "preds": preds, f"{prefix}_loss": loss, f"{prefix}_{self.val_metric}": metrics_tensor}
-
-    def save_metrics(self, latest_metrics, type_path) -> None:
-        self.metrics[type_path].append(latest_metrics)
-        save_json(self.metrics, self.metrics_save_path)
-
-    def calc_generative_metrics(self, preds, target) -> Dict:
-        return calculate_exact_match(preds, target)
-
-    def _generative_step(self, batch: dict) -> dict:
-        start_time = time.time()
-        generated_ids = self.model.generate(
-            batch["input_ids"],
-            retriever=self.retriever,
-            dedup=False,  # rag specific parameter
-            attention_mask=batch["attention_mask"],
-            use_cache=True,
-            min_length=1,
-            max_length=self.target_lens["val"],
-        )
-
-        gen_time = (time.time() - start_time) / batch["input_ids"].shape[0]
-        preds: List[str] = self.ids_to_clean_text(generated_ids)
-        target: List[str] = self.ids_to_clean_text(batch["decoder_input_ids"])
-        loss_tensors = self._step(batch)
-        base_metrics = {name: loss for name, loss in zip(self.loss_names, loss_tensors)}
-        gen_metrics: Dict = self.calc_generative_metrics(preds, target)
-
-        summ_len = np.mean(lmap(len, generated_ids))
-        base_metrics.update(gen_time=gen_time, gen_len=summ_len, preds=preds, target=target, **gen_metrics)
-        return base_metrics
-
-    def test_step(self, batch, batch_idx):
-        return self._generative_step(batch)
-
-    def test_epoch_end(self, outputs):
-        return self.validation_epoch_end(outputs, prefix="test")
-
-    def get_dataset(self, type_path) -> Seq2SeqDataset:
-        n_obs = self.n_obs[type_path]
-        max_target_length = self.target_lens[type_path]
-        dataset = Seq2SeqDataset(
-            self.tokenizer,
-            type_path=type_path,
-            n_obs=n_obs,
-            max_target_length=max_target_length,
-            **self.dataset_kwargs,
-        )
-        return dataset
-
-    def get_dataloader(self, type_path: str, batch_size: int, shuffle: bool = False) -> DataLoader:
-        dataset = self.get_dataset(type_path)
-        sampler = None
-        if self.hparams.sortish_sampler and type_path == "train":
-            assert self.hparams.gpus <= 1  # TODO: assert earlier
-            sampler = dataset.make_sortish_sampler(batch_size)
-            shuffle = False
-
-        dataloader = DataLoader(
-            dataset,
-            batch_size=batch_size,
-            collate_fn=dataset.collate_fn,
-            shuffle=shuffle,
-            num_workers=self.num_workers,
-            sampler=sampler,
-        )
-        return dataloader
-
-    def train_dataloader(self) -> DataLoader:
-        dataloader = self.get_dataloader("train", batch_size=self.hparams.train_batch_size, shuffle=True)
-        t_total = (
-            (len(dataloader.dataset) // (self.hparams.train_batch_size * max(1, self.hparams.gpus)))
-            // self.hparams.accumulate_grad_batches
-            * float(self.hparams.max_epochs)
-        )
-        scheduler = get_linear_schedule_with_warmup(
-            self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=t_total
-        )
-        if max(scheduler.get_last_lr()) > 0:
-            warnings.warn("All learning rates are 0")
-        self.lr_scheduler = scheduler
-        return dataloader
-
-    def val_dataloader(self) -> DataLoader:
-        return self.get_dataloader("val", batch_size=self.hparams.eval_batch_size)
-
-    def test_dataloader(self) -> DataLoader:
-        return self.get_dataloader("test", batch_size=self.hparams.eval_batch_size)
-
-    @pl.utilities.rank_zero_only
-    def on_save_checkpoint(self, checkpoint: Dict[str, Any]) -> None:
-        save_path = self.output_dir.joinpath("checkpoint{}".format(self.step_count))
-        self.model.config.save_step = self.step_count
-        self.model.save_pretrained(save_path)
-        self.tokenizer.save_pretrained(save_path)
-
-    @staticmethod
-    def add_model_specific_args(parser, root_dir):
-        BaseTransformer.add_model_specific_args(parser, root_dir)
-        add_generic_args(parser, root_dir)
-        parser.add_argument(
-            "--max_source_length",
-            default=128,
-            type=int,
-            help="The maximum total input sequence length after tokenization. Sequences longer "
-            "than this will be truncated, sequences shorter will be padded.",
-        )
-        parser.add_argument(
-            "--max_target_length",
-            default=25,
-            type=int,
-            help="The maximum total input sequence length after tokenization. Sequences longer "
-            "than this will be truncated, sequences shorter will be padded.",
-        )
-        parser.add_argument(
-            "--val_max_target_length",
-            default=25,
-            type=int,
-            help="The maximum total input sequence length after tokenization. Sequences longer "
-            "than this will be truncated, sequences shorter will be padded.",
-        )
-        parser.add_argument(
-            "--test_max_target_length",
-            default=25,
-            type=int,
-            help="The maximum total input sequence length after tokenization. Sequences longer "
-            "than this will be truncated, sequences shorter will be padded.",
-        )
-        parser.add_argument("--sortish_sampler", action="store_true", default=False)
-        parser.add_argument("--logger_name", type=str, choices=["default", "wandb", "wandb_shared"], default="default")
-        parser.add_argument("--n_train", type=int, default=-1, required=False, help="# examples. -1 means use all.")
-        parser.add_argument("--n_val", type=int, default=-1, required=False, help="# examples. -1 means use all.")
-        parser.add_argument("--n_test", type=int, default=-1, required=False, help="# examples. -1 means use all.")
-        parser.add_argument("--label_smoothing", type=float, default=0.0, required=False)
-        parser.add_argument(
-            "--prefix",
-            type=str,
-            default=None,
-            help="Prefix added at the beginning of each text, typically used with T5-based models.",
-        )
-        parser.add_argument(
-            "--early_stopping_patience",
-            type=int,
-            default=-1,
-            required=False,
-            help="-1 means never early stop. early_stopping_patience is measured in validation checks, not epochs. So val_check_interval will effect it.",
-        )
-        parser.add_argument(
-            "--distributed-port", type=int, default=-1, required=False, help="Port number for distributed training."
-        )
-        parser.add_argument(
-            "--model_type",
-            choices=["rag_sequence", "rag_token", "bart", "t5"],
-            type=str,
-            help="RAG model type: sequence or token, if none specified, the type is inferred from the model_name_or_path",
-        )
-        return parser
-
-
-def main(args, model=None) -> GenerativeQAModule:
-    Path(args.output_dir).mkdir(exist_ok=True)
-    if model is None:
-        model: GenerativeQAModule = GenerativeQAModule(args)
-
-    dataset = Path(args.data_dir).name
-    if (
-        args.logger_name == "default"
-        or args.fast_dev_run
-        or str(args.output_dir).startswith("/tmp")
-        or str(args.output_dir).startswith("/var")
-    ):
-        logger = True  # don't pollute wandb logs unnecessarily
-    elif args.logger_name == "wandb":
-        from pytorch_lightning.loggers import WandbLogger
-
-        project = os.environ.get("WANDB_PROJECT", dataset)
-        logger = WandbLogger(name=model.output_dir.name, project=project)
-
-    elif args.logger_name == "wandb_shared":
-        from pytorch_lightning.loggers import WandbLogger
-
-        logger = WandbLogger(name=model.output_dir.name, project=f"hf_{dataset}")
-
-    es_callback = (
-        get_early_stopping_callback(model.val_metric, args.early_stopping_patience)
-        if args.early_stopping_patience >= 0
-        else False
-    )
-    trainer: pl.Trainer = generic_train(
-        model,
-        args,
-        logging_callback=Seq2SeqLoggingCallback(),
-        checkpoint_callback=get_checkpoint_callback(args.output_dir, model.val_metric),
-        early_stopping_callback=es_callback,
-        logger=logger,
-    )
-    pickle_save(model.hparams, model.output_dir / "hparams.pkl")
-
-    if not args.do_predict:
-        return model
-
-    model.hparams.test_checkpoint = ""
-    checkpoints = list(sorted(glob.glob(os.path.join(args.output_dir, "*.ckpt"), recursive=True)))
-    if checkpoints:
-        model.hparams.test_checkpoint = checkpoints[-1]
-        trainer.resume_from_checkpoint = checkpoints[-1]  # best checkpoint
-    trainer.logger.log_hyperparams(model.hparams)
-
-    # test() without a model tests using the best checkpoint automatically
-    trainer.test()
-
-    return model
-
-
-if __name__ == "__main__":
-    parser = argparse.ArgumentParser()
-    parser = pl.Trainer.add_argparse_args(parser)
-    parser = GenerativeQAModule.add_model_specific_args(parser, os.getcwd())
-
-    args = parser.parse_args()
-
-    main(args)

examples/rag/finetune.sh

@@ -1,34 +0,0 @@
-# Add parent directory to python path to access lightning_base.py
-export PYTHONPATH="../":"${PYTHONPATH}"
-
-# A sample finetuning run, you need to specify data_dir, output_dir and model_name_or_path
-# run ./examples/rag/finetune.sh --help to see all the possible options
-
-python examples/rag/finetune.py \
-    --data_dir $DATA_DIR \
-    --output_dir $OUTPUT_DIR \
-    --model_name_or_path $MODLE_NAME_OR_PATH \
-    --model_type rag_sequence \
-    --fp16 \
-    --gpus 8 \
-    --do_train \
-    --do_predict \
-    --n_val -1 \
-    --val_check_interval 0.25 \
-    --train_batch_size 8 \
-    --eval_batch_size 1 \
-    --max_source_length 128 \
-    --max_target_length 25 \
-    --val_max_target_length 25 \
-    --test_max_target_length 25 \
-    --label_smoothing 0.1 \
-    --dropout 0.1 \
-    --attention_dropout 0.1 \
-    --weight_decay 0.001 \
-    --adam_epsilon 1e-08 \
-    --max_grad_norm 0.1 \
-    --lr_scheduler polynomial \
-    --learning_rate 3e-05 \
-    --num_train_epochs 100 \
-    --warmup_steps 500 \
-    --gradient_accumulation_steps 1

examples/rag/parse_dpr_relevance_data.py

@@ -1,47 +0,0 @@
-"""
-This script reads DPR retriever training data and parses each datapoint. We save a line per datapoint.
-Each line consists of the query followed by a tab-separated list of Wikipedia page titles constituting
-positive contexts for a given query.
-"""
-
-import argparse
-import json
-
-from tqdm import tqdm
-
-
-def main():
-    parser = argparse.ArgumentParser()
-
-    # Required parameters
-    parser.add_argument(
-        "--src_path",
-        type=str,
-        default="biencoder-nq-dev.json",
-        help="Path to raw DPR training data",
-    )
-    parser.add_argument(
-        "--evaluation_set",
-        type=str,
-        help="where to store parsed evaluation_set file",
-    )
-    parser.add_argument(
-        "--gold_data_path",
-        type=str,
-        help="where to store parsed gold_data_path file",
-    )
-    args = parser.parse_args()
-
-    with open(args.src_path, "r") as src_file, open(args.evaluation_set, "w") as eval_file, open(
-        args.gold_data_path, "w"
-    ) as gold_file:
-        dpr_records = json.load(src_file)
-        for dpr_record in tqdm(dpr_records):
-            question = dpr_record["question"]
-            contexts = [context["title"] for context in dpr_record["positive_ctxs"]]
-            eval_file.write(question + "\n")
-            gold_file.write("\t".join(contexts) + "\n")
-
-
-if __name__ == "__main__":
-    main()

examples/rag/utils.py

@@ -1,174 +0,0 @@
-import linecache
-import re
-import string
-from collections import Counter
-from logging import getLogger
-from pathlib import Path
-from typing import Dict, List
-
-import torch
-from torch.utils.data import Dataset
-
-from examples.seq2seq.utils import SortishSampler, trim_batch
-from transformers import BartTokenizer, T5Tokenizer
-
-
-def encode_line(tokenizer, line, max_length, padding_side, pad_to_max_length=True, return_tensors="pt"):
-    extra_kw = {"add_prefix_space": True} if isinstance(tokenizer, BartTokenizer) else {}
-    tokenizer.padding_side = padding_side
-    return tokenizer(
-        [line],
-        max_length=max_length,
-        padding="max_length" if pad_to_max_length else None,
-        truncation=True,
-        return_tensors=return_tensors,
-        add_special_tokens=True,
-        **extra_kw,
-    )
-
-
-class Seq2SeqDataset(Dataset):
-    def __init__(
-        self,
-        tokenizer,
-        data_dir,
-        max_source_length,
-        max_target_length,
-        type_path="train",
-        n_obs=None,
-        src_lang=None,
-        tgt_lang=None,
-        prefix="",
-    ):
-        super().__init__()
-        self.src_file = Path(data_dir).joinpath(type_path + ".source")
-        self.tgt_file = Path(data_dir).joinpath(type_path + ".target")
-        self.src_lens = self.get_char_lens(self.src_file)
-        self.max_source_length = max_source_length
-        self.max_target_length = max_target_length
-        assert min(self.src_lens) > 0, f"found empty line in {self.src_file}"
-        self.tokenizer = tokenizer
-        self.prefix = prefix
-        if n_obs is not None:
-            self.src_lens = self.src_lens[:n_obs]
-        self.pad_token_id = self.tokenizer.pad_token_id
-        self.src_lang = src_lang
-        self.tgt_lang = tgt_lang
-
-    def __len__(self):
-        return len(self.src_lens)
-
-    def __getitem__(self, index) -> Dict[str, torch.Tensor]:
-        index = index + 1  # linecache starts at 1
-        source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n")
-        tgt_line = linecache.getline(str(self.tgt_file), index).rstrip("\n")
-        assert source_line, f"empty source line for index {index}"
-        assert tgt_line, f"empty tgt line for index {index}"
-
-        # Need to add eos token manually for T5
-        if isinstance(self.tokenizer, T5Tokenizer):
-            source_line += self.tokenizer.eos_token
-            tgt_line += self.tokenizer.eos_token
-
-        # Pad source to the left and target to the right
-        source_inputs = encode_line(self.tokenizer, source_line, self.max_source_length, "right")  # "left")
-        target_inputs = encode_line(self.tokenizer, tgt_line, self.max_target_length, "right")
-
-        source_ids = source_inputs["input_ids"].squeeze()
-        target_ids = target_inputs["input_ids"].squeeze()
-        src_mask = source_inputs["attention_mask"].squeeze()
-        return {
-            "input_ids": source_ids,
-            "attention_mask": src_mask,
-            "decoder_input_ids": target_ids,
-        }
-
-    @staticmethod
-    def get_char_lens(data_file):
-        return [len(x) for x in Path(data_file).open().readlines()]
-
-    def collate_fn(self, batch) -> Dict[str, torch.Tensor]:
-        input_ids = torch.stack([x["input_ids"] for x in batch])
-        masks = torch.stack([x["attention_mask"] for x in batch])
-        target_ids = torch.stack([x["decoder_input_ids"] for x in batch])
-        pad_token_id = self.pad_token_id
-        y = trim_batch(target_ids, pad_token_id)
-        source_ids, source_mask = trim_batch(input_ids, pad_token_id, attention_mask=masks)
-        batch = {
-            "input_ids": source_ids,
-            "attention_mask": source_mask,
-            "decoder_input_ids": y,
-        }
-        return batch
-
-    def make_sortish_sampler(self, batch_size):
-        return SortishSampler(self.src_lens, batch_size)
-
-
-logger = getLogger(__name__)
-
-
-def normalize_answer(s):
-    """Lower text and remove punctuation, articles and extra whitespace."""
-
-    def remove_articles(text):
-        return re.sub(r"\b(a|an|the)\b", " ", text)
-
-    def white_space_fix(text):
-        return " ".join(text.split())
-
-    def remove_punc(text):
-        exclude = set(string.punctuation)
-        return "".join(ch for ch in text if ch not in exclude)
-
-    def lower(text):
-        return text.lower()
-
-    return white_space_fix(remove_articles(remove_punc(lower(s))))
-
-
-def f1_score(prediction, ground_truth):
-    prediction_tokens = normalize_answer(prediction).split()
-    ground_truth_tokens = normalize_answer(ground_truth).split()
-    common = Counter(prediction_tokens) & Counter(ground_truth_tokens)
-    num_same = sum(common.values())
-    if num_same == 0:
-        return 0
-    precision = 1.0 * num_same / len(prediction_tokens)
-    recall = 1.0 * num_same / len(ground_truth_tokens)
-    f1 = (2 * precision * recall) / (precision + recall)
-    return f1
-
-
-def exact_match_score(prediction, ground_truth):
-    return normalize_answer(prediction) == normalize_answer(ground_truth)
-
-
-def calculate_exact_match(output_lns: List[str], reference_lns: List[str]) -> Dict:
-    assert len(output_lns) == len(reference_lns)
-    em = 0
-    for hypo, pred in zip(output_lns, reference_lns):
-        em += exact_match_score(hypo, pred)
-    if len(output_lns) > 0:
-        em /= len(output_lns)
-    return {"em": em}
-
-
-def is_rag_model(model_prefix):
-    return model_prefix.startswith("rag")
-
-
-def set_extra_model_params(extra_params, hparams, config):
-    equivalent_param = {p: p for p in extra_params}
-    # T5 models don't have `dropout` param, they have `dropout_rate` instead
-    equivalent_param["dropout"] = "dropout_rate"
-    for p in extra_params:
-        if getattr(hparams, p, None):
-            if not hasattr(config, p) and not hasattr(config, equivalent_param[p]):
-                logger.info("config doesn't have a `{}` attribute".format(p))
-                delattr(hparams, p)
-                continue
-            set_p = p if hasattr(config, p) else equivalent_param[p]
-            setattr(config, set_p, getattr(hparams, p))
-            delattr(hparams, p)
-    return hparams, config

examples/requirements.txt

@@ -12,7 +12,7 @@ faiss
 streamlit
 elasticsearch
 pandas
-datasets
+nlp
 fire
 pytest
 conllu

examples/seq2seq/README.md

@@ -6,9 +6,8 @@ Please tag @sshleifer with any issues/unexpected behaviors, or send a PR!
 For `bertabs` instructions, see [`bertabs/README.md`](bertabs/README.md).
 
 
-## Datasets
-
-#### XSUM:
+### Data
+XSUM Data:
 ```bash
 cd examples/seq2seq
 wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/xsum.tar.gz
@@ -18,33 +17,23 @@ export XSUM_DIR=${PWD}/xsum
 this should make a directory called `xsum/` with files like `test.source`.
 To use your own data, copy that files format. Each article to be summarized is on its own line.
 
-#### CNN/DailyMail
+CNN/DailyMail data
 ```bash
 cd examples/seq2seq
-wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/cnn_dm_v2.tgz
-tar -xzvf cnn_dm_v2.tgz  # empty lines removed
-mv cnn_cln cnn_dm
+wget https://s3.amazonaws.com/datasets.huggingface.co/summarization/cnn_dm.tgz
+tar -xzvf cnn_dm.tgz
 export CNN_DIR=${PWD}/cnn_dm
+this should make a directory called `cnn_dm/` with files like `test.source`.
 ```
-this should make a directory called `cnn_dm/` with 6 files.
 
-#### WMT16 English-Romanian Translation Data:
+WMT16 English-Romanian Translation Data:
 download with this command:
 ```bash
 wget https://s3.amazonaws.com/datasets.huggingface.co/translation/wmt_en_ro.tar.gz
 tar -xzvf wmt_en_ro.tar.gz
 export ENRO_DIR=${PWD}/wmt_en_ro
+this should make a directory called `wmt_en_ro/` with files like `test.source`.
 ```
-this should make a directory called `wmt_en_ro/` with 6 files.
-
-#### WMT English-German:
-```bash
-wget https://s3.amazonaws.com/datasets.huggingface.co/translation/wmt_en_de.tgz
-tar -xzvf wmt_en_de.tgz
-export DATA_DIR=${PWD}/wmt_en_de
-```
-
-#### Private Data
 
 If you are using your own data, it must be formatted as one directory with 6 files: 
 ```
@@ -82,12 +71,11 @@ Summarization Tips:
 (It rarely makes sense to start from `bart-large` unless you are a researching finetuning methods).
 
 **Update 2018-07-18**
-Datasets: `LegacySeq2SeqDataset` will be used for all tokenizers without a `prepare_seq2seq_batch` method. Otherwise, `Seq2SeqDataset` will be used.
-Future work/help wanted: A new dataset to support multilingual tasks.
+Datasets: `Seq2SeqDataset` should be used for all tokenizers without a `prepare_seq2seq_batch` method. For those who do (like Marian, MBart), `TranslationDataset` should be used.**
+A new dataset is needed to support multilingual tasks.
 
 
-### Finetuning Scripts
-All finetuning bash scripts call finetune.py (or distillation.py) with reasonable command line arguments. They usually require extra command line arguments to work.
+### Command Line Options
 
 To see all the possible command line options, run:
 
@@ -122,8 +110,6 @@ The following command should work on a 16GB GPU:
     --model_name_or_path facebook/bart-large
 ```
 
-There is a starter finetuning script for pegasus at `finetune_pegasus_xsum.sh`.
-
 ### Translation Finetuning
 
 First, follow the wmt_en_ro download instructions.
@@ -136,8 +122,8 @@ Best performing command:
 export ENRO_DIR='wmt_en_ro' # Download instructions above
 # export WANDB_PROJECT="MT" # optional
 export MAX_LEN=128
-export BS=4
-./train_mbart_cc25_enro.sh --output_dir enro_finetune_baseline --label_smoothing 0.1 --fp16_opt_level=O1 --logger_name wandb --sortish_sampler
+export BS=8
+./train_mbart_cc25_enro.sh --output_dir enro_finetune_baseline_dropper --label_smoothing 0 --fp16_opt_level=O1 --logger_name wandb --sortish_sampler
 ```
 This should take < 6h/epoch on a 16GB v100 and achieve test BLEU above 26
 To get results in line with fairseq, you need to do some postprocessing. (see `romanian_postprocessing.md`)
@@ -155,7 +141,7 @@ export BS=4
 As you train, `output_dir` will be filled with files, that look kind of like this (comments are mine).
 Some of them are metrics, some of them are checkpoints, some of them are metadata. Here is a quick tour:
 
-```bash
+```
 output_dir
 ├── best_tfmr  # this is a huggingface checkpoint generated by save_pretrained. It is the same model as the PL .ckpt file below
 │   ├── config.json

examples/seq2seq/bertabs/configuration_bertabs.py

@@ -28,7 +28,7 @@ BERTABS_FINETUNED_CONFIG_MAP = {
 
 
 class BertAbsConfig(PretrainedConfig):
-    r"""Class to store the configuration of the BertAbs model.
+    r""" Class to store the configuration of the BertAbs model.
 
     Arguments:
         vocab_size: int

examples/seq2seq/bertabs/convert_bertabs_original_pytorch_checkpoint.py

@@ -62,7 +62,7 @@ BertAbsConfig = namedtuple(
 
 
 def convert_bertabs_checkpoints(path_to_checkpoints, dump_path):
-    """Copy/paste and tweak the pre-trained weights provided by the creators
+    """ Copy/paste and tweak the pre-trained weights provided by the creators
     of BertAbs for the internal architecture.
     """
 
@@ -164,22 +164,13 @@ def convert_bertabs_checkpoints(path_to_checkpoints, dump_path):
 if __name__ == "__main__":
     parser = argparse.ArgumentParser()
     parser.add_argument(
-        "--bertabs_checkpoint_path",
-        default=None,
-        type=str,
-        required=True,
-        help="Path the official PyTorch dump.",
+        "--bertabs_checkpoint_path", default=None, type=str, required=True, help="Path the official PyTorch dump.",
     )
     parser.add_argument(
-        "--pytorch_dump_folder_path",
-        default=None,
-        type=str,
-        required=True,
-        help="Path to the output PyTorch model.",
+        "--pytorch_dump_folder_path", default=None, type=str, required=True, help="Path to the output PyTorch model.",
     )
     args = parser.parse_args()
 
     convert_bertabs_checkpoints(
-        args.bertabs_checkpoint_path,
-        args.pytorch_dump_folder_path,
+        args.bertabs_checkpoint_path, args.pytorch_dump_folder_path,
     )

examples/seq2seq/bertabs/modeling_bertabs.py

@@ -105,17 +105,10 @@ class BertAbs(BertAbsPreTrainedModel):
                 p.data.zero_()
 
     def forward(
-        self,
-        encoder_input_ids,
-        decoder_input_ids,
-        token_type_ids,
-        encoder_attention_mask,
-        decoder_attention_mask,
+        self, encoder_input_ids, decoder_input_ids, token_type_ids, encoder_attention_mask, decoder_attention_mask,
     ):
         encoder_output = self.bert(
-            input_ids=encoder_input_ids,
-            token_type_ids=token_type_ids,
-            attention_mask=encoder_attention_mask,
+            input_ids=encoder_input_ids, token_type_ids=token_type_ids, attention_mask=encoder_attention_mask,
         )
         encoder_hidden_states = encoder_output[0]
         dec_state = self.decoder.init_decoder_state(encoder_input_ids, encoder_hidden_states)
@@ -124,7 +117,8 @@ class BertAbs(BertAbsPreTrainedModel):
 
 
 class Bert(nn.Module):
-    """This class is not really necessary and should probably disappear."""
+    """ This class is not really necessary and should probably disappear.
+    """
 
     def __init__(self):
         super().__init__()
@@ -313,14 +307,7 @@ class TransformerDecoderLayer(nn.Module):
         self.register_buffer("mask", mask)
 
     def forward(
-        self,
-        inputs,
-        memory_bank,
-        src_pad_mask,
-        tgt_pad_mask,
-        previous_input=None,
-        layer_cache=None,
-        step=None,
+        self, inputs, memory_bank, src_pad_mask, tgt_pad_mask, previous_input=None, layer_cache=None, step=None,
     ):
         """
         Args:
@@ -344,25 +331,13 @@ class TransformerDecoderLayer(nn.Module):
             all_input = torch.cat((previous_input, input_norm), dim=1)
             dec_mask = None
 
-        query = self.self_attn(
-            all_input,
-            all_input,
-            input_norm,
-            mask=dec_mask,
-            layer_cache=layer_cache,
-            type="self",
-        )
+        query = self.self_attn(all_input, all_input, input_norm, mask=dec_mask, layer_cache=layer_cache, type="self",)
 
         query = self.drop(query) + inputs
 
         query_norm = self.layer_norm_2(query)
         mid = self.context_attn(
-            memory_bank,
-            memory_bank,
-            query_norm,
-            mask=src_pad_mask,
-            layer_cache=layer_cache,
-            type="context",
+            memory_bank, memory_bank, query_norm, mask=src_pad_mask, layer_cache=layer_cache, type="context",
         )
         output = self.feed_forward(self.drop(mid) + query)
 
@@ -447,14 +422,7 @@ class MultiHeadedAttention(nn.Module):
             self.final_linear = nn.Linear(model_dim, model_dim)
 
     def forward(
-        self,
-        key,
-        value,
-        query,
-        mask=None,
-        layer_cache=None,
-        type=None,
-        predefined_graph_1=None,
+        self, key, value, query, mask=None, layer_cache=None, type=None, predefined_graph_1=None,
     ):
         """
         Compute the context vector and the attention vectors.
@@ -660,7 +628,7 @@ def gelu(x):
 
 
 class PositionwiseFeedForward(nn.Module):
-    """A two-layer Feed-Forward-Network with residual layer norm.
+    """ A two-layer Feed-Forward-Network with residual layer norm.
 
     Args:
         d_model (int): the size of input for the first-layer of the FFN.
@@ -802,7 +770,8 @@ class Translator(object):
         self.max_length = args.max_length
 
     def translate(self, batch, step, attn_debug=False):
-        """Generates summaries from one batch of data."""
+        """ Generates summaries from one batch of data.
+        """
         self.model.eval()
         with torch.no_grad():
             batch_data = self.translate_batch(batch)
@@ -829,7 +798,8 @@ class Translator(object):
     # Where the beam search lives
     # I have no idea why it is being called from the method above
     def _fast_translate_batch(self, batch, max_length, min_length=0):
-        """Beam Search using the encoder inputs contained in `batch`."""
+        """ Beam Search using the encoder inputs contained in `batch`.
+        """
 
         # The batch object is funny
         # Instead of just looking at the size of the arguments we encapsulate
@@ -1011,7 +981,7 @@ def tile(x, count, dim=0):
 
 
 class BertSumOptimizer(object):
-    """Specific optimizer for BertSum.
+    """ Specific optimizer for BertSum.
 
     As described in [1], the authors fine-tune BertSum for abstractive
     summarization using two Adam Optimizers with different warm-up steps and
@@ -1029,16 +999,10 @@ class BertSumOptimizer(object):
 
         self.optimizers = {
             "encoder": torch.optim.Adam(
-                model.encoder.parameters(),
-                lr=lr["encoder"],
-                betas=(beta_1, beta_2),
-                eps=eps,
+                model.encoder.parameters(), lr=lr["encoder"], betas=(beta_1, beta_2), eps=eps,
             ),
             "decoder": torch.optim.Adam(
-                model.decoder.parameters(),
-                lr=lr["decoder"],
-                betas=(beta_1, beta_2),
-                eps=eps,
+                model.decoder.parameters(), lr=lr["decoder"], betas=(beta_1, beta_2), eps=eps,
             ),
         }
 

examples/seq2seq/bertabs/run_summarization.py

@@ -44,9 +44,8 @@ def evaluate(args):
         reference_summaries = []
         generated_summaries = []
 
-        import nltk
-
         import rouge
+        import nltk
 
         nltk.download("punkt")
         rouge_evaluator = rouge.Rouge(
@@ -99,7 +98,7 @@ def evaluate(args):
 
 
 def save_summaries(summaries, path, original_document_name):
-    """Write the summaries in fies that are prefixed by the original
+    """ Write the summaries in fies that are prefixed by the original
     files' name with the `_summary` appended.
 
     Attributes:
@@ -125,7 +124,7 @@ def save_summaries(summaries, path, original_document_name):
 
 
 def format_summary(translation):
-    """Transforms the output of the `from_batch` function
+    """ Transforms the output of the `from_batch` function
     into nicely formatted summaries.
     """
     raw_summary, _, _ = translation
@@ -190,12 +189,7 @@ def build_data_iterator(args, tokenizer):
     def collate_fn(data):
         return collate(data, tokenizer, block_size=512, device=args.device)
 
-    iterator = DataLoader(
-        dataset,
-        sampler=sampler,
-        batch_size=args.batch_size,
-        collate_fn=collate_fn,
-    )
+    iterator = DataLoader(dataset, sampler=sampler, batch_size=args.batch_size, collate_fn=collate_fn,)
 
     return iterator
 
@@ -206,7 +200,7 @@ def load_and_cache_examples(args, tokenizer):
 
 
 def collate(data, tokenizer, block_size, device):
-    """Collate formats the data passed to the data loader.
+    """ Collate formats the data passed to the data loader.
 
     In particular we tokenize the data batch after batch to avoid keeping them
     all in memory. We output the data as a namedtuple to fit the original BertAbs's
@@ -236,7 +230,7 @@ def collate(data, tokenizer, block_size, device):
 
 
 def decode_summary(summary_tokens, tokenizer):
-    """Decode the summary and return it in a format
+    """ Decode the summary and return it in a format
     suitable for evaluation.
     """
     summary_tokens = summary_tokens.to("cpu").numpy()
@@ -247,7 +241,8 @@ def decode_summary(summary_tokens, tokenizer):
 
 
 def main():
-    """The main function defines the interface with the users."""
+    """ The main function defines the interface with the users.
+    """
     parser = argparse.ArgumentParser()
     parser.add_argument(
         "--documents_dir",
@@ -272,41 +267,23 @@ def main():
     )
     # EVALUATION options
     parser.add_argument(
-        "--no_cuda",
-        default=False,
-        type=bool,
-        help="Whether to force the execution on CPU.",
+        "--no_cuda", default=False, type=bool, help="Whether to force the execution on CPU.",
     )
     parser.add_argument(
-        "--batch_size",
-        default=4,
-        type=int,
-        help="Batch size per GPU/CPU for training.",
+        "--batch_size", default=4, type=int, help="Batch size per GPU/CPU for training.",
     )
     # BEAM SEARCH arguments
     parser.add_argument(
-        "--min_length",
-        default=50,
-        type=int,
-        help="Minimum number of tokens for the summaries.",
+        "--min_length", default=50, type=int, help="Minimum number of tokens for the summaries.",
     )
     parser.add_argument(
-        "--max_length",
-        default=200,
-        type=int,
-        help="Maixmum number of tokens for the summaries.",
+        "--max_length", default=200, type=int, help="Maixmum number of tokens for the summaries.",
     )
     parser.add_argument(
-        "--beam_size",
-        default=5,
-        type=int,
-        help="The number of beams to start with for each example.",
+        "--beam_size", default=5, type=int, help="The number of beams to start with for each example.",
     )
     parser.add_argument(
-        "--alpha",
-        default=0.95,
-        type=float,
-        help="The value of alpha for the length penalty in the beam search.",
+        "--alpha", default=0.95, type=float, help="The value of alpha for the length penalty in the beam search.",
     )
     parser.add_argument(
         "--block_trigram",

examples/seq2seq/bertabs/test_utils_summarization.py

@@ -43,7 +43,8 @@ class SummarizationDataProcessingTest(unittest.TestCase):
         self.assertEqual(truncate_or_pad(sequence, self.block_size, 0), expected_output)
 
     def test_process_story_no_highlights(self):
-        """Processing a story with no highlights returns an empty list for the summary."""
+        """ Processing a story with no highlights returns an empty list for the summary.
+        """
         raw_story = """It was the year of Our Lord one thousand seven hundred and
         seventy-five.\n\nSpiritual revelations were conceded to England at that
         favoured period, as at this."""
@@ -51,7 +52,8 @@ class SummarizationDataProcessingTest(unittest.TestCase):
         self.assertEqual(summary_lines, [])
 
     def test_process_empty_story(self):
-        """An empty story returns an empty collection of lines."""
+        """ An empty story returns an empty collection of lines.
+        """
         raw_story = ""
         story_lines, summary_lines = process_story(raw_story)
         self.assertEqual(story_lines, [])

examples/seq2seq/bertabs/utils_summarization.py

@@ -11,7 +11,7 @@ from torch.utils.data import Dataset
 
 
 class CNNDMDataset(Dataset):
-    """Abstracts the dataset used to train seq2seq models.
+    """ Abstracts the dataset used to train seq2seq models.
 
     The class will process the documents that are located in the specified
     folder. The preprocessing will work on any document that is reasonably
@@ -31,7 +31,7 @@ class CNNDMDataset(Dataset):
     """
 
     def __init__(self, path="", prefix="train"):
-        """We initialize the class by listing all the documents to summarize.
+        """ We initialize the class by listing all the documents to summarize.
         Files are not read in memory due to the size of some datasets (like CNN/DailyMail).
         """
         assert os.path.isdir(path)
@@ -60,7 +60,7 @@ class CNNDMDataset(Dataset):
 
 
 def process_story(raw_story):
-    """Extract the story and summary from a story file.
+    """ Extract the story and summary from a story file.
 
     Arguments:
         raw_story (str): content of the story file as an utf-8 encoded string.
@@ -108,7 +108,7 @@ def _add_missing_period(line):
 
 
 def truncate_or_pad(sequence, block_size, pad_token_id):
-    """Adapt the source and target sequences' lengths to the block size.
+    """ Adapt the source and target sequences' lengths to the block size.
     If the sequence is shorter we append padding token to the right of the sequence.
     """
     if len(sequence) > block_size:
@@ -119,8 +119,8 @@ def truncate_or_pad(sequence, block_size, pad_token_id):
 
 
 def build_mask(sequence, pad_token_id):
-    """Builds the mask. The attention mechanism will only attend to positions
-    with value 1."""
+    """ Builds the mask. The attention mechanism will only attend to positions
+    with value 1. """
     mask = torch.ones_like(sequence)
     idx_pad_tokens = sequence == pad_token_id
     mask[idx_pad_tokens] = 0
@@ -128,7 +128,7 @@ def build_mask(sequence, pad_token_id):
 
 
 def encode_for_summarization(story_lines, summary_lines, tokenizer):
-    """Encode the story and summary lines, and join them
+    """ Encode the story and summary lines, and join them
     as specified in [1] by using `[SEP] [CLS]` tokens to separate
     sentences.
     """
@@ -141,7 +141,7 @@ def encode_for_summarization(story_lines, summary_lines, tokenizer):
 
 
 def compute_token_type_ids(batch, separator_token_id):
-    """Segment embeddings as described in [1]
+    """ Segment embeddings as described in [1]
 
     The values {0,1} were found in the repository [2].
 

examples/seq2seq/callbacks.py

@@ -75,33 +75,26 @@ class Seq2SeqLoggingCallback(pl.Callback):
         return self._write_logs(trainer, pl_module, "test")
 
 
-def get_checkpoint_callback(output_dir, metric, save_top_k=1, lower_is_better=False):
+def get_checkpoint_callback(output_dir, metric):
     """Saves the best model by validation ROUGE2 score."""
     if metric == "rouge2":
         exp = "{val_avg_rouge2:.4f}-{step_count}"
     elif metric == "bleu":
         exp = "{val_avg_bleu:.4f}-{step_count}"
-    elif metric == "loss":
-        exp = "{val_avg_loss:.4f}-{step_count}"
     else:
         raise NotImplementedError(
-            f"seq2seq callbacks only support rouge2, bleu and loss, got {metric}, You can make your own by adding to this function."
+            f"seq2seq callbacks only support rouge2 and bleu, got {metric}, You can make your own by adding to this function."
         )
 
     checkpoint_callback = ModelCheckpoint(
         filepath=os.path.join(output_dir, exp),
         monitor=f"val_{metric}",
-        mode="min" if "loss" in metric else "max",
-        save_top_k=save_top_k,
+        mode="max",
+        save_top_k=1,
         period=0,  # maybe save a checkpoint every time val is run, not just end of epoch.
     )
     return checkpoint_callback
 
 
 def get_early_stopping_callback(metric, patience):
-    return EarlyStopping(
-        monitor=f"val_{metric}",  # does this need avg?
-        mode="min" if "loss" in metric else "max",
-        patience=patience,
-        verbose=True,
-    )
+    return EarlyStopping(monitor=f"val_{metric}", mode="max", patience=patience, verbose=True,)

examples/seq2seq/convert_pl_checkpoint_to_hf.py

@@ -1,72 +0,0 @@
-import os
-from pathlib import Path
-from typing import Dict, List
-
-import fire
-import torch
-
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from transformers.utils.logging import get_logger
-
-
-logger = get_logger(__name__)
-
-
-def remove_prefix(text: str, prefix: str):
-    if text.startswith(prefix):
-        return text[len(prefix) :]
-    return text  # or whatever
-
-
-def sanitize(sd):
-    return {remove_prefix(k, "model."): v for k, v in sd.items()}
-
-
-def average_state_dicts(state_dicts: List[Dict[str, torch.Tensor]]):
-    new_sd = {}
-    for k in state_dicts[0].keys():
-        tensors = [sd[k] for sd in state_dicts]
-        new_t = sum(tensors) / len(tensors)
-        assert isinstance(new_t, torch.Tensor)
-        new_sd[k] = new_t
-    return new_sd
-
-
-def convert_pl_to_hf(pl_ckpt_path: str, hf_src_model_dir: str, save_path: str) -> None:
-    """Cleanup a pytorch-lightning .ckpt file or experiment dir and save a huggingface model with that state dict.
-    Silently allows extra pl keys (like teacher.) Puts all ckpt models into CPU RAM at once!
-
-    Args:
-        pl_ckpt_path (:obj:`str`): Path to a .ckpt file saved by pytorch_lightning or dir containing ckpt files.
-            If a directory is passed, all .ckpt files inside it will be averaged!
-        hf_src_model_dir (:obj:`str`): Path to a directory containing a correctly shaped checkpoint
-        save_path (:obj:`str`): Directory to save the new model
-
-    """
-    hf_model = AutoModelForSeq2SeqLM.from_pretrained(hf_src_model_dir)
-    if os.path.isfile(pl_ckpt_path):
-        ckpt_files = [pl_ckpt_path]
-    else:
-        assert os.path.isdir(pl_ckpt_path)
-        ckpt_files = list(Path(pl_ckpt_path).glob("*.ckpt"))
-        assert ckpt_files, f"could not find any ckpt files inside the {pl_ckpt_path} directory"
-
-    if len(ckpt_files) > 1:
-        logger.info(f"averaging the weights of {ckpt_files}")
-
-    state_dicts = [sanitize(torch.load(x, map_location="cpu")["state_dict"]) for x in ckpt_files]
-    state_dict = average_state_dicts(state_dicts)
-
-    missing, unexpected = hf_model.load_state_dict(state_dict, strict=False)
-    assert not missing, f"missing keys: {missing}"
-    hf_model.save_pretrained(save_path)
-    try:
-        tok = AutoTokenizer.from_pretrained(hf_src_model_dir)
-        tok.save_pretrained(save_path)
-    except Exception:
-        pass
-        # dont copy tokenizer if cant
-
-
-if __name__ == "__main__":
-    fire.Fire(convert_pl_to_hf)

examples/seq2seq/distil_marian_enro_teacher.sh

@@ -1,20 +0,0 @@
-#!/usr/bin/env bash
-export PYTHONPATH="../":"${PYTHONPATH}"
-export WANDB_PROJECT=dmar
-# export MAX_LEN=128
-python distillation.py \
-  --learning_rate=3e-4 \
-  --do_train \
-  --fp16 \
-  --val_check_interval 0.25 \
-  --teacher Helsinki-NLP/opus-mt-en-ro \
-  --max_source_length $MAX_LEN --max_target_length $MAX_LEN --val_max_target_length $MAX_LEN --test_max_target_length $MAX_LEN \
-  --student_decoder_layers 3 --student_encoder_layers 6 \
-  --freeze_encoder --freeze_embeds \
-  --model_name_or_path IGNORED \
-  --alpha_hid=3. \
-  --train_batch_size=$BS --eval_batch_size=$BS \
-  --tokenizer_name Helsinki-NLP/opus-mt-en-ro \
-  --warmup_steps 500 --logger_name wandb \
-  --fp16_opt_level O1 --task translation --normalize_hidden --num_sanity_val_steps=0 \
-  "$@"

examples/seq2seq/distil_marian_no_teacher.sh

@@ -1,17 +0,0 @@
-#!/usr/bin/env bash
-export PYTHONPATH="../":"${PYTHONPATH}"
-export WANDB_PROJECT=dmar
-python distillation.py \
-  --learning_rate=3e-4 \
-  --do_train \
-  --do_predict \
-  --fp16 --no_teacher \
-  --val_check_interval 0.25 \
-  --data_dir $ENRO_DIR \
-  --max_source_length $MAX_LEN --max_target_length $MAX_LEN --val_max_target_length $MAX_LEN --test_max_target_length $MAX_LEN \
-  --freeze_encoder --freeze_embeds \
-  --train_batch_size=$BS --eval_batch_size=$BS \
-  --tokenizer_name $m --model_name_or_path $m \
-  --warmup_steps 500 --sortish_sampler --logger_name wandb \
-  --gpus 1 --fp16_opt_level=O1 --task translation --num_sanity_val_steps=0 \
-  "$@"

examples/seq2seq/distillation.py

@@ -1,7 +1,6 @@
 import argparse
 import gc
 import os
-import warnings
 from pathlib import Path
 from typing import List
 
@@ -11,41 +10,36 @@ from torch import nn
 from torch.nn import functional as F
 
 from lightning_base import generic_train
-from transformers import AutoModelForSeq2SeqLM, MBartTokenizer, T5Config, T5ForConditionalGeneration
-from transformers.modeling_bart import shift_tokens_right
+from transformers import BartConfig, BartForConditionalGeneration, MBartTokenizer, T5Config, T5ForConditionalGeneration
 
 
 try:
     from .finetune import SummarizationModule, TranslationModule
-    from .finetune import main as ft_main
-    from .initialization_utils import copy_layers, init_student
+    from .initialization_utils import init_student, copy_layers
     from .utils import (
-        any_requires_grad,
-        assert_all_frozen,
-        calculate_bleu,
-        freeze_params,
-        label_smoothed_nll_loss,
-        pickle_load,
         use_task_specific_params,
+        pickle_load,
+        freeze_params,
+        assert_all_frozen,
+        any_requires_grad,
+        calculate_bleu_score,
     )
+    from .finetune import main as ft_main
 except ImportError:
     from finetune import SummarizationModule, TranslationModule
     from finetune import main as ft_main
-    from initialization_utils import copy_layers, init_student
+    from initialization_utils import init_student, copy_layers
     from utils import (
-        any_requires_grad,
-        assert_all_frozen,
-        calculate_bleu,
-        freeze_params,
-        label_smoothed_nll_loss,
-        pickle_load,
         use_task_specific_params,
+        pickle_load,
+        freeze_params,
+        assert_all_frozen,
+        any_requires_grad,
+        calculate_bleu_score,
     )
 
 
 class BartSummarizationDistiller(SummarizationModule):
-    """Supports Bart, Pegasus and other models that inherit from Bart."""
-
     loss_names = ["loss", "ce_loss", "mlm_loss", "enc_mse_loss", "hid_loss_enc", "hid_loss_dec"]
 
     def __init__(self, hparams):
@@ -80,31 +74,22 @@ class BartSummarizationDistiller(SummarizationModule):
     def pre_init(self, hparams):
         self.output_dir = Path(hparams.output_dir)
         self.output_dir.mkdir(exist_ok=True)
-        teacher = AutoModelForSeq2SeqLM.from_pretrained(hparams.teacher).eval()
+        teacher = BartForConditionalGeneration.from_pretrained(hparams.teacher).eval()
         student_updates = {
             "decoder_layers": hparams.student_decoder_layers,
             "encoder_layers": hparams.student_encoder_layers,
         }
         if hparams.length_penalty != -1:
             student_updates["length_penalty"] = hparams.length_penalty
+        d_layers_to_copy = get_layers_to_copy(student_updates["decoder_layers"], teacher.config.decoder_layers)
         e_layers_to_copy: List = get_layers_to_copy(student_updates["encoder_layers"], teacher.config.encoder_layers)
-        hparams.e_layer_to_copy = e_layers_to_copy
-
-        d_layers_to_copy: List = get_layers_to_copy(student_updates["decoder_layers"], teacher.config.decoder_layers)
-
-        if hparams.supervise_forward:
-            hparams.d_matches = get_layers_to_supervise(
-                student_updates["decoder_layers"], teacher.config.decoder_layers
-            )
-        else:
-            hparams.d_matches = d_layers_to_copy
         hparams.d_layer_to_copy = d_layers_to_copy
-
+        hparams.e_layer_to_copy = e_layers_to_copy
         kw = teacher.config.to_diff_dict()
         kw.update(student_updates)
         # Copy weights
-        student_cfg = teacher.config_class(**kw)
-        student = type(teacher)(student_cfg)
+        student_cfg = BartConfig(**kw)
+        student = BartForConditionalGeneration(student_cfg)
         student, _ = init_student(student, teacher)
         save_dir = self.output_dir.joinpath("student")
         self.copy_to_student(d_layers_to_copy, e_layers_to_copy, hparams, student, teacher)
@@ -175,32 +160,22 @@ class BartSummarizationDistiller(SummarizationModule):
     def _step(self, batch):
         # assert is_frozen(self.teacher)
         pad_token_id = self.tokenizer.pad_token_id
-        input_ids, src_mask, tgt_ids = batch["input_ids"], batch["attention_mask"], batch["labels"]
-        decoder_input_ids = shift_tokens_right(tgt_ids, pad_token_id)
+        input_ids, src_mask, y = batch["input_ids"], batch["attention_mask"], batch["decoder_input_ids"]
+        decoder_input_ids = y[:, :-1].contiguous()
+        labels = y[:, 1:].clone()
+        labels[y[:, 1:] == pad_token_id] = -100
         # noinspection PyCallingNonCallable
-        lm_logits, dec_hidden, enc_outputs, enc_hidden_state = self(
+        sloss, slogits, dec_hidden, enc_outputs, enc_hidden_state = self(
             input_ids,
             attention_mask=src_mask,
             decoder_input_ids=decoder_input_ids,
+            labels=labels,
             output_hidden_states=True,
             output_attentions=False,
-            use_cache=False,
-        )  # TODO(@sshleifer): return_dict=True cleanup
-
-        # Same cross entropy vs. label smoothing logic as finetune.py
-        assert lm_logits.shape[-1] == self.model.config.vocab_size
-        if self.hparams.label_smoothing == 0:
-            # Same behavior as modeling_bart.py, besides ignoring pad_token_id
-            loss_fct = torch.nn.CrossEntropyLoss(ignore_index=pad_token_id)
-            student_lm_loss = loss_fct(lm_logits.view(-1, lm_logits.shape[-1]), tgt_ids.view(-1))
-        else:
-            lprobs = torch.nn.functional.log_softmax(lm_logits, dim=-1)
-            student_lm_loss, _ = label_smoothed_nll_loss(
-                lprobs, tgt_ids, self.hparams.label_smoothing, ignore_index=pad_token_id
-            )
+        )
 
         def zero_tensor():
-            return torch.tensor(0.0).type_as(student_lm_loss)
+            return torch.tensor(0.0).type_as(sloss)
 
         loss_encoder, hid_loss_enc, hid_loss_dec = zero_tensor(), zero_tensor(), zero_tensor()
         if self.different_encoder:
@@ -224,40 +199,41 @@ class BartSummarizationDistiller(SummarizationModule):
                 attention_mask=src_mask,
                 encoder_outputs=teacher_enc_outputs,
                 decoder_input_ids=decoder_input_ids,
-                lm_labels=tgt_ids,
+                lm_labels=labels,
                 output_hidden_states=True,
             )
         dec_mask = decoder_input_ids.ne(pad_token_id)
-        loss_ce, s_logits_slct, t_logits_slct = self.calc_ce_loss(dec_mask, lm_logits, tlogits)
+        loss_ce, s_logits_slct, t_logits_slct = self.calc_ce_loss(dec_mask, slogits, tlogits)
         if self.alpha_hid > 0:
-            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.d_matches)
+            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.d_layer_to_copy)
 
         blended_loss = (
             self.alpha_ce * loss_ce
-            + self.alpha_mlm * student_lm_loss
+            + self.alpha_mlm * sloss
             + self.hparams.alpha_encoder_loss * loss_encoder
             + self.hparams.alpha_hid * (hid_loss_enc + hid_loss_dec)
         )
-        return blended_loss, loss_ce, student_lm_loss, loss_encoder, hid_loss_enc, hid_loss_dec
+        return blended_loss, loss_ce, sloss, loss_encoder, hid_loss_enc, hid_loss_dec
 
     def calc_hidden_loss(self, attention_mask, hidden_states, hidden_states_T, matches):
-        msg = "expected list or tuple for hidden_states, got tensor of shape: "
-        assert not isinstance(hidden_states, torch.Tensor), f"{msg}{hidden_states.shape}"
-        assert not isinstance(hidden_states_T, torch.Tensor), f"{msg}{hidden_states_T.shape}"
+        assert not isinstance(
+            hidden_states, torch.Tensor
+        ), f"expected list or tuple for hidden_states, got tensor of shape {hidden_states.shape}"
+        assert not isinstance(
+            hidden_states_T, torch.Tensor
+        ), f"expected list or tuple for hidden_states_T, got tensor of shape {hidden_states_T.shape}"
         mask = attention_mask.to(hidden_states[0])
         valid_count = mask.sum() * hidden_states[0].size(-1)
-        student_states = torch.stack([hidden_states[i] for i in range(len(matches))])
-        teacher_states = torch.stack([hidden_states_T[j] for j in matches])
-        if self.hparams.normalize_hidden:
-            student_states = F.layer_norm(student_states, student_states.shape[1:])
-            teacher_states = F.layer_norm(teacher_states, teacher_states.shape[1:])
-        mse = F.mse_loss(student_states, teacher_states, reduction="none")
-        masked_mse = (mse * mask.unsqueeze(0).unsqueeze(-1)).sum() / valid_count
-        return masked_mse
+        hidden_losses = [
+            (F.mse_loss(hidden_states[i], hidden_states_T[j], reduction="none") * mask.unsqueeze(-1)).sum()
+            / valid_count
+            for i, j in enumerate(matches)
+        ]
+        return sum(hidden_losses)
 
 
 def add_distill_args(parser):
-    parser.add_argument("--teacher", type=str)
+    parser.add_argument("--teacher", default="facebook/bart-large-cnn", type=str)
     parser.add_argument("--alpha_ce", default=0.8, type=float)
     parser.add_argument("--alpha_mlm", default=0.2, type=float)
     parser.add_argument("--alpha_encoder_loss", default=0.0, type=float)
@@ -266,19 +242,17 @@ def add_distill_args(parser):
     parser.add_argument("--student_encoder_layers", default=12, type=int, required=False)
     parser.add_argument("--no_teacher", action="store_true", default=False)
     parser.add_argument("--length_penalty", type=float, default=-1)
-    parser.add_argument("--supervise_forward", action="store_true", default=False)
-    parser.add_argument("--normalize_hidden", action="store_true", default=False)
 
 
 class BartTranslationDistiller(BartSummarizationDistiller):
-    """Supports Mbart, Marian, other models that inherit from Bart."""
-
     mode = "translation"
+    loss_names = ["loss"]
     metric_names = ["bleu"]
-    default_val_metric = "bleu"
+    val_metric = "bleu"
 
     def __init__(self, hparams, **kwargs):
         super().__init__(hparams, **kwargs)
+        assert isinstance(self.tokenizer, MBartTokenizer)
         assert hparams.src_lang is not None
         assert hparams.tgt_lang is not None
         self.dataset_kwargs["src_lang"] = hparams.src_lang
@@ -287,7 +261,7 @@ class BartTranslationDistiller(BartSummarizationDistiller):
             self.decoder_start_token_id = self.tokenizer.lang_code_to_id[hparams.tgt_lang]
 
     def calc_generative_metrics(self, preds, target) -> dict:
-        return calculate_bleu(preds, target)
+        return calculate_bleu_score(preds, target)
 
     @staticmethod
     def add_model_specific_args(parser, root_dir):
@@ -374,10 +348,7 @@ class T5SummarizationDistiller(BartSummarizationDistiller):
         if self.different_encoder:
             with torch.no_grad():
                 teacher_enc_outputs, teacher_enc_hid = self.teacher.encoder(
-                    source_ids,
-                    attention_mask=source_mask,
-                    output_hidden_states=True,
-                    use_cache=False,
+                    source_ids, attention_mask=source_mask, output_hidden_states=True, use_cache=False,
                 )
             if self.hparams.alpha_encoder_loss > 0:
                 loss_encoder = self.calc_mse_loss(enc_outputs, teacher_enc_outputs, source_mask)
@@ -395,14 +366,14 @@ class T5SummarizationDistiller(BartSummarizationDistiller):
                 attention_mask=source_mask,
                 encoder_outputs=teacher_enc_outputs,
                 decoder_input_ids=decoder_input_ids,
-                labels=labels,
+                lm_labels=labels,
                 output_hidden_states=True,
                 use_cache=False,
             )
 
         loss_ce, s_logits_slct, t_logits_slct = self.calc_ce_loss(dec_mask, slogits, tlogits)
         if self.alpha_hid > 0:
-            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.d_matches)
+            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.d_layer_to_copy)
 
         blended_loss = (
             self.alpha_ce * loss_ce
@@ -429,7 +400,6 @@ def create_module(args):
 
 
 def evaluate_checkpoint(ckpt_path: Path, dest_dir=None):
-    # TODO(SS): DELETE? Better to convert_pl_ckpt_to_hf and run_eval.py
     exp_dir = ckpt_path.parent
     if dest_dir is None:
         dest_dir = exp_dir
@@ -452,53 +422,33 @@ def evaluate_checkpoint(ckpt_path: Path, dest_dir=None):
     trainer.test(model)
 
 
-LAYERS_TO_COPY = {
-    # maps  num layers in student -> which teacher layers to copy.
-    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
-    12: {
-        1: [0],
-        2: [0, 6],
-        3: [0, 6, 11],
-        4: [0, 4, 8, 11],
-        6: [0, 2, 4, 7, 9, 11],
-        9: [0, 1, 2, 4, 5, 7, 9, 10, 11],
-        12: list(range(12)),
-    },
-    16: {  # maps  num layers in student -> which teacher layers to copy
-        1: [0],
-        2: [0, 8],
-        3: [0, 8, 15],
-        4: [0, 5, 10, 15],
-        6: [0, 3, 6, 9, 12, 15],
-        8: [0, 2, 4, 6, 8, 10, 12, 15],
-        9: [0, 1, 3, 5, 7, 9, 11, 13, 15],
-        16: list(range(16)),
-    },
-    6: {1: [0], 2: [0, 5], 3: [0, 2, 5], 4: [0, 1, 3, 5], 6: list(range(6))},
-}
-LAYERS_TO_SUPERVISE = {
-    12: {1: [11], 2: [5, 11], 3: [3, 7, 11], 6: [1, 3, 5, 8, 10, 11]},
-    16: {1: [15], 4: [4, 9, 12, 15], 8: [1, 3, 5, 7, 9, 11, 13, 15]},
-    6: {1: [5], 2: [3, 5], 3: [1, 4, 5], 4: [1, 2, 4, 5]},
-    2: {1: [1], 2: [0, 1]},
-}
-
-
-def get_layers_to_supervise(n_student, n_teacher):
-    return LAYERS_TO_SUPERVISE[n_teacher][n_student]
-
-
-def get_layers_to_copy(n_student, n_teacher):
-    try:
-        val = LAYERS_TO_COPY[n_teacher][n_student]
-        assert len(LAYERS_TO_SUPERVISE[n_teacher][n_student]) == len(val) == n_student
-        return val
-    except KeyError:
-        if n_student != n_teacher:
-            warnings.warn(
-                f"no hardcoded layers to copy for teacher {n_teacher} -> student {n_student}, defaulting to first {n_student}"
-            )
-        return list(range(n_student))
+def get_layers_to_copy(n_to_get, tot):
+    all_layers = list(range(tot))
+    if tot == 12:  # Alternating for special cases
+        layers_to_copy = {  # maps  num layers in student -> which teacher layers to copy
+            1: [0],
+            2: [0, 6],
+            3: [0, 6, 11],
+            4: [0, 4, 8, 11],
+            6: [0, 2, 4, 7, 9, 11],
+            9: [0, 1, 2, 4, 5, 7, 9, 10, 11],
+            12: all_layers,
+        }
+        return layers_to_copy[n_to_get]
+    elif tot == 16:
+        layers_to_copy = {  # maps  num layers in student -> which teacher layers to copy
+            1: [0],
+            2: [0, 8],
+            3: [0, 8, 15],
+            4: [0, 5, 10, 15],
+            6: [0, 3, 6, 9, 12, 15],
+            8: [0, 2, 4, 6, 8, 10, 12, 15],
+            9: [0, 1, 3, 5, 7, 9, 11, 13, 15],
+            16: all_layers,
+        }
+        return layers_to_copy[n_to_get]
+    else:
+        return all_layers[:n_to_get]  # TODO: better version on theseus-bart branch
 
 
 def distill_main(args):

examples/seq2seq/download_wmt.py

@@ -5,25 +5,25 @@ from tqdm import tqdm
 
 
 def download_wmt_dataset(src_lang="ro", tgt_lang="en", dataset="wmt16", save_dir=None) -> None:
-    """Download a dataset using the datasets package and save it to the format expected by finetune.py
+    """Download a dataset using the nlp package and save it to the format expected by finetune.py
     Format of save_dir: train.source, train.target, val.source, val.target, test.source, test.target.
 
     Args:
         src_lang: <str> source language
         tgt_lang: <str> target language
-        dataset: <str> wmt16, wmt17, etc. wmt16 is a good start as it's small. To get the full list run `import datasets; print([d.id for d in datasets.list_datasets() if "wmt" in d.id])`
+        dataset: <str> wmt16, wmt17, etc. wmt16 is a good start as it's small. To get the full list run `import nlp; print([d.id for d in nlp.list_datasets() if "wmt" in d.id])`
         save_dir: <str>, where to save the datasets, defaults to f'{dataset}-{src_lang}-{tgt_lang}'
 
     Usage:
         >>> download_wmt_dataset('ro', 'en', dataset='wmt16') # saves to wmt16-ro-en
     """
     try:
-        import datasets
+        import nlp
     except (ModuleNotFoundError, ImportError):
-        raise ImportError("run pip install datasets")
+        raise ImportError("run pip install nlp")
     pair = f"{src_lang}-{tgt_lang}"
     print(f"Converting {dataset}-{pair}")
-    ds = datasets.load_dataset(dataset, pair)
+    ds = nlp.load_dataset(dataset, pair)
     if save_dir is None:
         save_dir = f"{dataset}-{pair}"
     save_dir = Path(save_dir)

examples/seq2seq/finetune.py

@@ -13,63 +13,65 @@ import torch
 from torch.utils.data import DataLoader
 
 from lightning_base import BaseTransformer, add_generic_args, generic_train
-from transformers import MBartTokenizer, T5ForConditionalGeneration
+from transformers import MarianTokenizer, MBartTokenizer, T5ForConditionalGeneration
 from transformers.modeling_bart import shift_tokens_right
 
 
 try:
-    from .callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
     from .utils import (
-        ROUGE_KEYS,
-        LegacySeq2SeqDataset,
-        Seq2SeqDataset,
         assert_all_frozen,
-        calculate_bleu,
-        calculate_rouge,
-        flatten_list,
-        freeze_params,
-        get_git_info,
-        label_smoothed_nll_loss,
+        use_task_specific_params,
         lmap,
+        flatten_list,
         pickle_save,
         save_git_info,
         save_json,
-        use_task_specific_params,
+        freeze_params,
+        calculate_rouge,
+        get_git_info,
+        ROUGE_KEYS,
+        calculate_bleu_score,
+        Seq2SeqDataset,
+        TranslationDataset,
+        label_smoothed_nll_loss,
     )
+
+    from .callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
+    from .loss_dropper import LossDropper
 except ImportError:
-    from callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
     from utils import (
-        ROUGE_KEYS,
-        LegacySeq2SeqDataset,
         Seq2SeqDataset,
+        TranslationDataset,
         assert_all_frozen,
-        calculate_bleu,
-        calculate_rouge,
-        flatten_list,
-        freeze_params,
-        get_git_info,
-        label_smoothed_nll_loss,
+        use_task_specific_params,
         lmap,
+        flatten_list,
         pickle_save,
         save_git_info,
         save_json,
-        use_task_specific_params,
+        freeze_params,
+        calculate_rouge,
+        get_git_info,
+        ROUGE_KEYS,
+        calculate_bleu_score,
+        label_smoothed_nll_loss,
     )
+    from callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
+    from loss_dropper import LossDropper
 
 logger = logging.getLogger(__name__)
 
 
 class SummarizationModule(BaseTransformer):
     mode = "summarization"
-    loss_names = ["loss"]
+    loss_names = ["loss", "dropper_mask_mean"]
     metric_names = ROUGE_KEYS
-    default_val_metric = "rouge2"
+    val_metric = "rouge2"
 
     def __init__(self, hparams, **kwargs):
-        if hparams.sortish_sampler and hparams.gpus > 1:
-            hparams.replace_sampler_ddp = False
         super().__init__(hparams, num_labels=None, mode=self.mode, **kwargs)
         use_task_specific_params(self.model, "summarization")
+        self.dropper = LossDropper(dropc=.05)
         save_git_info(self.hparams.output_dir)
         self.metrics_save_path = Path(self.output_dir) / "metrics.json"
         self.hparams_save_path = Path(self.output_dir) / "hparams.pkl"
@@ -105,20 +107,14 @@ class SummarizationModule(BaseTransformer):
 
         self.hparams.git_sha = get_git_info()["repo_sha"]
         self.num_workers = hparams.num_workers
-        self.decoder_start_token_id = None  # default to config
+        self.decoder_start_token_id = None
         if self.model.config.decoder_start_token_id is None and isinstance(self.tokenizer, MBartTokenizer):
             self.decoder_start_token_id = self.tokenizer.lang_code_to_id[hparams.tgt_lang]
             self.model.config.decoder_start_token_id = self.decoder_start_token_id
-        self.dataset_class = (
-            Seq2SeqDataset if hasattr(self.tokenizer, "prepare_seq2seq_batch") else LegacySeq2SeqDataset
-        )
-        self.eval_beams = self.model.config.num_beams if self.hparams.eval_beams is None else self.hparams.eval_beams
-        assert self.eval_beams >= 1, f"got self.eval_beams={self.eval_beams}. Need an integer > 1"
-        if self.hparams.eval_max_gen_length is not None:
-            self.eval_max_length = self.hparams.eval_max_gen_length
+        if isinstance(self.tokenizer, MBartTokenizer) or isinstance(self.tokenizer, MarianTokenizer):
+            self.dataset_class = TranslationDataset
         else:
-            self.eval_max_length = self.model.config.max_length
-        self.val_metric = self.default_val_metric if self.hparams.val_metric is None else self.hparams.val_metric
+            self.dataset_class = Seq2SeqDataset
 
     def freeze_embeds(self):
         """Freeze token embeddings and positional embeddings for bart, just token embeddings for t5."""
@@ -143,27 +139,45 @@ class SummarizationModule(BaseTransformer):
 
     def _step(self, batch: dict) -> Tuple:
         pad_token_id = self.tokenizer.pad_token_id
-        src_ids, src_mask = batch["input_ids"], batch["attention_mask"]
-        tgt_ids = batch["labels"]
-        if isinstance(self.model, T5ForConditionalGeneration):
-            decoder_input_ids = self.model._shift_right(tgt_ids)
+        source_ids, source_mask, target_ids = batch["input_ids"], batch["attention_mask"], batch["decoder_input_ids"]
+
+        if "labels" in batch:
+            lm_labels = batch["labels"]
+            decoder_input_ids = shift_tokens_right(lm_labels, pad_token_id)
+        elif isinstance(self.model, T5ForConditionalGeneration):
+            decoder_input_ids = self.model._shift_right(target_ids)
+            lm_labels = target_ids
         else:
-            decoder_input_ids = shift_tokens_right(tgt_ids, pad_token_id)
+            decoder_input_ids = target_ids[:, :-1].contiguous()  # Why this line?
+            lm_labels = target_ids[:, 1:].clone()  # why clone?
+
+        outputs = self(source_ids, attention_mask=source_mask, decoder_input_ids=decoder_input_ids, use_cache=False)
+        bs = source_ids.shape[0]
 
-        outputs = self(src_ids, attention_mask=src_mask, decoder_input_ids=decoder_input_ids, use_cache=False)
-        lm_logits = outputs[0]
         if self.hparams.label_smoothing == 0:
-            # Same behavior as modeling_bart.py, besides ignoring pad_token_id
-            ce_loss_fct = torch.nn.CrossEntropyLoss(ignore_index=pad_token_id)
 
+            # Same behavior as modeling_bart.py
+            loss_fct = torch.nn.CrossEntropyLoss(reduction='none', ignore_index=pad_token_id)
+            lm_logits = outputs[0]
             assert lm_logits.shape[-1] == self.model.config.vocab_size
-            loss = ce_loss_fct(lm_logits.view(-1, lm_logits.shape[-1]), tgt_ids.view(-1))
+
+            #loss_fct = torch.nn.NLLLoss(reduction='none', ignore_index=pad_token_id)
+            #logit_shape =
+            #weights = torch.ones(logit_shape
+            loss = loss_fct(lm_logits.view(-1, lm_logits.shape[-1]), lm_labels.view(-1))
+            loss = loss.view(-1, bs)
+            loss = loss.mean(dim=0)
+            mask = self.dropper(loss)
+            loss *= mask
+            loss = loss.mean()
+            return (loss, 1-mask.mean())
+            #loss = loss.view(-1, bs)
         else:
-            lprobs = torch.nn.functional.log_softmax(lm_logits, dim=-1)
+            lprobs = torch.nn.functional.log_softmax(outputs[0], dim=-1)
             loss, nll_loss = label_smoothed_nll_loss(
-                lprobs, tgt_ids, self.hparams.label_smoothing, ignore_index=pad_token_id
+                lprobs, lm_labels, self.hparams.label_smoothing, ignore_index=pad_token_id
             )
-        return (loss,)
+        return (loss,torch.tensor(1.))
 
     @property
     def pad(self) -> int:
@@ -174,7 +188,7 @@ class SummarizationModule(BaseTransformer):
 
         logs = {name: loss for name, loss in zip(self.loss_names, loss_tensors)}
         # tokens per batch
-        logs["tpb"] = batch["input_ids"].ne(self.pad).sum() + batch["labels"].ne(self.pad).sum()
+        logs["tpb"] = batch["input_ids"].ne(self.pad).sum() + batch["decoder_input_ids"].ne(self.pad).sum()
         return {"loss": loss_tensors[0], "log": logs}
 
     def validation_step(self, batch, batch_idx) -> Dict:
@@ -184,25 +198,15 @@ class SummarizationModule(BaseTransformer):
         self.step_count += 1
         losses = {k: torch.stack([x[k] for x in outputs]).mean() for k in self.loss_names}
         loss = losses["loss"]
-        generative_metrics = {
-            k: np.array([x[k] for x in outputs]).mean() for k in self.metric_names + ["gen_time", "gen_len"]
-        }
-        metric_val = (
-            generative_metrics[self.val_metric] if self.val_metric in generative_metrics else losses[self.val_metric]
-        )
-        metric_tensor: torch.FloatTensor = torch.tensor(metric_val).type_as(loss)
-        generative_metrics.update({k: v.item() for k, v in losses.items()})
-        losses.update(generative_metrics)
-        all_metrics = {f"{prefix}_avg_{k}": x for k, x in losses.items()}
-        all_metrics["step_count"] = self.step_count
-        self.save_metrics(all_metrics, prefix)  # writes to self.metrics_save_path
+        rouges = {k: np.array([x[k] for x in outputs]).mean() for k in self.metric_names + ["gen_time", "gen_len"]}
+        rouge_tensor: torch.FloatTensor = torch.tensor(rouges[self.val_metric]).type_as(loss)
+        rouges.update({k: v.item() for k, v in losses.items()})
+        losses.update(rouges)
+        metrics = {f"{prefix}_avg_{k}": x for k, x in losses.items()}
+        metrics["step_count"] = self.step_count
+        self.save_metrics(metrics, prefix)  # writes to self.metrics_save_path
         preds = flatten_list([x["preds"] for x in outputs])
-        return {
-            "log": all_metrics,
-            "preds": preds,
-            f"{prefix}_loss": loss,
-            f"{prefix}_{self.val_metric}": metric_tensor,
-        }
+        return {"log": metrics, "preds": preds, f"{prefix}_loss": loss, f"{prefix}_{self.val_metric}": rouge_tensor}
 
     def save_metrics(self, latest_metrics, type_path) -> None:
         self.metrics[type_path].append(latest_metrics)
@@ -213,19 +217,15 @@ class SummarizationModule(BaseTransformer):
 
     def _generative_step(self, batch: dict) -> dict:
         t0 = time.time()
-
-        # parser.add_argument('--eval_max_gen_length', type=int, default=None, help='never generate more than n tokens')
         generated_ids = self.model.generate(
             batch["input_ids"],
             attention_mask=batch["attention_mask"],
             use_cache=True,
             decoder_start_token_id=self.decoder_start_token_id,
-            num_beams=self.eval_beams,
-            max_length=self.eval_max_length,
         )
         gen_time = (time.time() - t0) / batch["input_ids"].shape[0]
         preds: List[str] = self.ids_to_clean_text(generated_ids)
-        target: List[str] = self.ids_to_clean_text(batch["labels"])
+        target: List[str] = self.ids_to_clean_text(batch["decoder_input_ids"])
         loss_tensors = self._step(batch)
         base_metrics = {name: loss for name, loss in zip(self.loss_names, loss_tensors)}
         rouge: Dict = self.calc_generative_metrics(preds, target)
@@ -255,7 +255,8 @@ class SummarizationModule(BaseTransformer):
         dataset = self.get_dataset(type_path)
         sampler = None
         if self.hparams.sortish_sampler and type_path == "train":
-            sampler = dataset.make_sortish_sampler(batch_size, distributed=self.hparams.gpus > 1)
+            assert self.hparams.gpus <= 1  # TODO: assert earlier
+            sampler = dataset.make_sortish_sampler(batch_size)
             shuffle = False
 
         dataloader = DataLoader(
@@ -321,14 +322,9 @@ class SummarizationModule(BaseTransformer):
             "--task", type=str, default="summarization", required=False, help="# examples. -1 means use all."
         )
         parser.add_argument("--label_smoothing", type=float, default=0.0, required=False)
+        parser.add_argument("--loss_dropper", type=float, default=0.0, required=False)
         parser.add_argument("--src_lang", type=str, default="", required=False)
         parser.add_argument("--tgt_lang", type=str, default="", required=False)
-        parser.add_argument("--eval_beams", type=int, default=None, required=False)
-        parser.add_argument(
-            "--val_metric", type=str, default=None, required=False, choices=["bleu", "rouge2", "loss", None]
-        )
-        parser.add_argument("--eval_max_gen_length", type=int, default=None, help="never generate more than n tokens")
-        parser.add_argument("--save_top_k", type=int, default=1, required=False, help="How many checkpoints to save")
         parser.add_argument(
             "--early_stopping_patience",
             type=int,
@@ -341,9 +337,8 @@ class SummarizationModule(BaseTransformer):
 
 class TranslationModule(SummarizationModule):
     mode = "translation"
-    loss_names = ["loss"]
     metric_names = ["bleu"]
-    default_val_metric = "bleu"
+    val_metric = "bleu"
 
     def __init__(self, hparams, **kwargs):
         super().__init__(hparams, **kwargs)
@@ -351,7 +346,7 @@ class TranslationModule(SummarizationModule):
         self.dataset_kwargs["tgt_lang"] = hparams.tgt_lang
 
     def calc_generative_metrics(self, preds, target) -> dict:
-        return calculate_bleu(preds, target)
+        return calculate_bleu_score(preds, target)
 
 
 def main(args, model=None) -> SummarizationModule:
@@ -359,10 +354,11 @@ def main(args, model=None) -> SummarizationModule:
     if len(os.listdir(args.output_dir)) > 3 and args.do_train:
         raise ValueError("Output directory ({}) already exists and is not empty.".format(args.output_dir))
     if model is None:
-        if "summarization" in args.task:
+        if args.task == "summarization":
             model: SummarizationModule = SummarizationModule(args)
         else:
             model: SummarizationModule = TranslationModule(args)
+
     dataset = Path(args.data_dir).name
     if (
         args.logger_name == "default"
@@ -386,17 +382,14 @@ def main(args, model=None) -> SummarizationModule:
         es_callback = get_early_stopping_callback(model.val_metric, args.early_stopping_patience)
     else:
         es_callback = False
-
-    lower_is_better = args.val_metric == "loss"
     trainer: pl.Trainer = generic_train(
         model,
         args,
         logging_callback=Seq2SeqLoggingCallback(),
-        checkpoint_callback=get_checkpoint_callback(
-            args.output_dir, model.val_metric, args.save_top_k, lower_is_better
-        ),
+        checkpoint_callback=get_checkpoint_callback(args.output_dir, model.val_metric),
         early_stopping_callback=es_callback,
         logger=logger,
+        # TODO: early stopping callback seems messed up
     )
     pickle_save(model.hparams, model.output_dir / "hparams.pkl")
     if not args.do_predict:

examples/seq2seq/finetune_pegasus_xsum.sh

@@ -10,5 +10,5 @@ python finetune.py \
     --n_val 1000 \
     --val_check_interval 0.25 \
     --max_source_length 512 --max_target_length 56 \
-    --freeze_embeds --label_smoothing 0.1 --adafactor --task summarization_xsum \
+    --freeze_embeds --max_target_length 56 --label_smoothing 0.1 \
     "$@"

examples/seq2seq/loss_dropper.py

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

examples/seq2seq/run_eval.py

@@ -1,10 +1,6 @@
 import argparse
 import json
-import time
-import warnings
-from logging import getLogger
 from pathlib import Path
-from typing import Dict, List
 
 import torch
 from tqdm import tqdm
@@ -12,12 +8,10 @@ from tqdm import tqdm
 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
 
 
-logger = getLogger(__name__)
-
 try:
-    from .utils import calculate_bleu, calculate_rouge, parse_numeric_cl_kwargs, use_task_specific_params
+    from .utils import calculate_rouge, use_task_specific_params, calculate_bleu_score, trim_batch
 except ImportError:
-    from utils import calculate_bleu, calculate_rouge, parse_numeric_cl_kwargs, use_task_specific_params
+    from utils import calculate_rouge, use_task_specific_params, calculate_bleu_score, trim_batch
 
 DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
 
@@ -29,47 +23,44 @@ def chunks(lst, n):
 
 
 def generate_summaries_or_translations(
-    examples: List[str],
+    examples: list,
     out_file: str,
     model_name: str,
     batch_size: int = 8,
     device: str = DEFAULT_DEVICE,
     fp16=False,
     task="summarization",
-    prefix=None,
-    **generate_kwargs,
-) -> Dict:
-    """Save model.generate results to <out_file>, and return how long it took."""
+    decoder_start_token_id=None,
+    **gen_kwargs,
+) -> None:
     fout = Path(out_file).open("w", encoding="utf-8")
     model_name = str(model_name)
     model = AutoModelForSeq2SeqLM.from_pretrained(model_name).to(device)
     if fp16:
         model = model.half()
+    if decoder_start_token_id is None:
+        decoder_start_token_id = gen_kwargs.pop("decoder_start_token_id", None)
 
     tokenizer = AutoTokenizer.from_pretrained(model_name)
-    logger.info(f"Inferred tokenizer type: {tokenizer.__class__}")  # if this is wrong, check config.model_type.
 
-    start_time = time.time()
-    # update config with task specific params
+    # update config with summarization specific params
     use_task_specific_params(model, task)
-    if prefix is None:
-        prefix = prefix or getattr(model.config, "prefix", "") or ""
-    for examples_chunk in tqdm(list(chunks(examples, batch_size))):
-        examples_chunk = [prefix + text for text in examples_chunk]
-        batch = tokenizer(examples_chunk, return_tensors="pt", truncation=True, padding="longest").to(device)
+
+    for batch in tqdm(list(chunks(examples, batch_size))):
+        if "t5" in model_name:
+            batch = [model.config.prefix + text for text in batch]
+        batch = tokenizer(batch, return_tensors="pt", truncation=True, padding="max_length").to(device)
+        input_ids, attention_mask = trim_batch(**batch, pad_token_id=tokenizer.pad_token_id)
         summaries = model.generate(
-            input_ids=batch.input_ids,
-            attention_mask=batch.attention_mask,
-            **generate_kwargs,
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            decoder_start_token_id=decoder_start_token_id,
+            **gen_kwargs,
         )
         dec = tokenizer.batch_decode(summaries, skip_special_tokens=True, clean_up_tokenization_spaces=False)
         for hypothesis in dec:
             fout.write(hypothesis + "\n")
             fout.flush()
-    fout.close()
-    runtime = int(time.time() - start_time)  # seconds
-    n_obs = len(examples)
-    return dict(n_obs=n_obs, runtime=runtime, seconds_per_sample=round(runtime / n_obs, 4))
 
 
 def run_generate():
@@ -77,30 +68,29 @@ def run_generate():
     parser.add_argument("model_name", type=str, help="like facebook/bart-large-cnn,t5-base, etc.")
     parser.add_argument("input_path", type=str, help="like cnn_dm/test.source")
     parser.add_argument("save_path", type=str, help="where to save summaries")
-    parser.add_argument("--reference_path", type=str, required=False, help="like cnn_dm/test.target")
-    parser.add_argument("--score_path", type=str, required=False, default="metrics.json", help="where to save metrics")
+
+    parser.add_argument("--reference_path", type=str, required=False, help="like cnn_dm/test_reference_summaries.txt")
+    parser.add_argument("--score_path", type=str, required=False, help="where to save the rouge score in json format")
     parser.add_argument("--device", type=str, required=False, default=DEFAULT_DEVICE, help="cuda, cuda:1, cpu etc.")
-    parser.add_argument(
-        "--prefix", type=str, required=False, default=None, help="will be added to the begininng of src examples"
-    )
-    parser.add_argument("--task", type=str, default="summarization", help="used for task_specific_params + metrics")
+    parser.add_argument("--task", type=str, default="summarization", help="typically translation or summarization")
     parser.add_argument("--bs", type=int, default=8, required=False, help="batch size")
+    parser.add_argument(
+        "--decoder_start_token_id",
+        type=int,
+        default=None,
+        required=False,
+        help="decoder_start_token_id (otherwise will look at config)",
+    )
     parser.add_argument(
         "--n_obs", type=int, default=-1, required=False, help="How many observations. Defaults to all."
     )
     parser.add_argument("--fp16", action="store_true")
-    # Unspecified args like --num_beams=2 --decoder_start_token_id=4 are passed to model.generate
-    args, rest = parser.parse_known_args()
-    parsed = parse_numeric_cl_kwargs(rest)
-    if parsed:
-        print(f"parsed the following generate kwargs: {parsed}")
+    args = parser.parse_args()
     examples = [" " + x.rstrip() if "t5" in args.model_name else x.rstrip() for x in open(args.input_path).readlines()]
     if args.n_obs > 0:
         examples = examples[: args.n_obs]
     Path(args.save_path).parent.mkdir(exist_ok=True)
-    if args.reference_path is None and Path(args.score_path).exists():
-        warnings.warn(f"score_path {args.score_path} will be overwritten unless you type ctrl-c.")
-    runtime_metrics = generate_summaries_or_translations(
+    generate_summaries_or_translations(
         examples,
         args.save_path,
         args.model_name,
@@ -108,24 +98,20 @@ def run_generate():
         device=args.device,
         fp16=args.fp16,
         task=args.task,
-        prefix=args.prefix,
-        **parsed,
+        decoder_start_token_id=args.decoder_start_token_id,
     )
     if args.reference_path is None:
         return
     # Compute scores
-    score_fn = calculate_bleu if "translation" in args.task else calculate_rouge
+    score_fn = calculate_bleu_score if "translation" in args.task else calculate_rouge
     output_lns = [x.rstrip() for x in open(args.save_path).readlines()]
     reference_lns = [x.rstrip() for x in open(args.reference_path).readlines()][: len(output_lns)]
     scores: dict = score_fn(output_lns, reference_lns)
-    scores.update(runtime_metrics)
     print(scores)
     if args.score_path is not None:
-        json.dump(scores, open(args.score_path, "w"))
+        json.dump(scores, open(args.score_path, "w+"))
     return scores
 
 
 if __name__ == "__main__":
-    # Usage for MT:
-    # python run_eval.py MODEL_NAME $DATA_DIR/test.source $save_dir/test_translations.txt --reference_path $DATA_DIR/test.target --score_path $save_dir/test_bleu.json  --task translation $@
     run_generate()

examples/seq2seq/test_bash_script.py

@@ -10,10 +10,9 @@ import pytorch_lightning as pl
 import timeout_decorator
 import torch
 
-from transformers import BartForConditionalGeneration, MarianMTModel
+from transformers import BartForConditionalGeneration
 from transformers.testing_utils import slow
 
-from .distillation import BartSummarizationDistiller, distill_main
 from .finetune import SummarizationModule, main
 from .test_seq2seq_examples import CUDA_AVAILABLE, MBART_TINY
 from .utils import load_json
@@ -21,7 +20,6 @@ from .utils import load_json
 
 MODEL_NAME = MBART_TINY
 # TODO(SS): MODEL_NAME = "sshleifer/student_mbart_en_ro_1_1"
-MARIAN_MODEL = "sshleifer/student_marian_en_ro_6_1"
 
 
 @slow
@@ -29,7 +27,6 @@ MARIAN_MODEL = "sshleifer/student_marian_en_ro_6_1"
 def test_model_download():
     """This warms up the cache so that we can time the next test without including download time, which varies between machines."""
     BartForConditionalGeneration.from_pretrained(MODEL_NAME)
-    MarianMTModel.from_pretrained(MARIAN_MODEL)
 
 
 @timeout_decorator.timeout(120)
@@ -38,30 +35,34 @@ def test_model_download():
 def test_train_mbart_cc25_enro_script():
     data_dir = "examples/seq2seq/test_data/wmt_en_ro"
     env_vars_to_replace = {
-        "--fp16_opt_level=O1": "",
-        "$MAX_LEN": 128,
+        "$MAX_LEN": 200,
         "$BS": 4,
         "$GAS": 1,
         "$ENRO_DIR": data_dir,
         "facebook/mbart-large-cc25": MODEL_NAME,
-        # Download is 120MB in previous test.
+        # 1 encoder and 1 decoder layer from finetuned mbart en-ro. Should be able to start >0 and improve quickly.
+        # Download is 600MB in previous test.
         "val_check_interval=0.25": "val_check_interval=1.0",
     }
 
     # Clean up bash script
     bash_script = Path("examples/seq2seq/train_mbart_cc25_enro.sh").open().read().split("finetune.py")[1].strip()
-    bash_script = bash_script.replace("\\\n", "").strip().replace('"$@"', "")
+    bash_script = bash_script.replace("\\\n", "").strip().replace("$@", "")
     for k, v in env_vars_to_replace.items():
         bash_script = bash_script.replace(k, str(v))
-    output_dir = tempfile.mkdtemp(prefix="output_mbart")
+    output_dir = tempfile.mkdtemp(prefix="output")
 
-    bash_script = bash_script.replace("--fp16 ", "")
+    if CUDA_AVAILABLE:
+        gpus = 1  # torch.cuda.device_count()
+    else:
+        gpus = 0
+    bash_script = bash_script.replace("--fp16", "")
     testargs = (
         ["finetune.py"]
         + bash_script.split()
         + [
             f"--output_dir={output_dir}",
-            "--gpus=1",
+            f"--gpus={gpus}",
             "--learning_rate=3e-1",
             "--warmup_steps=0",
             "--val_check_interval=1.0",
@@ -81,86 +82,7 @@ def test_train_mbart_cc25_enro_script():
     metrics = load_json(model.metrics_save_path)
     first_step_stats = metrics["val"][0]
     last_step_stats = metrics["val"][-1]
-    assert len(metrics["val"]) == (args.max_epochs / args.val_check_interval) + 1  # +1 accounts for val_sanity_check
-
-    assert last_step_stats["val_avg_gen_time"] >= 0.01
-
-    assert first_step_stats["val_avg_bleu"] < last_step_stats["val_avg_bleu"]  # model learned nothing
-    assert 1.0 >= last_step_stats["val_avg_gen_time"]  # model hanging on generate. Maybe bad config was saved.
-    assert isinstance(last_step_stats[f"val_avg_{model.val_metric}"], float)
-
-    # check lightning ckpt can be loaded and has a reasonable statedict
-    contents = os.listdir(output_dir)
-    ckpt_path = [x for x in contents if x.endswith(".ckpt")][0]
-    full_path = os.path.join(args.output_dir, ckpt_path)
-    ckpt = torch.load(full_path, map_location="cpu")
-    expected_key = "model.model.decoder.layers.0.encoder_attn_layer_norm.weight"
-    assert expected_key in ckpt["state_dict"]
-    assert ckpt["state_dict"]["model.model.decoder.layers.0.encoder_attn_layer_norm.weight"].dtype == torch.float32
-
-    # TODO(SS): turn on args.do_predict when PL bug fixed.
-    if args.do_predict:
-        contents = {os.path.basename(p) for p in contents}
-        assert "test_generations.txt" in contents
-        assert "test_results.txt" in contents
-        # assert len(metrics["val"]) ==  desired_n_evals
-        assert len(metrics["test"]) == 1
-
-
-@timeout_decorator.timeout(600)
-@slow
-@pytest.mark.skipif(not CUDA_AVAILABLE, reason="too slow to run on CPU")
-def test_opus_mt_distill_script():
-    data_dir = "examples/seq2seq/test_data/wmt_en_ro"
-    env_vars_to_replace = {
-        "--fp16_opt_level=O1": "",
-        "$MAX_LEN": 128,
-        "$BS": 16,
-        "$GAS": 1,
-        "$ENRO_DIR": data_dir,
-        "$m": "sshleifer/student_marian_en_ro_6_1",
-        "val_check_interval=0.25": "val_check_interval=1.0",
-    }
-
-    # Clean up bash script
-    bash_script = (
-        Path("examples/seq2seq/distil_marian_no_teacher.sh").open().read().split("distillation.py")[1].strip()
-    )
-    bash_script = bash_script.replace("\\\n", "").strip().replace('"$@"', "")
-    bash_script = bash_script.replace("--fp16 ", " ")
-
-    for k, v in env_vars_to_replace.items():
-        bash_script = bash_script.replace(k, str(v))
-    output_dir = tempfile.mkdtemp(prefix="marian_output")
-    bash_script = bash_script.replace("--fp16", "")
-    epochs = 6
-    testargs = (
-        ["distillation.py"]
-        + bash_script.split()
-        + [
-            f"--output_dir={output_dir}",
-            "--gpus=1",
-            "--learning_rate=1e-3",
-            f"--num_train_epochs={epochs}",
-            "--warmup_steps=10",
-            "--val_check_interval=1.0",
-        ]
-    )
-    with patch.object(sys, "argv", testargs):
-        parser = argparse.ArgumentParser()
-        parser = pl.Trainer.add_argparse_args(parser)
-        parser = BartSummarizationDistiller.add_model_specific_args(parser, os.getcwd())
-        args = parser.parse_args()
-        args.do_predict = False
-        # assert args.gpus == gpus THIS BREAKS for multigpu
-
-        model = distill_main(args)
-
-    # Check metrics
-    metrics = load_json(model.metrics_save_path)
-    first_step_stats = metrics["val"][0]
-    last_step_stats = metrics["val"][-1]
-    assert len(metrics["val"]) >= (args.max_epochs / args.val_check_interval)  # +1 accounts for val_sanity_check
+    assert len(metrics["val"]) == (args.max_epochs / args.val_check_interval)  # +1 accounts for val_sanity_check
 
     assert last_step_stats["val_avg_gen_time"] >= 0.01
 

examples/seq2seq/test_seq2seq_examples.py

@@ -10,20 +10,18 @@ from unittest.mock import patch
 import pytest
 import pytorch_lightning as pl
 import torch
+from pytest import param
 from torch.utils.data import DataLoader
 
 import lightning_base
-from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer
-from transformers.hf_api import HfApi
-from transformers.modeling_bart import shift_tokens_right
-from transformers.testing_utils import CaptureStderr, CaptureStdout, require_multigpu, require_torch_and_cuda, slow
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
+from transformers.testing_utils import CaptureStderr, CaptureStdout, require_multigpu
 
-from .convert_pl_checkpoint_to_hf import convert_pl_to_hf
 from .distillation import distill_main, evaluate_checkpoint
 from .finetune import SummarizationModule, main
 from .pack_dataset import pack_data_dir
 from .run_eval import generate_summaries_or_translations, run_generate
-from .utils import LegacySeq2SeqDataset, Seq2SeqDataset, label_smoothed_nll_loss, lmap, load_json
+from .utils import Seq2SeqDataset, TranslationDataset, label_smoothed_nll_loss, lmap, load_json
 
 
 logging.basicConfig(level=logging.DEBUG)
@@ -31,14 +29,7 @@ logging.basicConfig(level=logging.DEBUG)
 logger = logging.getLogger()
 CUDA_AVAILABLE = torch.cuda.is_available()
 CHEAP_ARGS = {
-    "supervise_forward": True,
-    "normalize_hidden": True,
     "label_smoothing": 0.2,
-    "eval_max_gen_length": None,
-    "eval_beams": 1,
-    "val_metric": "loss",
-    "save_top_k": 1,
-    "adafactor": True,
     "early_stopping_patience": 2,
     "logger_name": "default",
     "length_penalty": 0.5,
@@ -124,34 +115,10 @@ class TestSummarizationDistiller(unittest.TestCase):
         logging.disable(logging.CRITICAL)  # remove noisy download output from tracebacks
         return cls
 
-    @slow
-    @require_torch_and_cuda
-    def test_hub_configs(self):
-        """I put require_torch_and_cuda cause I only want this to run with self-scheduled."""
-
-        model_list = HfApi().model_list()
-        org = "sshleifer"
-        model_ids = [x.modelId for x in model_list if x.modelId.startswith(org)]
-        allowed_to_be_broken = ["sshleifer/blenderbot-3B", "sshleifer/blenderbot-90M"]
-        failures = []
-        for m in model_ids:
-            if m in allowed_to_be_broken:
-                continue
-            try:
-                AutoConfig.from_pretrained(m)
-            except Exception:
-                failures.append(m)
-        assert not failures, f"The following models could not be loaded through AutoConfig: {failures}"
-
     @require_multigpu
     def test_multigpu(self):
-        updates = dict(
-            no_teacher=True,
-            freeze_encoder=True,
-            gpus=2,
-            sortish_sampler=True,
-        )
-        self._test_distiller_cli(updates, check_contents=False)
+        updates = dict(no_teacher=True, freeze_encoder=True, gpus=2, sortish_sampler=False,)
+        self._test_distiller_cli(updates)
 
     def test_distill_no_teacher(self):
         updates = dict(student_encoder_layers=2, student_decoder_layers=1, no_teacher=True)
@@ -178,9 +145,6 @@ class TestSummarizationDistiller(unittest.TestCase):
         self.assertTrue(Path(out_path).exists())
 
         evaluate_checkpoint(ckpts[0], dest_dir=Path(tempfile.mkdtemp()))
-        out_path_new = tempfile.mkdtemp()
-        convert_pl_to_hf(ckpts[0], transformer_ckpts[0].parent, out_path_new)
-        assert os.path.exists(os.path.join(out_path_new, "pytorch_model.bin"))
 
     def test_loss_fn(self):
         model = AutoModelForSeq2SeqLM.from_pretrained(BART_TINY, return_dict=True)
@@ -217,7 +181,6 @@ class TestSummarizationDistiller(unittest.TestCase):
             tgt_lang="ro_RO",
         )
         model = self._test_distiller_cli(updates, check_contents=False)
-        assert model.model.config.model_type == "mbart"
 
         ckpts = list(Path(model.output_dir).glob("*.ckpt"))
         self.assertEqual(1, len(ckpts))
@@ -265,9 +228,9 @@ class TestSummarizationDistiller(unittest.TestCase):
         if not check_contents:
             return model
         contents = os.listdir(output_dir)
+        ckpt_name = "val_avg_rouge2=0.0000-step_count=2.ckpt"  # "val_avg_rouge2=0.0000-epoch=1.ckpt"  # "epoch=1-val_avg_rouge2=0.0000.ckpt"
         contents = {os.path.basename(p) for p in contents}
-        ckpt_files = [p for p in contents if p.endswith("ckpt")]
-        assert len(ckpt_files) > 0
+        self.assertIn(ckpt_name, contents)
 
         self.assertIn("test_generations.txt", contents)
         self.assertIn("test_results.txt", contents)
@@ -284,28 +247,13 @@ class TestSummarizationDistiller(unittest.TestCase):
 
 
 @pytest.mark.parametrize(["model"], [pytest.param(T5_TINY), pytest.param(BART_TINY), pytest.param(MBART_TINY)])
-def test_run_eval(model):
+def test_run_eval_bart(model):
     input_file_name = Path(tempfile.mkdtemp()) / "utest_input.source"
     output_file_name = input_file_name.parent / "utest_output.txt"
     assert not output_file_name.exists()
     articles = [" New York (CNN)When Liana Barrientos was 23 years old, she got married in Westchester County."]
     _dump_articles(input_file_name, articles)
-    score_path = str(Path(tempfile.mkdtemp()) / "scores.json")
-    task = "translation_en_to_de" if model == T5_TINY else "summarization"
-    testargs = [
-        "run_eval.py",
-        model,
-        str(input_file_name),
-        str(output_file_name),
-        "--score_path",
-        score_path,
-        "--task",
-        task,
-        "--num_beams",
-        "2",
-        "--length_penalty",
-        "2.0",
-    ]
+    testargs = ["run_eval.py", model, str(input_file_name), str(output_file_name)]  # TODO: test score_path
     with patch.object(sys, "argv", testargs):
         run_generate()
         assert Path(output_file_name).exists()
@@ -313,8 +261,7 @@ def test_run_eval(model):
 
 
 @pytest.mark.parametrize(
-    ["model"],
-    [pytest.param(T5_TINY), pytest.param(BART_TINY), pytest.param(MBART_TINY), pytest.param(MARIAN_TINY)],
+    ["model"], [pytest.param(T5_TINY), pytest.param(BART_TINY), pytest.param(MBART_TINY), pytest.param(MARIAN_TINY)],
 )
 def test_finetune(model):
     args_d: dict = CHEAP_ARGS.copy()
@@ -382,8 +329,7 @@ def test_finetune_extra_model_args():
     output_dir = tempfile.mkdtemp(prefix="output_1_")
     args_d1 = args_d.copy()
     args_d1.update(
-        model_name_or_path=model,
-        output_dir=output_dir,
+        model_name_or_path=model, output_dir=output_dir,
     )
     extra_model_params = ("encoder_layerdrop", "decoder_layerdrop", "dropout", "attention_dropout")
     for p in extra_model_params:
@@ -398,8 +344,7 @@ def test_finetune_extra_model_args():
     output_dir = tempfile.mkdtemp(prefix="output_2_")
     args_d2 = args_d.copy()
     args_d2.update(
-        model_name_or_path=model,
-        output_dir=output_dir,
+        model_name_or_path=model, output_dir=output_dir,
     )
     unsupported_param = "encoder_layerdrop"
     args_d2[unsupported_param] = 0.5
@@ -486,27 +431,18 @@ def test_pack_dataset():
     assert orig_paths == new_paths
 
 
-@pytest.mark.parametrize(
-    ["tok_name"],
-    [
-        pytest.param(MBART_TINY),
-        pytest.param(MARIAN_TINY),
-        pytest.param(T5_TINY),
-        pytest.param(BART_TINY),
-        pytest.param("google/pegasus-xsum"),
-    ],
-)
-def test_seq2seq_dataset_truncation(tok_name):
+@pytest.mark.parametrize(["tok_name"], [pytest.param(MBART_TINY), pytest.param(MARIAN_TINY)])
+def test_mbart_dataset_truncation(tok_name):
     tokenizer = AutoTokenizer.from_pretrained(tok_name)
     tmp_dir = make_test_data_dir()
     max_len_source = max(len(tokenizer.encode(a)) for a in ARTICLES)
     max_len_target = max(len(tokenizer.encode(a)) for a in SUMMARIES)
     max_src_len = 4
     max_tgt_len = 8
-    assert max_len_target > max_src_len  # Will be truncated
-    assert max_len_source > max_src_len  # Will be truncated
-    src_lang, tgt_lang = "ro_RO", "de_DE"  # ignored for all but mbart, but never causes error.
-    train_dataset = Seq2SeqDataset(
+    assert max_len_target > max_src_len  # Truncated
+    assert max_len_source > max_src_len
+    src_lang, tgt_lang = "ro_RO", "de_DE"  # NOT WHAT IT WAS TRAINED ON
+    train_dataset = TranslationDataset(
         tokenizer,
         data_dir=tmp_dir,
         type_path="train",
@@ -522,11 +458,10 @@ def test_seq2seq_dataset_truncation(tok_name):
         # show that articles were trimmed.
         assert batch["input_ids"].shape[1] == max_src_len
         # show that targets are the same len
-        assert batch["labels"].shape[1] == max_tgt_len
-        if tok_name != MBART_TINY:
+        assert batch["decoder_input_ids"].shape[1] == max_tgt_len
+        if tok_name == MARIAN_TINY:
             continue
         # check language codes in correct place
-        batch["decoder_input_ids"] = shift_tokens_right(batch["labels"], tokenizer.pad_token_id)
         assert batch["decoder_input_ids"][0, 0].item() == tokenizer.lang_code_to_id[tgt_lang]
         assert batch["decoder_input_ids"][0, -1].item() == tokenizer.eos_token_id
         assert batch["input_ids"][0, -2].item() == tokenizer.eos_token_id
@@ -535,19 +470,15 @@ def test_seq2seq_dataset_truncation(tok_name):
         break  # No need to test every batch
 
 
-@pytest.mark.parametrize(["tok"], [pytest.param(BART_TINY), pytest.param("bert-base-cased")])
-def test_legacy_dataset_truncation(tok):
+@pytest.mark.parametrize(["tok"], [pytest.param(T5_TINY), pytest.param(BART_TINY), param(MARIAN_TINY)])
+def test_summarization_dataset_truncation(tok):
     tokenizer = AutoTokenizer.from_pretrained(tok)
     tmp_dir = make_test_data_dir()
     max_len_source = max(len(tokenizer.encode(a)) for a in ARTICLES)
     max_len_target = max(len(tokenizer.encode(a)) for a in SUMMARIES)
     trunc_target = 4
-    train_dataset = LegacySeq2SeqDataset(
-        tokenizer,
-        data_dir=tmp_dir,
-        type_path="train",
-        max_source_length=20,
-        max_target_length=trunc_target,
+    train_dataset = Seq2SeqDataset(
+        tokenizer, data_dir=tmp_dir, type_path="train", max_source_length=20, max_target_length=trunc_target,
     )
     dataloader = DataLoader(train_dataset, batch_size=2, collate_fn=train_dataset.collate_fn)
     for batch in dataloader:
@@ -556,6 +487,6 @@ def test_legacy_dataset_truncation(tok):
         assert batch["input_ids"].shape[1] == max_len_source
         assert 20 >= batch["input_ids"].shape[1]  # trimmed significantly
         # show that targets were truncated
-        assert batch["labels"].shape[1] == trunc_target  # Truncated
+        assert batch["decoder_input_ids"].shape[1] == trunc_target  # Truncated
         assert max_len_target > trunc_target  # Truncated
         break  # No need to test every batch

examples/seq2seq/train_distilbart_cnn.sh

@@ -6,7 +6,7 @@ export GAS=1
 
 python finetune.py \
     --learning_rate=3e-5 \
-    --fp16 \
+    --fp16 --fp16_opt_level=O1 \
     --gpus 1 \
     --do_train \
     --do_predict \

examples/seq2seq/train_distilbart_xsum.sh

@@ -6,7 +6,7 @@ python distillation.py \
   --learning_rate=3e-4 \
   --do_train \
   --do_predict \
-  --fp16 \
+  --fp16 --fp16_opt_level=O1 \
   --val_check_interval 0.1 --n_val 1000 \
   --teacher facebook/bart-large-xsum --data_dir $XSUM_DIR \
   --max_target_length=60 --val_max_target_length=60 --test_max_target_length=100 \

examples/seq2seq/train_mbart_cc25_enro.sh

@@ -14,5 +14,4 @@ python finetune.py \
     --task translation \
     --warmup_steps 500 \
     --freeze_embeds \
-    --model_name_or_path=facebook/mbart-large-cc25 \
     "$@"

examples/seq2seq/utils.py

@@ -1,17 +1,16 @@
 import itertools
 import json
 import linecache
-import math
 import os
 import pickle
+import warnings
 from logging import getLogger
 from pathlib import Path
-from typing import Callable, Dict, Iterable, List, Union
+from typing import Callable, Dict, Iterable, List
 
 import git
 import numpy as np
 import torch
-import torch.distributed as dist
 from rouge_score import rouge_scorer, scoring
 from sacrebleu import corpus_bleu
 from torch import nn
@@ -42,7 +41,6 @@ def label_smoothed_nll_loss(lprobs, target, epsilon, ignore_index=-100):
 
 
 def encode_line(tokenizer, line, max_length, pad_to_max_length=True, return_tensors="pt"):
-    """Only used by LegacyDataset"""
     extra_kw = {"add_prefix_space": True} if isinstance(tokenizer, BartTokenizer) else {}
     return tokenizer(
         [line],
@@ -59,15 +57,13 @@ def lmap(f: Callable, x: Iterable) -> List:
     return list(map(f, x))
 
 
-def calculate_bleu(output_lns, refs_lns, **kwargs) -> dict:
+def calculate_bleu_score(output_lns, refs_lns, **kwargs) -> dict:
     """Uses sacrebleu's corpus_bleu implementation."""
-    return {"bleu": round(corpus_bleu(output_lns, [refs_lns], **kwargs).score, 4)}
+    return {"bleu": corpus_bleu(output_lns, [refs_lns], **kwargs).score}
 
 
 def trim_batch(
-    input_ids,
-    pad_token_id,
-    attention_mask=None,
+    input_ids, pad_token_id, attention_mask=None,
 ):
     """Remove columns that are populated exclusively by pad_token_id"""
     keep_column_mask = input_ids.ne(pad_token_id).any(dim=0)
@@ -77,7 +73,7 @@ def trim_batch(
         return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])
 
 
-class AbstractSeq2SeqDataset(Dataset):
+class Seq2SeqDataset(Dataset):
     def __init__(
         self,
         tokenizer,
@@ -104,31 +100,11 @@ class AbstractSeq2SeqDataset(Dataset):
         self.pad_token_id = self.tokenizer.pad_token_id
         self.src_lang = src_lang
         self.tgt_lang = tgt_lang
-        self.add_prefix_space = isinstance(self.tokenizer, BartTokenizer)
 
     def __len__(self):
         return len(self.src_lens)
 
-    @staticmethod
-    def get_char_lens(data_file):
-        return [len(x) for x in Path(data_file).open().readlines()]
-
-    def make_sortish_sampler(self, batch_size, distributed=False):
-        if distributed:
-            return DistributedSortishSampler(self, batch_size)
-        else:
-            return SortishSampler(self.src_lens, batch_size)
-
-    def __getitem__(self, item):
-        raise NotImplementedError("You must implement this")
-
-    def collate_fn(self, batch):
-        raise NotImplementedError("You must implement this")
-
-
-class LegacySeq2SeqDataset(AbstractSeq2SeqDataset):
     def __getitem__(self, index) -> Dict[str, torch.Tensor]:
-        """Call tokenizer on src and tgt_lines"""
         index = index + 1  # linecache starts at 1
         source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n")
         tgt_line = linecache.getline(str(self.tgt_file), index).rstrip("\n")
@@ -143,27 +119,42 @@ class LegacySeq2SeqDataset(AbstractSeq2SeqDataset):
         return {
             "input_ids": source_ids,
             "attention_mask": src_mask,
-            "labels": target_ids,
+            "decoder_input_ids": target_ids,
         }
 
+    @staticmethod
+    def get_char_lens(data_file):
+        return [len(x) for x in Path(data_file).open().readlines()]
+
     def collate_fn(self, batch) -> Dict[str, torch.Tensor]:
         input_ids = torch.stack([x["input_ids"] for x in batch])
         masks = torch.stack([x["attention_mask"] for x in batch])
-        target_ids = torch.stack([x["labels"] for x in batch])
+        target_ids = torch.stack([x["decoder_input_ids"] for x in batch])
         pad_token_id = self.pad_token_id
         y = trim_batch(target_ids, pad_token_id)
         source_ids, source_mask = trim_batch(input_ids, pad_token_id, attention_mask=masks)
         batch = {
             "input_ids": source_ids,
             "attention_mask": source_mask,
-            "labels": y,
+            "decoder_input_ids": y,
         }
         return batch
 
+    def make_sortish_sampler(self, batch_size):
+        return SortishSampler(self.src_lens, batch_size)
 
-class Seq2SeqDataset(AbstractSeq2SeqDataset):
+
+class TranslationDataset(Seq2SeqDataset):
     """A dataset that calls prepare_seq2seq_batch."""
 
+    def __init__(self, *args, **kwargs):
+        super().__init__(*args, **kwargs)
+        if self.max_source_length != self.max_target_length:
+            warnings.warn(
+                f"Mbart is using sequence lengths {self.max_source_length}, {self.max_target_length}. "
+                f"Imbalanced sequence lengths may be undesired for translation tasks"
+            )
+
     def __getitem__(self, index) -> Dict[str, str]:
         index = index + 1  # linecache starts at 1
         source_line = self.prefix + linecache.getline(str(self.src_file), index).rstrip("\n")
@@ -176,7 +167,6 @@ class Seq2SeqDataset(AbstractSeq2SeqDataset):
         }
 
     def collate_fn(self, batch) -> Dict[str, torch.Tensor]:
-        """Call prepare_seq2seq_batch."""
         batch_encoding = self.tokenizer.prepare_seq2seq_batch(
             [x["src_texts"] for x in batch],
             src_lang=self.src_lang,
@@ -184,8 +174,6 @@ class Seq2SeqDataset(AbstractSeq2SeqDataset):
             tgt_lang=self.tgt_lang,
             max_length=self.max_source_length,
             max_target_length=self.max_target_length,
-            return_tensors="pt",
-            add_prefix_space=self.add_prefix_space,
         )
         return batch_encoding.data
 
@@ -196,77 +184,24 @@ class SortishSampler(Sampler):
     def __init__(self, data, batch_size):
         self.data, self.bs = data, batch_size
 
+    def key(self, i):
+        return self.data[i]
+
     def __len__(self) -> int:
         return len(self.data)
 
     def __iter__(self):
-        return iter(sortish_sampler_indices(self.data, self.bs))
-
-
-def sortish_sampler_indices(data: List, bs: int) -> np.array:
-    "Go through the text data by order of src length with a bit of randomness. From fastai repo."
-
-    def key_fn(i):
-        return data[i]
-
-    idxs = np.random.permutation(len(data))
-    sz = bs * 50
-    ck_idx = [idxs[i : i + sz] for i in range(0, len(idxs), sz)]
-    sort_idx = np.concatenate([sorted(s, key=key_fn, reverse=True) for s in ck_idx])
-    sz = bs
-    ck_idx = [sort_idx[i : i + sz] for i in range(0, len(sort_idx), sz)]
-    max_ck = np.argmax([key_fn(ck[0]) for ck in ck_idx])  # find the chunk with the largest key,
-    ck_idx[0], ck_idx[max_ck] = ck_idx[max_ck], ck_idx[0]  # then make sure it goes first.
-    sort_idx = np.concatenate(np.random.permutation(ck_idx[1:])) if len(ck_idx) > 1 else np.array([], dtype=np.int)
-    sort_idx = np.concatenate((ck_idx[0], sort_idx))
-    return sort_idx
-
-
-class DistributedSortishSampler(Sampler):
-    """Copied from torch DistributedSampler"""
-
-    def __init__(self, dataset, batch_size, num_replicas=None, rank=None):
-        if num_replicas is None:
-            if not dist.is_available():
-                raise RuntimeError("Requires distributed package to be available")
-            num_replicas = dist.get_world_size()
-        if rank is None:
-            if not dist.is_available():
-                raise RuntimeError("Requires distributed package to be available")
-            rank = dist.get_rank()
-        self.dataset = dataset
-        self.num_replicas = num_replicas
-        self.rank = rank
-        self.epoch = 0
-        self.num_samples = int(math.ceil(len(self.dataset) * 1.0 / self.num_replicas))
-        self.total_size = self.num_samples * self.num_replicas
-        self.batch_size = batch_size
-
-    def __iter__(self) -> Iterable:
-        g = torch.Generator()
-        g.manual_seed(self.epoch)
-        available_indices = self.get_indices_for_rank()  # indices[self.rank: self.total_size: self.num_replicas]
-
-        sortish_data = [self.dataset.src_lens[i] for i in available_indices]
-        sortish_indices = sortish_sampler_indices(sortish_data, self.batch_size)
-        indices = [available_indices[i] for i in sortish_indices]
-        assert len(indices) == self.num_samples
-        return iter(indices)
-
-    def get_indices_for_rank(self) -> np.array:
-        indices = list(range(len(self.dataset)))
-        # add extra samples to make it evenly divisible
-        indices += indices[: (self.total_size - len(indices))]
-        assert len(indices) == self.total_size
-        # subsample
-        available_indices = indices[self.rank : self.total_size : self.num_replicas]
-        return available_indices
-
-    def __len__(self):
-        return self.num_samples
-
-    def set_epoch(self, epoch):
-        self.epoch = epoch
+        idxs = np.random.permutation(len(self.data))
+        sz = self.bs * 50
+        ck_idx = [idxs[i : i + sz] for i in range(0, len(idxs), sz)]
+        sort_idx = np.concatenate([sorted(s, key=self.key, reverse=True) for s in ck_idx])
+        sz = self.bs
+        ck_idx = [sort_idx[i : i + sz] for i in range(0, len(sort_idx), sz)]
+        max_ck = np.argmax([self.key(ck[0]) for ck in ck_idx])  # find the chunk with the largest key,
+        ck_idx[0], ck_idx[max_ck] = ck_idx[max_ck], ck_idx[0]  # then make sure it goes first.
+        sort_idx = np.concatenate(np.random.permutation(ck_idx[1:])) if len(ck_idx) > 1 else np.array([], dtype=np.int)
+        sort_idx = np.concatenate((ck_idx[0], sort_idx))
+        return iter(sort_idx)
 
 
 logger = getLogger(__name__)
@@ -336,14 +271,10 @@ def calculate_rouge(output_lns: List[str], reference_lns: List[str], use_stemmer
         aggregator.add_scores(scores)
 
     result = aggregator.aggregate()
-    return {k: round(v.mid.fmeasure * 100, 4) for k, v in result.items()}
-
-
-# Utilities for freezing parameters and checking whether they are frozen
+    return {k: v.mid.fmeasure * 100 for k, v in result.items()}
 
 
 def freeze_params(model: nn.Module):
-    """Set requires_grad=False for each of model.parameters()"""
     for par in model.parameters():
         par.requires_grad = False
 
@@ -367,23 +298,3 @@ def assert_not_all_frozen(model):
     model_grads: List[bool] = list(grad_status(model))
     npars = len(model_grads)
     assert any(model_grads), f"none of {npars} weights require grad"
-
-
-# CLI Parsing utils
-
-
-def parse_numeric_cl_kwargs(unparsed_args: List[str]) -> Dict[str, Union[int, float]]:
-    """Parse an argv list of unspecified command line args to a dict. Assumes all values are numeric."""
-    result = {}
-    assert len(unparsed_args) % 2 == 0, f"got odd number of unparsed args: {unparsed_args}"
-    num_pairs = len(unparsed_args) // 2
-    for pair_num in range(num_pairs):
-        i = 2 * pair_num
-        assert unparsed_args[i].startswith("--")
-        try:
-            value = int(unparsed_args[i + 1])
-        except ValueError:
-            value = float(unparsed_args[i + 1])  # this can raise another informative ValueError
-
-        result[unparsed_args[i][2:]] = value
-    return result

examples/test_examples.py

@@ -22,8 +22,7 @@ from unittest.mock import patch
 
 import torch
 
-from transformers.file_utils import is_apex_available
-from transformers.testing_utils import TestCasePlus, torch_device
+from transformers.testing_utils import TestCasePlus
 
 
 SRC_DIRS = [
@@ -36,8 +35,8 @@ sys.path.extend(SRC_DIRS)
 if SRC_DIRS is not None:
     import run_generation
     import run_glue
-    import run_language_modeling
     import run_pl_glue
+    import run_language_modeling
     import run_squad
 
 
@@ -53,11 +52,6 @@ def get_setup_file():
     return args.f
 
 
-def is_cuda_and_apex_available():
-    is_using_cuda = torch.cuda.is_available() and torch_device == "cuda"
-    return is_using_cuda and is_apex_available()
-
-
 class ExamplesTests(TestCasePlus):
     def test_run_glue(self):
         stream_handler = logging.StreamHandler(sys.stdout)
@@ -80,13 +74,7 @@ class ExamplesTests(TestCasePlus):
             --warmup_steps=2
             --seed=42
             --max_seq_length=128
-            """
-        output_dir = "./tests/fixtures/tests_samples/temp_dir_{}".format(hash(testargs))
-        testargs += "--output_dir " + output_dir
-        testargs = testargs.split()
-
-        if is_cuda_and_apex_available():
-            testargs.append("--fp16")
+            """.split()
 
         with patch.object(sys, "argv", testargs):
             result = run_glue.main()
@@ -114,9 +102,7 @@ class ExamplesTests(TestCasePlus):
             --max_seq_length=128
             """.split()
         if torch.cuda.is_available():
-            testargs += ["--gpus=1"]
-        if is_cuda_and_apex_available():
-            testargs.append("--fp16")
+            testargs += ["--fp16", "--gpus=1"]
 
         with patch.object(sys, "argv", testargs):
             result = run_pl_glue.main()
@@ -149,13 +135,8 @@ class ExamplesTests(TestCasePlus):
             --do_train
             --do_eval
             --num_train_epochs=1
-            """
-        output_dir = "./tests/fixtures/tests_samples/temp_dir_{}".format(hash(testargs))
-        testargs += "--output_dir " + output_dir
-        testargs = testargs.split()
-
-        if torch_device != "cuda":
-            testargs.append("--no_cuda")
+            --no_cuda
+            """.split()
 
         with patch.object(sys, "argv", testargs):
             result = run_language_modeling.main()
@@ -194,14 +175,7 @@ class ExamplesTests(TestCasePlus):
         logger.addHandler(stream_handler)
 
         testargs = ["run_generation.py", "--prompt=Hello", "--length=10", "--seed=42"]
-
-        if is_cuda_and_apex_available():
-            testargs.append("--fp16")
-
-        model_type, model_name = (
-            "--model_type=gpt2",
-            "--model_name_or_path=sshleifer/tiny-gpt2",
-        )
+        model_type, model_name = ("--model_type=gpt2", "--model_name_or_path=sshleifer/tiny-gpt2")
         with patch.object(sys, "argv", testargs + [model_type, model_name]):
             result = run_generation.main()
             self.assertGreaterEqual(len(result[0]), 10)

examples/text-classification/run_glue.py

@@ -13,7 +13,7 @@
 # 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.
-""" Finetuning the library models for sequence classification on GLUE."""
+""" Finetuning the library models for sequence classification on GLUE (Bert, XLM, XLNet, RoBERTa, Albert, XLM-RoBERTa)."""
 
 
 import dataclasses

examples/text-classification/run_pl_glue.py

@@ -153,11 +153,7 @@ class GLUETransformer(BaseTransformer):
         )
 
         parser.add_argument(
-            "--task",
-            default="",
-            type=str,
-            required=True,
-            help="The GLUE task to run",
+            "--task", default="", type=str, required=True, help="The GLUE task to run",
         )
         parser.add_argument(
             "--gpus",
@@ -181,10 +177,7 @@ def main():
 
     # If output_dir not provided, a folder will be generated in pwd
     if args.output_dir is None:
-        args.output_dir = os.path.join(
-            "./results",
-            f"{args.task}_{time.strftime('%Y%m%d_%H%M%S')}",
-        )
+        args.output_dir = os.path.join("./results", f"{args.task}_{time.strftime('%Y%m%d_%H%M%S')}",)
         os.makedirs(args.output_dir)
 
     model = GLUETransformer(args)

examples/text-classification/run_xnli.py

@@ -328,11 +328,7 @@ def load_and_cache_examples(args, task, tokenizer, evaluate=False):
             processor.get_test_examples(args.data_dir) if evaluate else processor.get_train_examples(args.data_dir)
         )
         features = convert_examples_to_features(
-            examples,
-            tokenizer,
-            max_length=args.max_seq_length,
-            label_list=label_list,
-            output_mode=output_mode,
+            examples, tokenizer, max_length=args.max_seq_length, label_list=label_list, output_mode=output_mode,
         )
         if args.local_rank in [-1, 0]:
             logger.info("Saving features into cached file %s", cached_features_file)

examples/text-generation/pplm/run_pplm.py

@@ -698,9 +698,7 @@ def run_pplm_example(
                 for word_id in pert_gen_tok_text.tolist()[0]:
                     if word_id in bow_word_ids:
                         pert_gen_text += "{}{}{}".format(
-                            colorama.Fore.RED,
-                            tokenizer.decode([word_id]),
-                            colorama.Style.RESET_ALL,
+                            colorama.Fore.RED, tokenizer.decode([word_id]), colorama.Style.RESET_ALL,
                         )
                     else:
                         pert_gen_text += tokenizer.decode([word_id])
@@ -731,10 +729,7 @@ if __name__ == "__main__":
     parser.add_argument("--cond_text", type=str, default="The lake", help="Prefix texts to condition on")
     parser.add_argument("--uncond", action="store_true", help="Generate from end-of-text as prefix")
     parser.add_argument(
-        "--num_samples",
-        type=int,
-        default=1,
-        help="Number of samples to generate from the modified latents",
+        "--num_samples", type=int, default=1, help="Number of samples to generate from the modified latents",
     )
     parser.add_argument(
         "--bag_of_words",
@@ -756,22 +751,13 @@ if __name__ == "__main__":
         help="Discriminator to use",
     )
     parser.add_argument(
-        "--discrim_weights",
-        type=str,
-        default=None,
-        help="Weights for the generic discriminator",
+        "--discrim_weights", type=str, default=None, help="Weights for the generic discriminator",
     )
     parser.add_argument(
-        "--discrim_meta",
-        type=str,
-        default=None,
-        help="Meta information for the generic discriminator",
+        "--discrim_meta", type=str, default=None, help="Meta information for the generic discriminator",
     )
     parser.add_argument(
-        "--class_label",
-        type=int,
-        default=-1,
-        help="Class label used for the discriminator",
+        "--class_label", type=int, default=-1, help="Class label used for the discriminator",
     )
     parser.add_argument("--length", type=int, default=100)
     parser.add_argument("--stepsize", type=float, default=0.02)
@@ -787,10 +773,7 @@ if __name__ == "__main__":
         help="Length of past which is being optimized; 0 corresponds to infinite window length",
     )
     parser.add_argument(
-        "--horizon_length",
-        type=int,
-        default=1,
-        help="Length of future to optimize over",
+        "--horizon_length", type=int, default=1, help="Length of future to optimize over",
     )
     parser.add_argument("--decay", action="store_true", help="whether to decay or not")
     parser.add_argument("--gamma", type=float, default=1.5)
@@ -800,10 +783,7 @@ if __name__ == "__main__":
     parser.add_argument("--no_cuda", action="store_true", help="no cuda")
     parser.add_argument("--colorama", action="store_true", help="colors keywords")
     parser.add_argument(
-        "--repetition_penalty",
-        type=float,
-        default=1.0,
-        help="Penalize repetition. More than 1.0 -> less repetition",
+        "--repetition_penalty", type=float, default=1.0, help="Penalize repetition. More than 1.0 -> less repetition",
     )
 
     args = parser.parse_args()

examples/text-generation/pplm/run_pplm_discrim_train.py

@@ -242,12 +242,7 @@ def train_discriminator(
 
         text = torchtext_data.Field()
         label = torchtext_data.Field(sequential=False)
-        train_data, val_data, test_data = datasets.SST.splits(
-            text,
-            label,
-            fine_grained=True,
-            train_subtrees=True,
-        )
+        train_data, val_data, test_data = datasets.SST.splits(text, label, fine_grained=True, train_subtrees=True,)
 
         x = []
         y = []

examples/text-generation/run_generation.py

@@ -41,9 +41,7 @@ from transformers import (
 
 
 logging.basicConfig(
-    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
-    datefmt="%m/%d/%Y %H:%M:%S",
-    level=logging.INFO,
+    format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s", datefmt="%m/%d/%Y %H:%M:%S", level=logging.INFO,
 )
 logger = logging.getLogger(__name__)
 
@@ -61,7 +59,7 @@ MODEL_CLASSES = {
 # Padding text to help Transformer-XL and XLNet with short prompts as proposed by Aman Rusia
 # in https://github.com/rusiaaman/XLNet-gen#methodology
 # and https://medium.com/@amanrusia/xlnet-speaks-comparison-to-gpt-2-ea1a4e9ba39e
-PREFIX = """In 1991, the remains of Russian Tsar Nicholas II and his family
+PADDING_TEXT = """In 1991, the remains of Russian Tsar Nicholas II and his family
 (except for Alexei and Maria) are discovered.
 The voice of Nicholas's young son, Tsarevich Alexei Nikolaevich, narrates the
 remainder of the story. 1883 Western Siberia,
@@ -122,14 +120,12 @@ def prepare_xlm_input(args, model, tokenizer, prompt_text):
 
 
 def prepare_xlnet_input(args, _, tokenizer, prompt_text):
-    prefix = args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
-    prompt_text = prefix + prompt_text
+    prompt_text = (args.padding_text if args.padding_text else PADDING_TEXT) + prompt_text
     return prompt_text
 
 
 def prepare_transfoxl_input(args, _, tokenizer, prompt_text):
-    prefix = args.prefix if args.prefix else args.padding_text if args.padding_text else PREFIX
-    prompt_text = prefix + prompt_text
+    prompt_text = (args.padding_text if args.padding_text else PADDING_TEXT) + prompt_text
     return prompt_text
 
 
@@ -184,30 +180,17 @@ def main():
     parser.add_argument("--k", type=int, default=0)
     parser.add_argument("--p", type=float, default=0.9)
 
-    parser.add_argument("--prefix", type=str, default="", help="Text added prior to input.")
-    parser.add_argument("--padding_text", type=str, default="", help="Deprecated, the use of `--prefix` is preferred.")
+    parser.add_argument("--padding_text", type=str, default="", help="Padding text for Transfo-XL and XLNet.")
     parser.add_argument("--xlm_language", type=str, default="", help="Optional language when used with the XLM model.")
 
     parser.add_argument("--seed", type=int, default=42, help="random seed for initialization")
     parser.add_argument("--no_cuda", action="store_true", help="Avoid using CUDA when available")
     parser.add_argument("--num_return_sequences", type=int, default=1, help="The number of samples to generate.")
-    parser.add_argument(
-        "--fp16",
-        action="store_true",
-        help="Whether to use 16-bit (mixed) precision (through NVIDIA apex) instead of 32-bit",
-    )
     args = parser.parse_args()
 
     args.device = torch.device("cuda" if torch.cuda.is_available() and not args.no_cuda else "cpu")
     args.n_gpu = 0 if args.no_cuda else torch.cuda.device_count()
 
-    logger.warning(
-        "device: %s, n_gpu: %s, 16-bits training: %s",
-        args.device,
-        args.n_gpu,
-        args.fp16,
-    )
-
     set_seed(args)
 
     # Initialize the model and tokenizer
@@ -221,9 +204,6 @@ def main():
     model = model_class.from_pretrained(args.model_name_or_path)
     model.to(args.device)
 
-    if args.fp16:
-        model.half()
-
     args.length = adjust_length_to_model(args.length, max_sequence_length=model.config.max_position_embeddings)
     logger.info(args)
 
@@ -244,8 +224,7 @@ def main():
             preprocessed_prompt_text, add_special_tokens=False, return_tensors="pt", **tokenizer_kwargs
         )
     else:
-        prefix = args.prefix if args.prefix else args.padding_text
-        encoded_prompt = tokenizer.encode(prefix + prompt_text, add_special_tokens=False, return_tensors="pt")
+        encoded_prompt = tokenizer.encode(prompt_text, add_special_tokens=False, return_tensors="pt")
     encoded_prompt = encoded_prompt.to(args.device)
 
     if encoded_prompt.size()[-1] == 0:

examples/token-classification/run.sh

@@ -1,11 +1,8 @@
-## The relevant files are currently on a shared Google
-## drive at https://drive.google.com/drive/folders/1kC0I2UGl2ltrluI9NqDjaQJGw5iliw_J
-## Monitor for changes and eventually migrate to nlp dataset
-curl -L 'https://drive.google.com/uc?export=download&id=1Jjhbal535VVz2ap4v4r_rN1UEHTdLK5P' \
+curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-train.tsv?attredirects=0&d=1' \
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > train.txt.tmp
-curl -L 'https://drive.google.com/uc?export=download&id=1ZfRcQThdtAR5PPRjIDtrVP7BtXSCUBbm' \
+curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-dev.tsv?attredirects=0&d=1' \
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > dev.txt.tmp
-curl -L 'https://drive.google.com/uc?export=download&id=1u9mb7kNJHWQCWyweMDRMuTFoOHOfeBTH' \
+curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-test.tsv?attredirects=0&d=1' \
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > test.txt.tmp
 
 export MAX_LENGTH=128

examples/token-classification/run_pl.sh

@@ -3,14 +3,11 @@
 # for seqeval metrics import
 pip install -r ../requirements.txt
 
-## The relevant files are currently on a shared Google
-## drive at https://drive.google.com/drive/folders/1kC0I2UGl2ltrluI9NqDjaQJGw5iliw_J
-## Monitor for changes and eventually migrate to nlp dataset
-curl -L 'https://drive.google.com/uc?export=download&id=1Jjhbal535VVz2ap4v4r_rN1UEHTdLK5P' \
+curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-train.tsv?attredirects=0&d=1' \
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > train.txt.tmp
-curl -L 'https://drive.google.com/uc?export=download&id=1ZfRcQThdtAR5PPRjIDtrVP7BtXSCUBbm' \
+curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-dev.tsv?attredirects=0&d=1' \
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > dev.txt.tmp
-curl -L 'https://drive.google.com/uc?export=download&id=1u9mb7kNJHWQCWyweMDRMuTFoOHOfeBTH' \
+curl -L 'https://sites.google.com/site/germeval2014ner/data/NER-de-test.tsv?attredirects=0&d=1' \
 | grep -v "^#" | cut -f 2,3 | tr '\t' ' ' > test.txt.tmp
 
 export MAX_LENGTH=128
@@ -32,6 +29,7 @@ mkdir -p $OUTPUT_DIR
 export PYTHONPATH="../":"${PYTHONPATH}"
 
 python3 run_pl_ner.py --data_dir ./ \
+--model_type bert \
 --labels ./labels.txt \
 --model_name_or_path $BERT_MODEL \
 --output_dir $OUTPUT_DIR \
@@ -39,6 +37,5 @@ python3 run_pl_ner.py --data_dir ./ \
 --num_train_epochs $NUM_EPOCHS \
 --train_batch_size $BATCH_SIZE \
 --seed $SEED \
---gpus 1 \
 --do_train \
 --do_predict

examples/token-classification/run_pl_ner.py

@@ -86,7 +86,7 @@ class NERTransformer(BaseTransformer):
                 logger.info("Saving features into cached file %s", cached_features_file)
                 torch.save(features, cached_features_file)
 
-    def get_dataloader(self, mode: int, batch_size: int, shuffle: bool = False) -> DataLoader:
+    def get_dataloader(self, mode: int, batch_size: int) -> DataLoader:
         "Load datasets. Called after prepare data."
         cached_features_file = self._feature_file(mode)
         logger.info("Loading features from cached file %s", cached_features_file)