boom boom

merge batch parity
Merge branch 'master' into dropper-celoss
2020-08-21 19:55:50 -04:00 · 2020-08-21 19:54:03 -04:00 · 2020-08-21 19:50:36 -04:00 · 2020-08-21 17:15:54 -04:00 · 2020-08-21 17:09:37 -04:00 · 2020-08-21 16:31:49 -04:00
536 changed files with 10094 additions and 43760 deletions
@@ -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,14 +235,15 @@ 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" }}
                  paths:
                      - '~/.cache/pip'
            - run: black --check --line-length 119 --target-version py35 examples templates tests src utils
-            - run: isort --check-only examples templates tests src utils
+            - run: isort --check-only --recursive examples templates tests src utils
            - run: flake8 examples templates tests src utils
            - run: python utils/check_repo.py
    check_repository_consistency:
@@ -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
@@ -1,2 +0,0 @@
-<!-- This line specifies which issue to close after the pull request is merged. -->
-Fixes #{issue number}
@@ -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: |
@@ -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"
@@ -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:

@@ -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"
@@ -22,125 +22,59 @@
 <p>State-of-the-art Natural Language Processing for PyTorch and TensorFlow 2.0
 </h3>

-🤗 Transformers provides thousands of pretrained models to perform tasks on texts such as classification, information extraction, question answering, summarization, translation, text generation, etc in 100+ languages. Its aim is to make cutting-edge NLP easier to use for everyone. 
-
-🤗 Transformers provides APIs to quickly download and use those pretrained models on a given text, fine-tune them on your own datasets then share them with the community on our [model hub](https://huggingface.co/models). At the same time, each python module defining an architecture can be used as a standalone and modified to enable quick research experiments. 
-
-🤗 Transformers is backed by the two most popular deep learning libraries, [PyTorch](https://pytorch.org/) and [TensorFlow](https://www.tensorflow.org/), with a seamless integration between them, allowing you to train your models with one then load it for inference with the other.
+🤗 Transformers (formerly known as `pytorch-transformers` and `pytorch-pretrained-bert`) provides state-of-the-art general-purpose architectures (BERT, GPT-2, RoBERTa, XLM, DistilBert, XLNet, T5, CTRL...) for Natural Language Understanding (NLU) and Natural Language Generation (NLG) with over thousands of pretrained models in 100+ languages and deep interoperability between PyTorch & TensorFlow 2.0.

 ### Recent contributors
 [![](https://sourcerer.io/fame/clmnt/huggingface/transformers/images/0)](https://sourcerer.io/fame/clmnt/huggingface/transformers/links/0)[![](https://sourcerer.io/fame/clmnt/huggingface/transformers/images/1)](https://sourcerer.io/fame/clmnt/huggingface/transformers/links/1)[![](https://sourcerer.io/fame/clmnt/huggingface/transformers/images/2)](https://sourcerer.io/fame/clmnt/huggingface/transformers/links/2)[![](https://sourcerer.io/fame/clmnt/huggingface/transformers/images/3)](https://sourcerer.io/fame/clmnt/huggingface/transformers/links/3)[![](https://sourcerer.io/fame/clmnt/huggingface/transformers/images/4)](https://sourcerer.io/fame/clmnt/huggingface/transformers/links/4)[![](https://sourcerer.io/fame/clmnt/huggingface/transformers/images/5)](https://sourcerer.io/fame/clmnt/huggingface/transformers/links/5)[![](https://sourcerer.io/fame/clmnt/huggingface/transformers/images/6)](https://sourcerer.io/fame/clmnt/huggingface/transformers/links/6)[![](https://sourcerer.io/fame/clmnt/huggingface/transformers/images/7)](https://sourcerer.io/fame/clmnt/huggingface/transformers/links/7)

-## Online demos
+### Features
+- High performance on NLU and NLG tasks
+- Low barrier to entry for educators and practitioners

-You can test most of our models directly on their pages from the [model hub](https://huggingface.co/models). We also offer an [inference API](https://huggingface.co/pricing) to use those models.
+State-of-the-art NLP for everyone
+- Deep learning researchers
+- Hands-on practitioners
+- AI/ML/NLP teachers and educators

-Here are a few examples: 
- [Masked word completion with BERT](https://huggingface.co/bert-base-uncased?text=Paris+is+the+%5BMASK%5D+of+France)
- [Name Entity Recognition with Electra](https://huggingface.co/dbmdz/electra-large-discriminator-finetuned-conll03-english?text=My+name+is+Sarah+and+I+live+in+London+city)
- [Text generation with GPT-2](https://huggingface.co/gpt2?text=A+long+time+ago%2C+)
- [Natural Langugage Inference with RoBERTa](https://huggingface.co/roberta-large-mnli?text=The+dog+was+lost.+Nobody+lost+any+animal)
- [Summarization with BART](https://huggingface.co/facebook/bart-large-cnn?text=The+tower+is+324+metres+%281%2C063+ft%29+tall%2C+about+the+same+height+as+an+81-storey+building%2C+and+the+tallest+structure+in+Paris.+Its+base+is+square%2C+measuring+125+metres+%28410+ft%29+on+each+side.+During+its+construction%2C+the+Eiffel+Tower+surpassed+the+Washington+Monument+to+become+the+tallest+man-made+structure+in+the+world%2C+a+title+it+held+for+41+years+until+the+Chrysler+Building+in+New+York+City+was+finished+in+1930.+It+was+the+first+structure+to+reach+a+height+of+300+metres.+Due+to+the+addition+of+a+broadcasting+aerial+at+the+top+of+the+tower+in+1957%2C+it+is+now+taller+than+the+Chrysler+Building+by+5.2+metres+%2817+ft%29.+Excluding+transmitters%2C+the+Eiffel+Tower+is+the+second+tallest+free-standing+structure+in+France+after+the+Millau+Viaduct)
- [Question answering with DistilBERT](https://huggingface.co/distilbert-base-uncased-distilled-squad?text=Which+name+is+also+used+to+describe+the+Amazon+rainforest+in+English%3F&context=The+Amazon+rainforest+%28Portuguese%3A+Floresta+Amaz%C3%B4nica+or+Amaz%C3%B4nia%3B+Spanish%3A+Selva+Amaz%C3%B3nica%2C+Amazon%C3%ADa+or+usually+Amazonia%3B+French%3A+For%C3%AAt+amazonienne%3B+Dutch%3A+Amazoneregenwoud%29%2C+also+known+in+English+as+Amazonia+or+the+Amazon+Jungle%2C+is+a+moist+broadleaf+forest+that+covers+most+of+the+Amazon+basin+of+South+America.+This+basin+encompasses+7%2C000%2C000+square+kilometres+%282%2C700%2C000+sq+mi%29%2C+of+which+5%2C500%2C000+square+kilometres+%282%2C100%2C000+sq+mi%29+are+covered+by+the+rainforest.+This+region+includes+territory+belonging+to+nine+nations.+The+majority+of+the+forest+is+contained+within+Brazil%2C+with+60%25+of+the+rainforest%2C+followed+by+Peru+with+13%25%2C+Colombia+with+10%25%2C+and+with+minor+amounts+in+Venezuela%2C+Ecuador%2C+Bolivia%2C+Guyana%2C+Suriname+and+French+Guiana.+States+or+departments+in+four+nations+contain+%22Amazonas%22+in+their+names.+The+Amazon+represents+over+half+of+the+planet%27s+remaining+rainforests%2C+and+comprises+the+largest+and+most+biodiverse+tract+of+tropical+rainforest+in+the+world%2C+with+an+estimated+390+billion+individual+trees+divided+into+16%2C000+species)
- [Translation with T5](https://huggingface.co/t5-base?text=My+name+is+Wolfgang+and+I+live+in+Berlin)
+Lower compute costs, smaller carbon footprint
+- Researchers can share trained models instead of always retraining
+- Practitioners can reduce compute time and production costs
+- Dozens of architectures with over 1,000 pretrained models, some in more than 100 languages

-**[Write With Transformer](https://transformer.huggingface.co)**, built by the Hugging Face team, is the official demo of this repo’s text generation capabilities.
+Choose the right framework for every part of a model's lifetime
+- Train state-of-the-art models in 3 lines of code
+- Deep interoperability between TensorFlow 2.0 and PyTorch models
+- Move a single model between TF2.0/PyTorch frameworks at will
+- Seamlessly pick the right framework for training, evaluation, production

-## Quick tour

-To immediately use a model on a given text, we provide the `pipeline` API. Pipelines group together a pretrained model with the preprocessing that was used during that model training. Here is how to quickly use a pipeline to classify positive versus negative texts 
-
-```python
->>> from transformers import pipeline
-
-# Allocate a pipeline for sentiment-analysis
->>> classifier = pipeline('sentiment-analysis')
->>> classifier('We are very happy to include pipeline into the transformers repository.')
-[{'label': 'POSITIVE', 'score': 0.9978193640708923}]
-```
-
-The second line of code downloads and caches the pretrained model used by the pipeline, the third line evaluates it on the given text. Here the answer is "positive" with a confidence of 99.8%. 
-
-This is another example of pipeline used for that can extract question answers from some context:
-
-``` python
->>> from transformers import pipeline
-
-# Allocate a pipeline for question-answering
->>> question_answerer = pipeline('question-answering')
->>> question_answerer({
-...     'question': 'What is the name of the repository ?',
-...     'context': 'Pipeline have been included in the huggingface/transformers repository'
-... })
-{'score': 0.5135612454720828, 'start': 35, 'end': 59, 'answer': 'huggingface/transformers'}
-
-```
-
-On top of the answer, the pretrained model used here returned its confidence score, along with the start position and its end position in the tokenized sentence. You can learn more about the tasks supported by the `pipeline` API in [this tutorial](https://huggingface.co/transformers/task_summary.html).
-
-To download and use any of the pretrained models on your given task, you just need to use those three lines of codes (PyTorch verison):
-```python
->>> from transformers import AutoTokenizer, AutoModel
-
->>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
->>> model = AutoModel.from_pretrained("bert-base-uncased")
-
->>> inputs = tokenizer("Hello world!", return_tensors="pt")
->>> outputs = model(**inputs)
-```
-or for TensorFlow:
-```python
->>> from transformers import AutoTokenizer, TFAutoModel
-
->>> tokenizer = AutoTokenizer.from_pretrained("bert-base-uncased")
->>> model = TFAutoModel.from_pretrained("bert-base-uncased")
-
->>> inputs = tokenizer("Hello world!", return_tensors="tf")
->>> outputs = model(**inputs)
-```
-
-The tokenizer is responsible for all the preprocessing the pretrained model expects, and can be called directly on one (or list) of texts (as we can see on the fourth line of both code examples). It will output a dictionary you can directly pass to your model (which is done on the fifth line).
-
-The model itself is a regular [Pytorch `nn.Module`](https://pytorch.org/docs/stable/nn.html#torch.nn.Module) or a [TensorFlow `tf.keras.Model`](https://www.tensorflow.org/api_docs/python/tf/keras/Model) (depending on your backend) which you can use normally. For instance, [this tutorial](https://huggingface.co/transformers/training.html) explains how to integrate such a model in classic PyTorch or TensorFlow training loop, or how to use our `Trainer` API to quickly fine-tune the on a new dataset.
-
-## Why should I use transformers?
-
-1. Easy-to-use state-of-the-art models:
-    - High performance on NLU and NLG tasks.
-    - Low barrier to entry for educators and practitioners.
-    - Few user-facing abastractions with just three classes to learn.
-    - A unified API for using all our pretrained models.
-
-1. Lower compute costs, smaller carbon footprint:
-    - Researchers can share trained models instead of always retraining.
-    - Practitioners can reduce compute time and production costs.
-    - Dozens of architectures with over 2,000 pretrained models, some in more than 100 languages.
-
-1. Choose the right framework for every part of a model's lifetime:
-    - Train state-of-the-art models in 3 lines of code.
-    - Move a single model between TF2.0/PyTorch frameworks at will.
-    - Seamlessly pick the right framework for training, evaluation, production.
-
-1. Easily customize a model or an example to your needs:
-    - Examples for each architecture to reproduce the results by the official authors of said architecture.
-    - Expose the models internal as consistently as possible.
-    - Model files can be used independently of the library for quick experiments. 
-
-## Why shouldn't I use transformers?
-
- This library is not a modular toolbox of building blocks for neural nets. The code in the model files is not refactored with additional abstractions on purpose, so that researchers can quickly iterate on each of the models without diving in additional abstractions/files.
- The training API is not intended to work on any model but is optimized to work with the models provided by the library. For generic machine learning loops, you should use another library.
- While we strive to present as many use cases as possible, the scripts in our [examples folder](https://github.com/huggingface/transformers/tree/master/examples) are just that: examples. It is expected that they won't work out-of-the box on your specific problem and that you will be required to change a few lines of code to adapt them to your needs.
+| Section | Description |
+|-|-|
+| [Installation](#installation) | How to install the package |
+| [Model architectures](#model-architectures) | Architectures (with pretrained weights) |
+| [Online demo](#online-demo) | Experimenting with this repo’s text generation capabilities |
+| [Quick tour: Usage](#quick-tour) | Tokenizers & models usage: Bert and GPT-2 |
+| [Quick tour: TF 2.0 and PyTorch ](#Quick-tour-TF-20-training-and-PyTorch-interoperability) | Train a TF 2.0 model in 10 lines of code, load it in PyTorch |
+| [Quick tour: pipelines](#quick-tour-of-pipelines) | Using Pipelines: Wrapper around tokenizer and models to use finetuned models |
+| [Quick tour: Fine-tuning/usage scripts](#quick-tour-of-the-fine-tuningusage-scripts) | Using provided scripts: GLUE, SQuAD and Text generation |
+| [Quick tour: Share your models ](#Quick-tour-of-model-sharing) | Upload and share your fine-tuned models with the community |
+| [Migrating from pytorch-transformers to transformers](#Migrating-from-pytorch-transformers-to-transformers) | Migrating your code from pytorch-transformers to transformers |
+| [Migrating from pytorch-pretrained-bert to pytorch-transformers](#Migrating-from-pytorch-pretrained-bert-to-transformers) | Migrating your code from pytorch-pretrained-bert to transformers |
+| [Documentation](https://huggingface.co/transformers/) | Full API documentation and more |

 ## Installation

-This repository is tested on Python 3.6+, PyTorch 1.0.0+ (PyTorch 1.3.1+ for [examples](https://github.com/huggingface/transformers/tree/master/examples)) and TensorFlow 2.0.
+This repo is tested on Python 3.6+, PyTorch 1.0.0+ (PyTorch 1.3.1+ for examples) and TensorFlow 2.0.

 You should install 🤗 Transformers in a [virtual environment](https://docs.python.org/3/library/venv.html). If you're unfamiliar with Python virtual environments, check out the [user guide](https://packaging.python.org/guides/installing-using-pip-and-virtual-environments/).

-First, create a virtual environment with the version of Python you're going to use and activate it.
+Create a virtual environment with the version of Python you're going to use and activate it.

-Then, you will need to install one of, or both, TensorFlow 2.0 and PyTorch.
+Now, if you want to use 🤗 Transformers, you can install it with pip. If you'd like to play with the examples, you must install it from source.
+
+### With pip
+
+First you need to install one of, or both, TensorFlow 2.0 and PyTorch.
 Please refer to [TensorFlow installation page](https://www.tensorflow.org/install/pip#tensorflow-2.0-rc-is-available) and/or [PyTorch installation page](https://pytorch.org/get-started/locally/#start-locally) regarding the specific install command for your platform.

 When TensorFlow 2.0 and/or PyTorch has been installed, 🤗 Transformers can be installed using pip as follows:
@@ -149,11 +83,68 @@ When TensorFlow 2.0 and/or PyTorch has been installed, 🤗 Transformers can be
 pip install transformers
 ```

-If you'd like to play with the examples, you must [install the library from source](https://huggingface.co/transformers/installation.html#installing-from-source).
+### From source

-## Models architectures
+Here also, you first need to install one of, or both, TensorFlow 2.0 and PyTorch.
+Please refer to [TensorFlow installation page](https://www.tensorflow.org/install/pip#tensorflow-2.0-rc-is-available) and/or [PyTorch installation page](https://pytorch.org/get-started/locally/#start-locally) regarding the specific install command for your platform.

-🤗 Transformers currently provides the following architectures (see [here](https://huggingface.co/transformers/model_summary.html) for a high-level summary of each them):
+When TensorFlow 2.0 and/or PyTorch has been installed, you can install from source by cloning the repository and running:
+
+```bash
+git clone https://github.com/huggingface/transformers
+cd transformers
+pip install .
+```
+
+When you update the repository, you should upgrade the transformers installation and its dependencies as follows:
+
+```bash
+git pull
+pip install --upgrade .
+```
+
+### Run the examples
+
+Examples are included in the repository but are not shipped with the library.
+
+Therefore, in order to run the latest versions of the examples, you need to install from source, as described above.
+
+Look at the [README](https://github.com/huggingface/transformers/blob/master/examples/README.md) for how to run examples.
+
+### Tests
+
+A series of tests are included for the library and for some example scripts. Library tests can be found in the [tests folder](https://github.com/huggingface/transformers/tree/master/tests) and examples tests in the [examples folder](https://github.com/huggingface/transformers/tree/master/examples).
+
+Depending on which framework is installed (TensorFlow 2.0 and/or PyTorch), the irrelevant tests will be skipped. Ensure that both frameworks are installed if you want to execute all tests.
+
+Here's the easiest way to run tests for the library:
+
+```bash
+pip install -e ".[testing]"
+make test
+```
+
+and for the examples:
+
+```bash
+pip install -e ".[testing]"
+pip install -r examples/requirements.txt
+make test-examples
+```
+
+For details, refer to the [contributing guide](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#tests).
+
+### Do you want to run a Transformer model on a mobile device?
+
+You should check out our [`swift-coreml-transformers`](https://github.com/huggingface/swift-coreml-transformers) repo.
+
+It contains a set of tools to convert PyTorch or TensorFlow 2.0 trained Transformer models (currently contains `GPT-2`, `DistilGPT-2`, `BERT`, and `DistilBERT`) to CoreML models that run on iOS devices.
+
+At some point in the future, you'll be able to seamlessly move from pre-training or fine-tuning models to productizing them in CoreML, or prototype a model or an app in CoreML then research its hyperparameters or architecture from TensorFlow 2.0 and/or PyTorch. Super exciting!
+
+## Model architectures
+
+🤗 Transformers currently provides the following NLU/NLG architectures:

 1. **[BERT](https://huggingface.co/transformers/model_doc/bert.html)** (from Google) released with the paper [BERT: Pre-training of Deep Bidirectional Transformers for Language Understanding](https://arxiv.org/abs/1810.04805) by Jacob Devlin, Ming-Wei Chang, Kenton Lee and Kristina Toutanova.
 2. **[GPT](https://huggingface.co/transformers/model_doc/gpt.html)** (from OpenAI) released with the paper [Improving Language Understanding by Generative Pre-Training](https://blog.openai.com/language-unsupervised/) by Alec Radford, Karthik Narasimhan, Tim Salimans and Ilya Sutskever.
@@ -181,33 +172,537 @@ 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. You can find more details on the performances in the Examples section of the [documentation](https://huggingface.co/transformers/examples.html).
+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).
+
+## Online demo
+
+You can test our inference API on most model pages from the model hub: https://huggingface.co/models
+
+For example: 
+- [Masked word completion with BERT](https://huggingface.co/bert-base-uncased?text=Paris+is+the+%5BMASK%5D+of+France)
+- [NER with Electra](https://huggingface.co/dbmdz/electra-large-discriminator-finetuned-conll03-english?text=My+name+is+Sarah+and+I+live+in+London+city)
+- [Text generation with GPT-2](https://huggingface.co/gpt2?text=A+long+time+ago%2C+)
+- [NLI with RoBERTa](https://huggingface.co/roberta-large-mnli?text=The+dog+was+lost.+Nobody+lost+any+animal)
+- [Summarization with BART](https://huggingface.co/facebook/bart-large-cnn?text=The+tower+is+324+metres+%281%2C063+ft%29+tall%2C+about+the+same+height+as+an+81-storey+building%2C+and+the+tallest+structure+in+Paris.+Its+base+is+square%2C+measuring+125+metres+%28410+ft%29+on+each+side.+During+its+construction%2C+the+Eiffel+Tower+surpassed+the+Washington+Monument+to+become+the+tallest+man-made+structure+in+the+world%2C+a+title+it+held+for+41+years+until+the+Chrysler+Building+in+New+York+City+was+finished+in+1930.+It+was+the+first+structure+to+reach+a+height+of+300+metres.+Due+to+the+addition+of+a+broadcasting+aerial+at+the+top+of+the+tower+in+1957%2C+it+is+now+taller+than+the+Chrysler+Building+by+5.2+metres+%2817+ft%29.+Excluding+transmitters%2C+the+Eiffel+Tower+is+the+second+tallest+free-standing+structure+in+France+after+the+Millau+Viaduct)
+- [Question answering with DistilBERT](https://huggingface.co/distilbert-base-uncased-distilled-squad?text=Which+name+is+also+used+to+describe+the+Amazon+rainforest+in+English%3F&context=The+Amazon+rainforest+%28Portuguese%3A+Floresta+Amaz%C3%B4nica+or+Amaz%C3%B4nia%3B+Spanish%3A+Selva+Amaz%C3%B3nica%2C+Amazon%C3%ADa+or+usually+Amazonia%3B+French%3A+For%C3%AAt+amazonienne%3B+Dutch%3A+Amazoneregenwoud%29%2C+also+known+in+English+as+Amazonia+or+the+Amazon+Jungle%2C+is+a+moist+broadleaf+forest+that+covers+most+of+the+Amazon+basin+of+South+America.+This+basin+encompasses+7%2C000%2C000+square+kilometres+%282%2C700%2C000+sq+mi%29%2C+of+which+5%2C500%2C000+square+kilometres+%282%2C100%2C000+sq+mi%29+are+covered+by+the+rainforest.+This+region+includes+territory+belonging+to+nine+nations.+The+majority+of+the+forest+is+contained+within+Brazil%2C+with+60%25+of+the+rainforest%2C+followed+by+Peru+with+13%25%2C+Colombia+with+10%25%2C+and+with+minor+amounts+in+Venezuela%2C+Ecuador%2C+Bolivia%2C+Guyana%2C+Suriname+and+French+Guiana.+States+or+departments+in+four+nations+contain+%22Amazonas%22+in+their+names.+The+Amazon+represents+over+half+of+the+planet%27s+remaining+rainforests%2C+and+comprises+the+largest+and+most+biodiverse+tract+of+tropical+rainforest+in+the+world%2C+with+an+estimated+390+billion+individual+trees+divided+into+16%2C000+species)
+- [Translation with T5](https://huggingface.co/t5-base?text=My+name+is+Wolfgang+and+I+live+in+Berlin)


-## Learn more
+**[Write With Transformer](https://transformer.huggingface.co)**, built by the Hugging Face team at transformer.huggingface.co, is the official demo of this repo’s text generation capabilities.

-| Section | Description |
-|-|-|
-| [Documentation](https://huggingface.co/transformers/) | Full API documentation and tutorials |
-| [Task summary](https://huggingface.co/transformers/task_summary.html) | Tasks supported by 🤗 Transformers |
-| [Preprocessing tutorial](https://huggingface.co/transformers/preprocessing.html) | Using the `Tokenizer` class to prepare data for the models |
-| [Training and fine-tuning](https://huggingface.co/transformers/training.html) | Using the models provided by 🤗 Transformers in a PyTorch/TensorFlow training loop and the `Trainer` API |
-| [Quick tour: Fine-tuning/usage scripts](https://github.com/huggingface/transformers/tree/master/examples) | Example scripts for fine-tuning models on a wide range of tasks |
-| [Model sharing and uploading](https://huggingface.co/transformers/model_sharing.html) | Upload and share your fine-tuned models with the community |
-| [Migration](https://huggingface.co/transformers/migration.html) | Migrate to 🤗 Transformers from `pytorch-transformers` or `pytorch-pretrained-bert` |
+## Quick tour
+
+Let's do a very quick overview of the model architectures in 🤗 Transformers. Detailed examples for each model architecture (Bert, GPT, GPT-2, Transformer-XL, XLNet and XLM) can be found in the [full documentation](https://huggingface.co/transformers/).
+
+```python
+import torch
+from transformers import *
+
+# Transformers has a unified API
+# for 10 transformer architectures and 30 pretrained weights.
+#          Model          | Tokenizer          | Pretrained weights shortcut
+MODELS = [(BertModel,       BertTokenizer,       'bert-base-uncased'),
+          (OpenAIGPTModel,  OpenAIGPTTokenizer,  'openai-gpt'),
+          (GPT2Model,       GPT2Tokenizer,       'gpt2'),
+          (CTRLModel,       CTRLTokenizer,       'ctrl'),
+          (TransfoXLModel,  TransfoXLTokenizer,  'transfo-xl-wt103'),
+          (XLNetModel,      XLNetTokenizer,      'xlnet-base-cased'),
+          (XLMModel,        XLMTokenizer,        'xlm-mlm-enfr-1024'),
+          (DistilBertModel, DistilBertTokenizer, 'distilbert-base-cased'),
+          (RobertaModel,    RobertaTokenizer,    'roberta-base'),
+          (XLMRobertaModel, XLMRobertaTokenizer, 'xlm-roberta-base'),
+         ]
+
+# To use TensorFlow 2.0 versions of the models, simply prefix the class names with 'TF', e.g. `TFRobertaModel` is the TF 2.0 counterpart of the PyTorch model `RobertaModel`
+
+# Let's encode some text in a sequence of hidden-states using each model:
+for model_class, tokenizer_class, pretrained_weights in MODELS:
+    # Load pretrained model/tokenizer
+    tokenizer = tokenizer_class.from_pretrained(pretrained_weights)
+    model = model_class.from_pretrained(pretrained_weights)
+
+    # Encode text
+    input_ids = torch.tensor([tokenizer.encode("Here is some text to encode", add_special_tokens=True)])  # Add special tokens takes care of adding [CLS], [SEP], <s>... tokens in the right way for each model.
+    with torch.no_grad():
+        last_hidden_states = model(input_ids)[0]  # Models outputs are now tuples
+
+# Each architecture is provided with several class for fine-tuning on down-stream tasks, e.g.
+BERT_MODEL_CLASSES = [BertModel, BertForPreTraining, BertForMaskedLM, BertForNextSentencePrediction,
+                      BertForSequenceClassification, BertForTokenClassification, BertForQuestionAnswering]
+
+# All the classes for an architecture can be initiated from pretrained weights for this architecture
+# Note that additional weights added for fine-tuning are only initialized
+# and need to be trained on the down-stream task
+pretrained_weights = 'bert-base-uncased'
+tokenizer = BertTokenizer.from_pretrained(pretrained_weights)
+for model_class in BERT_MODEL_CLASSES:
+    # Load pretrained model/tokenizer
+    model = model_class.from_pretrained(pretrained_weights)
+
+    # Models can return full list of hidden-states & attentions weights at each layer
+    model = model_class.from_pretrained(pretrained_weights,
+                                        output_hidden_states=True,
+                                        output_attentions=True)
+    input_ids = torch.tensor([tokenizer.encode("Let's see all hidden-states and attentions on this text")])
+    all_hidden_states, all_attentions = model(input_ids)[-2:]
+
+    # Models are compatible with Torchscript
+    model = model_class.from_pretrained(pretrained_weights, torchscript=True)
+    traced_model = torch.jit.trace(model, (input_ids,))
+
+    # Simple serialization for models and tokenizers
+    model.save_pretrained('./directory/to/save/')  # save
+    model = model_class.from_pretrained('./directory/to/save/')  # re-load
+    tokenizer.save_pretrained('./directory/to/save/')  # save
+    tokenizer = BertTokenizer.from_pretrained('./directory/to/save/')  # re-load
+
+    # SOTA examples for GLUE, SQUAD, text generation...
+```
+
+## Quick tour TF 2.0 training and PyTorch interoperability
+
+Let's do a quick example of how a TensorFlow 2.0 model can be trained in 12 lines of code with 🤗 Transformers and then loaded in PyTorch for fast inspection/tests.
+
+```python
+import tensorflow as tf
+import tensorflow_datasets
+from transformers import *
+
+# Load dataset, tokenizer, model from pretrained model/vocabulary
+tokenizer = BertTokenizer.from_pretrained('bert-base-cased')
+model = TFBertForSequenceClassification.from_pretrained('bert-base-cased')
+data = tensorflow_datasets.load('glue/mrpc')
+
+# Prepare dataset for GLUE as a tf.data.Dataset instance
+train_dataset = glue_convert_examples_to_features(data['train'], tokenizer, max_length=128, task='mrpc')
+valid_dataset = glue_convert_examples_to_features(data['validation'], tokenizer, max_length=128, task='mrpc')
+train_dataset = train_dataset.shuffle(100).batch(32).repeat(2)
+valid_dataset = valid_dataset.batch(64)
+
+# Prepare training: Compile tf.keras model with optimizer, loss and learning rate schedule
+optimizer = tf.keras.optimizers.Adam(learning_rate=3e-5, epsilon=1e-08, clipnorm=1.0)
+loss = tf.keras.losses.SparseCategoricalCrossentropy(from_logits=True)
+metric = tf.keras.metrics.SparseCategoricalAccuracy('accuracy')
+model.compile(optimizer=optimizer, loss=loss, metrics=[metric])
+
+# Train and evaluate using tf.keras.Model.fit()
+history = model.fit(train_dataset, epochs=2, steps_per_epoch=115,
+                    validation_data=valid_dataset, validation_steps=7)
+
+# Load the TensorFlow model in PyTorch for inspection
+model.save_pretrained('./save/')
+pytorch_model = BertForSequenceClassification.from_pretrained('./save/', from_tf=True)
+
+# Quickly test a few predictions - MRPC is a paraphrasing task, let's see if our model learned the task
+sentence_0 = "This research was consistent with his findings."
+sentence_1 = "His findings were compatible with this research."
+sentence_2 = "His findings were not compatible with this research."
+inputs_1 = tokenizer(sentence_0, sentence_1, add_special_tokens=True, return_tensors='pt')
+inputs_2 = tokenizer(sentence_0, sentence_2, add_special_tokens=True, return_tensors='pt')
+
+pred_1 = pytorch_model(inputs_1['input_ids'], token_type_ids=inputs_1['token_type_ids'])[0].argmax().item()
+pred_2 = pytorch_model(inputs_2['input_ids'], token_type_ids=inputs_2['token_type_ids'])[0].argmax().item()
+
+print("sentence_1 is", "a paraphrase" if pred_1 else "not a paraphrase", "of sentence_0")
+print("sentence_2 is", "a paraphrase" if pred_2 else "not a paraphrase", "of sentence_0")
+```
+
+## Quick tour of the fine-tuning/usage scripts
+
+**Important**
+Before running the fine-tuning scripts, please read the
+[instructions](#run-the-examples) on how to
+setup your environment to run the examples.
+
+The library comprises several example scripts with SOTA performances for NLU and NLG tasks:
+
+- `run_glue.py`: an example fine-tuning sequence classification models on nine different GLUE tasks (*sequence-level classification*)
+- `run_squad.py`: an example fine-tuning question answering models on the question answering dataset SQuAD 2.0 (*token-level classification*)
+- `run_ner.py`: an example fine-tuning token classification models on named entity recognition (*token-level classification*)
+- `run_generation.py`: an example using GPT, GPT-2, CTRL, Transformer-XL and XLNet for conditional language generation
+- other model-specific examples (see the documentation).
+
+Here are three quick usage examples for these scripts:
+
+### `run_glue.py`: Fine-tuning on GLUE tasks for sequence classification
+
+The [General Language Understanding Evaluation (GLUE) benchmark](https://gluebenchmark.com/) is a collection of nine sentence- or sentence-pair language understanding tasks for evaluating and analyzing natural language understanding systems.
+
+Before running any of these GLUE tasks you should download the
+[GLUE data](https://gluebenchmark.com/tasks) by running
+[this script](https://gist.github.com/W4ngatang/60c2bdb54d156a41194446737ce03e2e)
+and unpack it to some directory `$GLUE_DIR`.
+
+You should also install the additional packages required by the examples:
+
+```shell
+pip install -r ./examples/requirements.txt
+```
+
+```shell
+export GLUE_DIR=/path/to/glue
+export TASK_NAME=MRPC
+
+python ./examples/text-classification/run_glue.py \
+    --model_name_or_path bert-base-uncased \
+    --task_name $TASK_NAME \
+    --do_train \
+    --do_eval \
+    --data_dir $GLUE_DIR/$TASK_NAME \
+    --max_seq_length 128 \
+    --per_device_eval_batch_size=8   \
+    --per_device_train_batch_size=8   \
+    --learning_rate 2e-5 \
+    --num_train_epochs 3.0 \
+    --output_dir /tmp/$TASK_NAME/
+```
+
+where task name can be one of CoLA, SST-2, MRPC, STS-B, QQP, MNLI, QNLI, RTE, WNLI.
+
+The dev set results will be present within the text file 'eval_results.txt' in the specified output_dir. In case of MNLI, since there are two separate dev sets, matched and mismatched, there will be a separate output folder called '/tmp/MNLI-MM/' in addition to '/tmp/MNLI/'.
+
+#### Fine-tuning XLNet model on the STS-B regression task
+
+This example code fine-tunes XLNet on the STS-B corpus using parallel training on a server with 4 V100 GPUs.
+Parallel training is a simple way to use several GPUs (but is slower and less flexible than distributed training, see below).
+
+```shell
+export GLUE_DIR=/path/to/glue
+
+python ./examples/text-classification/run_glue.py \
+    --model_name_or_path xlnet-large-cased \
+    --do_train  \
+    --do_eval   \
+    --task_name=sts-b     \
+    --data_dir=${GLUE_DIR}/STS-B  \
+    --output_dir=./proc_data/sts-b-110   \
+    --max_seq_length=128   \
+    --per_device_eval_batch_size=8   \
+    --per_device_train_batch_size=8   \
+    --gradient_accumulation_steps=1 \
+    --max_steps=1200  \
+    --model_name=xlnet-large-cased   \
+    --overwrite_output_dir   \
+    --overwrite_cache \
+    --warmup_steps=120
+```
+
+On this machine we thus have a batch size of 32, please increase `gradient_accumulation_steps` to reach the same batch size if you have a smaller machine. These hyper-parameters should result in a Pearson correlation coefficient of `+0.917` on the development set.
+
+#### Fine-tuning Bert model on the MRPC classification task
+
+This example code fine-tunes the Bert Whole Word Masking model on the Microsoft Research Paraphrase Corpus (MRPC) corpus using distributed training on 8 V100 GPUs to reach a F1 > 92.
+
+```bash
+python -m torch.distributed.launch --nproc_per_node 8 ./examples/text-classification/run_glue.py   \
+    --model_name_or_path bert-large-uncased-whole-word-masking \
+    --task_name MRPC \
+    --do_train   \
+    --do_eval   \
+    --data_dir $GLUE_DIR/MRPC/   \
+    --max_seq_length 128   \
+    --per_device_eval_batch_size=8   \
+    --per_device_train_batch_size=8   \
+    --learning_rate 2e-5   \
+    --num_train_epochs 3.0  \
+    --output_dir /tmp/mrpc_output/ \
+    --overwrite_output_dir   \
+    --overwrite_cache \
+```
+
+Training with these hyper-parameters gave us the following results:
+
+```bash
+  acc = 0.8823529411764706
+  acc_and_f1 = 0.901702786377709
+  eval_loss = 0.3418912578906332
+  f1 = 0.9210526315789473
+  global_step = 174
+  loss = 0.07231863956341798
+```
+
+### `run_squad.py`: Fine-tuning on SQuAD for question-answering
+
+This example code fine-tunes BERT on the SQuAD dataset using distributed training on 8 V100 GPUs and Bert Whole Word Masking uncased model to reach a F1 > 93 on SQuAD:
+
+```bash
+python -m torch.distributed.launch --nproc_per_node=8 ./examples/question-answering/run_squad.py \
+    --model_type bert \
+    --model_name_or_path bert-large-uncased-whole-word-masking \
+    --do_train \
+    --do_eval \
+    --train_file $SQUAD_DIR/train-v1.1.json \
+    --predict_file $SQUAD_DIR/dev-v1.1.json \
+    --learning_rate 3e-5 \
+    --num_train_epochs 2 \
+    --max_seq_length 384 \
+    --doc_stride 128 \
+    --output_dir ../models/wwm_uncased_finetuned_squad/ \
+    --per_device_eval_batch_size=3   \
+    --per_device_train_batch_size=3   \
+```
+
+Training with these hyper-parameters gave us the following results:
+
+```bash
+python $SQUAD_DIR/evaluate-v1.1.py $SQUAD_DIR/dev-v1.1.json ../models/wwm_uncased_finetuned_squad/predictions.json
+{"exact_match": 86.91579943235573, "f1": 93.1532499015869}
+```
+
+This is the model provided as `bert-large-uncased-whole-word-masking-finetuned-squad`.
+
+### `run_generation.py`: Text generation with GPT, GPT-2, CTRL, Transformer-XL and XLNet
+
+A conditional generation script is also included to generate text from a prompt.
+The generation script includes the [tricks](https://github.com/rusiaaman/XLNet-gen#methodology) proposed by Aman Rusia to get high-quality generation with memory models like Transformer-XL and XLNet (include a predefined text to make short inputs longer).
+
+Here is how to run the script with the small version of OpenAI GPT-2 model:
+
+```shell
+python ./examples/text-generation/run_generation.py \
+    --model_type=gpt2 \
+    --length=20 \
+    --model_name_or_path=gpt2 \
+```
+
+and from the Salesforce CTRL model:
+```shell
+python ./examples/text-generation/run_generation.py \
+    --model_type=ctrl \
+    --length=20 \
+    --model_name_or_path=ctrl \
+    --temperature=0 \
+    --repetition_penalty=1.2 \
+```
+
+## Quick tour of model sharing
+
+Starting with `v2.2.2`, you can now upload and share your fine-tuned models with the community, using the <abbr title="Command-line interface">CLI</abbr> that's built-in to the library.
+
+**First, create an account on [https://huggingface.co/join](https://huggingface.co/join)**. Optionally, join an existing organization or create a new one. Then:
+
+```shell
+transformers-cli login
+# log in using the same credentials as on huggingface.co
+```
+Upload your model:
+```shell
+transformers-cli upload ./path/to/pretrained_model/
+
+# ^^ Upload folder containing weights/tokenizer/config
+# saved via `.save_pretrained()`
+
+transformers-cli upload ./config.json [--filename folder/foobar.json]
+
+# ^^ Upload a single file
+# (you can optionally override its filename, which can be nested inside a folder)
+```
+
+If you want your model to be namespaced by your organization name rather than your username, add the following flag to any command:
+```shell
+--organization organization_name
+```
+
+Your model will then be accessible through its identifier, a concatenation of your username (or organization name) and the folder name above:
+```python
+"username/pretrained_model"
+# or if an org:
+"organization_name/pretrained_model"
+```
+
+**Please add a README.md model card** to the repo under `model_cards/` with: model description, training params (dataset, preprocessing, hardware used, hyperparameters), evaluation results, intended uses & limitations, etc.
+
+Your model now has a page on huggingface.co/models 🔥
+
+Anyone can load it from code:
+```python
+tokenizer = AutoTokenizer.from_pretrained("namespace/pretrained_model")
+model = AutoModel.from_pretrained("namespace/pretrained_model")
+```
+
+List all your files on S3:
+```shell
+transformers-cli s3 ls
+```
+
+You can also delete unneeded files:
+
+```shell
+transformers-cli s3 rm …
+```
+
+## Quick tour of pipelines
+
+New in version `v2.3`: `Pipeline` are high-level objects which automatically handle tokenization, running your data through a transformers model
+and outputting the result in a structured object.
+
+You can create `Pipeline` objects for the following down-stream tasks:
+
+ - `feature-extraction`: Generates a tensor representation for the input sequence
+ - `ner`: Generates named entity mapping for each word in the input sequence.
+ - `sentiment-analysis`: Gives the polarity (positive / negative) of the whole input sequence.
+ - `text-classification`: Initialize a `TextClassificationPipeline` directly, or see `sentiment-analysis` for an example.
+ - `question-answering`: Provided some context and a question refering to the context, it will extract the answer to the question in the context.
+ - `fill-mask`: Takes an input sequence containing a masked token (e.g. `<mask>`) and return list of most probable filled sequences, with their probabilities.
+ - `summarization`
+ - `translation_xx_to_yy`
+
+```python
+>>> from transformers import pipeline
+
+# Allocate a pipeline for sentiment-analysis
+>>> nlp = pipeline('sentiment-analysis')
+>>> nlp('We are very happy to include pipeline into the transformers repository.')
+[{'label': 'POSITIVE', 'score': 0.9978193640708923}]
+
+# Allocate a pipeline for question-answering
+>>> nlp = pipeline('question-answering')
+>>> nlp({
+...     'question': 'What is the name of the repository ?',
+...     'context': 'Pipeline have been included in the huggingface/transformers repository'
+... })
+{'score': 0.5135612454720828, 'start': 35, 'end': 59, 'answer': 'huggingface/transformers'}
+
+```
+
+## Migrating from pytorch-transformers to transformers
+
+Here is a quick summary of what you should take care of when migrating from `pytorch-transformers` to `transformers`.
+
+### Positional order of some models' keywords inputs (`attention_mask`, `token_type_ids`...) changed
+
+To be able to use Torchscript (see #1010, #1204 and #1195) the specific order of some models **keywords inputs** (`attention_mask`, `token_type_ids`...) has been changed.
+
+If you used to call the models with keyword names for keyword arguments, e.g. `model(inputs_ids, attention_mask=attention_mask, token_type_ids=token_type_ids)`, this should not cause any change.
+
+If you used to call the models with positional inputs for keyword arguments, e.g. `model(inputs_ids, attention_mask, token_type_ids)`, you may have to double check the exact order of input arguments.
+
+
+## Migrating from pytorch-pretrained-bert to transformers
+
+Here is a quick summary of what you should take care of when migrating from `pytorch-pretrained-bert` to `transformers`.
+
+### Models always output `tuples`
+
+The main breaking change when migrating from `pytorch-pretrained-bert` to `transformers` is that every model's forward method always outputs a `tuple` with various elements depending on the model and the configuration parameters.
+
+The exact content of the tuples for each model is detailed in the models' docstrings and the [documentation](https://huggingface.co/transformers/).
+
+In pretty much every case, you will be fine by taking the first element of the output as the output you previously used in `pytorch-pretrained-bert`.
+
+Here is a `pytorch-pretrained-bert` to `transformers` conversion example for a `BertForSequenceClassification` classification model:
+
+```python
+# Let's load our model
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+
+# If you used to have this line in pytorch-pretrained-bert:
+loss = model(input_ids, labels=labels)
+
+# Now just use this line in transformers to extract the loss from the output tuple:
+outputs = model(input_ids, labels=labels)
+loss = outputs[0]
+
+# In transformers you can also have access to the logits:
+loss, logits = outputs[:2]
+
+# And even the attention weights if you configure the model to output them (and other outputs too, see the docstrings and documentation)
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased', output_attentions=True)
+outputs = model(input_ids, labels=labels)
+loss, logits, attentions = outputs
+```
+
+### Using hidden states
+
+By enabling the configuration option `output_hidden_states`, it was possible to retrieve the last hidden states of the encoder. In `pytorch-transformers` as well as `transformers` the return value has changed slightly: `all_hidden_states` now also includes the hidden state of the embeddings in addition to those of the encoding layers. This allows users to easily access the embeddings final state.
+
+### Serialization
+
+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 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.
+
+Here is an example:
+
+```python
+### Let's load a model and tokenizer
+model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+
+### Do some stuff to our model and tokenizer
+# Ex: add new tokens to the vocabulary and embeddings of our model
+tokenizer.add_tokens(['[SPECIAL_TOKEN_1]', '[SPECIAL_TOKEN_2]'])
+model.resize_token_embeddings(len(tokenizer))
+# Train our model
+train(model)
+
+### Now let's save our model and tokenizer to a directory
+model.save_pretrained('./my_saved_model_directory/')
+tokenizer.save_pretrained('./my_saved_model_directory/')
+
+### Reload the model and the tokenizer
+model = BertForSequenceClassification.from_pretrained('./my_saved_model_directory/')
+tokenizer = BertTokenizer.from_pretrained('./my_saved_model_directory/')
+```
+
+### Optimizers: BertAdam & OpenAIAdam are now AdamW, schedules are standard PyTorch schedules
+
+The two optimizers previously included, `BertAdam` and `OpenAIAdam`, have been replaced by a single `AdamW` optimizer which has a few differences:
+
+- it only implements weights decay correction,
+- schedules are now externals (see below),
+- gradient clipping is now also external (see below).
+
+The new optimizer `AdamW` matches PyTorch `Adam` optimizer API and let you use standard PyTorch or apex methods for the schedule and clipping.
+
+The schedules are now standard [PyTorch learning rate schedulers](https://pytorch.org/docs/stable/optim.html#how-to-adjust-learning-rate) and not part of the optimizer anymore.
+
+Here is a conversion examples from `BertAdam` with a linear warmup and decay schedule to `AdamW` and the same schedule:
+
+```python
+# Parameters:
+lr = 1e-3
+max_grad_norm = 1.0
+num_training_steps = 1000
+num_warmup_steps = 100
+warmup_proportion = float(num_warmup_steps) / float(num_training_steps)  # 0.1
+
+### Previously BertAdam optimizer was instantiated like this:
+optimizer = BertAdam(model.parameters(), lr=lr, schedule='warmup_linear', warmup=warmup_proportion, t_total=num_training_steps)
+### and used like this:
+for batch in train_data:
+    loss = model(batch)
+    loss.backward()
+    optimizer.step()
+
+### In Transformers, optimizer and schedules are splitted and instantiated like this:
+optimizer = AdamW(model.parameters(), lr=lr, correct_bias=False)  # To reproduce BertAdam specific behavior set correct_bias=False
+scheduler = get_linear_schedule_with_warmup(optimizer, num_warmup_steps=num_warmup_steps, num_training_steps=num_training_steps)  # PyTorch scheduler
+### and used like this:
+for batch in train_data:
+    model.train()
+    loss = model(batch)
+    loss.backward()
+    torch.nn.utils.clip_grad_norm_(model.parameters(), max_grad_norm)  # Gradient clipping is not in AdamW anymore (so you can use amp without issue)
+    optimizer.step()
+    scheduler.step()
+    optimizer.zero_grad()
+```

 ## 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}
@@ -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",
@@ -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 ---------------------------------------------------
@@ -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
@@ -3,9 +3,9 @@ Transformers

 State-of-the-art Natural Language Processing for Pytorch and TensorFlow 2.0.

-🤗 Transformers (formerly known as `pytorch-transformers` and `pytorch-pretrained-bert`) provides general-purpose
-architectures (BERT, GPT-2, RoBERTa, XLM, DistilBert, XLNet...) for Natural Language Understanding (NLU) and Natural
-Language Generation (NLG) with over 32+ pretrained models in 100+ languages and deep interoperability between
+🤗 Transformers (formerly known as `pytorch-transformers` and `pytorch-pretrained-bert`) provides general-purpose 
+architectures (BERT, GPT-2, RoBERTa, XLM, DistilBert, XLNet...) for Natural Language Understanding (NLU) and Natural 
+Language Generation (NLG) with over 32+ pretrained models in 100+ languages and deep interoperability between 
 TensorFlow 2.0 and PyTorch.

 This is the documentation of our repository `transformers <https://github.com/huggingface/transformers>`_.
@@ -127,17 +127,8 @@ conversion utilities for the following models:
 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. `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
+    Marjan Ghazvininejad, Mike Lewis, Luke Zettlemoyer.  
+25. `Other community models <https://huggingface.co/models>`_, contributed by the `community
    <https://huggingface.co/users>`_.

 .. toctree::
@@ -172,7 +163,6 @@ conversion utilities for the following models:
    converting_tensorflow_models
    migration
    contributing
-    testing
    serialization

 .. toctree::
@@ -195,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
@@ -223,10 +212,6 @@ conversion utilities for the following models:
    model_doc/dpr
    model_doc/pegasus
    model_doc/mbart
-    model_doc/fsmt
-    model_doc/funnel
-    model_doc/lxmert
-    model_doc/bertgeneration
    internal/modeling_utils
    internal/tokenization_utils
    internal/pipelines_utils
@@ -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:
@@ -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
@@ -13,11 +13,6 @@ The ``.optimization`` module provides:
 .. autoclass:: transformers.AdamW
    :members:

-``AdaFactor`` (PyTorch)
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.Adafactor
-
 ``AdamWeightDecay`` (TensorFlow)
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -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`
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -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`` 
 ~~~~~~~~~~~
@@ -1,131 +1,109 @@
-AutoClasses
+AutoModels
 -----------

 In many cases, the architecture you want to use can be guessed from the name or the path of the pretrained model you
-are supplying to the :obj:`from_pretrained()` method.
+are supplying to the ``from_pretrained`` method.
+
 AutoClasses are here to do this job for you so that you automatically retrieve the relevant model given the name/path
-to the pretrained weights/config/vocabulary.
+to the pretrained weights/config/vocabulary:

-Instantiating one of :class:`~transformers.AutoConfig`, :class:`~transformers.AutoModel`, and
-:class:`~transformers.AutoTokenizer` will directly create a class of the relevant architecture. For instance
+Instantiating one of ``AutoModel``, ``AutoConfig`` and ``AutoTokenizer`` will directly create a class of the relevant
+architecture (ex: ``model = AutoModel.from_pretrained('bert-base-cased')`` will create a instance of 
+:class:`~transformers.BertModel`).


-.. code-block:: python
-
-    model = AutoModel.from_pretrained('bert-base-cased')
-
-will create a model that is an instance of :class:`~transformers.BertModel`.
-
-There is one class of :obj:`AutoModel` for each task, and for each backend (PyTorch or TensorFlow).
-
-
-AutoConfig
-~~~~~~~~~~
+``AutoConfig``
+~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.AutoConfig
    :members:


-AutoTokenizer
-~~~~~~~~~~~~~
+``AutoTokenizer``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.AutoTokenizer
    :members:


-AutoModel
-~~~~~~~~~
+``AutoModel``
+~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.AutoModel
    :members:


-AutoModelForPreTraining
-~~~~~~~~~~~~~~~~~~~~~~~
+``AutoModelForPreTraining``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.AutoModelForPreTraining
    :members:


-AutoModelWithLMHead
-~~~~~~~~~~~~~~~~~~~
+``AutoModelWithLMHead``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.AutoModelWithLMHead
    :members:


-AutoModelForSequenceClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+``AutoModelForSequenceClassification``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.AutoModelForSequenceClassification
    :members:


-AutoModelForMultipleChoice
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.AutoModelForMultipleChoice
-    :members:
-
-
-AutoModelForTokenClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.AutoModelForTokenClassification
-    :members:
-
-
-AutoModelForQuestionAnswering
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+``AutoModelForQuestionAnswering``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.AutoModelForQuestionAnswering
    :members:


-TFAutoModel
-~~~~~~~~~~~
+``AutoModelForTokenClassification``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.AutoModelForTokenClassification
+    :members:
+
+``TFAutoModel``
+~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFAutoModel
    :members:


-TFAutoModelForPreTraining
-~~~~~~~~~~~~~~~~~~~~~~~~~
+``TFAutoModelForPreTraining``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFAutoModelForPreTraining
    :members:


-TFAutoModelWithLMHead
-~~~~~~~~~~~~~~~~~~~~~
+``TFAutoModelWithLMHead``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFAutoModelWithLMHead
    :members:


-TFAutoModelForSequenceClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+``TFAutoModelForSequenceClassification``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFAutoModelForSequenceClassification
    :members:


-TFAutoModelForMultipleChoice
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFAutoModelForMultipleChoice
-    :members:
-
-
-TFAutoModelForTokenClassification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.TFAutoModelForTokenClassification
-    :members:
-
-
-TFAutoModelForQuestionAnswering
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+``TFAutoModelForQuestionAnswering``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFAutoModelForQuestionAnswering
    :members:
+
+
+``TFAutoModelForTokenClassification``
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.TFAutoModelForTokenClassification
+    :members:
@@ -32,7 +32,7 @@ BartForConditionalGeneration
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.BartForConditionalGeneration
-    :members: forward
+    :members: generate, forward


 BartConfig
@@ -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:
@@ -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/1908.08345>`_ 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``
@@ -1,49 +0,0 @@
-FSMT
----------------------------------------------------
-**DISCLAIMER:** If you see something strange,
-file a `Github Issue <https://github.com/huggingface/transformers/issues/new?assignees=&labels=&template=bug-report.md&title>`__ and assign
-@stas00.
-
-Overview
-~~~~~~~~~~~~~~~~~~~~~
-
-FSMT (FairSeq MachineTranslation) models were introduced in "Facebook FAIR's WMT19 News Translation Task Submission" <this paper <https://arxiv.org/abs/1907.06616>__ by Nathan Ng, Kyra Yee, Alexei Baevski, Myle Ott, Michael Auli, Sergey Edunov.
-
-The abstract of the paper is the following:
-
-    This paper describes Facebook FAIR's submission to the WMT19 shared news translation task. We participate in two language pairs and four language directions, English <-> German and English <-> Russian. Following our submission from last year, our baseline systems are large BPE-based transformer models trained with the Fairseq sequence modeling toolkit which rely on sampled back-translations. This year we experiment with different bitext data filtering schemes, as well as with adding filtered back-translated data. We also ensemble and fine-tune our models on domain-specific data, then decode using noisy channel model reranking. Our submissions are ranked first in all four directions of the human evaluation campaign. On En->De, our system significantly outperforms other systems as well as human translations. This system improves upon our WMT'18 submission by 4.5 BLEU points.
-
-The original code can be found here <https://github.com/pytorch/fairseq/tree/master/examples/wmt19>__.
-
-Implementation Notes
-~~~~~~~~~~~~~~~~~~~~
-
- FSMT uses source and target vocab pair, that aren't combined into one. It doesn't share embed tokens either. Its tokenizer is very similar to `XLMTokenizer` and the main model is derived from `BartModel`.
-
-
-FSMTForConditionalGeneration
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FSMTForConditionalGeneration
-    :members: forward
-
-
-FSMTConfig
-~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FSMTConfig
-    :members:
-
-
-FSMTTokenizer
-~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FSMTTokenizer
-    :members:
-
-
-FSMTModel
-~~~~~~~~~~~~~
-
-.. autoclass:: transformers.FSMTModel
-    :members: forward
@@ -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:
@@ -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:
@@ -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
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
@@ -56,13 +56,6 @@ XLMRobertaModel
    :members:


-XLMRobertaForCausalLM
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-.. autoclass:: transformers.XLMRobertaForCausalLM
-    :members:
-
-
 XLMRobertaForMaskedLM
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

@@ -137,4 +130,4 @@ TFXLMRobertaForQuestionAnswering
 ~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~

 .. autoclass:: transformers.TFXLMRobertaForQuestionAnswering
-    :members:
+    :members:
@@ -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:

@@ -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
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^

@@ -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.                                                               |
+-------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
@@ -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::

@@ -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
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
@@ -1,943 +0,0 @@
-Testing
-==========
-
-
-Let's take a look at how 🤗 Transformer models are tested and how you can write new tests and improve the existing ones.
-
-There are 2 test suites in the repository:
-
-1. ``tests`` -- tests for the general API
-2. ``examples`` -- tests primarily for various applications that aren't part of the API
-
-How transformers are tested
---------------------------
-
-1. Once a PR is submitted it gets tested with 9 CircleCi jobs. Every new commit to that PR gets retested. These jobs are defined in this `config file <https://github.com/huggingface/transformers/blob/master/.circleci/config.yml>`__, so that if needed you can reproduce the same environment on your machine.
-   
-   These CI jobs don't run ``@slow`` tests.
-   
-2. There are 3 jobs run by `github actions <https://github.com/huggingface/transformers/actions>`__:
-
-   * `torch hub integration <https://github.com/huggingface/transformers/blob/master/.github/workflows/github-torch-hub.yml>`__:  checks whether torch hub integration works.
-
-   * `self-hosted (push) <https://github.com/huggingface/transformers/blob/master/.github/workflows/self-push.yml>`__: runs fast tests on GPU only on commits on ``master``. It only runs if a commit on ``master`` has updated the code in one of the following folders: ``src``, ``tests``, ``.github`` (to prevent running on added model cards, notebooks, etc.)
-     
-   * `self-hosted runner <https://github.com/huggingface/transformers/blob/master/.github/workflows/self-scheduled.yml>`__: runs slow tests on ``tests`` and ``examples``:
-
-   .. code-block:: bash
-
-    RUN_SLOW=1 USE_CUDA=1 pytest tests/
-    RUN_SLOW=1 USE_CUDA=1 pytest examples/
-
-   The results can be observed `here <https://github.com/huggingface/transformers/actions>`__.
-
-
-
-Running tests
-------------
-
-
-
-
-
-Choosing which tests to run
-~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-This document goes into many details of how tests can be run. If after reading everything, you need even more details you will find them `here <https://docs.pytest.org/en/latest/usage.html>`__.
-
-Here are some most useful ways of running tests.
-
-Run all:
-
-.. code-block:: console
-
-   pytest
-
-or:
-
-.. code-block:: bash
-
-   make test
-
-Note that the latter is defined as:
-
-.. code-block:: bash
-
-   python -m pytest -n auto --dist=loadfile -s -v ./tests/
-
-which tells pytest to:
-
-* run as many test processes as they are CPU cores (which could be too many if you don't have a ton of RAM!)
-* ensure that all tests from the same file will be run by the same test process
-* do not capture output
-* run in verbose mode
-
-
-
-Getting the list of all tests
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-All tests of the test suite:
-
-.. code-block:: bash
-
-   pytest --collect-only -q
-
-All tests of a given test file:
-
-.. code-block:: bash
-
-   pytest tests/test_optimization.py --collect-only -q
-
-
-   
-Run a specific test module
-~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-To run an individual test module:
-
-.. code-block:: bash
-
-   pytest tests/test_logging.py
-   
-
-Run specific tests
-~~~~~~~~~~~~~~~~~~
-
-Since unittest is used inside most of the tests, to run specific subtests you need to know the name of the unittest class containing those tests. For example, it could be:
-
-.. code-block:: bash
-
-   pytest tests/test_optimization.py::OptimizationTest::test_adam_w
-
-Here:
-
-* ``tests/test_optimization.py`` - the file with tests
-* ``OptimizationTest`` - the name of the class
-* ``test_adam_w`` - the name of the specific test function
-
-If the file contains multiple classes, you can choose to run only tests of a given class. For example:
-
-.. code-block:: bash
-
-   pytest tests/test_optimization.py::OptimizationTest
-
-
-will run all the tests inside that class.
-
-As mentioned earlier you can see what tests are contained inside the ``OptimizationTest`` class by running:
-
-.. code-block:: bash
-
-   pytest tests/test_optimization.py::OptimizationTest --collect-only -q
-
-  
-You can run tests by keyword expressions.
-
-To run only tests whose name contains ``adam``:
-
-.. code-block:: bash
-
-   pytest -k adam tests/test_optimization.py
-
-To run all tests except those whose name contains ``adam``:
-
-.. code-block:: bash
-
-   pytest -k "not adam" tests/test_optimization.py
-
-And you can combine the two patterns in one:
-
-
-.. code-block:: bash
-
-   pytest -k "ada and not adam" tests/test_optimization.py
-
-
-
-Run only modified tests
-~~~~~~~~~~~~~~~~~~~~~~~
-
-You can run the tests related to the unstaged files or the current branch (according to Git) by using `pytest-picked <https://github.com/anapaulagomes/pytest-picked>`__. This is a great way of quickly testing your changes didn't break anything, since it won't run the tests related to files you didn't touch.
-
-.. code-block:: bash
-
-    pip install pytest-picked
-
-.. code-block:: bash
-
-    pytest --picked
-
-All tests will be run from files and folders which are modified, but not
-yet committed.
-
-Automatically rerun failed tests on source modification
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-`pytest-xdist <https://github.com/pytest-dev/pytest-xdist>`__ provides a
-very useful feature of detecting all failed tests, and then waiting for
-you to modify files and continuously re-rerun those failing tests until
-they pass while you fix them. So that you don't need to re start pytest
-after you made the fix. This is repeated until all tests pass after
-which again a full run is performed.
-
-.. code-block:: bash
-
-    pip install pytest-xdist
-
-To enter the mode: ``pytest -f`` or ``pytest --looponfail``
-
-File changes are detected by looking at ``looponfailroots`` root
-directories and all of their contents (recursively). If the default for
-this value does not work for you, you can change it in your project by
-setting a configuration option in ``setup.cfg``:
-
-.. code-block:: ini
-
-    [tool:pytest]
-    looponfailroots = transformers tests
-
-or ``pytest.ini``/``tox.ini`` files:
-
-.. code-block:: ini
-
-    [pytest]
-    looponfailroots = transformers tests
-
-This would lead to only looking for file changes in the respective
-directories, specified relatively to the ini-file’s directory.
-
-`pytest-watch <https://github.com/joeyespo/pytest-watch>`__ is an
-alternative implementation of this functionality.
-
-
-Skip a test module
-~~~~~~~~~~~~~~~~~~
-
-If you want to run all test modules, except a few you can exclude them by giving an explicit list of tests to run. For example, to run all except ``test_modeling_*.py`` tests:
-
-.. code-block:: bash
-
-   pytest `ls -1 tests/*py | grep -v test_modeling`
-
-
-Clearing state
-~~~~~~~~~~~~~~
-
-CI builds and when isolation is important (against speed), cache should
-be cleared:
-
-.. code-block:: bash
-
-    pytest --cache-clear tests
-
-Running tests in parallel
-~~~~~~~~~~~~~~~~~~~~~~~~~
-
-As mentioned earlier ``make test`` runs tests in parallel via ``pytest-xdist`` plugin (``-n X`` argument, e.g. ``-n 2`` to run 2 parallel jobs).
-
-``pytest-xdist``'s ``--dist=`` option allows one to control how the tests are grouped. ``--dist=loadfile`` puts the tests located in one file onto the same process.
-
-Since the order of executed tests is different and unpredictable, if
-running the test suite with ``pytest-xdist`` produces failures (meaning
-we have some undetected coupled tests), use
-`pytest-replay <https://github.com/ESSS/pytest-replay>`__ to replay the
-tests in the same order, which should help with then somehow reducing
-that failing sequence to a minimum.
-
-Test order and repetition
-~~~~~~~~~~~~~~~~~~~~~~~~~
-
-It's good to repeat the tests several times, in sequence, randomly, or
-in sets, to detect any potential inter-dependency and state-related bugs
-(tear down). And the straightforward multiple repetition is just good to
-detect some problems that get uncovered by randomness of DL.
-
-
-Repeat tests
-^^^^^^^^^^^^
-
-* `pytest-flakefinder <https://github.com/dropbox/pytest-flakefinder>`__:
-
-.. code-block:: bash
-
-   pip install pytest-flakefinder
-
-And then run every test multiple times (50 by default):
-
-.. code-block:: bash
-
-   pytest --flake-finder --flake-runs=5 tests/test_failing_test.py
-   
-.. note::
-   This plugin doesn't work with ``-n`` flag from ``pytest-xdist``.
-   
-.. note::
-   There is another plugin ``pytest-repeat``, but it doesn't work with ``unittest``.
-
-
-Run tests in a random order
-^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-.. code-block:: bash
-
-    pip install pytest-random-order
-
-Important: the presence of ``pytest-random-order`` will automatically
-randomize tests, no configuration change or command line options is
-required.
-
-As explained earlier this allows detection of coupled tests - where one
-test's state affects the state of another. When ``pytest-random-order``
-is installed it will print the random seed it used for that session,
-e.g:
-
-.. code-block:: bash
-
-   pytest tests
-   [...]
-   Using --random-order-bucket=module
-   Using --random-order-seed=573663
-
-So that if the given particular sequence fails, you can reproduce it by
-adding that exact seed, e.g.:
-
-.. code-block:: bash
-
-   pytest --random-order-seed=573663
-   [...]
-   Using --random-order-bucket=module
-   Using --random-order-seed=573663
-
-It will only reproduce the exact order if you use the exact same list of
-tests (or no list at all). Once you start to manually narrowing
-down the list you can no longer rely on the seed, but have to list them
-manually in the exact order they failed and tell pytest to not randomize
-them instead using ``--random-order-bucket=none``, e.g.:
-
-.. code-block:: bash
-
-   pytest --random-order-bucket=none tests/test_a.py tests/test_c.py tests/test_b.py
-
-To disable the shuffling for all tests:
-
-.. code-block:: bash
-
-    pytest --random-order-bucket=none
-
-By default ``--random-order-bucket=module`` is implied, which will
-shuffle the files on the module levels. It can also shuffle on
-``class``, ``package``, ``global`` and ``none`` levels. For the complete
-details please see its `documentation <https://github.com/jbasko/pytest-random-order>`__.
-
-Another randomization alternative is: ``pytest-randomly`` <https://github.com/pytest-dev/pytest-randomly>`__. This module has a very similar functionality/interface, but it doesn't have the bucket modes available in ``pytest-random-order``. It has the same problem of imposing itself once installed.
-
-Look and feel variations
-~~~~~~~~~~~~~~~~~~~~~~~~
-
-pytest-sugar
-^^^^^^^^^^^^
-
-`pytest-sugar <https://github.com/Frozenball/pytest-sugar>`__ is a
-plugin that improves the look-n-feel, adds a progressbar, and show tests
-that fail and the assert instantly. It gets activated automatically upon
-installation.
-
-.. code-block:: bash
-                
-   pip install pytest-sugar
-
-To run tests without it, run:
-
-.. code-block:: bash
-
-    pytest -p no:sugar
-
-or uninstall it.
-
-
-
-Report each sub-test name and its progress
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-For a single or a group of tests via ``pytest`` (after
-``pip install pytest-pspec``):
-
-.. code-block:: bash
-
-   pytest --pspec tests/test_optimization.py 
-
-
-
-Instantly shows failed tests
-^^^^^^^^^^^^^^^^^^^^^^^^^^^^
-
-`pytest-instafail <https://github.com/pytest-dev/pytest-instafail>`__
-shows failures and errors instantly instead of waiting until the end of
-test session.
-
-.. code-block:: bash
-
-    pip install pytest-instafail
-
-.. code-block:: bash
-
-    pytest --instafail
-
-To GPU or not to GPU
-~~~~~~~~~~~~~~~~~~~~
-
-On a GPU-enabled setup, to test in CPU-only mode add ``CUDA_VISIBLE_DEVICES=""``:
-
-.. code-block:: bash
-                
-    CUDA_VISIBLE_DEVICES="" pytest tests/test_logging.py
-
-or if you have multiple gpus, you can tell which one to use in this test session, e.g. to use only the second gpu if you have gpus ``0`` and ``1``, you can run:
-
-.. code-block:: bash
-                
-    CUDA_VISIBLE_DEVICES="1" pytest tests/test_logging.py
-
-This is handy when you want to run different tasks on different GPUs.
-    
-And we have these decorators that require the condition described by the marker.
-
-```
-@require_torch
-@require_tf
-@require_multigpu
-@require_non_multigpu
-@require_torch_tpu
-@require_torch_and_cuda
-```
-
-This section will be expanded soon once our work in progress on those decorators is finished.
-
-Inside tests:
-
-* How many GPUs are available:
-
-.. code-block:: bash
-
-   torch.cuda.device_count()
-
-
-   
-
-
-Output capture
-~~~~~~~~~~~~~~
-
-During test execution any output sent to ``stdout`` and ``stderr`` is
-captured. If a test or a setup method fails, its according captured
-output will usually be shown along with the failure traceback.
-
-To disable output capturing and to get the ``stdout`` and ``stderr``
-normally, use ``-s`` or ``--capture=no``:
-
-.. code-block:: bash
-
-   pytest -s tests/test_logging.py
-
-To send test results to JUnit format output:
-
-.. code-block:: bash
-
-   py.test tests --junitxml=result.xml
-
-
-Color control
-~~~~~~~~~~~~~
-
-To have no color (e.g., yellow on white background is not readable):
-
-.. code-block:: bash
-
-   pytest --color=no tests/test_logging.py
-
-
-
-Sending test report to online pastebin service
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Creating a URL for each test failure:
-
-.. code-block:: bash
-
-   pytest --pastebin=failed tests/test_logging.py
-
-This will submit test run information to a remote Paste service and
-provide a URL for each failure. You may select tests as usual or add for
-example -x if you only want to send one particular failure.
-
-Creating a URL for a whole test session log:
-
-.. code-block:: bash
-
-   pytest --pastebin=all tests/test_logging.py
-
-
-
-Writing tests
-------------
-
-🤗 transformers tests are based on ``unittest``, but run by ``pytest``, so most of the time features from both systems can be used.
-
-You can read `here <https://docs.pytest.org/en/stable/unittest.html>`__ which features are supported, but the important thing to remember is that most ``pytest`` fixtures don't work. Neither parametrization, but we use the module ``parameterized`` that works in a similar way.
-
-
-Parametrization
-~~~~~~~~~~~~~~~
-
-Often, there is a need to run the same test multiple times, but with different arguments. It could be done from within the test, but then there is no way of running that test for just one set of arguments.
-
-.. code-block:: python
-                
-    # test_this1.py
-    import unittest
-    from parameterized import parameterized
-    class TestMathUnitTest(unittest.TestCase):
-        @parameterized.expand([
-            ("negative", -1.5, -2.0),
-            ("integer", 1, 1.0),
-            ("large fraction", 1.6, 1),
-        ])
-        def test_floor(self, name, input, expected):
-            assert_equal(math.floor(input), expected)
-
-Now, by default this test will be run 3 times, each time with the last 3 arguments of ``test_floor`` being assigned the corresponding arguments in the parameter list.
-
-and you could run just the ``negative`` and ``integer`` sets of params with:
-
-.. code-block:: bash
-
-   pytest -k "negative and integer" tests/test_mytest.py
-
-or all but ``negative`` sub-tests, with:
-
-.. code-block:: bash
-
-   pytest -k "not negative" tests/test_mytest.py
-
-Besides using the ``-k`` filter that was just mentioned, you can find out the exact name of each sub-test and run any or all of them using their exact names. 
-        
-.. code-block:: bash
-                
-    pytest test_this1.py --collect-only -q
-
-and it will list:
-                
-.. code-block:: bash
-
-    test_this1.py::TestMathUnitTest::test_floor_0_negative
-    test_this1.py::TestMathUnitTest::test_floor_1_integer
-    test_this1.py::TestMathUnitTest::test_floor_2_large_fraction
-
-So now you can run just 2 specific sub-tests:
-
-.. code-block:: bash
-
-    pytest test_this1.py::TestMathUnitTest::test_floor_0_negative  test_this1.py::TestMathUnitTest::test_floor_1_integer
-   
-The module `parameterized <https://pypi.org/project/parameterized/>`__ which is already in the developer dependencies of ``transformers`` works for both: ``unittests`` and ``pytest`` tests.
-
-If, however, the test is not a ``unittest``, you may use ``pytest.mark.parametrize`` (or you may see it being used in some existing tests, mostly under ``examples``).
-
-Here is the same example, this time using ``pytest``'s ``parametrize`` marker:
-
-.. code-block:: python
-
-    # test_this2.py
-    import pytest
-    @pytest.mark.parametrize(
-        "name, input, expected",
-        [
-            ("negative", -1.5, -2.0),
-            ("integer", 1, 1.0),
-            ("large fraction", 1.6, 1),
-        ],
-    )
-    def test_floor(name, input, expected):
-        assert_equal(math.floor(input), expected)
-
-Same as with ``parameterized``, with ``pytest.mark.parametrize`` you can have a fine control over which sub-tests are run, if the ``-k`` filter doesn't do the job. Except, this parametrization function creates a slightly different set of names for the sub-tests. Here is what they look like:
-        
-.. code-block:: bash
-                
-    pytest test_this2.py --collect-only -q
-
-and it will list:
-                
-.. code-block:: bash
-
-    test_this2.py::test_floor[integer-1-1.0]
-    test_this2.py::test_floor[negative--1.5--2.0]
-    test_this2.py::test_floor[large fraction-1.6-1]       
-
-So now you can run just the specific test:
-
-.. code-block:: bash
-
-    pytest test_this2.py::test_floor[negative--1.5--2.0] test_this2.py::test_floor[integer-1-1.0]
-
-as in the previous example.
-
-    
-
-Temporary files and directories
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-Using unique temporary files and directories are essential for parallel test running, so that the tests won't overwrite each other's data. Also we want to get the temp files and directories removed at the end of each test that created them. Therefore, using packages like ``tempfile``, which address these needs is essential.
-
-However, when debugging tests, you need to be able to see what goes into the temp file or directory and you want to know it's exact path and not having it randomized on every test re-run.
-
-A helper class :obj:`transformers.test_utils.TestCasePlus` is best used for such purposes. It's a sub-class of :obj:`unittest.TestCase`, so we can easily inherit from it in the test modules.
-
-Here is an example of its usage:
-
-.. code-block:: python
-
-    from transformers.testing_utils import TestCasePlus
-    class ExamplesTests(TestCasePlus):
-    def test_whatever(self):
-        tmp_dir = self.get_auto_remove_tmp_dir()
-
-This code creates a unique temporary directory, and sets :obj:`tmp_dir` to its location.
-
-In this and all the following scenarios the temporary directory will be auto-removed at the end of test, unless ``after=False`` is passed to the helper function.
-
-* Create a temporary directory of my choice and delete it at the end - useful for debugging when you want to monitor a specific directory:
-
-.. code-block:: python
-
-    def test_whatever(self):
-        tmp_dir = self.get_auto_remove_tmp_dir(tmp_dir="./tmp/run/test")
-
-* Create a temporary directory of my choice and do not delete it at the end---useful for when you want to look at the temp results:
-
-.. code-block:: python
-
-    def test_whatever(self):
-        tmp_dir = self.get_auto_remove_tmp_dir(tmp_dir="./tmp/run/test", after=False)
-
-* Create a temporary directory of my choice and ensure to delete it right away---useful for when you disabled deletion in the previous test run and want to make sure the that temporary directory is empty before the new test is run:
-
-.. code-block:: python
-
-   def test_whatever(self):
-        tmp_dir = self.get_auto_remove_tmp_dir(tmp_dir="./tmp/run/test", before=True)
-
-.. note::
-   In order to run the equivalent of ``rm -r`` safely, only subdirs of the project repository checkout are allowed if an explicit obj:`tmp_dir` is used, so that by mistake no ``/tmp`` or similar important part of the filesystem will get nuked. i.e. please always pass paths that start with ``./``.
-
-.. note::
-   Each test can register multiple temporary directories and they all will get auto-removed, unless requested otherwise.
-
-
-Skipping tests
-~~~~~~~~~~~~~~
-
-This is useful when a bug is found and a new test is written, yet the
-bug is not fixed yet. In order to be able to commit it to the main
-repository we need make sure it's skipped during ``make test``.
-
-Methods:
-
-  A **skip** means that you expect your test to pass only if some
-   conditions are met, otherwise pytest should skip running the test
-   altogether. Common examples are skipping windows-only tests on
-   non-windows platforms, or skipping tests that depend on an external
-   resource which is not available at the moment (for example a
-   database).
-
-  A **xfail** means that you expect a test to fail for some reason. A
-   common example is a test for a feature not yet implemented, or a bug
-   not yet fixed. When a test passes despite being expected to fail
-   (marked with pytest.mark.xfail), it’s an xpass and will be reported
-   in the test summary.
-
-One of the important differences between the two is that ``skip``
-doesn't run the test, and ``xfail`` does. So if the code that's buggy
-causes some bad state that will affect other tests, do not use
-``xfail``.
-
-Implementation
-^^^^^^^^^^^^^^
-
- Here is how to skip whole test unconditionally:
-
-.. code-block:: python
-
-    @unittest.skip("this bug needs to be fixed")
-    def test_feature_x():
-
-or via pytest:
-
-.. code-block:: python
-
-    @pytest.mark.skip(reason="this bug needs to be fixed")
-
-or the ``xfail`` way:
-
-.. code-block:: python
-
-    @pytest.mark.xfail
-    def test_feature_x():
-
-Here is how to skip a test based on some internal check inside the test:
-
-.. code-block:: python
-
-    def test_feature_x():
-        if not has_something():
-            pytest.skip("unsupported configuration")
-
-or the whole module:
-
-.. code-block:: python
-
-    import pytest
-    if not pytest.config.getoption("--custom-flag"):
-        pytest.skip("--custom-flag is missing, skipping tests", allow_module_level=True)
-
-or the ``xfail`` way:
-
-.. code-block:: python
-
-    def test_feature_x():
-        pytest.xfail("expected to fail until bug XYZ is fixed")
-
-Here is how to skip all tests in a module if some import is missing:
-
-.. code-block:: python
-
-    docutils = pytest.importorskip("docutils", minversion="0.3")
-
-  Skip a test based on a condition:
-
-.. code-block:: python
-
-    @pytest.mark.skipif(sys.version_info < (3,6), reason="requires python3.6 or higher")
-    def test_feature_x():
-
-or:
-
-.. code-block:: python
-
-    @unittest.skipIf(torch_device == "cpu", "Can't do half precision")
-    def test_feature_x():
-   
-or skip the whole module:
-
-.. code-block:: python
-
-    @pytest.mark.skipif(sys.platform == 'win32', reason="does not run on windows")
-    class TestClass():
-        def test_feature_x(self):
-
-More details, example and ways are `here <https://docs.pytest.org/en/latest/skipping.html>`__.
-
-Custom markers
-~~~~~~~~~~~~~~
-
-* Slow tests
-
-Tests that are too slow (e.g. once downloading huge model files) are marked with:
-
-.. code-block:: python
-
-    from transformers.testing_utils import slow
-    @slow
-    def test_integration_foo():
-
-To run such tests set ``RUN_SLOW=1`` env var, e.g.:
-
-.. code-block:: bash
-
-    RUN_SLOW=1 pytest tests
-    
-It's important that the decorator ``@slow`` appears last in the stack of decorators, as some decorators like ``parametrized`` may interfere with its normal functioning. Here is an example of the correct usage:
-
-.. code-block:: python
-
-    @parameterized.expand(...)
-    @slow
-    def test_integration_foo():
-
-Testing the stdout/stderr output
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-In order to test functions that write to ``stdout`` and/or ``stderr``,
-the test can access those streams using the ``pytest``'s `capsys
-system <https://docs.pytest.org/en/latest/capture.html>`__. Here is how
-this is accomplished:
-
-.. code-block:: python
-
-    import sys
-    def print_to_stdout(s): print(s)
-    def print_to_stderr(s): sys.stderr.write(s)
-    def test_result_and_stdout(capsys):
-        msg = "Hello"
-        print_to_stdout(msg)
-        print_to_stderr(msg)
-        out, err = capsys.readouterr() # consume the captured output streams
-        # optional: if you want to replay the consumed streams:
-        sys.stdout.write(out)
-        sys.stderr.write(err)
-        # test:
-        assert msg in out
-        assert msg in err
-
-And, of course, most of the time, ``stderr`` will come as a part of an
-exception, so try/except has to be used in such a case:
-
-.. code-block:: python
-
-    def raise_exception(msg): raise ValueError(msg)
-    def test_something_exception():
-        msg = "Not a good value"
-        error = ''
-        try:
-            raise_exception(msg)
-        except Exception as e:
-            error = str(e)
-            assert msg in error, f"{msg} is in the exception:\n{error}"
-
-Another approach to capturing stdout is via ``contextlib.redirect_stdout``:
-
-.. code-block:: python
-
-    from io import StringIO
-    from contextlib import redirect_stdout
-    def print_to_stdout(s): print(s)
-    def test_result_and_stdout():
-        msg = "Hello"
-        buffer = StringIO()
-        with redirect_stdout(buffer):
-            print_to_stdout(msg)
-        out = buffer.getvalue()
-        # optional: if you want to replay the consumed streams:
-        sys.stdout.write(out)
-        # test:
-        assert msg in out
-
-An important potential issue with capturing stdout is that it may
-contain ``\r`` characters that in normal ``print`` reset everything that
-has been printed so far. There is no problem with ``pytest``, but with
-``pytest -s`` these characters get included in the buffer, so to be able
-to have the test run with and without ``-s``, you have to make an extra
-cleanup to the captured output, using ``re.sub(r'~.*\r', '', buf, 0, re.M)``.
-
-But, then we have a helper context manager wrapper to automatically take
-care of it all, regardless of whether it has some ``\r``'s in it or
-not, so it's a simple:
-
-.. code-block:: python
-
-    from transformers.testing_utils import CaptureStdout
-    with CaptureStdout() as cs:
-        function_that_writes_to_stdout()
-    print(cs.out)
-
-Here is a full test example:
-
-.. code-block:: python
-
-    from transformers.testing_utils import CaptureStdout
-    msg = "Secret message\r"
-    final = "Hello World"
-    with CaptureStdout() as cs:
-        print(msg + final)
-    assert cs.out == final+"\n", f"captured: {cs.out}, expecting {final}"
-
-If you'd like to capture ``stderr`` use the :obj:`CaptureStderr` class
-instead:
-
-.. code-block:: python
-
-    from transformers.testing_utils import CaptureStderr
-    with CaptureStderr() as cs:
-        function_that_writes_to_stderr()
-    print(cs.err)
-
-If you need to capture both streams at once, use the parent
-:obj:`CaptureStd` class:
-
-.. code-block:: python
-
-    from transformers.testing_utils import CaptureStd
-    with CaptureStd() as cs:
-        function_that_writes_to_stdout_and_stderr()
-    print(cs.err, cs.out)
-
-
-
-Capturing logger stream
-~~~~~~~~~~~~~~~~~~~~~~~
-
-If you need to validate the output of a logger, you can use :obj:`CaptureLogger`:
-
-.. code-block:: python
-
-    from transformers import logging
-    from transformers.testing_utils import CaptureLogger
-
-    msg = "Testing 1, 2, 3"
-    logging.set_verbosity_info()
-    logger = logging.get_logger("transformers.tokenization_bart")
-    with CaptureLogger(logger) as cl:
-        logger.info(msg)
-    assert cl.out, msg+"\n"
-
-
-Testing with environment variables
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-If you want to test the impact of environment variables for a specific test you can use a helper decorator ``transformers.testing_utils.mockenv``
-
-.. code-block:: python
-
-    from transformers.testing_utils import mockenv
-    class HfArgumentParserTest(unittest.TestCase):
-        @mockenv(TRANSFORMERS_VERBOSITY="error")
-        def test_env_override(self):
-            env_level_str = os.getenv("TRANSFORMERS_VERBOSITY", None)
-
-
-Getting reproducible results
-~~~~~~~~~~~~~~~~~~~~~~~~~~~~
-
-In some situations you may want to remove randomness for your tests. To
-get identical reproducable results set, you will need to fix the seed:
-
-.. code-block:: python
-
-    seed = 42
-
-    # python RNG
-    import random
-    random.seed(seed)
-
-    # pytorch RNGs
-    import torch
-    torch.manual_seed(seed)
-    torch.backends.cudnn.deterministic = True
-    if torch.cuda.is_available(): torch.cuda.manual_seed_all(seed)
-
-    # numpy RNG
-    import numpy as np
-    np.random.seed(seed)
-
-    # tf RNG
-    tf.random.set_seed(seed)
-
-Debugging tests
-~~~~~~~~~~~~~~~
-
-To start a debugger at the point of the warning, do this:
-
-.. code-block:: bash
-
-    pytest tests/test_logging.py -W error::UserWarning --pdb
@@ -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``
@@ -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."}
@@ -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]

        """

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

        """

@@ -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)
@@ -721,7 +674,7 @@ def main():
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
            )
-
+            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        logger.info("Evaluate the following checkpoints: %s", checkpoints)

        for checkpoint in checkpoints:
@@ -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,
@@ -547,7 +547,7 @@ def main():
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
            )
-
+            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            global_step = checkpoint.split("-")[-1] if len(checkpoints) > 1 else ""
@@ -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],),
        ]
    )
@@ -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

@@ -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
@@ -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
    )
@@ -681,6 +679,7 @@ def main():
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
            )
+            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce model loading logs

        logger.info("Evaluate the following checkpoints: %s", checkpoints)

@@ -88,7 +88,7 @@ def main():
        )
    )

-    model.reset_memory_length(args.mem_len)
+    model.reset_length(args.tgt_len, args.ext_len, args.mem_len)
    if args.clamp_len > 0:
        model.clamp_len = args.clamp_len
    if args.same_length:
@@ -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)
@@ -677,7 +673,7 @@ def main():
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
            )
-
+            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            global_step = checkpoint.split("-")[-1] if len(checkpoints) > 1 else ""
@@ -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
@@ -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
@@ -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
@@ -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
@@ -842,6 +836,7 @@ def main():
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
            )
+            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce model loading logs

        logger.info("Evaluate the following checkpoints: %s", checkpoints)

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


@@ -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(
@@ -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,
@@ -119,7 +118,7 @@ class BaseTransformer(pl.LightningModule):
    def get_lr_scheduler(self):
        get_schedule_func = arg_to_scheduler[self.hparams.lr_scheduler]
        scheduler = get_schedule_func(
-            self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.total_steps()
+            self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.total_steps
        )
        scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
        return scheduler
@@ -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()
@@ -159,30 +150,29 @@ class BaseTransformer(pl.LightningModule):
    def test_epoch_end(self, outputs):
        return self.validation_end(outputs)

+    @property
    def total_steps(self) -> int:
        """The number of total training steps that will be run. Used for lr scheduler purposes."""
        num_devices = max(1, self.hparams.gpus)  # TODO: consider num_tpu_cores
        effective_batch_size = self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
-        return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
+        dataset_size = len(self.train_loader.dataset)
+        return (dataset_size / effective_batch_size) * self.hparams.max_epochs

    def setup(self, mode):
-        if mode == "test":
-            self.dataset_size = len(self.test_dataloader().dataset)
-        else:
+        if mode == "fit":
            self.train_loader = self.get_dataloader("train", self.hparams.train_batch_size, shuffle=True)
-            self.dataset_size = len(self.train_loader.dataset)

-    def get_dataloader(self, type_path: str, batch_size: int, shuffle: bool = False):
+    def get_dataloader(self, type_path, batch_size, shuffle=False):
        raise NotImplementedError("You must implement this for your task")

    def train_dataloader(self):
        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(
@@ -236,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(
@@ -261,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):
@@ -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).
@@ -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:", "")
@@ -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"],
@@ -1,5 +0,0 @@
-# LXMERT DEMO
-
-1. make a virtualenv: ``virtualenv venv`` and activate ``source venv/bin/activate``
-2. install reqs: ``pip install -r ./requirements.txt``
-3. usage is as shown in demo.ipynb
@@ -1,149 +0,0 @@
-import getopt
-import json
-import os
-
-# import numpy as np
-import sys
-from collections import OrderedDict
-
-import datasets
-import numpy as np
-import torch
-
-from modeling_frcnn import GeneralizedRCNN
-from processing_image import Preprocess
-from utils import Config
-
-
-"""
-USAGE:
-``python extracting_data.py -i <img_dir> -o <dataset_file>.datasets <batch_size>``
-"""
-
-
-TEST = False
-CONFIG = Config.from_pretrained("unc-nlp/frcnn-vg-finetuned")
-DEFAULT_SCHEMA = datasets.Features(
-    OrderedDict(
-        {
-            "attr_ids": datasets.Sequence(length=CONFIG.MAX_DETECTIONS, feature=datasets.Value("float32")),
-            "attr_probs": datasets.Sequence(length=CONFIG.MAX_DETECTIONS, feature=datasets.Value("float32")),
-            "boxes": datasets.Array2D((CONFIG.MAX_DETECTIONS, 4), dtype="float32"),
-            "img_id": datasets.Value("int32"),
-            "obj_ids": datasets.Sequence(length=CONFIG.MAX_DETECTIONS, feature=datasets.Value("float32")),
-            "obj_probs": datasets.Sequence(length=CONFIG.MAX_DETECTIONS, feature=datasets.Value("float32")),
-            "roi_features": datasets.Array2D((CONFIG.MAX_DETECTIONS, 2048), dtype="float32"),
-            "sizes": datasets.Sequence(length=2, feature=datasets.Value("float32")),
-            "preds_per_image": datasets.Value(dtype="int32"),
-        }
-    )
-)
-
-
-class Extract:
-    def __init__(self, argv=sys.argv[1:]):
-        inputdir = None
-        outputfile = None
-        subset_list = None
-        batch_size = 1
-        opts, args = getopt.getopt(argv, "i:o:b:s", ["inputdir=", "outfile=", "batch_size=", "subset_list="])
-        for opt, arg in opts:
-            if opt in ("-i", "--inputdir"):
-                inputdir = arg
-            elif opt in ("-o", "--outfile"):
-                outputfile = arg
-            elif opt in ("-b", "--batch_size"):
-                batch_size = int(arg)
-            elif opt in ("-s", "--subset_list"):
-                subset_list = arg
-
-        assert inputdir is not None  # and os.path.isdir(inputdir), f"{inputdir}"
-        assert outputfile is not None and not os.path.isfile(outputfile), f"{outputfile}"
-        if subset_list is not None:
-            with open(os.path.realpath(subset_list)) as f:
-                self.subset_list = set(map(lambda x: self._vqa_file_split()[0], tryload(f)))
-        else:
-            self.subset_list = None
-
-        self.config = CONFIG
-        if torch.cuda.is_available():
-            self.config.model.device = "cuda"
-        self.inputdir = os.path.realpath(inputdir)
-        self.outputfile = os.path.realpath(outputfile)
-        self.preprocess = Preprocess(self.config)
-        self.model = GeneralizedRCNN.from_pretrained("unc-nlp/frcnn-vg-finetuned", config=self.config)
-        self.batch = batch_size if batch_size != 0 else 1
-        self.schema = DEFAULT_SCHEMA
-
-    def _vqa_file_split(self, file):
-        img_id = int(file.split(".")[0].split("_")[-1])
-        filepath = os.path.join(self.inputdir, file)
-        return (img_id, filepath)
-
-    @property
-    def file_generator(self):
-        batch = []
-        for i, file in enumerate(os.listdir(self.inputdir)):
-            if self.subset_list is not None and i not in self.subset_list:
-                continue
-            batch.append(self._vqa_file_split(file))
-            if len(batch) == self.batch:
-                temp = batch
-                batch = []
-                yield list(map(list, zip(*temp)))
-
-        for i in range(1):
-            yield list(map(list, zip(*batch)))
-
-    def __call__(self):
-        # make writer
-        if not TEST:
-            writer = datasets.ArrowWriter(features=self.schema, path=self.outputfile)
-        # do file generator
-        for i, (img_ids, filepaths) in enumerate(self.file_generator):
-            images, sizes, scales_yx = self.preprocess(filepaths)
-            output_dict = self.model(
-                images,
-                sizes,
-                scales_yx=scales_yx,
-                padding="max_detections",
-                max_detections=self.config.MAX_DETECTIONS,
-                pad_value=0,
-                return_tensors="np",
-                location="cpu",
-            )
-            output_dict["boxes"] = output_dict.pop("normalized_boxes")
-            if not TEST:
-                output_dict["img_id"] = np.array(img_ids)
-                batch = self.schema.encode_batch(output_dict)
-                writer.write_batch(batch)
-            if TEST:
-                break
-            # finalizer the writer
-        if not TEST:
-            num_examples, num_bytes = writer.finalize()
-            print(f"Success! You wrote {num_examples} entry(s) and {num_bytes >> 20} mb")
-
-
-def tryload(stream):
-    try:
-        data = json.load(stream)
-        try:
-            data = list(data.keys())
-        except Exception:
-            data = [d["img_id"] for d in data]
-    except Exception:
-        try:
-            data = eval(stream.read())
-        except Exception:
-            data = stream.read().split("\n")
-    return data
-
-
-if __name__ == "__main__":
-    extract = Extract(sys.argv[1:])
-    extract()
-    if not TEST:
-        dataset = datasets.Dataset.from_file(extract.outputfile)
-        # wala!
-        # print(np.array(dataset[0:2]["roi_features"]).shape)
@@ -1,147 +0,0 @@
-"""
- coding=utf-8
- Copyright 2018, Antonio Mendoza Hao Tan, Mohit Bansal
- Adapted From Facebook Inc, Detectron2
-
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
-
-     http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.import copy
- """
-import sys
-from typing import Tuple
-
-import numpy as np
-import torch
-import torch.nn.functional as F
-from PIL import Image
-
-from utils import img_tensorize
-
-
-class ResizeShortestEdge:
-    def __init__(self, short_edge_length, max_size=sys.maxsize):
-        """
-        Args:
-            short_edge_length (list[min, max])
-            max_size (int): maximum allowed longest edge length.
-        """
-        self.interp_method = "bilinear"
-        self.max_size = max_size
-        self.short_edge_length = short_edge_length
-
-    def __call__(self, imgs):
-        img_augs = []
-        for img in imgs:
-            h, w = img.shape[:2]
-            # later: provide list and randomly choose index for resize
-            size = np.random.randint(self.short_edge_length[0], self.short_edge_length[1] + 1)
-            if size == 0:
-                return img
-            scale = size * 1.0 / min(h, w)
-            if h < w:
-                newh, neww = size, scale * w
-            else:
-                newh, neww = scale * h, size
-            if max(newh, neww) > self.max_size:
-                scale = self.max_size * 1.0 / max(newh, neww)
-                newh = newh * scale
-                neww = neww * scale
-            neww = int(neww + 0.5)
-            newh = int(newh + 0.5)
-
-            if img.dtype == np.uint8:
-                pil_image = Image.fromarray(img)
-                pil_image = pil_image.resize((neww, newh), Image.BILINEAR)
-                img = np.asarray(pil_image)
-            else:
-                img = img.permute(2, 0, 1).unsqueeze(0)  # 3, 0, 1)  # hw(c) -> nchw
-                img = F.interpolate(img, (newh, neww), mode=self.interp_method, align_corners=False).squeeze(0)
-            img_augs.append(img)
-
-        return img_augs
-
-
-class Preprocess:
-    def __init__(self, cfg):
-        self.aug = ResizeShortestEdge([cfg.INPUT.MIN_SIZE_TEST, cfg.INPUT.MIN_SIZE_TEST], cfg.INPUT.MAX_SIZE_TEST)
-        self.input_format = cfg.INPUT.FORMAT
-        self.size_divisibility = cfg.SIZE_DIVISIBILITY
-        self.pad_value = cfg.PAD_VALUE
-        self.max_image_size = cfg.INPUT.MAX_SIZE_TEST
-        self.device = cfg.MODEL.DEVICE
-        self.pixel_std = torch.tensor(cfg.MODEL.PIXEL_STD).to(self.device).view(len(cfg.MODEL.PIXEL_STD), 1, 1)
-        self.pixel_mean = torch.tensor(cfg.MODEL.PIXEL_MEAN).to(self.device).view(len(cfg.MODEL.PIXEL_STD), 1, 1)
-        self.normalizer = lambda x: (x - self.pixel_mean) / self.pixel_std
-
-    def pad(self, images):
-        max_size = tuple(max(s) for s in zip(*[img.shape for img in images]))
-        image_sizes = [im.shape[-2:] for im in images]
-        images = [
-            F.pad(
-                im,
-                [0, max_size[-1] - size[1], 0, max_size[-2] - size[0]],
-                value=self.pad_value,
-            )
-            for size, im in zip(image_sizes, images)
-        ]
-
-        return torch.stack(images), torch.tensor(image_sizes)
-
-    def __call__(self, images, single_image=False):
-        with torch.no_grad():
-            if not isinstance(images, list):
-                images = [images]
-            if single_image:
-                assert len(images) == 1
-            for i in range(len(images)):
-                if isinstance(images[i], torch.Tensor):
-                    images.insert(i, images.pop(i).to(self.device).float())
-                elif not isinstance(images[i], torch.Tensor):
-                    images.insert(
-                        i,
-                        torch.as_tensor(img_tensorize(images.pop(i), input_format=self.input_format))
-                        .to(self.device)
-                        .float(),
-                    )
-            # resize smallest edge
-            raw_sizes = torch.tensor([im.shape[:2] for im in images])
-            images = self.aug(images)
-            # transpose images and convert to torch tensors
-            # images = [torch.as_tensor(i.astype("float32")).permute(2, 0, 1).to(self.device) for i in images]
-            # now normalize before pad to aoid useless arithmatic
-            images = [self.normalizer(x) for x in images]
-            # now pad them to do the following operations
-            images, sizes = self.pad(images)
-            # Normalize
-
-            if self.size_divisibility > 0:
-                raise NotImplementedError()
-            # pad
-            scales_yx = torch.true_divide(raw_sizes, sizes)
-            if single_image:
-                return images[0], sizes[0], scales_yx[0]
-            else:
-                return images, sizes, scales_yx
-
-
-def _scale_box(boxes, scale_yx):
-    boxes[:, 0::2] *= scale_yx[:, 1]
-    boxes[:, 1::2] *= scale_yx[:, 0]
-    return boxes
-
-
-def _clip_box(tensor, box_size: Tuple[int, int]):
-    assert torch.isfinite(tensor).all(), "Box tensor contains infinite or NaN!"
-    h, w = box_size
-    tensor[:, 0].clamp_(min=0, max=w)
-    tensor[:, 1].clamp_(min=0, max=h)
-    tensor[:, 2].clamp_(min=0, max=w)
-    tensor[:, 3].clamp_(min=0, max=h)
@@ -1,99 +0,0 @@
-appdirs==1.4.3
-argon2-cffi==20.1.0
-async-generator==1.10
-attrs==20.2.0
-backcall==0.2.0
-bleach==3.1.5
-CacheControl==0.12.6
-certifi==2020.6.20
-cffi==1.14.2
-chardet==3.0.4
-click==7.1.2
-colorama==0.4.3
-contextlib2==0.6.0
-cycler==0.10.0
-datasets==1.0.0
-decorator==4.4.2
-defusedxml==0.6.0
-dill==0.3.2
-distlib==0.3.0
-distro==1.4.0
-entrypoints==0.3
-filelock==3.0.12
-future==0.18.2
-html5lib==1.0.1
-idna==2.8
-ipaddr==2.2.0
-ipykernel==5.3.4
-ipython
-ipython-genutils==0.2.0
-ipywidgets==7.5.1
-jedi==0.17.2
-Jinja2==2.11.2
-joblib==0.16.0
-jsonschema==3.2.0
-jupyter==1.0.0
-jupyter-client==6.1.7
-jupyter-console==6.2.0
-jupyter-core==4.6.3
-jupyterlab-pygments==0.1.1
-kiwisolver==1.2.0
-lockfile==0.12.2
-MarkupSafe==1.1.1
-matplotlib==3.3.1
-mistune==0.8.4
-msgpack==0.6.2
-nbclient==0.5.0
-nbconvert==6.0.1
-nbformat==5.0.7
-nest-asyncio==1.4.0
-notebook==6.1.4
-numpy==1.19.2
-opencv-python==4.4.0.42
-packaging==20.3
-pandas==1.1.2
-pandocfilters==1.4.2
-parso==0.7.1
-pep517==0.8.2
-pexpect==4.8.0
-pickleshare==0.7.5
-Pillow==7.2.0
-progress==1.5
-prometheus-client==0.8.0
-prompt-toolkit==3.0.7
-ptyprocess==0.6.0
-pyaml==20.4.0
-pyarrow==1.0.1
-pycparser==2.20
-Pygments==2.6.1
-pyparsing==2.4.6
-pyrsistent==0.16.0
-python-dateutil==2.8.1
-pytoml==0.1.21
-pytz==2020.1
-PyYAML==5.3.1
-pyzmq==19.0.2
-qtconsole==4.7.7
-QtPy==1.9.0
-regex==2020.7.14
-requests==2.22.0
-retrying==1.3.3
-sacremoses==0.0.43
-Send2Trash==1.5.0
-sentencepiece==0.1.91
-six==1.14.0
-terminado==0.8.3
-testpath==0.4.4
-tokenizers==0.8.1rc2
-torch==1.6.0
-torchvision==0.7.0
-tornado==6.0.4
-tqdm==4.48.2
-traitlets
-git+https://github.com/huggingface/transformers.git
-urllib3==1.25.8
-wcwidth==0.2.5
-webencodings==0.5.1
-wget==3.2
-widgetsnbextension==3.5.1
-xxhash==2.0.0
@@ -1,559 +0,0 @@
-"""
- coding=utf-8
- Copyright 2018, Antonio Mendoza Hao Tan, Mohit Bansal, Huggingface team :)
- Adapted From Facebook Inc, Detectron2
-
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
-
-     http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.import copy
- """
-
-import copy
-import fnmatch
-import json
-import os
-import pickle as pkl
-import shutil
-import sys
-import tarfile
-import tempfile
-from collections import OrderedDict
-from contextlib import contextmanager
-from functools import partial
-from hashlib import sha256
-from io import BytesIO
-from pathlib import Path
-from urllib.parse import urlparse
-from zipfile import ZipFile, is_zipfile
-
-import numpy as np
-from PIL import Image
-from tqdm.auto import tqdm
-
-import cv2
-import requests
-import wget
-from filelock import FileLock
-from yaml import Loader, dump, load
-
-
-try:
-    import torch
-
-    _torch_available = True
-except ImportError:
-    _torch_available = False
-
-
-try:
-    from torch.hub import _get_torch_home
-
-    torch_cache_home = _get_torch_home()
-except ImportError:
-    torch_cache_home = os.path.expanduser(
-        os.getenv("TORCH_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch"))
-    )
-
-default_cache_path = os.path.join(torch_cache_home, "transformers")
-
-CLOUDFRONT_DISTRIB_PREFIX = "https://cdn.huggingface.co"
-S3_BUCKET_PREFIX = "https://s3.amazonaws.com/models.huggingface.co/bert"
-PATH = "/".join(str(Path(__file__).resolve()).split("/")[:-1])
-CONFIG = os.path.join(PATH, "config.yaml")
-ATTRIBUTES = os.path.join(PATH, "attributes.txt")
-OBJECTS = os.path.join(PATH, "objects.txt")
-PYTORCH_PRETRAINED_BERT_CACHE = os.getenv("PYTORCH_PRETRAINED_BERT_CACHE", default_cache_path)
-PYTORCH_TRANSFORMERS_CACHE = os.getenv("PYTORCH_TRANSFORMERS_CACHE", PYTORCH_PRETRAINED_BERT_CACHE)
-TRANSFORMERS_CACHE = os.getenv("TRANSFORMERS_CACHE", PYTORCH_TRANSFORMERS_CACHE)
-WEIGHTS_NAME = "pytorch_model.bin"
-CONFIG_NAME = "config.yaml"
-
-
-def load_labels(objs=OBJECTS, attrs=ATTRIBUTES):
-    vg_classes = []
-    with open(objs) as f:
-        for object in f.readlines():
-            vg_classes.append(object.split(",")[0].lower().strip())
-
-    vg_attrs = []
-    with open(attrs) as f:
-        for object in f.readlines():
-            vg_attrs.append(object.split(",")[0].lower().strip())
-    return vg_classes, vg_attrs
-
-
-def load_checkpoint(ckp):
-    r = OrderedDict()
-    with open(ckp, "rb") as f:
-        ckp = pkl.load(f)["model"]
-    for k in copy.deepcopy(list(ckp.keys())):
-        v = ckp.pop(k)
-        if isinstance(v, np.ndarray):
-            v = torch.tensor(v)
-        else:
-            assert isinstance(v, torch.tensor), type(v)
-        r[k] = v
-    return r
-
-
-class Config:
-    _pointer = {}
-
-    def __init__(self, dictionary: dict, name: str = "root", level=0):
-        self._name = name
-        self._level = level
-        d = {}
-        for k, v in dictionary.items():
-            if v is None:
-                raise ValueError()
-            k = copy.deepcopy(k)
-            v = copy.deepcopy(v)
-            if isinstance(v, dict):
-                v = Config(v, name=k, level=level + 1)
-            d[k] = v
-            setattr(self, k, v)
-
-        self._pointer = d
-
-    def __repr__(self):
-        return str(list((self._pointer.keys())))
-
-    def __setattr__(self, key, val):
-        self.__dict__[key] = val
-        self.__dict__[key.upper()] = val
-        levels = key.split(".")
-        last_level = len(levels) - 1
-        pointer = self._pointer
-        if len(levels) > 1:
-            for i, l in enumerate(levels):
-                if hasattr(self, l) and isinstance(getattr(self, l), Config):
-                    setattr(getattr(self, l), ".".join(levels[i:]), val)
-                if l == last_level:
-                    pointer[l] = val
-                else:
-                    pointer = pointer[l]
-
-    def to_dict(self):
-        return self._pointer
-
-    def dump_yaml(self, data, file_name):
-        with open(f"{file_name}", "w") as stream:
-            dump(data, stream)
-
-    def dump_json(self, data, file_name):
-        with open(f"{file_name}", "w") as stream:
-            json.dump(data, stream)
-
-    @staticmethod
-    def load_yaml(config):
-        with open(config) as stream:
-            data = load(stream, Loader=Loader)
-        return data
-
-    def __str__(self):
-        t = "    "
-        if self._name != "root":
-            r = f"{t * (self._level-1)}{self._name}:\n"
-        else:
-            r = ""
-        level = self._level
-        for i, (k, v) in enumerate(self._pointer.items()):
-            if isinstance(v, Config):
-                r += f"{t * (self._level)}{v}\n"
-                self._level += 1
-            else:
-                r += f"{t * (self._level)}{k}: {v} ({type(v).__name__})\n"
-            self._level = level
-        return r[:-1]
-
-    @classmethod
-    def from_pretrained(cls, pretrained_model_name_or_path: str, **kwargs):
-        config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
-        return cls(config_dict)
-
-    @classmethod
-    def get_config_dict(cls, pretrained_model_name_or_path: str, **kwargs):
-
-        cache_dir = kwargs.pop("cache_dir", None)
-        force_download = kwargs.pop("force_download", False)
-        resume_download = kwargs.pop("resume_download", False)
-        proxies = kwargs.pop("proxies", None)
-        local_files_only = kwargs.pop("local_files_only", False)
-
-        if os.path.isdir(pretrained_model_name_or_path):
-            config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
-        elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
-            config_file = pretrained_model_name_or_path
-        else:
-            config_file = hf_bucket_url(pretrained_model_name_or_path, filename=CONFIG_NAME, use_cdn=False)
-
-        try:
-            # Load from URL or cache if already cached
-            resolved_config_file = cached_path(
-                config_file,
-                cache_dir=cache_dir,
-                force_download=force_download,
-                proxies=proxies,
-                resume_download=resume_download,
-                local_files_only=local_files_only,
-            )
-            # Load config dict
-            if resolved_config_file is None:
-                raise EnvironmentError
-
-            config_file = Config.load_yaml(resolved_config_file)
-
-        except EnvironmentError:
-            msg = "Can't load config for"
-            raise EnvironmentError(msg)
-
-        if resolved_config_file == config_file:
-            print("loading configuration file from path")
-        else:
-            print("loading configuration file cache")
-
-        return Config.load_yaml(resolved_config_file), kwargs
-
-
-# quick compare tensors
-def compare(in_tensor):
-
-    out_tensor = torch.load("dump.pt", map_location=in_tensor.device)
-    n1 = in_tensor.numpy()
-    n2 = out_tensor.numpy()[0]
-    print(n1.shape, n1[0, 0, :5])
-    print(n2.shape, n2[0, 0, :5])
-    assert np.allclose(
-        n1, n2, rtol=0.01, atol=0.1
-    ), f"{sum([1 for x in np.isclose(n1, n2, rtol=0.01, atol=0.1).flatten() if x == False])/len(n1.flatten())*100:.4f} % element-wise mismatch"
-    raise Exception("tensors are all good")
-
-    # Hugging face functiions below
-
-
-def is_remote_url(url_or_filename):
-    parsed = urlparse(url_or_filename)
-    return parsed.scheme in ("http", "https")
-
-
-def hf_bucket_url(model_id: str, filename: str, use_cdn=True) -> str:
-    endpoint = CLOUDFRONT_DISTRIB_PREFIX if use_cdn else S3_BUCKET_PREFIX
-    legacy_format = "/" not in model_id
-    if legacy_format:
-        return f"{endpoint}/{model_id}-{filename}"
-    else:
-        return f"{endpoint}/{model_id}/{filename}"
-
-
-def http_get(
-    url,
-    temp_file,
-    proxies=None,
-    resume_size=0,
-    user_agent=None,
-):
-    ua = "python/{}".format(sys.version.split()[0])
-    if _torch_available:
-        ua += "; torch/{}".format(torch.__version__)
-    if isinstance(user_agent, dict):
-        ua += "; " + "; ".join("{}/{}".format(k, v) for k, v in user_agent.items())
-    elif isinstance(user_agent, str):
-        ua += "; " + user_agent
-    headers = {"user-agent": ua}
-    if resume_size > 0:
-        headers["Range"] = "bytes=%d-" % (resume_size,)
-    response = requests.get(url, stream=True, proxies=proxies, headers=headers)
-    if response.status_code == 416:  # Range not satisfiable
-        return
-    content_length = response.headers.get("Content-Length")
-    total = resume_size + int(content_length) if content_length is not None else None
-    progress = tqdm(
-        unit="B",
-        unit_scale=True,
-        total=total,
-        initial=resume_size,
-        desc="Downloading",
-    )
-    for chunk in response.iter_content(chunk_size=1024):
-        if chunk:  # filter out keep-alive new chunks
-            progress.update(len(chunk))
-            temp_file.write(chunk)
-    progress.close()
-
-
-def get_from_cache(
-    url,
-    cache_dir=None,
-    force_download=False,
-    proxies=None,
-    etag_timeout=10,
-    resume_download=False,
-    user_agent=None,
-    local_files_only=False,
-):
-
-    if cache_dir is None:
-        cache_dir = TRANSFORMERS_CACHE
-    if isinstance(cache_dir, Path):
-        cache_dir = str(cache_dir)
-
-    os.makedirs(cache_dir, exist_ok=True)
-
-    etag = None
-    if not local_files_only:
-        try:
-            response = requests.head(url, allow_redirects=True, proxies=proxies, timeout=etag_timeout)
-            if response.status_code == 200:
-                etag = response.headers.get("ETag")
-        except (EnvironmentError, requests.exceptions.Timeout):
-            # etag is already None
-            pass
-
-    filename = url_to_filename(url, etag)
-
-    # get cache path to put the file
-    cache_path = os.path.join(cache_dir, filename)
-
-    # etag is None = we don't have a connection, or url doesn't exist, or is otherwise inaccessible.
-    # try to get the last downloaded one
-    if etag is None:
-        if os.path.exists(cache_path):
-            return cache_path
-        else:
-            matching_files = [
-                file
-                for file in fnmatch.filter(os.listdir(cache_dir), filename + ".*")
-                if not file.endswith(".json") and not file.endswith(".lock")
-            ]
-            if len(matching_files) > 0:
-                return os.path.join(cache_dir, matching_files[-1])
-            else:
-                # If files cannot be found and local_files_only=True,
-                # the models might've been found if local_files_only=False
-                # Notify the user about that
-                if local_files_only:
-                    raise ValueError(
-                        "Cannot find the requested files in the cached path and outgoing traffic has been"
-                        " disabled. To enable model look-ups and downloads online, set 'local_files_only'"
-                        " to False."
-                    )
-                return None
-
-    # From now on, etag is not None.
-    if os.path.exists(cache_path) and not force_download:
-        return cache_path
-
-    # Prevent parallel downloads of the same file with a lock.
-    lock_path = cache_path + ".lock"
-    with FileLock(lock_path):
-
-        # If the download just completed while the lock was activated.
-        if os.path.exists(cache_path) and not force_download:
-            # Even if returning early like here, the lock will be released.
-            return cache_path
-
-        if resume_download:
-            incomplete_path = cache_path + ".incomplete"
-
-            @contextmanager
-            def _resumable_file_manager():
-                with open(incomplete_path, "a+b") as f:
-                    yield f
-
-            temp_file_manager = _resumable_file_manager
-            if os.path.exists(incomplete_path):
-                resume_size = os.stat(incomplete_path).st_size
-            else:
-                resume_size = 0
-        else:
-            temp_file_manager = partial(tempfile.NamedTemporaryFile, dir=cache_dir, delete=False)
-            resume_size = 0
-
-        # Download to temporary file, then copy to cache dir once finished.
-        # Otherwise you get corrupt cache entries if the download gets interrupted.
-        with temp_file_manager() as temp_file:
-            print(
-                "%s not found in cache or force_download set to True, downloading to %s",
-                url,
-                temp_file.name,
-            )
-
-            http_get(
-                url,
-                temp_file,
-                proxies=proxies,
-                resume_size=resume_size,
-                user_agent=user_agent,
-            )
-
-        os.replace(temp_file.name, cache_path)
-
-        meta = {"url": url, "etag": etag}
-        meta_path = cache_path + ".json"
-        with open(meta_path, "w") as meta_file:
-            json.dump(meta, meta_file)
-
-    return cache_path
-
-
-def url_to_filename(url, etag=None):
-
-    url_bytes = url.encode("utf-8")
-    url_hash = sha256(url_bytes)
-    filename = url_hash.hexdigest()
-
-    if etag:
-        etag_bytes = etag.encode("utf-8")
-        etag_hash = sha256(etag_bytes)
-        filename += "." + etag_hash.hexdigest()
-
-    if url.endswith(".h5"):
-        filename += ".h5"
-
-    return filename
-
-
-def cached_path(
-    url_or_filename,
-    cache_dir=None,
-    force_download=False,
-    proxies=None,
-    resume_download=False,
-    user_agent=None,
-    extract_compressed_file=False,
-    force_extract=False,
-    local_files_only=False,
-):
-    if cache_dir is None:
-        cache_dir = TRANSFORMERS_CACHE
-    if isinstance(url_or_filename, Path):
-        url_or_filename = str(url_or_filename)
-    if isinstance(cache_dir, Path):
-        cache_dir = str(cache_dir)
-
-    if is_remote_url(url_or_filename):
-        # URL, so get it from the cache (downloading if necessary)
-        output_path = get_from_cache(
-            url_or_filename,
-            cache_dir=cache_dir,
-            force_download=force_download,
-            proxies=proxies,
-            resume_download=resume_download,
-            user_agent=user_agent,
-            local_files_only=local_files_only,
-        )
-    elif os.path.exists(url_or_filename):
-        # File, and it exists.
-        output_path = url_or_filename
-    elif urlparse(url_or_filename).scheme == "":
-        # File, but it doesn't exist.
-        raise EnvironmentError("file {} not found".format(url_or_filename))
-    else:
-        # Something unknown
-        raise ValueError("unable to parse {} as a URL or as a local path".format(url_or_filename))
-
-    if extract_compressed_file:
-        if not is_zipfile(output_path) and not tarfile.is_tarfile(output_path):
-            return output_path
-
-        # Path where we extract compressed archives
-        # We avoid '.' in dir name and add "-extracted" at the end: "./model.zip" => "./model-zip-extracted/"
-        output_dir, output_file = os.path.split(output_path)
-        output_extract_dir_name = output_file.replace(".", "-") + "-extracted"
-        output_path_extracted = os.path.join(output_dir, output_extract_dir_name)
-
-        if os.path.isdir(output_path_extracted) and os.listdir(output_path_extracted) and not force_extract:
-            return output_path_extracted
-
-        # Prevent parallel extractions
-        lock_path = output_path + ".lock"
-        with FileLock(lock_path):
-            shutil.rmtree(output_path_extracted, ignore_errors=True)
-            os.makedirs(output_path_extracted)
-            if is_zipfile(output_path):
-                with ZipFile(output_path, "r") as zip_file:
-                    zip_file.extractall(output_path_extracted)
-                    zip_file.close()
-            elif tarfile.is_tarfile(output_path):
-                tar_file = tarfile.open(output_path)
-                tar_file.extractall(output_path_extracted)
-                tar_file.close()
-            else:
-                raise EnvironmentError("Archive format of {} could not be identified".format(output_path))
-
-        return output_path_extracted
-
-    return output_path
-
-
-def get_data(query, delim=","):
-    assert isinstance(query, str)
-    if os.path.isfile(query):
-        with open(query) as f:
-            data = eval(f.read())
-    else:
-        req = requests.get(query)
-        try:
-            data = requests.json()
-        except Exception:
-            data = req.content.decode()
-            assert data is not None, "could not connect"
-            try:
-                data = eval(data)
-            except Exception:
-                data = data.split("\n")
-        req.close()
-    return data
-
-
-def get_image_from_url(url):
-    response = requests.get(url)
-    img = np.array(Image.open(BytesIO(response.content)))
-    return img
-
-
-# to load legace frcnn checkpoint from detectron
-def load_frcnn_pkl_from_url(url):
-    fn = url.split("/")[-1]
-    if fn not in os.listdir(os.getcwd()):
-        wget.download(url)
-    with open(fn, "rb") as stream:
-        weights = pkl.load(stream)
-    model = weights.pop("model")
-    new = {}
-    for k, v in model.items():
-        new[k] = torch.from_numpy(v)
-        if "running_var" in k:
-            zero = torch.Tensor([0])
-            k2 = k.replace("running_var", "num_batches_tracked")
-            new[k2] = zero
-    return new
-
-
-def get_demo_path():
-    print(f"{os.path.abspath(os.path.join(PATH, os.pardir))}/demo.ipynb")
-
-
-def img_tensorize(im, input_format="RGB"):
-    assert isinstance(im, str)
-    if os.path.isfile(im):
-        img = cv2.imread(im)
-    else:
-        img = get_image_from_url(im)
-        assert img is not None, f"could not connect to: {im}"
-    img = cv2.cvtColor(img, cv2.COLOR_BGR2RGB)
-    if input_format == "RGB":
-        img = img[:, :, ::-1]
-    return img
-
-
-def chunk(images, batch=1):
-    return (images[i : i + batch] for i in range(0, len(images), batch))
@@ -1,499 +0,0 @@
-"""
- coding=utf-8
- Copyright 2018, Antonio Mendoza Hao Tan, Mohit Bansal
- Adapted From Facebook Inc, Detectron2
-
- Licensed under the Apache License, Version 2.0 (the "License");
- you may not use this file except in compliance with the License.
- You may obtain a copy of the License at
-
-     http://www.apache.org/licenses/LICENSE-2.0
-
- Unless required by applicable law or agreed to in writing, software
- distributed under the License is distributed on an "AS IS" BASIS,
- WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
- See the License for the specific language governing permissions and
- limitations under the License.import copy
- """
-import colorsys
-import io
-
-import matplotlib as mpl
-import matplotlib.colors as mplc
-import matplotlib.figure as mplfigure
-import numpy as np
-import torch
-from matplotlib.backends.backend_agg import FigureCanvasAgg
-
-import cv2
-from utils import img_tensorize
-
-
-_SMALL_OBJ = 1000
-
-
-class SingleImageViz:
-    def __init__(
-        self,
-        img,
-        scale=1.2,
-        edgecolor="g",
-        alpha=0.5,
-        linestyle="-",
-        saveas="test_out.jpg",
-        rgb=True,
-        pynb=False,
-        id2obj=None,
-        id2attr=None,
-        pad=0.7,
-    ):
-        """
-        img: an RGB image of shape (H, W, 3).
-        """
-        if isinstance(img, torch.Tensor):
-            img = img.numpy().astype("np.uint8")
-        if isinstance(img, str):
-            img = img_tensorize(img)
-        assert isinstance(img, np.ndarray)
-
-        width, height = img.shape[1], img.shape[0]
-        fig = mplfigure.Figure(frameon=False)
-        dpi = fig.get_dpi()
-        width_in = (width * scale + 1e-2) / dpi
-        height_in = (height * scale + 1e-2) / dpi
-        fig.set_size_inches(width_in, height_in)
-        ax = fig.add_axes([0.0, 0.0, 1.0, 1.0])
-        ax.axis("off")
-        ax.set_xlim(0.0, width)
-        ax.set_ylim(height)
-
-        self.saveas = saveas
-        self.rgb = rgb
-        self.pynb = pynb
-        self.img = img
-        self.edgecolor = edgecolor
-        self.alpha = 0.5
-        self.linestyle = linestyle
-        self.font_size = int(np.sqrt(min(height, width)) * scale // 3)
-        self.width = width
-        self.height = height
-        self.scale = scale
-        self.fig = fig
-        self.ax = ax
-        self.pad = pad
-        self.id2obj = id2obj
-        self.id2attr = id2attr
-        self.canvas = FigureCanvasAgg(fig)
-
-    def add_box(self, box, color=None):
-        if color is None:
-            color = self.edgecolor
-        (x0, y0, x1, y1) = box
-        width = x1 - x0
-        height = y1 - y0
-        self.ax.add_patch(
-            mpl.patches.Rectangle(
-                (x0, y0),
-                width,
-                height,
-                fill=False,
-                edgecolor=color,
-                linewidth=self.font_size // 3,
-                alpha=self.alpha,
-                linestyle=self.linestyle,
-            )
-        )
-
-    def draw_boxes(self, boxes, obj_ids=None, obj_scores=None, attr_ids=None, attr_scores=None):
-        if len(boxes.shape) > 2:
-            boxes = boxes[0]
-        if len(obj_ids.shape) > 1:
-            obj_ids = obj_ids[0]
-        if len(obj_scores.shape) > 1:
-            obj_scores = obj_scores[0]
-        if len(attr_ids.shape) > 1:
-            attr_ids = attr_ids[0]
-        if len(attr_scores.shape) > 1:
-            attr_scores = attr_scores[0]
-        if isinstance(boxes, torch.Tensor):
-            boxes = boxes.numpy()
-        if isinstance(boxes, list):
-            boxes = np.array(boxes)
-        assert isinstance(boxes, np.ndarray)
-        areas = np.prod(boxes[:, 2:] - boxes[:, :2], axis=1)
-        sorted_idxs = np.argsort(-areas).tolist()
-        boxes = boxes[sorted_idxs] if boxes is not None else None
-        obj_ids = obj_ids[sorted_idxs] if obj_ids is not None else None
-        obj_scores = obj_scores[sorted_idxs] if obj_scores is not None else None
-        attr_ids = attr_ids[sorted_idxs] if attr_ids is not None else None
-        attr_scores = attr_scores[sorted_idxs] if attr_scores is not None else None
-
-        assigned_colors = [self._random_color(maximum=1) for _ in range(len(boxes))]
-        assigned_colors = [assigned_colors[idx] for idx in sorted_idxs]
-        if obj_ids is not None:
-            labels = self._create_text_labels_attr(obj_ids, obj_scores, attr_ids, attr_scores)
-            for i in range(len(boxes)):
-                color = assigned_colors[i]
-                self.add_box(boxes[i], color)
-                self.draw_labels(labels[i], boxes[i], color)
-
-    def draw_labels(self, label, box, color):
-        x0, y0, x1, y1 = box
-        text_pos = (x0, y0)
-        instance_area = (y1 - y0) * (x1 - x0)
-        small = _SMALL_OBJ * self.scale
-        if instance_area < small or y1 - y0 < 40 * self.scale:
-            if y1 >= self.height - 5:
-                text_pos = (x1, y0)
-            else:
-                text_pos = (x0, y1)
-
-        height_ratio = (y1 - y0) / np.sqrt(self.height * self.width)
-        lighter_color = self._change_color_brightness(color, brightness_factor=0.7)
-        font_size = np.clip((height_ratio - 0.02) / 0.08 + 1, 1.2, 2)
-        font_size *= 0.75 * self.font_size
-
-        self.draw_text(
-            text=label,
-            position=text_pos,
-            color=lighter_color,
-        )
-
-    def draw_text(
-        self,
-        text,
-        position,
-        color="g",
-        ha="left",
-    ):
-        rotation = 0
-        font_size = self.font_size
-        color = np.maximum(list(mplc.to_rgb(color)), 0.2)
-        color[np.argmax(color)] = max(0.8, np.max(color))
-        bbox = {
-            "facecolor": "black",
-            "alpha": self.alpha,
-            "pad": self.pad,
-            "edgecolor": "none",
-        }
-        x, y = position
-        self.ax.text(
-            x,
-            y,
-            text,
-            size=font_size * self.scale,
-            family="sans-serif",
-            bbox=bbox,
-            verticalalignment="top",
-            horizontalalignment=ha,
-            color=color,
-            zorder=10,
-            rotation=rotation,
-        )
-
-    def save(self, saveas=None):
-        if saveas is None:
-            saveas = self.saveas
-        if saveas.lower().endswith(".jpg") or saveas.lower().endswith(".png"):
-            cv2.imwrite(
-                saveas,
-                self._get_buffer()[:, :, ::-1],
-            )
-        else:
-            self.fig.savefig(saveas)
-
-    def _create_text_labels_attr(self, classes, scores, attr_classes, attr_scores):
-        labels = [self.id2obj[i] for i in classes]
-        attr_labels = [self.id2attr[i] for i in attr_classes]
-        labels = [
-            f"{label} {score:.2f} {attr} {attr_score:.2f}"
-            for label, score, attr, attr_score in zip(labels, scores, attr_labels, attr_scores)
-        ]
-        return labels
-
-    def _create_text_labels(self, classes, scores):
-        labels = [self.id2obj[i] for i in classes]
-        if scores is not None:
-            if labels is None:
-                labels = ["{:.0f}%".format(s * 100) for s in scores]
-            else:
-                labels = ["{} {:.0f}%".format(li, s * 100) for li, s in zip(labels, scores)]
-        return labels
-
-    def _random_color(self, maximum=255):
-        idx = np.random.randint(0, len(_COLORS))
-        ret = _COLORS[idx] * maximum
-        if not self.rgb:
-            ret = ret[::-1]
-        return ret
-
-    def _get_buffer(self):
-        if not self.pynb:
-            s, (width, height) = self.canvas.print_to_buffer()
-            if (width, height) != (self.width, self.height):
-                img = cv2.resize(self.img, (width, height))
-            else:
-                img = self.img
-        else:
-            buf = io.BytesIO()  # works for cairo backend
-            self.canvas.print_rgba(buf)
-            width, height = self.width, self.height
-            s = buf.getvalue()
-            img = self.img
-
-        buffer = np.frombuffer(s, dtype="uint8")
-        img_rgba = buffer.reshape(height, width, 4)
-        rgb, alpha = np.split(img_rgba, [3], axis=2)
-
-        try:
-            import numexpr as ne  # fuse them with numexpr
-
-            visualized_image = ne.evaluate("img * (1 - alpha / 255.0) + rgb * (alpha / 255.0)")
-        except ImportError:
-            alpha = alpha.astype("float32") / 255.0
-            visualized_image = img * (1 - alpha) + rgb * alpha
-
-        return visualized_image.astype("uint8")
-
-    def _change_color_brightness(self, color, brightness_factor):
-        assert brightness_factor >= -1.0 and brightness_factor <= 1.0
-        color = mplc.to_rgb(color)
-        polygon_color = colorsys.rgb_to_hls(*mplc.to_rgb(color))
-        modified_lightness = polygon_color[1] + (brightness_factor * polygon_color[1])
-        modified_lightness = 0.0 if modified_lightness < 0.0 else modified_lightness
-        modified_lightness = 1.0 if modified_lightness > 1.0 else modified_lightness
-        modified_color = colorsys.hls_to_rgb(polygon_color[0], modified_lightness, polygon_color[2])
-        return modified_color
-
-
-# Color map
-_COLORS = (
-    np.array(
-        [
-            0.000,
-            0.447,
-            0.741,
-            0.850,
-            0.325,
-            0.098,
-            0.929,
-            0.694,
-            0.125,
-            0.494,
-            0.184,
-            0.556,
-            0.466,
-            0.674,
-            0.188,
-            0.301,
-            0.745,
-            0.933,
-            0.635,
-            0.078,
-            0.184,
-            0.300,
-            0.300,
-            0.300,
-            0.600,
-            0.600,
-            0.600,
-            1.000,
-            0.000,
-            0.000,
-            1.000,
-            0.500,
-            0.000,
-            0.749,
-            0.749,
-            0.000,
-            0.000,
-            1.000,
-            0.000,
-            0.000,
-            0.000,
-            1.000,
-            0.667,
-            0.000,
-            1.000,
-            0.333,
-            0.333,
-            0.000,
-            0.333,
-            0.667,
-            0.000,
-            0.333,
-            1.000,
-            0.000,
-            0.667,
-            0.333,
-            0.000,
-            0.667,
-            0.667,
-            0.000,
-            0.667,
-            1.000,
-            0.000,
-            1.000,
-            0.333,
-            0.000,
-            1.000,
-            0.667,
-            0.000,
-            1.000,
-            1.000,
-            0.000,
-            0.000,
-            0.333,
-            0.500,
-            0.000,
-            0.667,
-            0.500,
-            0.000,
-            1.000,
-            0.500,
-            0.333,
-            0.000,
-            0.500,
-            0.333,
-            0.333,
-            0.500,
-            0.333,
-            0.667,
-            0.500,
-            0.333,
-            1.000,
-            0.500,
-            0.667,
-            0.000,
-            0.500,
-            0.667,
-            0.333,
-            0.500,
-            0.667,
-            0.667,
-            0.500,
-            0.667,
-            1.000,
-            0.500,
-            1.000,
-            0.000,
-            0.500,
-            1.000,
-            0.333,
-            0.500,
-            1.000,
-            0.667,
-            0.500,
-            1.000,
-            1.000,
-            0.500,
-            0.000,
-            0.333,
-            1.000,
-            0.000,
-            0.667,
-            1.000,
-            0.000,
-            1.000,
-            1.000,
-            0.333,
-            0.000,
-            1.000,
-            0.333,
-            0.333,
-            1.000,
-            0.333,
-            0.667,
-            1.000,
-            0.333,
-            1.000,
-            1.000,
-            0.667,
-            0.000,
-            1.000,
-            0.667,
-            0.333,
-            1.000,
-            0.667,
-            0.667,
-            1.000,
-            0.667,
-            1.000,
-            1.000,
-            1.000,
-            0.000,
-            1.000,
-            1.000,
-            0.333,
-            1.000,
-            1.000,
-            0.667,
-            1.000,
-            0.333,
-            0.000,
-            0.000,
-            0.500,
-            0.000,
-            0.000,
-            0.667,
-            0.000,
-            0.000,
-            0.833,
-            0.000,
-            0.000,
-            1.000,
-            0.000,
-            0.000,
-            0.000,
-            0.167,
-            0.000,
-            0.000,
-            0.333,
-            0.000,
-            0.000,
-            0.500,
-            0.000,
-            0.000,
-            0.667,
-            0.000,
-            0.000,
-            0.833,
-            0.000,
-            0.000,
-            1.000,
-            0.000,
-            0.000,
-            0.000,
-            0.167,
-            0.000,
-            0.000,
-            0.333,
-            0.000,
-            0.000,
-            0.500,
-            0.000,
-            0.000,
-            0.667,
-            0.000,
-            0.000,
-            0.833,
-            0.000,
-            0.000,
-            1.000,
-            0.000,
-            0.000,
-            0.000,
-            0.143,
-            0.143,
-            0.143,
-            0.857,
-            0.857,
-            0.857,
-            1.000,
-            1.000,
-            1.000,
-        ]
-    )
-    .astype(np.float32)
-    .reshape(-1, 3)
-)
@@ -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:
@@ -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")

@@ -934,7 +901,7 @@ def main():
            checkpoints = list(
                os.path.dirname(c) for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
            )
-
+            logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce logging
        logger.info("Evaluate the following checkpoints: %s", checkpoints)
        for checkpoint in checkpoints:
            global_step = checkpoint.split("-")[-1] if len(checkpoints) > 1 else ""
@@ -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",
@@ -1098,7 +1086,7 @@ def main():
                    os.path.dirname(c)
                    for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
                )
-
+                logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce model loading logs
        else:
            logger.info("Loading checkpoint %s for evaluation", args.model_name_or_path)
            checkpoints = [args.model_name_or_path]
@@ -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`
@@ -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]
@@ -792,7 +792,7 @@ def main():
                    os.path.dirname(c)
                    for c in sorted(glob.glob(args.output_dir + "/**/" + WEIGHTS_NAME, recursive=True))
                )
-
+                logging.getLogger("transformers.modeling_utils").setLevel(logging.WARN)  # Reduce model loading logs
        else:
            logger.info("Loading checkpoint %s for evaluation", args.model_name_or_path)
            checkpoints = [args.model_name_or_path]
@@ -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:
@@ -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:
@@ -5,14 +5,14 @@ psutil
 sacrebleu
 rouge-score
 tensorflow_datasets
-pytorch-lightning==0.9.0
+pytorch-lightning==0.8.5
 matplotlib
 git-python==1.0.3
 faiss
 streamlit
 elasticsearch
 pandas
-datasets
+nlp
 fire
 pytest
-conllu
+conllu
@@ -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,35 +17,25 @@ 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 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
+this should make a directory called `wmt_en_ro/` with files like `test.source`.
 ```

-#### Private Data
-
-If you are using your own data, it must be formatted as one directory with 6 files:
+If you are using your own data, it must be formatted as one directory with 6 files: 
 ```
 train.source
 train.target
@@ -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
@@ -229,83 +215,6 @@ python run_eval.py sshleifer/distilbart-cnn-12-6 $DATA_DIR/val.source dbart_val_
 ```


-### Multi-GPU Evaluation
-here is a command to run xsum evaluation on 8 GPUS. It is more than linearly faster than run_eval.py in some cases 
-because it uses SortishSampler to minimize padding. You can also use it on 1 GPU. `data_dir` must have 
-`{type_path}.source` and `{type_path}.target`. Run `python run_distributed_eval.py --help` for all clargs.
-
-```bash
-python -m torch.distributed.launch --nproc_per_node=8  run_distributed_eval.py \
-    --model_name sshleifer/distilbart-large-xsum-12-3  \
-    --save_dir xsum_generations \
-    --data_dir xsum \
-    --fp16  # you can pass generate kwargs like num_beams here, just like run_eval.py
-```
-
-Contributions that implement this command for other distributed hardware setups are welcome!
-
-#### run_eval tips and tricks
-
-When using `run_eval.py`, the following features can be useful:
-
-* if you running the script multiple times and want to make it easier to track what arguments produced that output, use `--dump-args`. Along with the results it will also dump any custom params that were passed to the script. For example if you used: `--num_beams 8 --early_stopping true`, the output will be:
-   ```
-   {'bleu': 26.887, 'n_obs': 10, 'runtime': 1, 'seconds_per_sample': 0.1, 'num_beams': 8, 'early_stopping': True}
-   ```
-
-   `--info` is an additional argument available for the same purpose of tracking the conditions of the experiment. It's useful to pass things that weren't in the argument list, e.g. a language pair `--info "lang:en-ru"`. But also if you pass `--info` without a value it will fallback to the current date/time string, e.g. `2020-09-13 18:44:43`.
-
-   If using `--dump-args --info`, the output will be:
-   
-   ```
-   {'bleu': 26.887, 'n_obs': 10, 'runtime': 1, 'seconds_per_sample': 0.1, 'num_beams': 8, 'early_stopping': True, 'info': '2020-09-13 18:44:43'}
-   ```
-
-   If using `--dump-args --info "pair:en-ru chkpt=best`, the output will be:
-   
-   ```
-   {'bleu': 26.887, 'n_obs': 10, 'runtime': 1, 'seconds_per_sample': 0.1, 'num_beams': 8, 'early_stopping': True, 'info': 'pair=en-ru chkpt=best'}
-   ```
-      
-
-* if you need to perform a parametric search in order to find the best ones that lead to the highest BLEU score, let `run_eval_search.py` to do the searching for you.
-
-   The script accepts the exact same arguments as `run_eval.py`, plus an additional argument `--search`. The value of `--search` is parsed, reformatted and fed to ``run_eval.py`` as additional args.
-
-   The format for the `--search` value is a simple string with hparams and colon separated values to try, e.g.:
-   ```
-    --search "num_beams=5:10 length_penalty=0.8:1.0:1.2 early_stopping=true:false"
-   ```
-   which will generate `12` `(2*3*2)` searches for a product of each hparam. For example the example that was just used will invoke `run_eval.py` repeatedly with:
-   
-   ```
-    --num_beams 5 --length_penalty 0.8 --early_stopping true
-    --num_beams 5 --length_penalty 0.8 --early_stopping false
-    [...]
-    --num_beams 10 --length_penalty 1.2 --early_stopping false
-   ```
-   
-   On completion, this function prints a markdown table of the results sorted by the best BLEU score and the winning arguments.
-
-```
-bleu  | num_beams | length_penalty | early_stopping
----- | --------- | -------------- | --------------
-26.71 |         5 |            1.1 |              1
-26.66 |         5 |            0.9 |              1
-26.66 |         5 |            0.9 |              0
-26.41 |         5 |            1.1 |              0
-21.94 |         1 |            0.9 |              1
-21.94 |         1 |            0.9 |              0
-21.94 |         1 |            1.1 |              1
-21.94 |         1 |            1.1 |              0
-
-Best score args:
-stas/wmt19-en-ru data/en-ru/val.source data/en-ru/test_translations.txt --reference_path data/en-ru/val.target --score_path data/en-ru/test_bleu.json --bs 8 --task translation --num_beams 5 --length_penalty 1.1 --early_stopping True
-```
-
-If you pass `--info "some experiment-specific info"` it will get printed before the results table - this is useful for scripting and multiple runs, so one can tell the different sets of results from each other.
-
-
 ### DistilBART
 ![DBART](https://huggingface.co/front/thumbnails/distilbart_large.png)

@@ -354,43 +263,3 @@ runtime: 13H on V-100 16GB GPU.
 ```bash
 pytest examples/seq2seq/
 ```
-
-### Converting pytorch-lightning checkpoints
-pytorch lightning ``-do_predict`` often fails, after you are done training, the best way to evaluate your model is to convert it.
-
-This should be done for you, with a file called `{save_dir}/best_tfmr`. 
-
-If that file doesn't exist but you have a lightning `.ckpt` file, you can run
-```bash
-python convert_pl_checkpoint_to_hf.py PATH_TO_CKPT  randomly_initialized_hf_model_path save_dir/best_tfmr
-```
-Then either `run_eval` or `run_distributed_eval` with `save_dir/best_tfmr` (see previous sections)
-
-
-## Experimental Features 
-These features are harder to use and not always useful.
-
-###  Dynamic Batch Size for MT
-`finetune.py` has a command line arg `--max_tokens_per_batch` that allows batches to be dynamically sized.
-This feature can only be used:
- with fairseq installed
- on 1 GPU
- without sortish sampler
- after calling `python save_len_file.py $tok $data_dir`
-
-For example, 
-```bash
-python save_len_file.py Helsinki-NLP/opus-mt-en-ro  wmt_en_ro
-./dynamic_bs_example.sh --max_tokens_per_batch=2000 --output_dir benchmark_dynamic_bs
-```
-splits `wmt_en_ro/train` into 11,197 uneven lengthed batches and can finish 1 epoch in 8 minutes on a v100.
-
-For comparison,
-```bash
-./dynamic_bs_example.sh --sortish_sampler --train_batch_size 48
-```
-uses 12,723 batches of length 48 and takes slightly more time 9.5 minutes.
-
-The feature is still experimental, because:
-+ we can make it much more robust if we have memory mapped/preprocessed datasets.
-+ The speedup over sortish sampler is not that large at the moment.
@@ -1,5 +0,0 @@
-import os
-import sys
-
-
-sys.path.insert(1, os.path.dirname(os.path.realpath(__file__)))
@@ -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
@@ -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,
    )
@@ -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,
            ),
        }

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

@@ -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,)
@@ -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)
@@ -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 \
-  "$@"
@@ -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 \
-  "$@"
@@ -1,7 +1,6 @@
 import argparse
 import gc
 import os
-import warnings
 from pathlib import Path
 from typing import List

@@ -10,24 +9,37 @@ import torch
 from torch import nn
 from torch.nn import functional as F

-from finetune import SummarizationModule, TranslationModule
-from finetune import main as ft_main
-from initialization_utils import copy_layers, init_student
-from transformers import AutoModelForSeq2SeqLM, MBartTokenizer, T5Config, T5ForConditionalGeneration
-from transformers.modeling_bart import shift_tokens_right
-from utils import (
-    any_requires_grad,
-    assert_all_frozen,
-    calculate_bleu,
-    freeze_params,
-    label_smoothed_nll_loss,
-    use_task_specific_params,
-)
+from lightning_base import generic_train
+from transformers import BartConfig, BartForConditionalGeneration, MBartTokenizer, T5Config, T5ForConditionalGeneration
+
+
+try:
+    from .finetune import SummarizationModule, TranslationModule
+    from .initialization_utils import init_student, copy_layers
+    from .utils import (
+        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 init_student, copy_layers
+    from utils import (
+        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):
@@ -62,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)
@@ -157,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:
@@ -206,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)
@@ -248,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
@@ -269,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):
@@ -356,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)
@@ -377,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
@@ -410,54 +399,56 @@ def create_module(args):
    return 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],
-        12: [0, 1, 2, 3, 4, 5, 6, 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 evaluate_checkpoint(ckpt_path: Path, dest_dir=None):
+    exp_dir = ckpt_path.parent
+    if dest_dir is None:
+        dest_dir = exp_dir
+    clash = list(dest_dir.glob("test_generations*"))
+    if clash:
+        print(f"SKIPPING to avoid overwriting {clash}")
+    ckpt = torch.load(ckpt_path, map_location="cpu")
+    if "hparams" in ckpt:
+        args = argparse.Namespace(**ckpt["hparams"])
+    else:
+        args = argparse.Namespace(**pickle_load(exp_dir / "hparams.pkl"))
+    args.resume_from_checkpoint = str(ckpt_path)
+    args.do_train = False
+    args.output_dir = str(dest_dir)
+    args.n_gpu = 1
+    args.eval_batch_size = 16
+    Path(args.output_dir).mkdir(exist_ok=True)
+    model = create_module(args)
+    trainer: pl.Trainer = generic_train(model, args, early_stopping_callback=False)
+    trainer.test(model)


-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):
@@ -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)
@@ -1,17 +0,0 @@
-#!/usr/bin/env bash
-export PYTHONPATH="../":"${PYTHONPATH}"
-export WANDB_PROJECT=dmar
-export MAX_LEN=128
-export m=sshleifer/student_marian_en_ro_6_1
-python finetune.py \
-  --learning_rate=3e-4 \
-  --do_train \
-  --fp16 \
-  --data_dir wmt_en_ro \
-  --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=48 --eval_batch_size=64 \
-  --tokenizer_name $m --model_name_or_path $m --num_train_epochs=1 \
-  --warmup_steps 500 --logger_name wandb --gpus 1 \
-  --fp16_opt_level=O1 --task translation \
-  "$@"
@@ -12,49 +12,66 @@ import pytorch_lightning as pl
 import torch
 from torch.utils.data import DataLoader

-from callbacks import Seq2SeqLoggingCallback, get_checkpoint_callback, get_early_stopping_callback
 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
-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,
-    lmap,
-    pickle_save,
-    save_git_info,
-    save_json,
-    use_task_specific_params,
-)


+try:
+    from .utils import (
+        assert_all_frozen,
+        use_task_specific_params,
+        lmap,
+        flatten_list,
+        pickle_save,
+        save_git_info,
+        save_json,
+        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 utils import (
+        Seq2SeqDataset,
+        TranslationDataset,
+        assert_all_frozen,
+        use_task_specific_params,
+        lmap,
+        flatten_list,
+        pickle_save,
+        save_git_info,
+        save_json,
+        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__)
-from pytorch_lightning.utilities import rank_zero_only
+

 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
-        elif hparams.max_tokens_per_batch is not None:
-            if hparams.gpus > 1:
-                raise NotImplementedError("Dynamic Batch size does not work for multi-gpu training")
-            if hparams.sortish_sampler:
-                raise ValueError("--sortish_sampler and --max_tokens_per_batch may not be used simultaneously")
-
        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"
@@ -81,6 +98,7 @@ class SummarizationModule(BaseTransformer):
        }
        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}"
+
        if self.hparams.freeze_embeds:
            self.freeze_embeds()
        if self.hparams.freeze_encoder:
@@ -89,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."""
@@ -127,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:
@@ -158,11 +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["bs"] = batch["input_ids"].shape[0]
-        logs["src_pad_tok"] = batch["input_ids"].eq(self.pad).sum()
-        logs["src_pad_frac"] = batch["input_ids"].eq(self.pad).float().mean()
-        # TODO(SS): make a wandb summary metric for this
+        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:
@@ -172,39 +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
-
-        def get_date_str(seconds=True) -> str:
-            """Returns 2019-09-25-10:02:07, for example."""
-            if seconds:
-                return time.strftime('%Y-%m-%d-%H:%M:%S')
-            else:
-                return time.strftime('%Y-%m-%d-%H:%M')
-
-        all_metrics['Time'] = get_date_str(seconds=True)
-        #all_metrics['n_obs'] =
-
-
-        all_metrics['rank'] = getattr(self.train_dataloader(), '_rank', -1.)
-
-        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)
@@ -215,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,39 +253,21 @@ class SummarizationModule(BaseTransformer):

    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

-        if self.hparams.sortish_sampler and type_path != "test":
-            sampler = dataset.make_sortish_sampler(batch_size, distributed=self.hparams.gpus > 1)
-            return DataLoader(
-                dataset,
-                batch_size=batch_size,
-                collate_fn=dataset.collate_fn,
-                shuffle=False,
-                num_workers=self.num_workers,
-                sampler=sampler,
-            )
-
-        elif self.hparams.max_tokens_per_batch is not None and type_path != "test":
-            batch_sampler = dataset.make_dynamic_sampler(
-                self.hparams.max_tokens_per_batch, distributed=self.hparams.gpus > 1
-            )
-            return DataLoader(
-                dataset,
-                batch_sampler=batch_sampler,
-                collate_fn=dataset.collate_fn,
-                # shuffle=False,
-                num_workers=self.num_workers,
-                # batch_size=None,
-            )
-        else:
-            return DataLoader(
-                dataset,
-                batch_size=batch_size,
-                collate_fn=dataset.collate_fn,
-                shuffle=shuffle,
-                num_workers=self.num_workers,
-                sampler=None,
-            )
+        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)
@@ -334,7 +314,6 @@ class SummarizationModule(BaseTransformer):
        parser.add_argument("--freeze_encoder", action="store_true")
        parser.add_argument("--freeze_embeds", action="store_true")
        parser.add_argument("--sortish_sampler", action="store_true", default=False)
-        parser.add_argument("--max_tokens_per_batch", type=int, default=None)
        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=500, required=False, help="# examples. -1 means use all.")
@@ -343,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,
@@ -363,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)
@@ -373,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:
@@ -381,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"
@@ -408,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:
@@ -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 \
    "$@"
@@ -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
@@ -12,7 +12,7 @@ Note: You need to have your test_generations.txt before you start this process.
 cd $HOME
 git clone git@github.com:moses-smt/mosesdecoder.git
 cd mosesdecoder  
-git clone git@github.com:rsennrich/wmt16-scripts.git
+git@github.com:rsennrich/wmt16-scripts.git
 ```

 (2) define a function for post processing.
@@ -1,210 +0,0 @@
-import argparse
-import shutil
-import time
-from json import JSONDecodeError
-from logging import getLogger
-from pathlib import Path
-from typing import Dict, List
-
-import torch
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-
-from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from utils import (
-    Seq2SeqDataset,
-    calculate_bleu,
-    calculate_rouge,
-    lmap,
-    load_json,
-    parse_numeric_n_bool_cl_kwargs,
-    save_json,
-    use_task_specific_params,
-    write_txt_file,
-)
-
-
-logger = getLogger(__name__)
-
-
-def eval_data_dir(
-    data_dir,
-    save_dir: str,
-    model_name: str,
-    bs: int = 8,
-    max_source_length: int = 1024,
-    type_path="val",
-    n_obs=None,
-    fp16=False,
-    task="summarization",
-    local_rank=None,
-    **generate_kwargs,
-) -> Dict:
-    """Run evaluation on part of the data for one gpu and save to {save_dir}/rank_{rank}_output.json"""
-    model_name = str(model_name)
-    assert local_rank is not None
-    torch.distributed.init_process_group(backend="nccl", rank=local_rank)
-
-    save_dir = Path(save_dir)
-    save_path = save_dir.joinpath(f"rank_{local_rank}_output.json")
-    torch.cuda.set_device(local_rank)
-    model = AutoModelForSeq2SeqLM.from_pretrained(model_name).cuda()
-    if fp16:
-        model = model.half()
-
-    tokenizer = AutoTokenizer.from_pretrained(model_name)
-    logger.info(f"Inferred tokenizer type: {tokenizer.__class__}")  # if this is wrong, check config.model_type.
-    use_task_specific_params(model, task)  # update config with task specific params
-    if max_source_length is None:
-        max_source_length = tokenizer.model_max_length
-    ds = Seq2SeqDataset(
-        tokenizer,
-        data_dir,
-        max_source_length,
-        max_target_length=1024,
-        type_path=type_path,
-        n_obs=n_obs,
-        prefix=model.config.prefix,
-    )
-    # I set shuffle=True for a more accurate progress bar.
-    # If all the longest samples are first, the prog bar estimate is too high at the beginning.
-    sampler = ds.make_sortish_sampler(bs, distributed=True, add_extra_examples=False, shuffle=True)
-    data_loader = DataLoader(ds, sampler=sampler, batch_size=bs, collate_fn=ds.collate_fn)
-    results = []
-    for batch in tqdm(data_loader):
-        summaries = model.generate(
-            input_ids=batch["input_ids"].to(model.device),
-            attention_mask=batch["attention_mask"].to(model.device),
-            **generate_kwargs,
-        )
-        preds = tokenizer.batch_decode(summaries, skip_special_tokens=True, clean_up_tokenization_spaces=False)
-        ids = batch["ids"]
-        for i, pred in enumerate(preds):
-            results.append(dict(pred=pred, id=ids[i].item()))
-    save_json(results, save_path)
-    return results, sampler.num_replicas
-
-
-def run_generate():
-    parser = argparse.ArgumentParser(
-        epilog="Unspecified args like --num_beams=2 --decoder_start_token_id=4 are passed to model.generate"
-    )
-    parser.add_argument("--data_dir", type=str, help="like cnn_dm/test.source")
-    parser.add_argument(
-        "--model_name",
-        type=str,
-        help="like facebook/bart-large-cnn,t5-base, etc.",
-        default="sshleifer/distilbart-xsum-12-3",
-    )
-    parser.add_argument("--save_dir", type=str, help="where to save", default="tmp_gen")
-    parser.add_argument("--max_source_length", type=int, default=None)
-    parser.add_argument(
-        "--type_path", type=str, default="test", help="which subset to evaluate typically train/val/test"
-    )
-    parser.add_argument("--reference_path", type=str, required=False, help="like cnn_dm/test.target")
-    parser.add_argument("--task", type=str, default="summarization", help="used for task_specific_params + metrics")
-    parser.add_argument("--bs", type=int, default=8, required=False, help="batch size")
-    parser.add_argument(
-        "--local_rank", type=int, default=-1, required=False, help="should be passed by distributed.launch"
-    )
-
-    parser.add_argument(
-        "--n_obs", type=int, default=None, required=False, help="How many observations. Defaults to all."
-    )
-    parser.add_argument(
-        "--sync_timeout",
-        type=int,
-        default=600,
-        required=False,
-        help="How long should master process wait for other processes to finish.",
-    )
-    parser.add_argument("--fp16", action="store_true")
-    parser.add_argument("--debug", action="store_true")
-    start_time = time.time()
-    args, rest = parser.parse_known_args()
-    generate_kwargs = parse_numeric_n_bool_cl_kwargs(rest)
-    if generate_kwargs and args.local_rank <= 0:
-        print(f"parsed the following generate kwargs: {generate_kwargs}")
-    json_save_dir = Path(args.save_dir + "_tmp")
-    Path(json_save_dir).mkdir(exist_ok=True)  # this handles locking.
-    intermediate_files = list(json_save_dir.glob("rank_*.json"))
-    if intermediate_files:
-        raise ValueError(f"Found files at {json_save_dir} please move or remove them.")
-        # In theory, a node could finish and save before another node hits this. If this happens, we can address later.
-
-    Path(args.save_dir).mkdir(exist_ok=True)
-    results, num_replicas = eval_data_dir(
-        args.data_dir,
-        json_save_dir,
-        args.model_name,
-        type_path=args.type_path,
-        bs=args.bs,
-        fp16=args.fp16,
-        task=args.task,
-        local_rank=args.local_rank,
-        n_obs=args.n_obs,
-        max_source_length=args.max_source_length,
-        **generate_kwargs,
-    )
-
-    if args.local_rank <= 0:
-        save_dir = Path(args.save_dir)
-        save_dir.mkdir(exist_ok=True)
-        partial_results = gather_results_from_each_node(num_replicas, json_save_dir, args.sync_timeout)
-        preds = combine_partial_results(partial_results)
-        tgt_file = Path(args.data_dir).joinpath(args.type_path + ".target")
-        labels = [x.rstrip() for x in open(tgt_file).readlines()][: len(preds)]
-
-        # Calculate metrics, save metrics,  and save _generations.txt
-        calc_bleu = "translation" in args.task
-        score_fn = calculate_bleu if calc_bleu else calculate_rouge
-        metric_name = "bleu" if calc_bleu else "rouge"
-        metrics: Dict = score_fn(preds, labels)
-        metrics["n_obs"] = len(preds)
-        runtime = time.time() - start_time
-        metrics["seconds_per_sample"] = round(runtime / metrics["n_obs"], 4)
-        metrics["n_gpus"] = num_replicas
-        # TODO(@stas00): add whatever metadata to metrics
-        metrics_save_path = save_dir.joinpath(f"{args.type_path}_{metric_name}.json")
-        save_json(metrics, metrics_save_path, indent=None)
-        print(metrics)
-        write_txt_file(preds, save_dir.joinpath(f"{args.type_path}_generations.txt"))
-        if args.debug:
-            write_txt_file(labels, save_dir.joinpath(f"{args.type_path}.target"))
-        else:
-            shutil.rmtree(json_save_dir)
-
-
-def combine_partial_results(partial_results) -> List:
-    """Concatenate partial results into one file, then sort it by id."""
-    records = []
-    for partial_result in partial_results:
-        records.extend(partial_result)
-    records = list(sorted(records, key=lambda x: x["id"]))
-    preds = [x["pred"] for x in records]
-    return preds
-
-
-def gather_results_from_each_node(num_replicas, save_dir, timeout) -> List[Dict[str, List]]:
-    # WAIT FOR lots of .json files
-    start_wait = time.time()
-    logger.info("waiting for all nodes to finish")
-    json_data = None
-    while (time.time() - start_wait) < timeout:
-        json_files = list(save_dir.glob("rank_*.json"))
-        if len(json_files) < num_replicas:
-            continue
-        try:
-            # make sure all json files are fully saved
-            json_data = lmap(load_json, json_files)
-            return json_data
-        except JSONDecodeError:
-            continue
-    else:
-        raise TimeoutError("Rank 0 gave up on waiting for other processes")
-    # Unreachable
-
-
-if __name__ == "__main__":
-    # Usage for MT:
-    run_generate()
@@ -1,21 +1,17 @@
 import argparse
-import datetime
 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

 from transformers import AutoModelForSeq2SeqLM, AutoTokenizer
-from utils import calculate_bleu, calculate_rouge, parse_numeric_n_bool_cl_kwargs, use_task_specific_params


-logger = getLogger(__name__)
-
+try:
+    from .utils import calculate_rouge, use_task_specific_params, calculate_bleu_score, trim_batch
+except ImportError:
+    from utils import calculate_rouge, use_task_specific_params, calculate_bleu_score, trim_batch

 DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"

@@ -27,105 +23,74 @@ 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 datetime_now():
-    return datetime.datetime.now().strftime("%Y-%m-%d %H:%M:%S")
-
-
-def run_generate(verbose=True):
-    """
-
-    Takes input text, generates output, and then using reference calculates the BLEU scores.
-
-    The results are saved to a file and returned to the caller, and printed out unless ``verbose=False`` is passed.
-
-    Args:
-        verbose (:obj:`bool`, `optional`, defaults to :obj:`True`): print results to stdout
-
-    Returns:
-        a tuple: ``(scores, params}``
-        - ``scores``: a dict of scores data ``{'bleu': 39.6501, 'n_obs': 2000, 'runtime': 186, 'seconds_per_sample': 0.093}``
-        - ``params``: a dict of custom params, e.g. ``{'num_beams': 5, 'length_penalty': 0.8}``
-    """
-
+def run_generate():
    parser = argparse.ArgumentParser()
    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")
-    parser.add_argument("--dump-args", action="store_true", help="print the custom hparams with the results")
-    parser.add_argument(
-        "--info",
-        nargs="?",
-        type=str,
-        const=datetime_now(),
-        help="use in conjunction w/ --dump-args to print with the results whatever other info you'd like, e.g. lang=en-ru. If no value is passed, the current datetime string will be used.",
-    )
-    # Unspecified args like --num_beams=2 --decoder_start_token_id=4 are passed to model.generate
-    args, rest = parser.parse_known_args()
-    parsed_args = parse_numeric_n_bool_cl_kwargs(rest)
-    if parsed_args and verbose:
-        print(f"parsed the following generate kwargs: {parsed_args}")
+    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,
@@ -133,36 +98,20 @@ def run_generate(verbose=True):
        device=args.device,
        fp16=args.fp16,
        task=args.task,
-        prefix=args.prefix,
-        **parsed_args,
+        decoder_start_token_id=args.decoder_start_token_id,
    )
-
    if args.reference_path is None:
-        return {}
-
+        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)
-
-    if args.dump_args:
-        scores.update(parsed_args)
-    if args.info:
-        scores["info"] = args.info
-
-    if verbose:
-        print(scores)
-
+    print(scores)
    if args.score_path is not None:
-        path = args.score_path
-        json.dump(scores, open(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(verbose=True)
+    run_generate()
@@ -1,135 +0,0 @@
-import argparse
-import itertools
-import operator
-import sys
-from collections import OrderedDict
-
-from run_eval import datetime_now, run_generate
-
-
-# A table of supported tasks and the list of scores in the order of importance to be sorted by.
-# To add a new task, simply list the score names that `run_eval.run_generate()` returns
-task_score_names = {
-    "translation": ["bleu"],
-    "summarization": ["rouge1", "rouge2", "rougeL"],
-}
-
-
-def parse_search_arg(search):
-    groups = search.split()
-    entries = {k: vs for k, vs in (g.split("=") for g in groups)}
-    entry_names = list(entries.keys())
-    sets = [list((f"--{k} {v}") for v in vs.split(":")) for k, vs in entries.items()]
-    matrix = [list(x) for x in itertools.product(*sets)]
-    return matrix, entry_names
-
-
-def run_search():
-    """
-     Run parametric search over the desired hparam space with help of ``run_eval.py``.
-
-     All the arguments except ``--search`` are passed to ``run_eval.py`` as is. The values inside of "--search" are parsed, reformatted and fed to ``run_eval.py`` as additional args.
-
-    The format for the ``--search`` value is a simple string with hparams and colon separated values to try, e.g.:
-    ```
-     --search "num_beams=5:10 length_penalty=0.8:1.0:1.2 early_stopping=true:false"
-    ```
-    which will generate ``12`` ``(2*3*2)`` searches for a product of each hparam. For example the example that was just used will invoke ``run_eval.py`` repeatedly with:
-
-    ```
-     --num_beams 5 --length_penalty 0.8 --early_stopping true
-     --num_beams 5 --length_penalty 0.8 --early_stopping false
-     [...]
-     --num_beams 10 --length_penalty 1.2 --early_stopping false
-    ```
-
-    On completion, this function prints a markdown table of the results sorted by the best BLEU score and the winning arguments.
-
-
-    """
-    prog = sys.argv[0]
-
-    parser = argparse.ArgumentParser(
-        usage="\n\nImportant: this script accepts all arguments `run_eval.py` accepts and then a few extra, therefore refer to `run_eval.py -h` for the complete list."
-    )
-    parser.add_argument(
-        "--search",
-        type=str,
-        required=False,
-        help='param space to search, e.g. "num_beams=5:10 length_penalty=0.8:1.0:1.2"',
-    )
-    parser.add_argument(
-        "--bs", type=int, default=8, required=False, help="initial batch size (may get reduced if it's too big)"
-    )
-    parser.add_argument("--task", type=str, help="used for task_specific_params + metrics")
-    parser.add_argument(
-        "--info",
-        nargs="?",
-        type=str,
-        const=datetime_now(),
-        help="add custom notes to be printed before the results table. If no value is passed, the current datetime string will be used.",
-    )
-    args, args_main = parser.parse_known_args()
-    # we share some of the args
-    args_main.extend(["--task", args.task])
-    args_normal = [prog] + args_main
-
-    # to support variations like translation_en_to_de"
-    task = "translation" if "translation" in args.task else "summarization"
-
-    matrix, col_names = parse_search_arg(args.search)
-    col_names[0:0] = task_score_names[task]  # score cols first
-    col_widths = {col: len(str(col)) for col in col_names}
-    results = []
-    for r in matrix:
-        hparams = {k: v for k, v in (x.replace("--", "").split() for x in r)}
-        args_exp = " ".join(r).split()
-        args_exp.extend(["--bs", str(args.bs)])  # in case we need to reduce its size due to CUDA OOM
-        sys.argv = args_normal + args_exp
-
-        # XXX: need to trap CUDA OOM and lower args.bs if that happens and retry
-
-        scores = run_generate(verbose=False)
-        # make sure scores are first in the table
-        result = OrderedDict()
-        for score in task_score_names[task]:
-            result[score] = scores[score]
-        result.update(hparams)
-        results.append(result)
-
-        # find widest entries
-        for k, v in result.items():
-            l = len(str(v))
-            if l > col_widths[k]:
-                col_widths[k] = l
-
-    results_sorted = sorted(results, key=operator.itemgetter(*task_score_names[task]), reverse=True)
-    print(" | ".join([f"{col:{col_widths[col]}}" for col in col_names]))
-    print(" | ".join([f"{'-'*col_widths[col]}" for col in col_names]))
-    for row in results_sorted:
-        print(" | ".join([f"{row[col]:{col_widths[col]}}" for col in col_names]))
-
-    best = results_sorted[0]
-    for score in task_score_names[task]:
-        del best[score]
-    best_args = [f"--{k} {v}" for k, v in best.items()]
-    dyn_args = ["--bs", str(args.bs)]
-    if args.info:
-        print(f"\nInfo: {args.info}")
-    print("\nBest score args:")
-    print(" ".join(args_main + best_args + dyn_args))
-
-    return results_sorted
-
-
-if __name__ == "__main__":
-    # Usage:
-    # [normal-run_eval_search.py cmd plus] \
-    # --search="num_beams=1:5:10 length_penalty=0.8:1:1.2 early_stopping=true:false"
-    #
-    # Example:
-    # PYTHONPATH="src:examples/seq2seq" python examples/seq2seq/run_eval_search.py $MODEL_NAME \
-    # $DATA_DIR/val.source $SAVE_DIR/test_translations.txt --reference_path $DATA_DIR/val.target \
-    # --score_path $SAVE_DIR/test_bleu.json --bs $BS --task translation \
-    # --search="num_beams=1:5:10 length_penalty=0.8:1:1.2 early_stopping=true:false"
-    run_search()
@@ -1,46 +0,0 @@
-import fire
-from torch.utils.data import DataLoader
-from tqdm import tqdm
-
-from transformers import AutoTokenizer
-
-
-try:
-    from .utils import Seq2SeqDataset, pickle_save
-except ImportError:
-    from utils import Seq2SeqDataset, pickle_save
-
-
-def save_len_file(
-    tokenizer_name, data_dir, max_source_length=1024, max_target_length=1024, consider_target=False, **kwargs
-):
-    """Save max(src_len, tgt_len) for each example to allow dynamic batching."""
-    tok = AutoTokenizer.from_pretrained(tokenizer_name)
-    train_ds = Seq2SeqDataset(tok, data_dir, max_source_length, max_target_length, type_path="train", **kwargs)
-    pad = tok.pad_token_id
-
-    def get_lens(ds):
-        dl = tqdm(
-            DataLoader(ds, batch_size=512, num_workers=8, shuffle=False, collate_fn=ds.collate_fn),
-            desc=str(ds.len_file),
-        )
-        max_lens = []
-        for batch in dl:
-            src_lens = batch["input_ids"].ne(pad).sum(1).tolist()
-            tgt_lens = batch["labels"].ne(pad).sum(1).tolist()
-            if consider_target:
-                for src, tgt in zip(src_lens, tgt_lens):
-                    max_lens.append(max(src, tgt))
-            else:
-                max_lens.extend(src_lens)
-        return max_lens
-
-    train_lens = get_lens(train_ds)
-    val_ds = Seq2SeqDataset(tok, data_dir, max_source_length, max_target_length, type_path="val", **kwargs)
-    val_lens = get_lens(val_ds)
-    pickle_save(train_lens, train_ds.len_file)
-    pickle_save(val_lens, val_ds.len_file)
-
-
-if __name__ == "__main__":
-    fire.Fire(save_len_file)
@@ -10,17 +10,16 @@ import pytorch_lightning as pl
 import timeout_decorator
 import torch

-from distillation import BartSummarizationDistiller, distill_main
-from finetune import SummarizationModule, main
-from test_seq2seq_examples import CUDA_AVAILABLE, MBART_TINY
-from transformers import BartForConditionalGeneration, MarianMTModel
+from transformers import BartForConditionalGeneration
 from transformers.testing_utils import slow
-from utils import load_json
+
+from .finetune import SummarizationModule, main
+from .test_seq2seq_examples import CUDA_AVAILABLE, MBART_TINY
+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
@@ -28,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)
@@ -37,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",
@@ -80,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

@@ -1,33 +0,0 @@
-#!/usr/bin/env python
-
-import io
-import json
-import subprocess
-
-
-pairs = [
-    ["en", "ru"],
-    ["ru", "en"],
-    ["en", "de"],
-    ["de", "en"],
-]
-
-n_objs = 8
-
-
-def get_all_data(pairs, n_objs):
-    text = {}
-    for src, tgt in pairs:
-        pair = f"{src}-{tgt}"
-        cmd = f"sacrebleu -t wmt19 -l {pair} --echo src".split()
-        src_lines = subprocess.run(cmd, stdout=subprocess.PIPE).stdout.decode("utf-8").splitlines()
-        cmd = f"sacrebleu -t wmt19 -l {pair} --echo ref".split()
-        tgt_lines = subprocess.run(cmd, stdout=subprocess.PIPE).stdout.decode("utf-8").splitlines()
-        text[pair] = {"src": src_lines[:n_objs], "tgt": tgt_lines[:n_objs]}
-    return text
-
-
-text = get_all_data(pairs, n_objs)
-filename = "./fsmt_val_data.json"
-with io.open(filename, "w", encoding="utf-8") as f:
-    bleu_data = json.dump(text, f, indent=2, ensure_ascii=False)
@@ -1,90 +0,0 @@
-{
-  "en-ru": {
-    "src": [
-      "Welsh AMs worried about 'looking like muppets'",
-      "There is consternation among some AMs at a suggestion their title should change to MWPs (Member of the Welsh Parliament).",
-      "It has arisen because of plans to change the name of the assembly to the Welsh Parliament.",
-      "AMs across the political spectrum are worried it could invite ridicule.",
-      "One Labour AM said his group was concerned \"it rhymes with Twp and Pwp.\"",
-      "For readers outside of Wales: In Welsh twp means daft and pwp means poo.",
-      "A Plaid AM said the group as a whole was \"not happy\" and has suggested alternatives.",
-      "A Welsh Conservative said his group was \"open minded\" about the name change, but noted it was a short verbal hop from MWP to Muppet."
-    ],
-    "tgt": [
-      "Члены Национальной ассамблеи Уэльса обеспокоены, что \"выглядят как куклы\"",
-      "Некоторые члены Национальной ассамблеи Уэльса в ужасе от предложения о том, что их наименование должно измениться на MPW (члены Парламента Уэльса).",
-      "Этот вопрос был поднят в связи с планами по переименованию ассамблеи в Парламент Уэльса.",
-      "Члены Национальной ассамблеи Уэльса всего политического спектра обеспокоены, что это может породить насмешки.",
-      "Один из лейбористских членов Национальной ассамблеи Уэльса сказал, что его партия обеспокоена тем, что \"это рифмуется с Twp и Pwp\".",
-      "Для читателей за предлами Уэльса: по-валлийски twp означает \"глупый\", а pwp означает \"какашка\".",
-      "Член Национальной ассамблеи от Плайд сказал, что эта партия в целом \"не счастлива\" и предложил альтернативы.",
-      "Представитель Консервативной партии Уэльса сказал, что его партия \"открыта\" к переименованию, но отметил, что между WMP и Muppet небольшая разница в произношении."
-    ]
-  },
-  "ru-en": {
-    "src": [
-      "Названо число готовящихся к отправке в Донбасс новобранцев из Украины",
-      "Официальный представитель Народной милиции самопровозглашенной Луганской Народной Республики (ЛНР) Андрей Марочко заявил, что зимой 2018-2019 года Украина направит в Донбасс не менее 3 тыс. новобранцев.",
-      "По его словам, таким образом Киев планирует \"хоть как-то доукомплектовать подразделения\".",
-      "\"Нежелание граждан Украины проходить службу в рядах ВС Украины, массовые увольнения привели к низкой укомплектованности подразделений\", - рассказал Марочко, которого цитирует \"РИА Новости\".",
-      "Он также не исключил, что реальные цифры призванных в армию украинцев могут быть увеличены в случае необходимости.",
-      "В 2014-2017 годах Киев начал так называемую антитеррористическую операцию (АТО), которую позже сменили на операцию объединенных сил (ООС).",
-      "Предполагалось, что эта мера приведет к усилению роли украинских силовиков в урегулировании ситуации.",
-      "В конце августа 2018 года ситуация в Донбассе обострилась из-за убийства главы ДНР Александра Захарченко."
-    ],
-    "tgt": [
-      "The number of new Ukrainian recruits ready to go to Donbass has become public",
-      "Official representative of the peoples’ militia of the self-proclaimed Lugansk People’s Republic Andrey Marochko claimed that Ukrainian will send at least 3 thousand new recruits to Donbass in winter 2018-2019.",
-      "This is how Kyiv tries “at least somehow to staff the units,” he said.",
-      "“The unwillingness of Ukrainian citizens to serve in the Ukraine’s military forces, mass resignments lead to low understaffing,” said Marochko cited by RIA Novosti.",
-      "Also, he doesn’t exclude that the real numbers of conscripts in the Ukrainian army can be raised is necessary.",
-      "In 2014-2017, Kyiv started so-called antiterrorist operation, that ws later changed to the united forces operation.",
-      "This measure was supposed to strengthen the role of the Ukrainian military in settling the situation.",
-      "In the late August 2018, the situation in Donbass escalated as the DNR head Aleksandr Zakharchenko was killed."
-    ]
-  },
-  "en-de": {
-    "src": [
-      "Welsh AMs worried about 'looking like muppets'",
-      "There is consternation among some AMs at a suggestion their title should change to MWPs (Member of the Welsh Parliament).",
-      "It has arisen because of plans to change the name of the assembly to the Welsh Parliament.",
-      "AMs across the political spectrum are worried it could invite ridicule.",
-      "One Labour AM said his group was concerned \"it rhymes with Twp and Pwp.\"",
-      "For readers outside of Wales: In Welsh twp means daft and pwp means poo.",
-      "A Plaid AM said the group as a whole was \"not happy\" and has suggested alternatives.",
-      "A Welsh Conservative said his group was \"open minded\" about the name change, but noted it was a short verbal hop from MWP to Muppet."
-    ],
-    "tgt": [
-      "Walisische Ageordnete sorgen sich \"wie Dödel auszusehen\"",
-      "Es herrscht Bestürzung unter einigen Mitgliedern der Versammlung über einen Vorschlag, der ihren Titel zu MWPs (Mitglied der walisischen Parlament) ändern soll.",
-      "Der Grund dafür waren Pläne, den Namen der Nationalversammlung in Walisisches Parlament zu ändern.",
-      "Mitglieder aller Parteien der Nationalversammlung haben Bedenken, dass sie sich dadurch Spott aussetzen könnten.",
-      "Ein Labour-Abgeordneter sagte, dass seine Gruppe \"sich mit Twp und Pwp reimt\".",
-      "Hinweis für den Leser: „twp“ im Walisischen bedeutet „bescheuert“ und „pwp“ bedeutet „Kacke“.",
-      "Ein Versammlungsmitglied von Plaid Cymru sagte, die Gruppe als Ganzes sei \"nicht glücklich\" und hat Alternativen vorgeschlagen.",
-      "Ein walisischer Konservativer sagte, seine Gruppe wäre „offen“ für eine Namensänderung, wies aber darauf hin, dass es von „MWP“ (Mitglied des Walisischen Parlaments) nur ein kurzer verbaler Sprung zu „Muppet“ ist."
-    ]
-  },
-  "de-en": {
-    "src": [
-      "Schöne Münchnerin 2018: Schöne Münchnerin 2018 in Hvar: Neun Dates",
-      "Von az, aktualisiert am 04.05.2018 um 11:11",
-      "Ja, sie will...",
-      "\"Schöne Münchnerin\" 2018 werden!",
-      "Am Nachmittag wartet erneut eine Überraschung auf unsere Kandidatinnen: sie werden das romantische Candlelight-Shooting vor der MY SOLARIS nicht alleine bestreiten, sondern an der Seite von Male-Model Fabian!",
-      "Hvar - Flirten, kokettieren, verführen - keine einfachen Aufgaben für unsere Mädchen.",
-      "Insbesondere dann, wenn in Deutschland ein Freund wartet.",
-      "Dennoch liefern die neun \"Schöne Münchnerin\"-Kandidatinnen beim Shooting mit People-Fotograf Tuan ab und trotzen Wind, Gischt und Regen wie echte Profis."
-    ],
-    "tgt": [
-      "The Beauty of Munich 2018: the Beauty of Munich 2018 in Hvar: Nine dates",
-      "From A-Z, updated on 04/05/2018 at 11:11",
-      "Yes, she wants to...",
-      "to become \"The Beauty of Munich\" in 2018!",
-      "In the afternoon there is another surprise waiting for our contestants: they will be competing for the romantic candlelight photo shoot at MY SOLARIS not alone, but together with a male-model Fabian!",
-      "Hvar with its flirting, coquetting, and seduction is not an easy task for our girls.",
-      "Especially when there is a boyfriend waiting in Germany.",
-      "Despite dealing with wind, sprays and rain, the nine contestants of \"The Beauty of Munich\" behaved like real professionals at the photo shoot with People-photographer Tuan."
-    ]
-  }
-}
@@ -1,11 +1,8 @@
-Corrections to votes and voting intentions: see Minutes Assignment conferred on a Member: see Minutes Membership of committees and delegations: see Minutes Decisions concerning certain documents: see Minutes Forwarding of texts adopted during the sitting: see Minutes Dates for next sittings: see Minutes
 Membership of Parliament: see Minutes Approval of Minutes of previous sitting: see Minutes Membership of Parliament: see Minutes Verification of credentials: see Minutes Documents received: see Minutes Written statements and oral questions (tabling): see Minutes Petitions: see Minutes Texts of agreements forwarded by the Council: see Minutes Action taken on Parliament's resolutions: see Minutes Agenda for next sitting: see Minutes Closure of sitting (The sitting was closed at 7.45 p.m.)
 Election of Vice-Presidents of the European Parliament (deadline for submitting nominations): see Minutes (The sitting was suspended at 12.40 p.m. and resumed at 3.00 p.m.) Election of Quaestors of the European Parliament (deadline for submitting nominations): see Minutes (The sitting was suspended at 3.25 p.m. and resumed at 6.00 p.m.) Agenda for next sitting: see Minutes Closure of sitting (The sitting was closed at 6.15 p.m.) Opening of the sitting (The sitting was opened at 9.35 a.m.) Documents received: see Minutes Approval of Minutes of previous sitting: see Minutes Membership of Parliament: see Minutes
 Membership of committees (deadline for tabling amendments): see Minutes (The sitting was suspended at 7 p.m. and resumed at 9 p.m.) Agenda for next sitting: see Minutes Closure of sitting (The sitting was suspended at 23.25 p.m.) Documents received: see Minutes Communication of Council common positions: see Minutes (The sitting was suspended at 11.35 a.m. and resumed for voting time at noon) Approval of Minutes of previous sitting: see Minutes Committee of Inquiry into the crisis of the Equitable Life Assurance Society (extension of mandate): see Minutes
 Announcement by the President: see Minutes 1. Membership of committees (vote) 2. Amendment of the ACP-EC Partnership Agreement (vote) 4. Certification of train drivers operating locomotives and trains on the railway system in the Community (vote) 6. Law applicable to non-contractual obligations ("ROME II") (vote) 8. Seventh and eighth annual reports on arms exports (vote) Corrections to votes and voting intentions: see Minutes Membership of committees and delegations: see Minutes Request for waiver of parliamentary immunity: see Minutes Decisions concerning certain documents: see Minutes
-Written statements for entry
 Written statements for entry in the register (Rule 116): see Minutes Forwarding of texts adopted during the sitting: see Minutes Dates for next sittings: see Minutes Adjournment of the session I declare the session of the European Parliament adjourned. (The sitting was closed at 1 p.m.) Approval of Minutes of previous sitting: see Minutes Membership of Parliament: see Minutes Request for the defence of parliamentary immunity: see Minutes Appointments to committees (proposal by the Conference of Presidents): see Minutes Documents received: see Minutes Texts of agreements forwarded by the Council: see Minutes
 Action taken on Parliament's resolutions: see Minutes Oral questions and written statements (tabling): see Minutes Written statements (Rule 116): see Minutes Agenda: see Minutes 1. Appointments to parliamentary committees (vote): see Minutes Voting time Agenda for next sitting: see Minutes Closure of sitting (The sitting was closed at 12 midnight) Opening of the sitting (The sitting was opened at 09.05) Documents received: see Minutes Approval of Minutes of previous sitting: see Minutes 1. Protection of passengers against displaced luggage (vote) 2.
 Approval of motor vehicles with regard to the forward field of vision of the driver (vote) 3. EC-Korea Agreement on scientific and technological cooperation (vote) 4. Mainstreaming sustainability in development cooperation policies (vote) 5. Draft Amending Budget No 1/2007 (vote) 7. EC-Gabon Fisheries Partnership (vote) 10. Limitation periods in cross-border disputes involving personal injuries and fatal accidents (vote) 12. Strategy for a strengthened partnership with the Pacific Islands (vote) 13. The European private company statute (vote) That concludes the vote.
-Corrections to votes and voting intentions: see Minutes Assignment conferred on a Member: see Minutes Membership of committees and delegations: see Minutes Decisions concerning certain documents: see Minutes Forwarding of texts adopted during the sitting: see Minutes Dates for next sittings: see Minutes
-Written statements for entry
+Corrections to votes and voting intentions: see Minutes Assignment conferred on a Member: see Minutes Membership of committees and delegations: see Minutes Decisions concerning certain documents: see Minutes Forwarding of texts adopted during the sitting: see Minutes Dates for next sittings: see Minutes
@@ -1,11 +1,8 @@
-Corectările voturilor şi intenţiile de vot: a se vedea procesul-verbal Misiune încredinţată unui deputat: consultaţi procesul-verbal Componenţa comisiilor şi a delegaţiilor: a se vedea procesul-verbal Decizii privind anumite documente: a se vedea procesul-verbal Transmiterea textelor adoptate în cursul prezentei şedinţe: a se vedea procesul-verbal Calendarul următoarelor şedinţe: a se vedea procesul-verbal
 Componenţa Parlamentului: a se vedea procesul-verbal Aprobarea procesului-verbal al şedinţei precedente: a se vedea procesul-verbal Componenţa Parlamentului: a se vedea procesul-verbal Verificarea prerogativelor: a se vedea procesul-verbal Depunere de documente: a se vedea procesul-verbal Declaraţii scrise şi întrebări orale (depunere): consultaţi procesul-verbal Petiţii: a se vedea procesul-verbal Transmiterea de către Consiliu a textelor acordurilor: a se vedea procesul-verbal Cursul dat rezoluţiilor Parlamentului: a se vedea procesul-verbal Ordinea de zi a următoarei şedinţe: a se vedea procesul-verbal Ridicarea şedinţei (Se levanta la sesión a las 19.45 horas)
 Alegerea vicepreşedinţilor Parlamentului European (termenul de depunere a candidaturilor): consultaţi procesul-verbal (Die Sitzung wird um 12.40 Uhr unterbrochen und um 15.00 Uhr wiederaufgenommen). Alegerea chestorilor Parlamentului European (termenul de depunere a candidaturilor): consultaţi procesul-verbal (Die Sitzung wird um 15.25 Uhr unterbrochen und um 18.00 Uhr wiederaufgenommen). Ordinea de zi a următoarei şedinţe: a se vedea procesul-verbal Ridicarea şedinţei (Die Sitzung wird um 18.15 Uhr geschlossen.) Deschiderea şedinţei (Die Sitzung wird um 9.35 Uhr eröffnet.) Depunerea documentelor: a se vedea procesul-verbal Aprobarea procesului-verbal al şedinţei precedente: a se vedea procesul-verbal Componenţa Parlamentului: a se vedea procesul-verbal
 Componenţa comisiilor (termenul de depunere a amendamentelor): consultaţi procesul-verbal (La seduta, sospesa alle 19.00, è ripresa alle 21.00) Ordinea de zi a următoarei şedinţe: a se vedea procesul-verbal Ridicarea şedinţei (Die Sitzung wird um 23.25 Uhr geschlossen.) Depunerea documentelor: a se vedea procesul-verbal Comunicarea poziţiilor comune ale Parlamentului: a se vedea procesul-verbal (La séance, suspendue à 11h35 dans l'attente de l'Heure des votes, est reprise à midi) Aprobarea procesului-verbal al şedinţei precedente: a se vedea procesul-verbal Comisia de anchetă privind criza societăţii de asigurări "Equitable Life” (prelungirea mandatului): consultaţi procesul-verbal
 Comunicarea Preşedintelui: consultaţi procesul-verbal 1. Componenţa comisiilor (vot) 2. Modificarea Acordului de parteneriat ACP-CE ("Acordul de la Cotonou”) (vot) 4. Certificarea mecanicilor de locomotivă care conduc locomotive şi trenuri în sistemul feroviar comunitar (vot) 6. Legea aplicabilă obligaţiilor necontractuale ("Roma II”) (vot) 8. Al şaptelea şi al optulea raport anual privind exportul de armament (vot) Corectările voturilor şi intenţiile de vot: a se vedea procesul-verbal Componenţa comisiilor şi a delegaţiilor: a se vedea procesul-verbal Cerere de ridicare a imunităţii parlamentare: consultaţi procesul-verbal Decizii privind anumite documente: a se vedea procesul-verbal
-Declaraţii scrise înscrise
 Declaraţii scrise înscrise în registru (articolul 116 din Regulamentul de procedură): a se vedea procesul-verbal Transmiterea textelor adoptate în cursul prezentei şedinţe: a se vedea procesul-verbal Calendarul următoarelor şedinţe: a se vedea procesul-verbal Întreruperea sesiunii Dichiaro interrotta la sessione del Parlamento europeo. (La seduta è tolta alle 13.00) Aprobarea procesului-verbal al şedinţei precedente: a se vedea procesul-verbal Componenţa Parlamentului: a se vedea procesul-verbal Cerere de apărare a imunităţii parlamentare: consultaţi procesul-verbal Numiri în comisii (propunerea Conferinţei preşedinţilor): consultaţi procesul-verbal Depunerea documentelor: a se vedea procesul-verbal Transmiterea de către Consiliu a textelor acordurilor: a se vedea procesul-verbal
 Continuări ale rezoluţiilor Parlamentului: consultaţi procesul-verbal Declaraţii scrise şi întrebări orale (depunere): consultaţi procesul-verbal Declaraţii scrise (articolul 116 din Regulamentul de procedură) Ordinea de zi: a se vedea procesul-verbal 1. Numiri în comisiile parlamentare (vot): consultaţi procesul-verbal Timpul afectat votului Ordinea de zi a următoarei şedinţe: a se vedea procesul-verbal Ridicarea şedinţei (La seduta è tolta alle 24.00) Deschiderea şedinţei (The sitting was opened at 09.05) Depunerea documentelor: a se vedea procesul-verbal Aprobarea procesului-verbal al şedinţei precedente: a se vedea procesul-verbal 1. Protecţia pasagerilor împotriva deplasării bagajelor (vot) 2.
 Omologarea vehiculelor cu motor cu privire la câmpul de vizibilitate înainte al conducătorului auto (vot) 3. Acordul CE-Coreea de cooperare ştiinţifică şi tehnologică (vot) 4. Integrarea durabilităţii în politicile de cooperare pentru dezvoltare (vot) 5. Proiect de buget rectificativ nr.1/2007 (vot) 7. Acordul de parteneriat în domeniul pescuitului între Comunitatea Europeană şi Republica Gaboneză (vot) 10. Termenele de prescripţie aplicabile în cadrul litigiilor transfrontaliere cu privire la vătămările corporale şi accidentele mortale (vot) 12. Relaţiile UE cu insulele din Pacific: Strategie pentru un parteneriat consolidat (vot) 13. Statutul societăţii private europene (vot) Damit ist die Abstimmungsstunde beendet.
-Corectările voturilor şi intenţiile de vot: a se vedea procesul-verbal Misiune încredinţată unui deputat: consultaţi procesul-verbal Componenţa comisiilor şi a delegaţiilor: a se vedea procesul-verbal Decizii privind anumite documente: a se vedea procesul-verbal Transmiterea textelor adoptate în cursul prezentei şedinţe: a se vedea procesul-verbal Calendarul următoarelor şedinţe: a se vedea procesul-verbal
-Declaraţii scrise înscrise
+Corectările voturilor şi intenţiile de vot: a se vedea procesul-verbal Misiune încredinţată unui deputat: consultaţi procesul-verbal Componenţa comisiilor şi a delegaţiilor: a se vedea procesul-verbal Decizii privind anumite documente: a se vedea procesul-verbal Transmiterea textelor adoptate în cursul prezentei şedinţe: a se vedea procesul-verbal Calendarul următoarelor şedinţe: a se vedea procesul-verbal
--- a/Show More
+++ b/Show More
Author	SHA1	Message	Date
Sam Shleifer	c2788f19aa	boom boom	2020-08-21 19:55:50 -04:00
Sam Shleifer	6aaac683c5	merge batch parity	2020-08-21 19:54:03 -04:00
Sam Shleifer	d68c671132	Merge branch 'master' into dropper-celoss	2020-08-21 19:50:36 -04:00
Sam Shleifer	f5d68bb13b	fix tests	2020-08-21 17:15:54 -04:00
Sam Shleifer	0da45503e8	split out bart tokenizer tests	2020-08-21 17:09:37 -04:00
Sam Shleifer	aebbe11a67	Merge branch 'master' into batch-parity-cleaner	2020-08-21 16:31:49 -04:00
Sam Shleifer	f9613aaf24	Merge remote-tracking branch 'upstream/master' into batch-parity-cleaner	2020-08-20 14:37:55 -04:00
Sam Shleifer	b2c039bfad	boom boom	2020-08-20 11:37:02 -04:00
Sam Shleifer	f328679d93	boom boom	2020-08-20 11:31:40 -04:00
Sam Shleifer	2207e5d8cb	tests pass	2020-08-19 22:41:38 -04:00
Sam Shleifer	75978bf5e9	batch parity	2020-08-19 22:39:03 -04:00
Sam Shleifer	018c1bb3da	broken test	2020-08-19 14:33:04 -04:00
Sam Shleifer	b9ca11e3ae	Dropc 0.5	2020-08-18 22:03:35 -04:00
Sam Shleifer	6bdf998dff	dropc03	2020-08-18 10:17:15 -04:00
Sam Shleifer	fcb96a2c1e	test dropc=0, ce loss	2020-08-17 23:48:58 -04:00
Sam Shleifer	e5a3f09cd4	dropc_zero	2020-08-17 23:44:48 -04:00
Sam Shleifer	7d5bdb82eb	Merge branch 'master' into dropper	2020-08-17 23:42:36 -04:00
Sam Shleifer	556fdea687	boom boom	2020-08-16 22:50:05 -04:00
Sam Shleifer	9ba04a609e	Merge branch 'master' into dropper	2020-08-16 22:50:02 -04:00
Sam Shleifer	d63f811af5	Doesnt break	2020-08-16 22:43:36 -04:00
Sam Shleifer	faac0718cb	boom boom	2020-08-16 21:36:09 -04:00
Sam Shleifer	7618c08e55	merge master	2020-08-16 21:06:19 -04:00
Sam Shleifer	1734ba169b	asked for help	2020-07-22 10:54:11 -04:00