Release: v3.3.1

Fix Trainer tests in a multiGPU env (#7458 )
Catch import datasets common errors (#7456 )
2020-09-29 14:17:34 -04:00 · 2020-09-29 14:06:41 -04:00 · 2020-09-29 13:42:09 -04:00 · 2020-09-29 13:38:47 -04:00 · 2020-09-29 12:26:26 -04:00 · 2020-09-29 10:41:18 -04:00
115 changed files with 3429 additions and 960 deletions
@@ -49,4 +49,5 @@ deploy_doc "10d7239" v2.10.0
 deploy_doc "b42586e" v2.11.0
 deploy_doc "7fb8bdf" v3.0.2
 deploy_doc "4b3ee9c" v3.1.0
-deploy_doc "3ebb1b3" # v3.2.0 Latest stable release
+deploy_doc "3ebb1b3" v3.2.0
+deploy_doc "0613f05" # v3.3.0 Latest stable release
@@ -1,2 +1,61 @@
-<!-- This line specifies which issue to close after the pull request is merged. -->
-Fixes #{issue number}
+# What does this PR do?
+
+<!--
+Congratulations! You've made it this far! You're not quite done yet though.
+
+Once merged, your PR is going to appear in the release notes with the title you set, so make sure it's a great title that fully reflects the extent of your awesome contribution.
+
+Then, please replace this with a description of the change and which issue is fixed (if applicable). Please also include relevant motivation and context. List any dependencies (if any) that are required for this change.
+
+Once you're done, someone will review your PR shortly (see the section "Who can review?" below to tag some potential reviewers). They may suggest changes to make the code even better. If no one reviewed your PR after a week has passed, don't hesitate to post a new comment @-mentioning the same persons---sometimes notifications get lost.
+-->
+
+<!-- Remove if not applicable -->
+
+Fixes # (issue)
+
+
+## Before submitting
+- [ ] This PR fixes a typo or improves the docs (you can dimiss the other checks if that's the case).
+- [ ] Did you read the [contributor guideline](https://github.com/huggingface/transformers/blob/master/CONTRIBUTING.md#start-contributing-pull-requests), 
+      Pull Request section?
+- [ ] Was this discussed/approved via a Github issue or the [forum](https://discuss.huggingface.co/)? Please add a link
+      to the it if that's the case.
+- [ ] Did you make sure to update the documentation with your changes? Here are the
+      [documentation guidelines](https://github.com/huggingface/transformers/tree/master/docs), and
+      [here are tips on formatting docstrings](https://github.com/huggingface/transformers/tree/master/docs#writing-source-documentation).
+- [ ] Did you write any new necessary tests? 
+
+
+## Who can review?
+
+Anyone in the community is free to review the PR once the tests have passed. Feel free to tag
+members/contributors which may be interested in your PR.
+
+<!-- Your PR will be replied to more quickly if you can figure out the right person to tag with @
+
+ If you know how to use git blame, that is the easiest way, otherwise, here is a rough guide of **who to tag**.
+ Please tag fewer than 3 people.
+
+ albert, bert, GPT2, XLM: @LysandreJik 
+ tokenizers: @mfuntowicz
+ Trainer: @sgugger
+ Speed and Memory Benchmarks: @patrickvonplaten
+ Model Cards: @julien-c
+ Translation: @sshleifer
+ Summarization: @sshleifer
+ TextGeneration: @TevenLeScao 
+ examples/distillation: @VictorSanh
+ nlp datasets: [different repo](https://github.com/huggingface/nlp)
+ rust tokenizers: [different repo](https://github.com/huggingface/tokenizers)
+ Text Generation: @TevenLeScao
+ Blenderbot, Bart, Marian, Pegasus: @sshleifer
+ T5: @patrickvonplaten
+ Longformer/Reformer: @patrickvonplaten
+ TransfoXL/XLNet: @TevenLeScao 
+ examples/seq2seq: @sshleifer
+ examples/bert-loses-patience: @JetRunner
+ tensorflow: @jplu
+ examples/token-classification: @stefan-it
+ documentation: @sgugger
+ -->
@@ -0,0 +1,129 @@
+
+# Contributor Covenant Code of Conduct
+
+## Our Pledge
+
+We as members, contributors, and leaders pledge to make participation in our
+community a harassment-free experience for everyone, regardless of age, body
+size, visible or invisible disability, ethnicity, sex characteristics, gender
+identity and expression, level of experience, education, socio-economic status,
+nationality, personal appearance, race, religion, or sexual identity
+and orientation.
+
+We pledge to act and interact in ways that contribute to an open, welcoming,
+diverse, inclusive, and healthy community.
+
+## Our Standards
+
+Examples of behavior that contributes to a positive environment for our
+community include:
+
+* Demonstrating empathy and kindness toward other people
+* Being respectful of differing opinions, viewpoints, and experiences
+* Giving and gracefully accepting constructive feedback
+* Accepting responsibility and apologizing to those affected by our mistakes,
+  and learning from the experience
+* Focusing on what is best not just for us as individuals, but for the
+  overall community
+
+Examples of unacceptable behavior include:
+
+* The use of sexualized language or imagery, and sexual attention or
+  advances of any kind
+* Trolling, insulting or derogatory comments, and personal or political attacks
+* Public or private harassment
+* Publishing others' private information, such as a physical or email
+  address, without their explicit permission
+* Other conduct which could reasonably be considered inappropriate in a
+  professional setting
+
+## Enforcement Responsibilities
+
+Community leaders are responsible for clarifying and enforcing our standards of
+acceptable behavior and will take appropriate and fair corrective action in
+response to any behavior that they deem inappropriate, threatening, offensive,
+or harmful.
+
+Community leaders have the right and responsibility to remove, edit, or reject
+comments, commits, code, wiki edits, issues, and other contributions that are
+not aligned to this Code of Conduct, and will communicate reasons for moderation
+decisions when appropriate.
+
+## Scope
+
+This Code of Conduct applies within all community spaces, and also applies when
+an individual is officially representing the community in public spaces.
+Examples of representing our community include using an official e-mail address,
+posting via an official social media account, or acting as an appointed
+representative at an online or offline event.
+
+## Enforcement
+
+Instances of abusive, harassing, or otherwise unacceptable behavior may be
+reported to the community leaders responsible for enforcement at
+feedback@huggingface.co.
+All complaints will be reviewed and investigated promptly and fairly.
+
+All community leaders are obligated to respect the privacy and security of the
+reporter of any incident.
+
+## Enforcement Guidelines
+
+Community leaders will follow these Community Impact Guidelines in determining
+the consequences for any action they deem in violation of this Code of Conduct:
+
+### 1. Correction
+
+**Community Impact**: Use of inappropriate language or other behavior deemed
+unprofessional or unwelcome in the community.
+
+**Consequence**: A private, written warning from community leaders, providing
+clarity around the nature of the violation and an explanation of why the
+behavior was inappropriate. A public apology may be requested.
+
+### 2. Warning
+
+**Community Impact**: A violation through a single incident or series
+of actions.
+
+**Consequence**: A warning with consequences for continued behavior. No
+interaction with the people involved, including unsolicited interaction with
+those enforcing the Code of Conduct, for a specified period of time. This
+includes avoiding interactions in community spaces as well as external channels
+like social media. Violating these terms may lead to a temporary or
+permanent ban.
+
+### 3. Temporary Ban
+
+**Community Impact**: A serious violation of community standards, including
+sustained inappropriate behavior.
+
+**Consequence**: A temporary ban from any sort of interaction or public
+communication with the community for a specified period of time. No public or
+private interaction with the people involved, including unsolicited interaction
+with those enforcing the Code of Conduct, is allowed during this period.
+Violating these terms may lead to a permanent ban.
+
+### 4. Permanent Ban
+
+**Community Impact**: Demonstrating a pattern of violation of community
+standards, including sustained inappropriate behavior,  harassment of an
+individual, or aggression toward or disparagement of classes of individuals.
+
+**Consequence**: A permanent ban from any sort of public interaction within
+the community.
+
+## Attribution
+
+This Code of Conduct is adapted from the [Contributor Covenant][homepage],
+version 2.0, available at
+https://www.contributor-covenant.org/version/2/0/code_of_conduct.html.
+
+Community Impact Guidelines were inspired by [Mozilla's code of conduct
+enforcement ladder](https://github.com/mozilla/diversity).
+
+[homepage]: https://www.contributor-covenant.org
+
+For answers to common questions about this code of conduct, see the FAQ at
+https://www.contributor-covenant.org/faq. Translations are available at
+https://www.contributor-covenant.org/translations.
@@ -9,6 +9,9 @@ It also helps us if you spread the word: reference the library from blog posts
 on the awesome projects it made possible, shout out on Twitter every time it has
 helped you, or simply star the repo to say "thank you".

+Whichever way you choose to contribute, please be mindful to respect our
+[code of conduct](https://github.com/huggingface/transformers/blob/master/CODE_OF_CONDUCT.md).
+
 ## You can contribute in so many ways!

 There are 4 ways you can contribute to transformers:
@@ -170,12 +173,19 @@ Follow these steps to start contributing:
   $ make style
   ```

-   `transformers` also uses `flake8` to check for coding mistakes. Quality
+   `transformers` also uses `flake8` and a few custom scripts to check for coding mistakes. Quality
   control runs in CI, however you can also run the same checks with:

   ```bash
   $ make quality
   ```
+   You can do the automatic style corrections and code verifications that can't be automated in one go:
+
+   ```bash
+   $ make fixup
+   ```
+
+   This target is also optimized to only work with files modified by the PR you're working on.

   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
@@ -1,19 +1,46 @@
-.PHONY: quality style test test-examples docs
+.PHONY: modified_only_fixup extra_quality_checks quality style fixup fix-copies test test-examples docs
+
+
+check_dirs := examples templates tests src utils
+
+# get modified files since the branch was made
+fork_point_sha := $(shell git merge-base --fork-point master)
+joined_dirs    := $(shell echo $(check_dirs) | tr " " "|")
+modified_files := $(shell git diff --name-only $(fork_point_sha) | egrep '^($(joined_dirs))')
+#$(info modified files are: $(modified_files))
+
+modified_only_fixup:
+	@if [ -n "$(modified_files)" ]; then \
+		echo "Checking/fixing $(modified_files)"; \
+		black $(modified_files); \
+		isort $(modified_files); \
+		flake8 $(modified_files); \
+	else \
+		echo "No relevant files were modified"; \
+	fi

 # Check that source code meets quality standards

-quality:
-	black --check examples templates tests src utils
-	isort --check-only examples templates tests src utils
-	flake8 examples templates tests src utils
+extra_quality_checks:
 	python utils/check_copies.py
 	python utils/check_repo.py

-# Format source code automatically
+# this target runs checks on all files
+quality:
+	black --check $(check_dirs)
+	isort --check-only $(check_dirs)
+	flake8 $(check_dirs)
+	${MAKE} extra_quality_checks
+
+# Format source code automatically and check is there are any problems left that need manual fixing

 style:
-	black examples templates tests src utils
-	isort examples templates tests src utils
+	black $(check_dirs)
+	isort $(check_dirs)
+
+# Super fast fix and check target that only works on relevant modified files since the branch was made
+
+fixup: modified_only_fixup extra_quality_checks

 # Make marked copies of snippets of codes conform to the original

@@ -16,6 +16,9 @@
    <a href="https://github.com/huggingface/transformers/releases">
        <img alt="GitHub release" src="https://img.shields.io/github/release/huggingface/transformers.svg">
    </a>
+    <a href="https://github.com/huggingface/transformers/blob/master/CODE_OF_CONDUCT.md">
+        <img alt="Contributor Covenant" src="https://img.shields.io/badge/Contributor%20Covenant-v2.0%20adopted-ff69b4.svg">
+    </a>
 </p>

 <h3 align="center">
@@ -125,6 +125,12 @@ a.copybtn {
    background-color: #6670FF;
 }

+/* The section headers in the toc tree */
+.wy-menu-vertical p.caption{
+    background-color: #4d59ff;
+    line-height: 40px;
+}
+
 /* The selected items in the toc tree */
 .wy-menu-vertical li.current{
    background-color: #A6B0FF;
@@ -1,10 +1,11 @@
 // These two things need to be updated at each release for the version selector.
 // Last stable version
-const stableVersion = "v3.2.0"
+const stableVersion = "v3.3.0"
 // Dictionary doc folder to label
 const versionMapping = {
    "master": "master",
-    "": "v3.2.0",
+    "": "v3.3.0",
+    "v3.2.0": "v3.2.0",
    "v3.1.0": "v3.1.0 (stable)",
    "v3.0.2": "v3.0.0/v3.0.1/v3.0.2",
    "v2.11.0": "v2.11.0",
@@ -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.2.0'
+release = u'3.3.1'


 # -- General configuration ---------------------------------------------------
@@ -46,7 +46,10 @@ The documentation is organized in five parts:
 - **ADVANCED GUIDES** contains more advanced guides that are more specific to a given script or part of the library.
 - **RESEARCH** focuses on tutorials that have less to do with how to use the library but more about general resarch in
  transformers model
- **PACKAGE REFERENCE** contains the documentation of each public class and function.
+- The three last section contain the documentation of each public class and function, grouped in:
+    - **MAIN CLASSES** for the main classes exposing the important APIs of the library.
+    - **MODELS** for the classes and functions related to each model implemented in the library.
+    - **INTERNAL HELPERS** for the classes and functions we use internally.

 The library currently contains PyTorch and Tensorflow implementations, pre-trained model weights, usage scripts and
 conversion utilities for the following models:
@@ -188,50 +191,60 @@ conversion utilities for the following models:

 .. toctree::
    :maxdepth: 2
-    :caption: Package Reference
+    :caption: Main Classes

    main_classes/configuration
-    main_classes/output
-    main_classes/model
-    main_classes/tokenizer
-    main_classes/pipelines
-    main_classes/trainer
-    main_classes/optimizer_schedules
-    main_classes/processors
    main_classes/logging
-    model_doc/auto
-    model_doc/encoderdecoder
-    model_doc/bert
-    model_doc/gpt
-    model_doc/transformerxl
-    model_doc/gpt2
-    model_doc/xlm
-    model_doc/xlnet
-    model_doc/roberta
-    model_doc/distilbert
-    model_doc/ctrl
-    model_doc/camembert
+    main_classes/model
+    main_classes/optimizer_schedules
+    main_classes/output
+    main_classes/pipelines
+    main_classes/processors
+    main_classes/tokenizer
+    main_classes/trainer
+
+.. toctree::
+    :maxdepth: 2
+    :caption: Models
+
    model_doc/albert
-    model_doc/xlmroberta
-    model_doc/flaubert
+    model_doc/auto
    model_doc/bart
-    model_doc/t5
-    model_doc/electra
+    model_doc/bert
+    model_doc/bertgeneration
+    model_doc/camembert
+    model_doc/ctrl
    model_doc/dialogpt
-    model_doc/reformer
-    model_doc/marian
-    model_doc/longformer
-    model_doc/retribert
-    model_doc/mobilebert
+    model_doc/distilbert
    model_doc/dpr
-    model_doc/pegasus
-    model_doc/mbart
+    model_doc/electra
+    model_doc/encoderdecoder
+    model_doc/flaubert
    model_doc/fsmt
    model_doc/funnel
-    model_doc/lxmert
-    model_doc/bertgeneration
    model_doc/layoutlm
+    model_doc/longformer
+    model_doc/lxmert
+    model_doc/marian
+    model_doc/mbart
+    model_doc/mobilebert
+    model_doc/gpt
+    model_doc/gpt2
+    model_doc/pegasus
    model_doc/rag
+    model_doc/reformer
+    model_doc/retribert
+    model_doc/roberta
+    model_doc/t5
+    model_doc/transformerxl
+    model_doc/xlm
+    model_doc/xlmroberta
+    model_doc/xlnet
+
+.. toctree::
+    :maxdepth: 2
+    :caption: Internal Helpers
+
    internal/modeling_utils
-    internal/tokenization_utils
    internal/pipelines_utils
+    internal/tokenization_utils
@@ -52,3 +52,7 @@ Other functions
 .. autofunction:: transformers.logging.set_verbosity

 .. autofunction:: transformers.logging.get_logger
+
+.. autofunction:: transformers.logging.enable_explicit_format
+
+.. autofunction:: transformers.logging.reset_format
@@ -112,8 +112,7 @@ Make sure there are no garbage files in the directory you'll upload. It should o
 - a `tf_model.h5` file, which is the TensorFlow checkpoint (unless you can't have it for some reason) ;
 - a `special_tokens_map.json`, which is part of your :doc:`tokenizer <main_classes/tokenizer>` save;
 - a `tokenizer_config.json`, which is part of your :doc:`tokenizer <main_classes/tokenizer>` save;
- a `vocab.txt`, which is the vocabulary of your tokenizer, part of your :doc:`tokenizer <main_classes/tokenizer>`
-  save;
+- files named `vocab.json`, `vocab.txt`, `merges.txt`, or similar, which contain the vocabulary of your tokenizer, part of your :doc:`tokenizer <main_classes/tokenizer>` save;
 - maybe a `added_tokens.json`, which is part of your :doc:`tokenizer <main_classes/tokenizer>` save.

 Other files can safely be deleted.
@@ -221,4 +220,3 @@ You can also delete unneeded files with
 .. code-block::

    transformers-cli s3 rm awesome-name-you-picked/filename
-
@@ -23,11 +23,11 @@ This is still a work-in-progress – in particular documentation is still sparse
 | [**`multiple-choice`**](https://github.com/huggingface/transformers/tree/master/examples/multiple-choice)           | SWAG, RACE, ARC | ✅ | ✅ | -  | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/ViktorAlm/notebooks/blob/master/MPC_GPU_Demo_for_TF_and_PT.ipynb)
 | [**`question-answering`**](https://github.com/huggingface/transformers/tree/master/examples/question-answering)     | SQuAD           | ✅ | ✅ | -  | -
 | [**`text-generation`**](https://github.com/huggingface/transformers/tree/master/examples/text-generation)           | -               | n/a | n/a | n/a | [![Open In Colab](https://colab.research.google.com/assets/colab-badge.svg)](https://colab.research.google.com/github/huggingface/blog/blob/master/notebooks/02_how_to_generate.ipynb)
-| [**`distillation`**](https://github.com/huggingface/transformers/tree/master/examples/distillation)       | All               | -  | -  | -  | -
-| [**`summarization`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)     | CNN/Daily Mail    | -  | -  | ✅  | -
-| [**`translation`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)         | WMT               | -  | -  | ✅  | -
-| [**`bertology`**](https://github.com/huggingface/transformers/tree/master/examples/bertology)             | -                 | -  | -  | -  | -
-| [**`adversarial`**](https://github.com/huggingface/transformers/tree/master/examples/adversarial)         | HANS              | ✅ | -  | -  | -
+| [**`distillation`**](https://github.com/huggingface/transformers/tree/master/examples/distillation)                 | All             | - | -  | -  | -
+| [**`summarization`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)                     | CNN/Daily Mail  | ✅  | -  | ✅  | -
+| [**`translation`**](https://github.com/huggingface/transformers/tree/master/examples/seq2seq)                       | WMT             | ✅  | -  | ✅  | -
+| [**`bertology`**](https://github.com/huggingface/transformers/tree/master/examples/bertology)                       | -               | - | -  | -  | -
+| [**`adversarial`**](https://github.com/huggingface/transformers/tree/master/examples/adversarial)                   | HANS            | ✅ | -  | -  | -


 <br>
@@ -1,24 +1,28 @@
 # Intro
-RAG is a seq2seq model which encapsulates two core components: a question encoder and a generator.
+Aimed at tackling the knowledge-intensive NLP tasks (think tasks a human wouldn't be expected to solve without access to external knowledge sources), RAG models are seq2seq models with access to a retrieval mechanism providing relevant context documents at training and evaluation time.
+
+A RAG model encapsulates two core components: a question encoder and a generator.
 During a forward pass, we encode the input with the question encoder and pass it
 to the retriever to extract relevant context documents. The documents are then prepended to the input.
-Such contextualized inputs is passed to the generator.
-
-The question encoder can be any `autoencoding` model, preferably :obj:`~transformers.DPRQuestionEncoder`, and the generator can be any `seq2seq` model, preferably :obj:`~transformers.BartForConditionalGeneration`.
-
-The model can be initialized with a :obj:`~transformers.RagRetriever` for end-to-end generation or used in combination with the outputs of a retriever in multiple steps - see examples for more details.
-The model is compatible any `autoencoding` model as the ``question_encoder`` and any `seq2seq` model with language model head as the ``generator``.
-The model has been tested with :class:`~transformers.DPRQuestionEncoder` as the ``question_encoder`` and :class:`~transformers.BartForConditionalGeneration` or :class:`~transformers.T5ForConditionalGeneration` as the ``generator``.
-
-RAG models were released with the paper `Retrieval-Augmented Generation for
-Knowledge-Intensive NLP Tasks <https://arxiv.org/abs/2005.11401>`_ by Patrick Lewis, Ethan Perez, Aleksandra Piktus et al.
-
+Such contextualized inputs are passed to the generator.

+Read more about RAG  at https://arxiv.org/abs/2005.11401.
 # Finetuning
-Our finetuning logic is based on scripts from [`examples/seq2seq`](https://github.com/huggingface/transformers/tree/master/examples/seq2seq).
-Follow instructions there regarding data preprocessing. A sample finetuning command:

+
+Our finetuning logic is based on scripts from [`examples/seq2seq`](https://github.com/huggingface/transformers/tree/master/examples/seq2seq). We accept training data in the same format as specified there - we expect a directory consisting of 6 text files:
+```bash
+train.source
+train.target
+val.source
+val.target
+test.source
+test.target
 ```
+
+A sample finetuning command (run ` ./examples/rag/finetune.py --help` to list all available options):
+
+```bash
 python examples/rag/finetune.py \
    --data_dir $DATA_DIR \
    --output_dir $OUTPUT_DIR \
@@ -27,62 +31,85 @@ python examples/rag/finetune.py \
    --fp16 \
    --gpus 8
 ```
+We publish two `base` models which can serve as a starting point for finetuning on downstream tasks (use them as `model_name_or_path`):
+- [`facebook/rag-sequence-base`](https://huggingface.co/facebook/rag-sequence-base) - a base for finetuning `RagSequenceForGeneration` models,
+- [`facebook/rag-token-base`](https://huggingface.co/facebook/rag-token-base) - a base for finetuning `RagTokenForGeneration` models.
+
+The `base` models initialize the question encoder with [`facebook/dpr-question_encoder-single-nq-base`](https://huggingface.co/facebook/dpr-question_encoder-single-nq-base) and the generator with [`facebook/bart-large`](https://huggingface.co/facebook/bart-large).
+
+If you would like to initialize finetuning with a base model using different question encoder and generator architectures, you can build it with a consolidation script, e.g.:
+```
+python examples/rag/consolidate_rag_checkpoint.py \
+    --model_type rag_sequence \
+    --generator_name_or_path facebook/bart-large-cnn \
+    --question_encoder_name_or_path facebook/dpr-question_encoder-single-nq-base \
+    --dest path/to/checkpoint
+```
+You will then be able to pass `path/to/checkpoint` as `model_name_or_path` to the `finetune.py` script.


 # Evaluation
-Apart from the parameters specifying the model to evaluate and some extra parameters, the evaluation script expects paths to two files:
- `evaluation_set` - a path to a file specifying the evaluation dataset, a single datapoint per line, e.g.
-```who is the owner of reading football club```
- `gold_data_path` - a path to a file contaning ground truth answers for datapoints from the `evaluation_set`.
+Our evaluation script enables two modes of evaluation (controlled by the `eval_mode` argument): `e2e` - end2end evaluation, returns EM (exact match) and F1 scores calculated for the downstream task and `retrieval` - which returns precision@k of the documents retrieved for provided inputs.

-We expect the following formats of the gold data file:
+The evaluation script expects paths to two files:
+- `evaluation_set` - a path to a file specifying the evaluation dataset, a single input per line.
+- `gold_data_path` - a path to a file contaning ground truth answers for datapoints from the `evaluation_set`, a single output per line. Check below for expected formats of the gold data files.

- for e2e evaluation, we support two formats of the gold file:
-    - `qa` - where a single line in the following format: input [tab] output_list, e.g.:
-    ```
-    who is the owner of reading football club	['Xiu Li Dai', 'Dai Yongge', 'Dai Xiuli', 'Yongge Dai']
-    ```
-    - `ans` - where a single line of the gold file contains the expected output string, e.g.:
-    ```
-    Xiu Li Dai
-    ```

- for retrieval evaluation, we expect a tab-separated list of Wikipedia page titles constituting positive contexts for a given query, e.g. given a question `who sings does he love me with reba`, a line with ground truth retrieval data could look as follows:
+## Retrieval evaluation
+For `retrieval` evaluation, we expect a gold data file where each line will consist of a tab-separated list of document titles constituting positive contexts for respective datapoints from the `evaluation_set`. E.g. given a question `who sings does he love me with reba` in the `evaluation_set`, a respective ground truth line could look as follows:
 ```
 Does He Love You	Does He Love You	Red Sandy Spika dress of Reba McEntire	Greatest Hits Volume Two (Reba McEntire album)	Shoot for the Moon (album)
 ```

-## Retrieval evaluation
-
 We demonstrate how to evaluate retrieval against DPR evaluation data. You can download respective files from links listed [here](https://github.com/facebookresearch/DPR/blob/master/data/download_data.py#L39-L45).

 1. Download and unzip the gold data file. We use the `biencoder-nq-dev` from https://dl.fbaipublicfiles.com/dpr/data/retriever/biencoder-nq-dev.json.gz.
 2. Parse the unziped file using the `parse_dpr_relevance_data.py`
-```
-python examples/rag/parse_dpr_relevance_data.py --src_path path/to/unziped/biencoder-nq-dev.json --evaluation_set path/to/output/biencoder-nq-dev.questions --gold_data_path path/to/output/biencoder-nq-dev.pages
-```
+    ```bash
+    python examples/rag/parse_dpr_relevance_data.py \
+        --src_path path/to/unziped/biencoder-nq-dev.json \
+        --evaluation_set path/to/output/biencoder-nq-dev.questions \
+        --gold_data_path path/to/output/biencoder-nq-dev.pages
+    ```
 3. Run evaluation:
-```
-python examples/rag/eval_rag.py \
-    --model_name_or_path $MODEL_NAME_OR_PATH \ # model name or path of the model we're evaluating
-    --model_type rag_sequence \ # RAG model type (rag_token or rag_sequence)
-    --evaluation_set path/to/output/biencoder-nq-dev.questions \ # an input dataset for evaluation
-    --gold_data_path path/to/output/biencoder-nq-dev.pages \ # a dataset containing ground truth answers for samples from the evaluation_set
-    --predictions_path path/to/retrieval_preds.tsv  \ # name of file in which predictions will be stored
-    --eval_mode retrieval \ # indicates whether we're performing retrieval evaluation or e2e evaluation
-    --recalculate # if predictions_filename already exists, and this option is set - we regenerate the answers, otherwise we reuse the predicsion file to calculate metrics.
-```
+    ```bash
+    python examples/rag/eval_rag.py \
+        --model_name_or_path facebook/rag-sequence-nq \ # model name or path of the model we're evaluating
+        --model_type rag_sequence \ # RAG model type (rag_token or rag_sequence)
+        --evaluation_set path/to/output/biencoder-nq-dev.questions \ # an input dataset for evaluation
+        --gold_data_path path/to/output/biencoder-nq-dev.pages \ # a dataset containing ground truth answers for samples from the evaluation_set
+        --predictions_path path/to/retrieval_preds.tsv  \ # name of file where predictions will be stored
+        --eval_mode retrieval \ # indicates whether we're performing retrieval evaluation or e2e evaluation
+        --k 1 # parameter k for the precision@k metric
+    ```


 ## End-to-end evaluation
+
+We support two formats of the gold data file (controlled by the `gold_data_mode` parameter):
+- `qa` - where a single line has the following format: `input [tab] output_list`, e.g.:
 ```
+who is the owner of reading football club	['Xiu Li Dai', 'Dai Yongge', 'Dai Xiuli', 'Yongge Dai']
+```
+- `ans` - where a single line contains a single expected answer, e.g.:
+```
+Xiu Li Dai
+```
+
+Predictions of the model for the samples from the `evaluation_set` will be saved under the path specified by the `predictions_path` parameter. If this path already exists, the script will use saved predictions to calculate metrics. Add `--recalculate` parameter to force the script to perform inference from scratch.
+
+An example e2e evaluation run could look as follows:
+```bash
 python examples/rag/eval_rag.py \
-	--model_name_or_path $MODEL_NAME_OR_PATH \
+    --model_name_or_path facebook/rag-sequence-nq \
    --model_type rag_sequence \
    --evaluation_set path/to/test.source \
    --gold_data_path path/to/gold_data \
    --predictions_path path/to/e2e_preds.txt \
-    --eval_mode e2e \ # indicates whether we're performing retrieval evaluation or e2e evaluation (default)
+    --eval_mode e2e \
+    --gold_data_mode qa \
    --n_docs 5 \ # You can experiment with retrieving different number of documents at evaluation time
-    --print_predictions
+    --print_predictions \
+    --recalculate \ # adding this parameter will force recalculating predictions even if predictions_path already exists
 ```
@@ -1,7 +1,20 @@
 import logging
 import os
+from pathlib import Path

-from pytorch_lightning.callbacks import ModelCheckpoint
+import numpy as np
+import pytorch_lightning as pl
+import torch
+from pytorch_lightning.callbacks import EarlyStopping, ModelCheckpoint
+from pytorch_lightning.utilities import rank_zero_only
+
+from utils import save_json
+
+
+def count_trainable_parameters(model):
+    model_parameters = filter(lambda p: p.requires_grad, model.parameters())
+    params = sum([np.prod(p.size()) for p in model_parameters])
+    return params


 logger = logging.getLogger(__name__)
@@ -28,3 +41,76 @@ def get_checkpoint_callback(output_dir, metric):
        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,
+    )
+
+
+class Seq2SeqLoggingCallback(pl.Callback):
+    def on_batch_end(self, trainer, pl_module):
+        lrs = {f"lr_group_{i}": param["lr"] for i, param in enumerate(pl_module.trainer.optimizers[0].param_groups)}
+        pl_module.logger.log_metrics(lrs)
+
+    @rank_zero_only
+    def _write_logs(
+        self, trainer: pl.Trainer, pl_module: pl.LightningModule, type_path: str, save_generations=True
+    ) -> None:
+        logger.info(f"***** {type_path} results at step {trainer.global_step:05d} *****")
+        metrics = trainer.callback_metrics
+        trainer.logger.log_metrics({k: v for k, v in metrics.items() if k not in ["log", "progress_bar", "preds"]})
+        # Log results
+        od = Path(pl_module.hparams.output_dir)
+        if type_path == "test":
+            results_file = od / "test_results.txt"
+            generations_file = od / "test_generations.txt"
+        else:
+            # this never gets hit. I prefer not to save intermediate generations, and results are in metrics.json
+            # If people want this it will be easy enough to add back.
+            results_file = od / f"{type_path}_results/{trainer.global_step:05d}.txt"
+            generations_file = od / f"{type_path}_generations/{trainer.global_step:05d}.txt"
+            results_file.parent.mkdir(exist_ok=True)
+            generations_file.parent.mkdir(exist_ok=True)
+        with open(results_file, "a+") as writer:
+            for key in sorted(metrics):
+                if key in ["log", "progress_bar", "preds"]:
+                    continue
+                val = metrics[key]
+                if isinstance(val, torch.Tensor):
+                    val = val.item()
+                msg = f"{key}: {val:.6f}\n"
+                writer.write(msg)
+
+        if not save_generations:
+            return
+
+        if "preds" in metrics:
+            content = "\n".join(metrics["preds"])
+            generations_file.open("w+").write(content)
+
+    @rank_zero_only
+    def on_train_start(self, trainer, pl_module):
+        try:
+            npars = pl_module.model.model.num_parameters()
+        except AttributeError:
+            npars = pl_module.model.num_parameters()
+
+        n_trainable_pars = count_trainable_parameters(pl_module)
+        # mp stands for million parameters
+        trainer.logger.log_metrics({"n_params": npars, "mp": npars / 1e6, "grad_mp": n_trainable_pars / 1e6})
+
+    @rank_zero_only
+    def on_test_end(self, trainer: pl.Trainer, pl_module: pl.LightningModule):
+        save_json(pl_module.metrics, pl_module.metrics_save_path)
+        return self._write_logs(trainer, pl_module, "test")
+
+    @rank_zero_only
+    def on_validation_end(self, trainer: pl.Trainer, pl_module):
+        save_json(pl_module.metrics, pl_module.metrics_save_path)
+        # Uncommenting this will save val generations
+        # return self._write_logs(trainer, pl_module, "valid")
@@ -0,0 +1,99 @@
+"""
+A script creating a RAG checkpoint from a generator and a question encoder checkpoints.
+"""
+
+import argparse
+from pathlib import Path
+
+from transformers import AutoConfig, AutoTokenizer, RagConfig, RagSequenceForGeneration, RagTokenForGeneration
+
+
+def consolidate(
+    model_type,
+    generator_name_or_path: str,
+    question_encoder_name_or_path: str,
+    dest_dir: Path,
+    config_name_or_path: str = None,
+    generator_tokenizer_name_or_path: str = None,
+    question_encoder_tokenizer_name_or_path: str = None,
+):
+
+    if config_name_or_path is None:
+        config_name_or_path = "facebook/rag-token-base" if model_type == "rag_token" else "facebook/rag-sequence-base"
+
+    if generator_tokenizer_name_or_path is None:
+        generator_tokenizer_name_or_path = generator_name_or_path
+
+    if question_encoder_tokenizer_name_or_path is None:
+        question_encoder_tokenizer_name_or_path = question_encoder_name_or_path
+
+    model_class = RagTokenForGeneration if model_type == "rag_token" else RagSequenceForGeneration
+
+    # Save model.
+    rag_config = RagConfig.from_pretrained(config_name_or_path)
+    gen_config = AutoConfig.from_pretrained(generator_name_or_path)
+    question_encoder_config = AutoConfig.from_pretrained(question_encoder_name_or_path)
+
+    rag_config.generator = gen_config
+    rag_config.question_encoder = question_encoder_config
+
+    rag_model = model_class.from_pretrained_question_encoder_generator(
+        question_encoder_name_or_path, generator_name_or_path, config=rag_config
+    )
+    rag_model.save_pretrained(dest_dir)
+
+    # Sanity check.
+    model_class.from_pretrained(dest_dir)
+
+    # Save tokenizers.
+    gen_tokenizer = AutoTokenizer.from_pretrained(generator_tokenizer_name_or_path)
+    gen_tokenizer.save_pretrained(dest_dir / "generator_tokenizer/")
+    question_encoder_tokenizer = AutoTokenizer.from_pretrained(question_encoder_tokenizer_name_or_path)
+    question_encoder_tokenizer.save_pretrained(dest_dir / "question_encoder_tokenizer/")
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument(
+        "--model_type",
+        choices=["rag_sequence", "rag_token"],
+        required=True,
+        type=str,
+        help="RAG model type: rag_sequence, rag_token",
+    )
+    parser.add_argument("--dest", type=str, required=True, help="Path to the output checkpoint directory.")
+    parser.add_argument("--generator_name_or_path", type=str, required=True, help="Generator model identifier")
+    parser.add_argument(
+        "--question_encoder_name_or_path", type=str, required=True, help="Question encoder model identifier"
+    )
+
+    parser.add_argument(
+        "--generator_tokenizer_name_or_path",
+        type=str,
+        help="Generator tokenizer identifier, if not specified, resolves to ``generator_name_or_path``",
+    )
+    parser.add_argument(
+        "--question_encoder_tokenizer_name_or_path",
+        type=str,
+        help="Question encoder tokenizer identifier, if not specified, resolves to ``question_encoder_name_or_path``",
+    )
+    parser.add_argument(
+        "--config_name_or_path",
+        type=str,
+        help="Identifier of the model config to use, if not provided, resolves to a base config for a given ``model_type``",
+    )
+
+    args = parser.parse_args()
+
+    dest_dir = Path(args.dest)
+    dest_dir.mkdir(exist_ok=True)
+
+    consolidate(
+        args.model_type,
+        args.generator_name_or_path,
+        args.question_encoder_name_or_path,
+        dest_dir,
+        args.config_name_or_path,
+        args.generator_tokenizer_name_or_path,
+        args.question_encoder_tokenizer_name_or_path,
+    )
@@ -115,11 +115,15 @@ def evaluate_batch_retrieval(args, rag_model, questions):

 def evaluate_batch_e2e(args, rag_model, questions):
    with torch.no_grad():
-        input_ids = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus(
+        inputs_dict = rag_model.retriever.question_encoder_tokenizer.batch_encode_plus(
            questions, return_tensors="pt", padding=True, truncation=True
-        )["input_ids"].to(args.device)
+        )
+
+        input_ids = inputs_dict.input_ids.to(args.device)
+        attention_mask = inputs_dict.attention_mask.to(args.device)
        outputs = rag_model.generate(  # rag_model overwrites generate
            input_ids,
+            attention_mask=attention_mask,
            num_beams=args.num_beams,
            min_length=args.min_length,
            max_length=args.max_length,
@@ -34,22 +34,23 @@ from transformers import logging as transformers_logging
 sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip

 from examples.lightning_base import BaseTransformer, add_generic_args, generic_train  # noqa: E402 # isort:skip
-from examples.rag.callbacks import get_checkpoint_callback  # noqa: E402 # isort:skip
+from examples.rag.callbacks import (  # noqa: E402 # isort:skip
+    get_checkpoint_callback,
+    get_early_stopping_callback,
+    Seq2SeqLoggingCallback,
+)
 from examples.rag.distributed_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
 from examples.rag.utils import (  # noqa: E402 # isort:skip
-    Seq2SeqDataset,
    calculate_exact_match,
-    is_rag_model,
-    set_extra_model_params,
-)
-from examples.seq2seq.callbacks import Seq2SeqLoggingCallback, get_early_stopping_callback  # noqa: E402 # isort:skip
-from examples.seq2seq.utils import (  # noqa: E402 # isort:skip
    flatten_list,
    get_git_info,
+    is_rag_model,
    lmap,
    pickle_save,
    save_git_info,
    save_json,
+    set_extra_model_params,
+    Seq2SeqDataset,
 )

 logging.basicConfig(level=logging.INFO)
@@ -265,6 +266,7 @@ class GenerativeQAModule(BaseTransformer):
        start_time = time.time()
        generated_ids = self.model.generate(
            batch["input_ids"],
+            attention_mask=batch["attention_mask"],
            do_deduplication=False,  # rag specific parameter
            use_cache=True,
            min_length=1,
@@ -302,11 +304,6 @@ class GenerativeQAModule(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

        dataloader = DataLoader(
            dataset,
@@ -314,7 +311,6 @@ class GenerativeQAModule(BaseTransformer):
            collate_fn=dataset.collate_fn,
            shuffle=shuffle,
            num_workers=self.num_workers,
-            sampler=sampler,
        )
        return dataloader

@@ -378,7 +374,6 @@ class GenerativeQAModule(BaseTransformer):
            help="The maximum total input sequence length after tokenization. Sequences longer "
            "than this will be truncated, sequences shorter will be padded.",
        )
-        parser.add_argument("--sortish_sampler", action="store_true", default=False)
        parser.add_argument("--logger_name", type=str, choices=["default", "wandb", "wandb_shared"], default="default")
        parser.add_argument("--n_train", type=int, default=-1, required=False, help="# examples. -1 means use all.")
        parser.add_argument("--n_val", type=int, default=-1, required=False, help="# examples. -1 means use all.")
@@ -153,4 +153,4 @@ class RagRetrieverTest(TestCase):
        self.assertEqual(len(doc_dicts[0]["id"]), n_docs)
        self.assertEqual(doc_dicts[0]["id"][0], "1")  # max inner product is reached with second doc
        self.assertEqual(doc_dicts[1]["id"][0], "0")  # max inner product is reached with first doc
-        self.assertListEqual(list(doc_ids), [1, 0])
+        self.assertListEqual(doc_ids.tolist(), [[1], [0]])
@@ -1,15 +1,20 @@
+import itertools
+import json
 import linecache
+import os
+import pickle
 import re
+import socket
 import string
 from collections import Counter
 from logging import getLogger
 from pathlib import Path
-from typing import Dict, List
+from typing import Callable, Dict, Iterable, List

+import git
 import torch
 from torch.utils.data import Dataset

-from examples.seq2seq.utils import SortishSampler, trim_batch
 from transformers import BartTokenizer, RagTokenizer, T5Tokenizer


@@ -27,6 +32,19 @@ def encode_line(tokenizer, line, max_length, padding_side, pad_to_max_length=Tru
    )


+def trim_batch(
+    input_ids,
+    pad_token_id,
+    attention_mask=None,
+):
+    """Remove columns that are populated exclusively by pad_token_id"""
+    keep_column_mask = input_ids.ne(pad_token_id).any(dim=0)
+    if attention_mask is None:
+        return input_ids[:, keep_column_mask]
+    else:
+        return (input_ids[:, keep_column_mask], attention_mask[:, keep_column_mask])
+
+
 class Seq2SeqDataset(Dataset):
    def __init__(
        self,
@@ -114,13 +132,52 @@ class Seq2SeqDataset(Dataset):
        }
        return batch

-    def make_sortish_sampler(self, batch_size):
-        return SortishSampler(self.src_lens, batch_size)
-

 logger = getLogger(__name__)


+def flatten_list(summary_ids: List[List]):
+    return [x for x in itertools.chain.from_iterable(summary_ids)]
+
+
+def save_git_info(folder_path: str) -> None:
+    """Save git information to output_dir/git_log.json"""
+    repo_infos = get_git_info()
+    save_json(repo_infos, os.path.join(folder_path, "git_log.json"))
+
+
+def save_json(content, path, indent=4, **json_dump_kwargs):
+    with open(path, "w") as f:
+        json.dump(content, f, indent=indent, **json_dump_kwargs)
+
+
+def load_json(path):
+    with open(path) as f:
+        return json.load(f)
+
+
+def get_git_info():
+    repo = git.Repo(search_parent_directories=True)
+    repo_infos = {
+        "repo_id": str(repo),
+        "repo_sha": str(repo.head.object.hexsha),
+        "repo_branch": str(repo.active_branch),
+        "hostname": str(socket.gethostname()),
+    }
+    return repo_infos
+
+
+def lmap(f: Callable, x: Iterable) -> List:
+    """list(map(f, x))"""
+    return list(map(f, x))
+
+
+def pickle_save(obj, path):
+    """pickle.dump(obj, path)"""
+    with open(path, "wb") as f:
+        return pickle.dump(obj, f)
+
+
 def normalize_answer(s):
    """Lower text and remove punctuation, articles and extra whitespace."""

@@ -11,6 +11,7 @@ git-python==1.0.3
 faiss-cpu
 streamlit
 elasticsearch
+nltk
 pandas
 datasets
 fire
@@ -100,7 +100,7 @@ All finetuning bash scripts call finetune.py (or distillation.py) with reasonabl
 To see all the possible command line options, run:

 ```bash
-./finetune.sh --help  # this calls python finetune.py --help
+ ./finetune.py --help 
 ```

 ### Finetuning Training Params
@@ -189,6 +189,36 @@ from transformers import AutoModelForSeq2SeqLM
 model = AutoModelForSeq2SeqLM.from_pretrained(f'{output_dir}/best_tfmr')
 ```

+### Fine-tuning using Seq2SeqTrainer
+To use `Seq2SeqTrainer` for fine-tuning you should use the `finetune_trainer.py` script. It subclasses `Trainer` to extend it for seq2seq training. Except the `Trainer` releated `TrainingArguments`, it shares the same argument names as that of `finetune.py` file. One notable difference is that, calculating generative metrics (BLEU, ROUGE) is optional and is controlled using the `--predict_with_generate` argument, set this argument to calculate BLEU and ROUGE metrics.
+
+With PyTorch 1.6+ it'll automatically use `native AMP` when `--fp16` is set. 
+
+To see all the possible command line options, run:
+
+```bash
+./builtin_trainer/finetune.sh --help # This calls python finetune_trainer.py --help
+```
+
+**At the moment, `Seq2SeqTrainer` does not support *with teacher* distillation.**
+
+All `Seq2SeqTrainer` based fine-tuning scripts are included in the `builtin_trainer` directory.
+
+#### TPU Training
+`Seq2SeqTrainer` supports TPU training with few caveats
+1. As `generate` method does not work on TPU at the moment, `predict_with_generate` can not be used. You should use `--prediction_loss_only` to only calculate loss, and do not set `--do_predict` and `--predict_with_generate`.
+2. All sequences should be padded to be of equal length otherwise it leads to extremely slow training. (`finetune_trainer.py` does this automatically when running on TPU.)
+
+We provide a very simple launcher script named `xla_spawn.py` that lets you run our example scripts on multiple TPU cores without any boilerplate. Just pass a --num_cores flag to this script, then your regular training script with its arguments (this is similar to the torch.distributed.launch helper for torch.distributed).
+
+`builtin_trainer/finetune_tpu.sh` script provides minimal arguments needed for TPU training.
+
+Following command fine-tunes `sshleifer/student_marian_en_ro_6_3` on TPU V3-8 and should complete one epoch in ~5-6 mins.
+
+```bash
+./builtin_trainer/train_distil_marian_enro_tpu.sh
+```
+
 ### Evaluation Commands

 To create summaries for each article in dataset, we use `run_eval.py`, here are a few commands that run eval for different tasks and models.
@@ -197,7 +227,7 @@ If 'translation' is in your task name, the computed metric will be BLEU. Otherwi
 For t5, you need to specify --task translation_{src}_to_{tgt} as follows:
 ```bash
 export DATA_DIR=wmt_en_ro
-python run_eval.py t5-base \
+./run_eval.py t5-base \
    $DATA_DIR/val.source t5_val_generations.txt \
    --reference_path $DATA_DIR/val.target \
    --score_path enro_bleu.json \
@@ -211,7 +241,7 @@ python run_eval.py t5-base \
 This command works for MBART, although the BLEU score is suspiciously low.
 ```bash
 export DATA_DIR=wmt_en_ro
-python run_eval.py facebook/mbart-large-en-ro $DATA_DIR/val.source mbart_val_generations.txt \
+./run_eval.py facebook/mbart-large-en-ro $DATA_DIR/val.source mbart_val_generations.txt \
    --reference_path $DATA_DIR/val.target \
    --score_path enro_bleu.json \
    --task translation \
@@ -224,7 +254,7 @@ python run_eval.py facebook/mbart-large-en-ro $DATA_DIR/val.source mbart_val_gen
 Summarization (xsum will be very similar):
 ```bash
 export DATA_DIR=cnn_dm
-python run_eval.py sshleifer/distilbart-cnn-12-6 $DATA_DIR/val.source dbart_val_generations.txt \
+./run_eval.py sshleifer/distilbart-cnn-12-6 $DATA_DIR/val.source dbart_val_generations.txt \
    --reference_path $DATA_DIR/val.target \
    --score_path cnn_rouge.json \
    --task summarization \
@@ -235,10 +265,10 @@ python run_eval.py sshleifer/distilbart-cnn-12-6 $DATA_DIR/val.source dbart_val_
    --fp16 \
    --bs 32
 ```
-### Multi-GPU Evalulation
+### 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.
+`{type_path}.source` and `{type_path}.target`. Run `./run_distributed_eval.py --help` for all clargs.

 ```bash
 python -m torch.distributed.launch --nproc_per_node=8  run_distributed_eval.py \
@@ -250,7 +280,7 @@ python -m torch.distributed.launch --nproc_per_node=8  run_distributed_eval.py \

 Contributions that implement this command for other distributed hardware setups are welcome!

-#### run_eval tips and tricks
+#### Single-GPU Eval: Tips and Tricks

 When using `run_eval.py`, the following features can be useful:

@@ -339,7 +369,7 @@ runtime: 6H on NVIDIA RTX 24GB GPU
 If you are using `wandb` and comparing the two distillation methods, using this entry point will make your logs consistent,
 because you will have the same hyperparameters logged in every run.

-#### With a teacher
+#### With a teacher (Intermediate Supervision)
 *Note* only BART variants are supported

 In this method, we use try to enforce that the student and teacher produce similar encoder_outputs, logits, and hidden_states using `BartSummarizationDistiller`.
@@ -348,7 +378,7 @@ This is how `sshleifer/distilbart-xsum*` checkpoints were produced.
 The command that produced `sshleifer/distilbart-xsum-12-6` is:

 ```bash
-./train_distilbart_xsum.sh
+./train_distilbart_xsum.sh --logger_name wandb --gpus 1
 ```

 runtime: 13H on V-100 16GB GPU.
@@ -371,11 +401,11 @@ 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`
+- after calling `./save_len_file.py $tok $data_dir`

 For example, 
 ```bash
-python save_len_file.py Helsinki-NLP/opus-mt-en-ro  wmt_en_ro
+./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.
@@ -0,0 +1,9 @@
+# the proper usage is documented in the README, you need to specify data_dir, output_dir and model_name_or_path
+# run ./builtin_trainer/finetune.sh --help to see all the possible options
+python finetune_trainer.py \
+    --learning_rate=3e-5 \
+    --fp16 \
+    --do_train --do_eval --do_predict --evaluate_during_training \
+    --predict_with_generate \
+    --n_val 1000 \
+    "$@"
@@ -0,0 +1,11 @@
+export TPU_NUM_CORES=8
+
+# the proper usage is documented in the README, you need to specify data_dir, output_dir and model_name_or_path
+# run ./builtin_trainer/finetune_tpu.sh --help to see all the possible options
+python xla_spawn.py --num_cores $TPU_NUM_CORES \
+    finetune_trainer.py \
+    --learning_rate=3e-5 \
+    --do_train --do_eval --evaluate_during_training \
+    --prediction_loss_only \
+    --n_val 1000 \
+    "$@"
@@ -0,0 +1,23 @@
+export WANDB_PROJECT=distil-marian
+export BS=64
+export GAS=1
+export m=sshleifer/student_marian_en_ro_6_3
+export MAX_LEN=128
+python finetune_trainer.py \
+    --tokenizer_name $m --model_name_or_path $m \
+    --data_dir $ENRO_DIR \
+    --output_dir marian_en_ro_6_3 --overwrite_output_dir \
+    --learning_rate=3e-4 \
+    --warmup_steps 500 --sortish_sampler \
+    --fp16 \
+    --gradient_accumulation_steps=$GAS \
+    --per_device_train_batch_size=$BS --per_device_eval_batch_size=$BS \
+    --freeze_encoder --freeze_embeds \
+    --num_train_epochs=6 \
+    --save_steps 3000 --eval_steps 3000 \
+    --max_source_length $MAX_LEN --max_target_length $MAX_LEN --val_max_target_length $MAX_LEN --test_max_target_length $MAX_LEN \
+    --do_train --do_eval --do_predict --evaluate_during_training\
+    --predict_with_generate --logging_first_step \
+    --task translation --label_smoothing 0.1 \
+    --run_name marian_en_ro_6_3 \
+    "$@"
@@ -0,0 +1,24 @@
+export WANDB_PROJECT=distil-marian
+export BS=64
+export m=sshleifer/student_marian_en_ro_6_3
+export MAX_LEN=128
+export TPU_NUM_CORES=8
+
+python xla_spawn.py --num_cores $TPU_NUM_CORES \
+    finetune_trainer.py \
+    --tokenizer_name $m --model_name_or_path $m \
+    --data_dir $ENRO_DIR \
+    --output_dir marian_en_ro_6_3 --overwrite_output_dir \
+    --learning_rate=3e-4 \
+    --warmup_steps 500 \
+    --per_device_train_batch_size=$BS --per_device_eval_batch_size=$BS \
+    --freeze_encoder --freeze_embeds \
+    --num_train_epochs=6 \
+    --save_steps 500 --eval_steps 500 \
+    --logging_first_step --logging_steps 200 \
+    --max_source_length $MAX_LEN --max_target_length $MAX_LEN --val_max_target_length $MAX_LEN --test_max_target_length $MAX_LEN \
+    --do_train --do_eval --evaluate_during_training \
+    --prediction_loss_only \
+    --task translation --label_smoothing 0.1 \
+    --run_name marian_en_ro_6_3 \
+    "$@"
@@ -0,0 +1,26 @@
+export WANDB_PROJECT=distilbart-cnn
+export BS=32
+export GAS=1
+export m=sshleifer/student_cnn_12_6
+export tok=facebook/bart-large
+export MAX_TGT_LEN=142
+
+python finetune_trainer.py \
+    --model_name_or_path $m --tokenizer_name $tok \ 
+    --data_dir $CNN_DIR \
+    --output_dir distilbart-cnn-12-6 --overwrite_output_dir \
+    --learning_rate=3e-5 \
+    --warmup_steps 500 --sortish_sampler \
+    --fp16 \
+    --n_val 500 \
+    --gradient_accumulation_steps=$GAS \
+    --per_device_train_batch_size=$BS --per_device_eval_batch_size=$BS \
+    --freeze_encoder --freeze_embeds \
+    --num_train_epochs=2 \
+    --save_steps 3000 --eval_steps 3000 \
+    --logging_first_step \
+    --max_target_length $MAX_TGT_LEN --val_max_target_length $MAX_TGT_LEN --test_max_target_length $MAX_TGT_LEN \
+    --do_train --do_eval --do_predict --evaluate_during_training \
+    --predict_with_generate \
+    --run_name distilbart-cnn-12-6 \
+    "$@"
@@ -0,0 +1,22 @@
+python finetune_trainer.py \
+    --model_name_or_path=facebook/mbart-large-cc25 \
+    --data_dir $ENRO_DIR \
+    --output_dir mbart_cc25_enro --overwrite_output_dir \
+    --learning_rate=3e-5 \
+    --warmup_steps 500 \ 
+    --fp16 \
+    --label_smoothing 0.1 \
+    --adam_eps 1e-06 \
+    --src_lang en_XX --tgt_lang ro_RO \
+    --freeze_embeds \
+    --per_device_train_batch_size=4 --per_device_eval_batch_size=4 \
+    --max_source_length 128 --max_target_length 128 \
+    --val_max_target_length 128 --test_max_target_length 128 \
+    --sortish_sampler \
+    --num_train_epochs 6 \
+    --save_steps 25000 --eval_steps 25000 --logging_steps 1000 \
+    --do_train --do_eval --do_predict --evaluate_during_training \
+    --predict_with_generate --logging_first_step 
+    --task translation \
+    --run_name mbart_en_ro \
+    "$@"
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 from typing import Union

 import fire
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import os
 from pathlib import Path
 from typing import Dict, List
@@ -1,7 +1,9 @@
+#!/usr/bin/env python
+
 import argparse
 import gc
 import os
-import warnings
+import sys
 from pathlib import Path
 from typing import List

@@ -12,19 +14,15 @@ 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 lightning_base import generic_train
-from transformers import AutoModelForSeq2SeqLM, MBartTokenizer, T5Config, T5ForConditionalGeneration
+from make_student import create_student_by_copying_alternating_layers, get_layers_to_supervise
+from transformers import AutoModelForSeq2SeqLM, MBartTokenizer, 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,
-    pickle_load,
-    use_task_specific_params,
-)
+from utils import calculate_bleu, freeze_params, label_smoothed_nll_loss, pickle_load, use_task_specific_params
+
+
+# need the parent dir module
+sys.path.insert(2, str(Path(__file__).resolve().parents[1]))
+from lightning_base import generic_train  # noqa


 class BartSummarizationDistiller(SummarizationModule):
@@ -34,87 +32,50 @@ class BartSummarizationDistiller(SummarizationModule):

    def __init__(self, hparams):
        assert Path(hparams.data_dir).exists()
-        student, student_cfg, teacher = self.pre_init(hparams)
+        self.output_dir = Path(hparams.output_dir)
+        self.output_dir.mkdir(exist_ok=True)

-        super().__init__(hparams, model=student, config=student_cfg)
+        save_dir = self.output_dir.joinpath("student")
+
+        hparams.model_name_or_path = str(save_dir)  # Tell lightning we are training the student
+        teacher = AutoModelForSeq2SeqLM.from_pretrained(hparams.teacher).eval()
+        use_task_specific_params(teacher, hparams.task)  # We copy good generation parameters to student by default
+        student, e_layer_ids, d_layer_ids = create_student_by_copying_alternating_layers(
+            teacher, e=hparams.student_encoder_layers, d=hparams.student_decoder_layers, save_path=save_dir
+        )
+        if hparams.length_penalty != -1:
+            student.config.length_penalty = hparams.length_penalty
+        super().__init__(hparams, model=student, config=student.config)
+        self.e_layer_ids, self.d_layer_ids = e_layer_ids, d_layer_ids  # type: List[int], List[int]
+        self.different_encoder = hparams.student_encoder_layers != teacher.config.encoder_layers
+        self.different_decoder = hparams.student_decoder_layers != teacher.config.decoder_layers
        self.teacher = teacher
-        use_task_specific_params(self.teacher, "summarization")
        freeze_params(self.teacher)
-        self.sanity_check_gradients()
+
+        if not self.different_encoder:  # To save RAM, delete teacher encoder and freeze student encoder.
+            try:
+                del self.teacher.model.encoder
+            except AttributeError:  # T5
+                del self.teacher.encoder
+        # Intermediate supervision: Decide which layers to supervise
+        if hparams.supervise_forward:
+            self.d_matches = get_layers_to_supervise(
+                n_student=len(self.d_layer_ids), n_teacher=self.teacher.config.decoder_layers
+            )
+        else:
+            self.d_matches = self.d_layer_ids
        self.ce_loss_fct = nn.KLDivLoss(reduction="batchmean")
        self.temperature = 2.0
        self.alpha_mlm = hparams.alpha_mlm
        self.alpha_ce = hparams.alpha_ce
        self.alpha_hid = hparams.alpha_hid
-        # self.alpha_cos = hparams.alpha_cos
-        self.alpha_encoder_loss = self.hparams.alpha_encoder_loss
+        self.alpha_encoder_loss = hparams.alpha_encoder_loss
        gc.collect()
        torch.cuda.empty_cache()

-    def sanity_check_gradients(self):
-        assert_all_frozen(self.teacher)
-        assert_all_frozen(self.model.model.decoder.embed_tokens)
-        assert_all_frozen(self.model.model.encoder.embed_tokens)
-        if self.different_encoder:
-            assert any_requires_grad(self.model.model.encoder)
-        else:
-            freeze_params(self.model.model.encoder)
-            del self.teacher.model.encoder
-
-    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()
-        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
-        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
-
-        kw = teacher.config.to_diff_dict()
-        kw.update(student_updates)
-        # Copy weights
-        student_cfg = teacher.config_class(**kw)
-        student = type(teacher)(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)
-        student.save_pretrained(save_dir)
-        hparams.model_name_or_path = str(save_dir)
-        return student, student_cfg, teacher
-
-    def copy_to_student(self, d_layers_to_copy, e_layers_to_copy, hparams, student, teacher):
-        if teacher.config.model_type == "t5":
-            return self.copy_t5_to_student(d_layers_to_copy, e_layers_to_copy, hparams, student, teacher)
-        self.different_encoder: bool = hparams.student_encoder_layers != teacher.config.encoder_layers
-        self.different_decoder = hparams.student_decoder_layers != teacher.config.decoder_layers
-        if self.different_decoder:
-            copy_layers(teacher.model.decoder.layers, student.model.decoder.layers, d_layers_to_copy)
-        if self.different_encoder:
-            copy_layers(teacher.model.encoder.layers, student.model.encoder.layers, e_layers_to_copy)
-
-    def copy_t5_to_student(self, d_layers_to_copy, e_layers_to_copy, hparams, student, teacher):
-        self.different_encoder: bool = hparams.student_encoder_layers != teacher.config.num_layers
-        self.different_decoder = hparams.student_decoder_layers != teacher.config.num_layers
-        if self.different_decoder:
-            copy_layers(teacher.decoder.block, student.decoder.block, d_layers_to_copy)
-        if self.different_encoder:
-            copy_layers(teacher.encoder.block, student.encoder.block, e_layers_to_copy)
-
    def calc_mse_loss(self, teacher_outputs: torch.Tensor, student_outputs: torch.Tensor, mask) -> torch.FloatTensor:
+        """Supervise MSE(teacher.encoder_outputs, student.encoder_outputs)."""
+        # raise NotImplementedError()
        if mask is not None:
            # mask has False at padding_idx
            sel_mask = mask[:, :, None].expand_as(student_outputs).bool()
@@ -126,20 +87,15 @@ class BartSummarizationDistiller(SummarizationModule):
        return F.mse_loss(s_logits_slct, t_logits_slct)

    def calc_ce_loss(self, mask, s_logits, t_logits):
-        if mask is not None:
-            # mask has False at padding_idx
-            sel_mask = mask[:, :, None].expand_as(s_logits)
-            s_logits_slct = torch.masked_select(
-                s_logits, sel_mask
-            )  # (bs * seq_length * voc_size) modulo the 1s in mask
-            t_logits_slct = torch.masked_select(
-                t_logits, sel_mask
-            )  # (bs * seq_length * voc_size) modulo the 1s in mask
-        else:
-            t_logits_slct = t_logits
-            s_logits_slct = s_logits  # (bs * seq_length * voc_size) modulo the 1s in mask
-        s_logits_slct = s_logits_slct.view(-1, s_logits.size(-1))  # (bs * seq_length, voc_size) modulo the 1s in mask
-        t_logits_slct = t_logits_slct.view(-1, s_logits.size(-1))  # (bs * seq_length, voc_size) modulo the 1s in mask
+        """Copy pasted from distillbert (transformers/examples/distillation/)"""
+
+        # mask has False at padding_idx
+        sel_mask = mask[:, :, None].expand_as(s_logits)
+        vocab_size = s_logits.size(-1)
+        s_logits_slct = torch.masked_select(s_logits, sel_mask)  # (bs * seq_length * voc_size) modulo the 1s in mask
+        t_logits_slct = torch.masked_select(t_logits, sel_mask)  # (bs * seq_length * voc_size) modulo the 1s in mask
+        s_logits_slct = s_logits_slct.view(-1, vocab_size)  # (bs * seq_length, voc_size) modulo the 1s in mask
+        t_logits_slct = t_logits_slct.view(-1, vocab_size)  # (bs * seq_length, voc_size) modulo the 1s in mask
        assert t_logits_slct.size() == s_logits_slct.size()
        loss_ce = (
            self.ce_loss_fct(
@@ -148,7 +104,7 @@ class BartSummarizationDistiller(SummarizationModule):
            )
            * (self.temperature) ** 2
        )
-        return loss_ce, s_logits_slct, t_logits_slct
+        return loss_ce

    @staticmethod
    def add_model_specific_args(parser, root_dir):
@@ -157,10 +113,14 @@ class BartSummarizationDistiller(SummarizationModule):
        return parser

    def _step(self, batch):
-        # assert is_frozen(self.teacher)
+        # assert is_frozen(self.teacher) copied_decoder_layers
        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, labels = batch["input_ids"], batch["attention_mask"], batch["labels"]
+        if isinstance(self.model, T5ForConditionalGeneration):
+            decoder_input_ids = self.model._shift_right(labels)
+        else:
+            decoder_input_ids = shift_tokens_right(labels, pad_token_id)
+
        # noinspection PyCallingNonCallable
        lm_logits, dec_hidden, enc_outputs, enc_hidden_state = self(
            input_ids,
@@ -176,11 +136,11 @@ class BartSummarizationDistiller(SummarizationModule):
        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))
+            student_lm_loss = loss_fct(lm_logits.view(-1, lm_logits.shape[-1]), labels.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
+                lprobs, labels, self.hparams.label_smoothing, ignore_index=pad_token_id
            )

        def zero_tensor():
@@ -189,15 +149,14 @@ class BartSummarizationDistiller(SummarizationModule):
        loss_encoder, hid_loss_enc, hid_loss_dec = zero_tensor(), zero_tensor(), zero_tensor()
        if self.different_encoder:
            with torch.no_grad():
-                teacher_enc_outputs, teacher_enc_hid, _ = self.teacher.model.encoder(
+                teacher_enc_outputs, teacher_enc_hid, _ = self.teacher.get_encoder()(
                    input_ids, attention_mask=src_mask, output_hidden_states=True
                )
+            # DEPRECATE THIS
            if self.hparams.alpha_encoder_loss > 0:
                loss_encoder = self.calc_mse_loss(enc_outputs, teacher_enc_outputs, src_mask)

-            hid_loss_enc = self.calc_hidden_loss(
-                src_mask, enc_hidden_state, teacher_enc_hid, self.hparams.e_layer_to_copy
-            )
+            hid_loss_enc = self.calc_hidden_loss(src_mask, enc_hidden_state, teacher_enc_hid, self.e_layer_ids)

        teacher_enc_outputs = (enc_outputs,)
        assert isinstance(teacher_enc_outputs, tuple), type(teacher_enc_outputs)
@@ -208,13 +167,15 @@ 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)
-        if self.alpha_hid > 0:
-            hid_loss_dec = self.calc_hidden_loss(dec_mask, dec_hidden, tdec_hidden, self.hparams.d_matches)
+        loss_ce = self.calc_ce_loss(dec_mask, lm_logits, tlogits)
+        if self.alpha_hid > 0:  # Intermediate supervision of decoder hidden states
+            hid_loss_dec = self.calc_hidden_loss(
+                dec_mask, dec_hidden, tdec_hidden, self.d_matches, normalize_hidden=self.hparams.normalize_hidden
+            )

        blended_loss = (
            self.alpha_ce * loss_ce
@@ -224,7 +185,9 @@ class BartSummarizationDistiller(SummarizationModule):
        )
        return blended_loss, loss_ce, student_lm_loss, loss_encoder, hid_loss_enc, hid_loss_dec

-    def calc_hidden_loss(self, attention_mask, hidden_states, hidden_states_T, matches):
+    @staticmethod
+    def calc_hidden_loss(attention_mask, hidden_states, hidden_states_T, matches, normalize_hidden):
+        """MSE(student_hid, teacher_hid[matches]). Called "Intermediate supervision" in paper. Inspired by TinyBERT."""
        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}"
@@ -232,7 +195,7 @@ class BartSummarizationDistiller(SummarizationModule):
        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:
+        if 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")
@@ -280,130 +243,9 @@ class BartTranslationDistiller(BartSummarizationDistiller):
        return parser


-class T5SummarizationDistiller(BartSummarizationDistiller):
-    def pre_init(self, hparams):
-        raise NotImplementedError("T5 Distillation does not work yet")
-        self.output_dir = Path(hparams.output_dir)
-        self.output_dir.mkdir(exist_ok=True)
-        teacher = T5ForConditionalGeneration.from_pretrained(hparams.teacher)
-        n_layer = hparams.student_decoder_layers
-        assert n_layer == hparams.student_encoder_layers  # TODO(SS): relax this constraint so that we can do 12-6.
-        d_layers_to_copy = get_layers_to_copy(n_layer, len(teacher.decoder.block))
-        e_layers_to_copy: List = get_layers_to_copy(n_layer, len(teacher.encoder.block))
-        student_updates = {"num_layers": n_layer}
-        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 = T5Config(**kw)
-        student = T5ForConditionalGeneration(student_cfg)
-        student, _ = init_student(student, teacher)
-        self.copy_to_student(d_layers_to_copy, e_layers_to_copy, hparams, student, teacher)
-        Path(hparams.output_dir).mkdir(exist_ok=True)
-        task_specific_params = student.config.task_specific_params
-        if task_specific_params is not None:
-            student.config.update(task_specific_params.get("summarization", {}))  # TODO: dont hardcode
-        save_dir = self.output_dir.joinpath("student")
-        save_dir.mkdir(exist_ok=True)
-
-        student.save_pretrained(save_dir)
-        hparams.model_name_or_path = str(save_dir)
-        return student, student_cfg, teacher
-
-    def freeze_embeds(self):
-        freeze_params(self.model.shared)
-        for d in [self.model.encoder, self.model.decoder]:
-            freeze_params(d.embed_tokens)
-
-    def sanity_check_gradients(self):
-        """T5"""
-        assert_all_frozen(self.teacher)
-        assert_all_frozen(self.model.decoder.embed_tokens)
-        assert_all_frozen(self.model.encoder.embed_tokens)
-        if self.different_encoder:
-            assert any_requires_grad(self.model.encoder)
-        else:
-            freeze_params(self.model.encoder)
-            del self.teacher.model.encoder
-        if self.different_decoder:
-            assert any_requires_grad(self.model.decoder)
-        else:
-            freeze_params(self.model.decoder)  # TODO(SS): very suspicious
-
-    def _step(self, batch):
-        pad_token_id = self.tokenizer.pad_token_id
-        source_ids, source_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
-        dec_mask = decoder_input_ids.ne(pad_token_id)
-
-        sloss, slogits, dec_hidden, enc_outputs, enc_hidden_state = self(
-            source_ids,
-            attention_mask=source_mask,
-            decoder_input_ids=decoder_input_ids,
-            labels=labels,
-            output_hidden_states=True,
-            output_attentions=False,
-            use_cache=False,
-        )
-
-        def zero_tensor():
-            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:
-            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,
-                )
-            if self.hparams.alpha_encoder_loss > 0:
-                loss_encoder = self.calc_mse_loss(enc_outputs, teacher_enc_outputs, source_mask)
-
-            hid_loss_enc = self.calc_hidden_loss(
-                source_mask, enc_hidden_state, teacher_enc_hid, self.hparams.e_layer_to_copy
-            )
-
-        teacher_enc_outputs = (enc_outputs,)
-        assert isinstance(teacher_enc_outputs, tuple), type(teacher_enc_outputs)
-
-        with torch.no_grad():
-            tloss, tlogits, tdec_hidden, _ = self.teacher(
-                source_ids,
-                attention_mask=source_mask,
-                encoder_outputs=teacher_enc_outputs,
-                decoder_input_ids=decoder_input_ids,
-                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)
-
-        blended_loss = (
-            self.alpha_ce * loss_ce
-            + 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, sloss, loss_encoder, hid_loss_enc, hid_loss_dec
-
-
 def create_module(args):
-    t5 = "t5" in args.model_name_or_path
    if args.no_teacher:
        module_cls = TranslationModule if "translation" in args.task else SummarizationModule
-    elif t5:  # DISTILL T5 WITH TEACHER FOR SUMMARIZATION
-        assert "translation" not in args.task, "t5 translation distillation not supported"
-        module_cls = T5SummarizationDistiller
    else:  # DISTILL WITH TEACHER
        module_cls = BartTranslationDistiller if "translation" in args.task else BartSummarizationDistiller
    args.setup_cls: str = module_cls.__name__
@@ -436,56 +278,6 @@ def evaluate_checkpoint(ckpt_path: Path, dest_dir=None):
    trainer.test(model)


-LAYERS_TO_COPY = {
-    # maps  num layers in student -> which teacher layers to copy.
-    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
-    12: {
-        1: [0],
-        2: [0, 6],
-        3: [0, 6, 11],
-        4: [0, 4, 8, 11],
-        6: [0, 2, 4, 7, 9, 11],
-        9: [0, 1, 2, 4, 5, 7, 9, 10, 11],
-        12: list(range(12)),
-    },
-    16: {  # maps  num layers in student -> which teacher layers to copy
-        1: [0],
-        2: [0, 8],
-        3: [0, 8, 15],
-        4: [0, 5, 10, 15],
-        6: [0, 3, 6, 9, 12, 15],
-        8: [0, 2, 4, 6, 8, 10, 12, 15],
-        9: [0, 1, 3, 5, 7, 9, 11, 13, 15],
-        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 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 distill_main(args):
    Path(args.output_dir).mkdir(exist_ok=True)
    if len(os.listdir(args.output_dir)) > 3 and args.do_train:
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 from pathlib import Path

 import fire
@@ -1,7 +1,10 @@
+#!/usr/bin/env python
+
 import argparse
 import glob
 import logging
 import os
+import sys
 import time
 from collections import defaultdict
 from pathlib import Path
@@ -13,7 +16,6 @@ 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.modeling_bart import shift_tokens_right
 from utils import (
@@ -34,6 +36,11 @@ from utils import (
 )


+# need the parent dir module
+sys.path.insert(2, str(Path(__file__).resolve().parents[1]))
+from lightning_base import BaseTransformer, add_generic_args, generic_train  # noqa
+
+
 logger = logging.getLogger(__name__)


@@ -1,6 +1,3 @@
-# Add parent directory to python path to access lightning_base.py
-export PYTHONPATH="../":"${PYTHONPATH}"
-
 # the proper usage is documented in the README, you need to specify data_dir, output_dir and model_name_or_path
 # run ./finetune.sh --help to see all the possible options
 python finetune.py \
@@ -0,0 +1,442 @@
+import json
+import logging
+import os
+import sys
+from dataclasses import dataclass, field
+from typing import Callable, Dict, List, Optional, Tuple
+
+import numpy as np
+import torch
+
+from seq2seq_trainer import Seq2SeqTrainer
+from transformers import (
+    AutoConfig,
+    AutoModelForSeq2SeqLM,
+    AutoTokenizer,
+    BartTokenizer,
+    EvalPrediction,
+    HfArgumentParser,
+    MBartTokenizer,
+    T5Tokenizer,
+    TrainingArguments,
+    set_seed,
+)
+from transformers.modeling_bart import shift_tokens_right
+from utils import (
+    LegacySeq2SeqDataset,
+    Seq2SeqDataset,
+    assert_all_frozen,
+    calculate_bleu,
+    calculate_rouge,
+    freeze_params,
+    lmap,
+    trim_batch,
+    use_task_specific_params,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+class Seq2SeqDataCollator:
+    def __init__(self, tokenizer, data_args, tpu_num_cores=None):
+        self.tokenizer = tokenizer
+        self.pad_token_id = tokenizer.pad_token_id
+        self.data_args = data_args
+        self.tpu_num_cores = tpu_num_cores
+        self.add_prefix_space = isinstance(tokenizer, BartTokenizer)
+
+    def __call__(self, batch) -> Dict[str, torch.Tensor]:
+        if hasattr(self.tokenizer, "prepare_seq2seq_batch"):
+            batch = self._encode(batch)
+            input_ids, attention_mask, labels = (
+                batch["input_ids"],
+                batch["attention_mask"],
+                batch["labels"],
+            )
+        else:
+            input_ids = torch.stack([x["input_ids"] for x in batch])
+            attention_mask = torch.stack([x["attention_mask"] for x in batch])
+            labels = torch.stack([x["labels"] for x in batch])
+
+            labels = trim_batch(labels, self.pad_token_id)
+            input_ids, attention_mask = trim_batch(input_ids, self.pad_token_id, attention_mask=attention_mask)
+
+        if isinstance(self.tokenizer, T5Tokenizer):
+            decoder_input_ids = self._shift_right_t5(labels)
+            labels = labels
+        else:
+            decoder_input_ids = shift_tokens_right(labels, self.pad_token_id)
+            labels = labels
+
+        batch = {
+            "input_ids": input_ids,
+            "attention_mask": attention_mask,
+            "decoder_input_ids": decoder_input_ids,
+            "labels": labels,
+        }
+        return batch
+
+    def _shift_right_t5(self, input_ids):
+        decoder_start_token_id = self.pad_token_id
+
+        assert (
+            decoder_start_token_id is not None
+        ), "self.model.config.decoder_start_token_id has to be defined. In T5 it is usually set to the pad_token_id. See T5 docs for more information"
+
+        # shift inputs to the right
+        shifted_input_ids = input_ids.new_zeros(input_ids.shape)
+        shifted_input_ids[..., 1:] = input_ids[..., :-1].clone()
+        shifted_input_ids[..., 0] = decoder_start_token_id
+
+        return shifted_input_ids
+
+    def _encode(self, batch) -> Dict[str, torch.Tensor]:
+        batch_encoding = self.tokenizer.prepare_seq2seq_batch(
+            [x["src_texts"] for x in batch],
+            src_lang=self.data_args.src_lang,
+            tgt_texts=[x["tgt_texts"] for x in batch],
+            tgt_lang=self.data_args.tgt_lang,
+            max_length=self.data_args.max_source_length,
+            max_target_length=self.data_args.max_target_length,
+            padding="max_length" if self.tpu_num_cores is not None else "longest",  # TPU hack
+            return_tensors="pt",
+            add_prefix_space=self.add_prefix_space,
+        )
+        return batch_encoding.data
+
+
+@dataclass
+class Seq2SeqTrainingArguments(TrainingArguments):
+    """
+    Parameters:
+        label_smoothing (:obj:`float`, `optional`, defaults to 0):
+            The label smoothing epsilon to apply (if not zero).
+        sortish_sampler (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to SortishSamler or not. It sorts the inputs according to lenghts in-order to minimizing the padding size.
+        predict_with_generate (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether to use generate to calculate generative metrics (ROUGE, BLEU).
+    """
+
+    label_smoothing: Optional[float] = field(
+        default=0.0, metadata={"help": "The label smoothing epsilon to apply (if not zero)."}
+    )
+    sortish_sampler: bool = field(default=False, metadata={"help": "Whether to SortishSamler or not."})
+    predict_with_generate: bool = field(
+        default=False, metadata={"help": "Whether to use generate to calculate generative metrics (ROUGE, BLEU)."}
+    )
+
+
+@dataclass
+class ModelArguments:
+    """
+    Arguments pertaining to which model/config/tokenizer we are going to fine-tune from.
+    """
+
+    model_name_or_path: str = field(
+        metadata={"help": "Path to pretrained model or model identifier from huggingface.co/models"}
+    )
+    config_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained config name or path if not the same as model_name"}
+    )
+    tokenizer_name: Optional[str] = field(
+        default=None, metadata={"help": "Pretrained tokenizer name or path if not the same as model_name"}
+    )
+    cache_dir: Optional[str] = field(
+        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
+    )
+    freeze_encoder: bool = field(default=False, metadata={"help": "Whether tp freeze the encoder."})
+    freeze_embeds: bool = field(default=False, metadata={"help": "Whether  to freeze the embeddings."})
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+    """
+
+    data_dir: str = field(
+        metadata={"help": "The input data dir. Should contain the .tsv files (or other data files) for the task."}
+    )
+    task: Optional[str] = field(
+        default="summarization",
+        metadata={"help": "Task name, summarization (or summarization_{dataset} for pegasus) or translation"},
+    )
+    max_source_length: Optional[int] = field(
+        default=1024,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    max_target_length: Optional[int] = field(
+        default=128,
+        metadata={
+            "help": "The maximum total sequence length for target text after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    val_max_target_length: Optional[int] = field(
+        default=142,
+        metadata={
+            "help": "The maximum total sequence length for validation target text after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    test_max_target_length: Optional[int] = field(
+        default=142,
+        metadata={
+            "help": "The maximum total sequence length for test target text after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    n_train: Optional[int] = field(default=-1, metadata={"help": "# training examples. -1 means use all."})
+    n_val: Optional[int] = field(default=-1, metadata={"help": "# validation examples. -1 means use all."})
+    n_test: Optional[int] = field(default=-1, metadata={"help": "# test examples. -1 means use all."})
+    src_lang: Optional[str] = field(default=None, metadata={"help": "Source language id for translation."})
+    tgt_lang: Optional[str] = field(default=None, metadata={"help": "Target language id for translation."})
+    eval_beams: Optional[int] = field(default=None, metadata={"help": "# num_beams to use for evaluation."})
+
+
+def main():
+    # See all possible arguments in src/transformers/training_args.py
+    # or by passing the --help flag to this script.
+    # We now keep distinct sets of args, for a cleaner separation of concerns.
+
+    parser = HfArgumentParser((ModelArguments, DataTrainingArguments, Seq2SeqTrainingArguments))
+
+    if len(sys.argv) == 2 and sys.argv[1].endswith(".json"):
+        # If we pass only one argument to the script and it's the path to a json file,
+        # let's parse it to get our arguments.
+        model_args, data_args, training_args = parser.parse_json_file(json_file=os.path.abspath(sys.argv[1]))
+    else:
+        model_args, data_args, training_args = parser.parse_args_into_dataclasses()
+
+    if (
+        os.path.exists(training_args.output_dir)
+        and os.listdir(training_args.output_dir)
+        and training_args.do_train
+        and not training_args.overwrite_output_dir
+    ):
+        raise ValueError(
+            f"Output directory ({training_args.output_dir}) already exists and is not empty. Use --overwrite_output_dir to overcome."
+        )
+
+    # Setup logging
+    logging.basicConfig(
+        format="%(asctime)s - %(levelname)s - %(name)s -   %(message)s",
+        datefmt="%m/%d/%Y %H:%M:%S",
+        level=logging.INFO if training_args.local_rank in [-1, 0] else logging.WARN,
+    )
+    logger.warning(
+        "Process rank: %s, device: %s, n_gpu: %s, distributed training: %s, 16-bits training: %s",
+        training_args.local_rank,
+        training_args.device,
+        training_args.n_gpu,
+        bool(training_args.local_rank != -1),
+        training_args.fp16,
+    )
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Set seed
+    set_seed(training_args.seed)
+
+    # Load pretrained model and tokenizer
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+
+    config = AutoConfig.from_pretrained(
+        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+    )
+    tokenizer = AutoTokenizer.from_pretrained(
+        model_args.tokenizer_name if model_args.tokenizer_name else model_args.model_name_or_path,
+        cache_dir=model_args.cache_dir,
+    )
+    model = AutoModelForSeq2SeqLM.from_pretrained(
+        model_args.model_name_or_path,
+        from_tf=".ckpt" in model_args.model_name_or_path,
+        config=config,
+        cache_dir=model_args.cache_dir,
+    )
+
+    # use task specific params
+    use_task_specific_params(model, data_args.task)
+
+    # set num_beams for evaluation
+    if data_args.eval_beams is not None:
+        model.config.num_beams = data_args.eval_beams
+    assert model.config.num_beams >= 1, f"got eval_beams={model.config.num_beams}. Need an integer >= 1"
+
+    # set max length for generation
+    model.config.max_generate_length = data_args.val_max_target_length
+
+    # set decoder_start_token_id for MBart
+    if model.config.decoder_start_token_id is None and isinstance(tokenizer, MBartTokenizer):
+        decoder_start_token_id = tokenizer.lang_code_to_id[data_args.tgt_lang]
+        model.config.decoder_start_token_id = decoder_start_token_id
+
+    def build_compute_metrics_fn(task_name: str) -> Callable[[EvalPrediction], Dict]:
+        def non_pad_len(tokens: np.ndarray) -> int:
+            return np.count_nonzero(tokens != tokenizer.pad_token_id)
+
+        def decode_pred(pred: EvalPrediction) -> Tuple[List[str], List[str]]:
+            pred_str = tokenizer.batch_decode(pred.predictions, skip_special_tokens=True)
+            label_str = tokenizer.batch_decode(pred.label_ids, skip_special_tokens=True)
+            pred_str = lmap(str.strip, pred_str)
+            label_str = lmap(str.strip, label_str)
+            return pred_str, label_str
+
+        def summarization_metrics(pred: EvalPrediction) -> Dict:
+            pred_str, label_str = decode_pred(pred)
+            rouge: Dict = calculate_rouge(pred_str, label_str)
+            summ_len = np.mean(lmap(non_pad_len, pred.predictions))
+            rouge.update({"gen_len": summ_len})
+            return rouge
+
+        def translation_metrics(pred: EvalPrediction) -> Dict:
+            pred_str, label_str = decode_pred(pred)
+            bleu: Dict = calculate_bleu(pred_str, label_str)
+            gen_len = np.mean(lmap(non_pad_len, pred.predictions))
+            bleu.update({"gen_len": gen_len})
+            return bleu
+
+        compute_metrics_fn = summarization_metrics if "summarization" in task_name else translation_metrics
+        return compute_metrics_fn
+
+    def freeze_embeds(model: torch.nn.Module):
+        """Freeze token embeddings and positional embeddings for bart, just token embeddings for t5."""
+        try:
+            freeze_params(model.model.shared)
+            for d in [model.model.encoder, model.model.decoder]:
+                freeze_params(d.embed_positions)
+                freeze_params(d.embed_tokens)
+        except AttributeError:
+            freeze_params(model.shared)
+            for d in [model.encoder, model.decoder]:
+                freeze_params(d.embed_tokens)
+
+    if model_args.freeze_embeds:
+        freeze_embeds(model)
+    if model_args.freeze_encoder:
+        freeze_params(model.get_encoder())
+        assert_all_frozen(model.get_encoder())
+
+    dataset_class = Seq2SeqDataset if hasattr(tokenizer, "prepare_seq2seq_batch") else LegacySeq2SeqDataset
+
+    # Get datasets
+    train_dataset = (
+        dataset_class(
+            tokenizer,
+            type_path="train",
+            data_dir=data_args.data_dir,
+            n_obs=data_args.n_train,
+            max_target_length=data_args.max_target_length,
+            max_source_length=data_args.max_source_length,
+            prefix=model.config.prefix or "",
+        )
+        if training_args.do_train
+        else None
+    )
+    eval_dataset = (
+        dataset_class(
+            tokenizer,
+            type_path="val",
+            data_dir=data_args.data_dir,
+            n_obs=data_args.n_val,
+            max_target_length=data_args.val_max_target_length,
+            max_source_length=data_args.max_source_length,
+            prefix=model.config.prefix or "",
+        )
+        if training_args.do_eval
+        else None
+    )
+    test_dataset = (
+        dataset_class(
+            tokenizer,
+            type_path="test",
+            data_dir=data_args.data_dir,
+            n_obs=data_args.n_test,
+            max_target_length=data_args.test_max_target_length,
+            max_source_length=data_args.max_source_length,
+            prefix=model.config.prefix or "",
+        )
+        if training_args.do_predict
+        else None
+    )
+
+    # Initialize our Trainer
+    trainer = Seq2SeqTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        data_collator=Seq2SeqDataCollator(tokenizer, data_args, training_args.tpu_num_cores),
+        compute_metrics=build_compute_metrics_fn(data_args.task) if training_args.predict_with_generate else None,
+    )
+
+    # Training
+    if training_args.do_train:
+        trainer.train(
+            model_path=model_args.model_name_or_path if os.path.isdir(model_args.model_name_or_path) else None
+        )
+        trainer.save_model()
+        # For convenience, we also re-save the tokenizer to the same directory,
+        # so that you can share your model easily on huggingface.co/models =)
+        if trainer.is_world_process_zero():
+            tokenizer.save_pretrained(training_args.output_dir)
+
+    # Evaluation
+    eval_results = {}
+    if training_args.do_eval:
+        logger.info("*** Evaluate ***")
+
+        result = trainer.evaluate()
+
+        output_eval_file = os.path.join(training_args.output_dir, "eval_results.json")
+        if trainer.is_world_process_zero():
+            logger.info("***** Eval results *****")
+            for key, value in result.items():
+                logger.info("  %s = %s", key, value)
+
+            with open(output_eval_file, "w") as f:
+                json.dump(result, f)
+
+            eval_results.update(result)
+
+    if training_args.do_predict:
+        logging.info("*** Test ***")
+
+        test_output = trainer.predict(test_dataset=test_dataset)
+        test_metrics = test_output.metrics
+        test_metrics = {k.replace("eval", "test"): v for k, v in test_metrics.items()}
+
+        output_test_file = os.path.join(training_args.output_dir, "test_results.json")
+
+        if trainer.is_world_process_zero():
+            logger.info("***** Test results *****")
+            for key, value in test_metrics.items():
+                logger.info("  %s = %s", key, value)
+
+            with open(output_test_file, "w") as f:
+                json.dump(test_metrics, f)
+
+            if training_args.predict_with_generate:
+                test_preds = tokenizer.batch_decode(test_output.predictions, skip_special_tokens=True)
+                test_preds = lmap(str.strip, test_preds)
+                output_test_pred_file = os.path.join(training_args.output_dir, "test_generations.txt")
+                with open(output_test_pred_file, "w") as f:
+                    f.write("\n".join(test_preds))
+
+    return eval_results
+
+
+def _mp_fn(index):
+    # For xla_spawn (TPUs)
+    main()
+
+
+if __name__ == "__main__":
+    main()
@@ -1,20 +0,0 @@
-from typing import List
-
-from torch import nn
-
-
-def init_student(student, teacher):
-    teacher_state_dict = teacher.state_dict()
-    info = student.load_state_dict(teacher_state_dict, strict=False)
-    assert info.missing_keys == [], info.missing_keys
-    return student, info
-
-
-def copy_decoder_layers(teacher, student, l2copy=[0, 2, 4, 7, 9, 11]):
-    copy_layers(teacher.model.decoder.layers, student.model.decoder.layers, l2copy)
-
-
-def copy_layers(teacher_layers: nn.ModuleList, student_layers: nn.ModuleList, layers_to_copy: List) -> None:
-    layers_to_copy = nn.ModuleList([l for i, l in enumerate(teacher_layers) if i in layers_to_copy])
-    assert len(student_layers) == len(layers_to_copy), f"{len(student_layers)} != {len(layers_to_copy)}"
-    student_layers.load_state_dict(layers_to_copy.state_dict())
@@ -0,0 +1,169 @@
+import warnings
+from pathlib import Path
+from typing import List, Tuple, Union
+
+import fire
+from torch import nn
+
+from transformers import AutoModelForSeq2SeqLM, AutoTokenizer, PreTrainedModel
+from transformers.utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+def copy_layers(src_layers: nn.ModuleList, dest_layers: nn.ModuleList, layers_to_copy: List[int]) -> None:
+    layers_to_copy = nn.ModuleList([l for i, l in enumerate(src_layers) if i in layers_to_copy])
+    assert len(dest_layers) == len(layers_to_copy), f"{len(dest_layers)} != {len(layers_to_copy)}"
+    dest_layers.load_state_dict(layers_to_copy.state_dict())
+
+
+LAYERS_TO_COPY = {
+    # maps  num layers in teacher -> num_layers in student -> which teacher layers to copy.
+    # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP
+    12: {
+        1: [0],  # This says that if the teacher has 12 layers and the student has 1, copy layer 0 of the teacher
+        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 = {
+    # maps  num layers in student -> which teacher layers to copy.
+    6: {1: [5], 2: [3, 5], 3: [1, 4, 5], 4: [1, 2, 4, 5]},
+    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]},
+}
+
+
+def pick_layers_to_copy(n_student, n_teacher):
+    try:
+        val = LAYERS_TO_COPY[n_teacher][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_supervise(n_student, n_teacher) -> List[int]:
+    """Used or the --supervise_forward kwarg"""
+    if n_student > n_teacher:
+        raise ValueError(f"Cannot perform intermediate supervision for student {n_student} > teacher {n_teacher}")
+    elif n_teacher == n_student:
+        return list(range(n_teacher))
+    elif n_student == 1:
+        return [n_teacher - 1]
+    else:
+        return LAYERS_TO_SUPERVISE[n_teacher][n_student]
+
+
+def create_student_by_copying_alternating_layers(
+    teacher: Union[str, PreTrainedModel],
+    save_path: Union[str, Path] = "student",
+    e: Union[int, None] = None,
+    d: Union[int, None] = None,
+    copy_first_teacher_layers=False,
+    **extra_config_kwargs
+) -> Tuple[PreTrainedModel, List[int], List[int]]:
+    """Make a student by copying alternating layers from a teacher, save it to save_path.
+    Args:
+        teacher: str or PreTrainedModel if str, this will call AutoModelForSeq2SeqLM.from_pretrained(teacher) before
+        copying layers
+        save_path: where to save the student, defaults to student directory.
+        e: how many Encoder layers should the student have, default is fully copy of teacher
+        d: how many Decoder layers should the student have, default is fully copy of teacher
+        copy_first_teacher_layers: [bool] dont copy alternating layers, just the first e/d.
+        **extra_config_kwargs: extra kwargs to pass to the student, by default the teacher config is used.
+
+    Returns:
+        student: new, smaller model.  (Also saves it to save_path)
+        e_layers_to_copy: list of which teacher encoder layers were used
+        d_layers_to_copy: list of which teacher decoder layers were used
+    """
+    _msg = "encoder_layers and decoder_layers cannot be both None-- you would just have an identical teacher."
+    assert (e is not None) or (d is not None), _msg
+    if isinstance(teacher, str):
+        AutoTokenizer.from_pretrained(teacher).save_pretrained(save_path)  # purely for convenience
+        teacher = AutoModelForSeq2SeqLM.from_pretrained(teacher).eval()
+    else:
+
+        assert isinstance(teacher, PreTrainedModel), f"teacher must be a model or string got type {type(teacher)}"
+    init_kwargs = teacher.config.to_diff_dict()
+
+    try:
+        teacher_e, teacher_d = teacher.config.encoder_layers, teacher.config.decoder_layers
+        if e is None:
+            e = teacher_e
+        if d is None:
+            d = teacher_d
+        init_kwargs.update({"encoder_layers": e, "decoder_layers": d})
+    except AttributeError:  # T5
+        teacher_e, teacher_d = teacher.config.num_layers, teacher.config.num_decoder_layers
+        if e is None:
+            e = teacher_e
+        if d is None:
+            d = teacher_d
+        init_kwargs.update({"num_layers": e, "num_decoder_layers": d})
+
+    # Kwargs to instantiate student: teacher kwargs with updated layer numbers + **extra_config_kwargs
+    init_kwargs.update(extra_config_kwargs)
+
+    # Copy weights
+    student_cfg = teacher.config_class(**init_kwargs)
+    student = AutoModelForSeq2SeqLM.from_config(student_cfg)
+    # Start by copying the full teacher state dict this will copy the first N teacher layers to the student.
+    info = student.load_state_dict(teacher.state_dict(), strict=False)
+    assert info.missing_keys == [], info.missing_keys  # every student key should have a teacher keys.
+
+    if copy_first_teacher_layers:  # Our copying is done. We just log and save
+        e_layers_to_copy, d_layers_to_copy = list(range(e)), list(range(d))
+        logger.info(
+            f"Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to {save_path}"
+        )
+        student.save_pretrained(save_path)
+        return student, e_layers_to_copy, d_layers_to_copy
+
+    # Decide which layers of the teacher to copy. Not exactly alternating -- we try to keep first and last layer.
+    e_layers_to_copy: List[int] = pick_layers_to_copy(e, teacher_e)
+    d_layers_to_copy: List[int] = pick_layers_to_copy(d, teacher_d)
+
+    try:
+        copy_layers(teacher.model.encoder.layers, student.model.encoder.layers, e_layers_to_copy)
+        copy_layers(teacher.model.decoder.layers, student.model.decoder.layers, d_layers_to_copy)
+    except AttributeError:  # For t5, student.model.encoder.layers is called student.encoder.block
+        copy_layers(teacher.encoder.block, student.encoder.block, e_layers_to_copy)
+        copy_layers(teacher.decoder.block, student.decoder.block, d_layers_to_copy)
+    logger.info(
+        f"Copied encoder layers {e_layers_to_copy} and decoder layers {d_layers_to_copy}. Saving them to {save_path}"
+    )
+    student.config.init_metadata = dict(
+        teacher_type=teacher.config.model_type,
+        copied_encoder_layers=e_layers_to_copy,
+        copied_decoder_layers=d_layers_to_copy,
+    )
+    student.save_pretrained(save_path)
+    # Save information about copying for easier reproducibility
+
+    return student, e_layers_to_copy, d_layers_to_copy
+
+
+if __name__ == "__main__":
+    fire.Fire(create_student_by_copying_alternating_layers)
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 from pathlib import Path

 import fire
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 """Fill examples with bitext up to max_tokens without breaking up examples.
 [['I went', 'yo fui'],
 ['to the store', 'a la tienda']
@@ -0,0 +1,17 @@
+import fire
+
+from utils import calculate_rouge, save_json
+
+
+def calculate_rouge_path(pred_path, tgt_path, save_path=None, **kwargs):
+    """Kwargs will be passed to calculate_rouge"""
+    pred_lns = [x.strip() for x in open(pred_path).readlines()]
+    tgt_lns = [x.strip() for x in open(tgt_path).readlines()][: len(pred_lns)]
+    metrics = calculate_rouge(pred_lns, tgt_lns, **kwargs)
+    if save_path is not None:
+        save_json(metrics, save_path, indent=None)
+    return metrics  # these print nicely
+
+
+if __name__ == "__main__":
+    fire.Fire(calculate_rouge_path)
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import argparse
 import shutil
 import time
@@ -15,6 +17,7 @@ from utils import (
    Seq2SeqDataset,
    calculate_bleu,
    calculate_rouge,
+    chunks,
    lmap,
    load_json,
    parse_numeric_n_bool_cl_kwargs,
@@ -38,6 +41,7 @@ def eval_data_dir(
    fp16=False,
    task="summarization",
    local_rank=None,
+    num_return_sequences=1,
    src_lang=None,
    tgt_lang=None,
    prefix="",
@@ -54,10 +58,15 @@ def eval_data_dir(
    model = AutoModelForSeq2SeqLM.from_pretrained(model_name).cuda()
    if fp16:
        model = model.half()
+    # determine if we need to increase num_beams
+    use_task_specific_params(model, task)  # update config with task specific params
+    num_beams = generate_kwargs.pop("num_beams", model.config.num_beams)  # AttributeError risk?
+    if num_return_sequences > num_beams:
+        num_beams = num_return_sequences

    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
    if prefix is None:
@@ -82,10 +91,14 @@ def eval_data_dir(
        summaries = model.generate(
            input_ids=batch["input_ids"].to(model.device),
            attention_mask=batch["attention_mask"].to(model.device),
+            num_return_sequences=num_return_sequences,
+            num_beams=num_beams,
            **generate_kwargs,
        )
        preds = tokenizer.batch_decode(summaries, skip_special_tokens=True, clean_up_tokenization_spaces=False)
        ids = batch["ids"]
+        if num_return_sequences > 1:
+            preds = chunks(preds, num_return_sequences)  # batch size chunks, each of size num_return_seq
        for i, pred in enumerate(preds):
            results.append(dict(pred=pred, id=ids[i].item()))
    save_json(results, save_path)
@@ -108,7 +121,6 @@ def run_generate():
    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(
@@ -118,6 +130,9 @@ def run_generate():
    parser.add_argument(
        "--n_obs", type=int, default=None, required=False, help="How many observations. Defaults to all."
    )
+    parser.add_argument(
+        "--num_return_sequences", type=int, default=1, required=False, help="How many sequences to return"
+    )
    parser.add_argument(
        "--sync_timeout",
        type=int,
@@ -156,6 +171,7 @@ def run_generate():
        local_rank=args.local_rank,
        n_obs=args.n_obs,
        max_source_length=args.max_source_length,
+        num_return_sequences=args.num_return_sequences,
        prefix=args.prefix,
        src_lang=args.src_lang,
        tgt_lang=args.tgt_lang,
@@ -167,6 +183,11 @@ def run_generate():
        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)
+        if args.num_return_sequences > 1:
+            save_path = save_dir.joinpath("pseudolabel_results.json")
+            print(f"Saving aggregated results at {save_path}, intermediate in {json_save_dir}/")
+            save_json(preds, save_path)
+            return
        tgt_file = Path(args.data_dir).joinpath(args.type_path + ".target")
        labels = [x.rstrip() for x in open(tgt_file).readlines()][: len(preds)]

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import argparse
 import datetime
 import json
@@ -11,7 +13,7 @@ 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
+from utils import calculate_bleu, calculate_rouge, chunks, parse_numeric_n_bool_cl_kwargs, use_task_specific_params


 logger = getLogger(__name__)
@@ -20,12 +22,6 @@ logger = getLogger(__name__)
 DEFAULT_DEVICE = "cuda" if torch.cuda.is_available() else "cpu"


-def chunks(lst, n):
-    """Yield successive n-sized chunks from lst."""
-    for i in range(0, len(lst), n):
-        yield lst[i : i + n]
-
-
 def generate_summaries_or_translations(
    examples: List[str],
    out_file: str,
@@ -156,8 +152,7 @@ def run_generate(verbose=True):
        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

@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import argparse
 import itertools
 import operator
@@ -5,13 +7,14 @@ import sys
 from collections import OrderedDict

 from run_eval import datetime_now, run_generate
+from utils import ROUGE_KEYS


 # 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"],
+    "summarization": ROUGE_KEYS,
 }


@@ -1,14 +1,11 @@
+#!/usr/bin/env python
+
 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
+from utils import Seq2SeqDataset, pickle_save


 def save_len_file(
@@ -0,0 +1,26 @@
+#!/usr/bin/env python
+
+import fire
+
+from transformers import AutoConfig, AutoModelForSeq2SeqLM, AutoTokenizer
+
+
+def save_randomly_initialized_version(config_name: str, save_dir: str, **config_kwargs):
+    """Save a randomly initialized version of a model using a pretrained config.
+    Args:
+        config_name: which config to use
+        save_dir: where to save the resulting model and tokenizer
+        config_kwargs: Passed to AutoConfig
+
+    Usage::
+        save_randomly_initialized_version("facebook/bart-large-cnn", "distilbart_random_cnn_6_3", encoder_layers=6, decoder_layers=3, num_beams=3)
+    """
+    cfg = AutoConfig.from_pretrained(config_name, **config_kwargs)
+    model = AutoModelForSeq2SeqLM.from_config(cfg)
+    model.save_pretrained(save_dir)
+    AutoTokenizer.from_pretrained(config_name).save_pretrained(save_dir)
+    return model
+
+
+if __name__ == "__main__":
+    fire.Fire(save_randomly_initialized_version)
@@ -0,0 +1,21 @@
+import re
+
+
+try:
+    import nltk
+
+    NLTK_AVAILABLE = True
+except (ImportError, ModuleNotFoundError):
+    NLTK_AVAILABLE = False
+
+if NLTK_AVAILABLE:
+    try:
+        nltk.download("punkt", quiet=True)
+    except FileExistsError:  # multiprocessing race condition
+        pass
+
+
+def add_newline_to_end_of_each_sentence(x: str) -> str:
+    re.sub("<n>", "", x)  # remove pegasus newline char
+    assert NLTK_AVAILABLE, "nltk must be installed to separate newlines betwee sentences. (pip install nltk)"
+    return "\n".join(nltk.sent_tokenize(x))
@@ -0,0 +1,126 @@
+import logging
+from typing import Any, Dict, Optional, Tuple, Union
+
+import torch
+from torch import nn
+from torch.utils.data import DistributedSampler, RandomSampler
+
+from transformers import Trainer
+from transformers.file_utils import is_torch_tpu_available
+from transformers.trainer import get_tpu_sampler
+
+
+try:
+    from .utils import label_smoothed_nll_loss
+except ImportError:
+    from utils import label_smoothed_nll_loss
+
+
+logger = logging.getLogger(__name__)
+
+
+class Seq2SeqTrainer(Trainer):
+    def _get_train_sampler(self) -> Optional[torch.utils.data.sampler.Sampler]:
+        if isinstance(self.train_dataset, torch.utils.data.IterableDataset):
+            return None
+        elif is_torch_tpu_available():
+            return get_tpu_sampler(self.train_dataset)
+        else:
+            if self.args.sortish_sampler:
+                self.train_dataset.make_sortish_sampler(
+                    self.args.per_device_train_batch_size, distributed=self.args.n_gpu > 1
+                )
+
+            return (
+                RandomSampler(self.train_dataset)
+                if self.args.local_rank == -1
+                else DistributedSampler(self.train_dataset)
+            )
+
+    def compute_loss(self, model, inputs):
+        labels = inputs.pop("labels")
+        outputs = model(**inputs, use_cache=False)
+        logits = outputs[0]
+        return self._compute_loss(logits, labels, ignore_index=model.config.pad_token_id)
+
+    def _compute_loss(self, logits, labels, ignore_index):
+        if self.args.label_smoothing == 0:
+            # Same behavior as modeling_bart.py
+            loss_fct = torch.nn.CrossEntropyLoss(ignore_index=ignore_index)
+            assert logits.shape[-1] == self.model.config.vocab_size
+            loss = loss_fct(logits.view(-1, logits.shape[-1]), labels.view(-1))
+        else:
+            lprobs = torch.nn.functional.log_softmax(logits, dim=-1)
+            loss, nll_loss = label_smoothed_nll_loss(
+                lprobs, labels, self.args.label_smoothing, ignore_index=ignore_index
+            )
+        return loss
+
+    def prediction_step(
+        self, model: nn.Module, inputs: Dict[str, Union[torch.Tensor, Any]], prediction_loss_only: bool
+    ) -> Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
+        """
+        Perform an evaluation step on :obj:`model` using obj:`inputs`.
+
+        Subclass and override to inject custom behavior.
+
+        Args:
+            model (:obj:`nn.Module`):
+                The model to evaluate.
+            inputs (:obj:`Dict[str, Union[torch.Tensor, Any]]`):
+                The inputs and targets of the model.
+
+                The dictionary will be unpacked before being fed to the model. Most models expect the targets under the
+                argument :obj:`labels`. Check your model's documentation for all accepted arguments.
+            prediction_loss_only (:obj:`bool`):
+                Whether or not to return the loss only.
+
+        Return:
+            Tuple[Optional[float], Optional[torch.Tensor], Optional[torch.Tensor]]:
+            A tuple with the loss, logits and labels (each being optional).
+        """
+        inputs = self._prepare_inputs(inputs)
+
+        max_length = (
+            model.config.max_generate_length
+            if hasattr(model.config, "max_generate_length")
+            else model.config.max_position_embeddings
+        )
+
+        with torch.no_grad():
+            if self.args.predict_with_generate and not self.args.prediction_loss_only:
+                generated_tokens = model.generate(
+                    inputs["input_ids"],
+                    attention_mask=inputs["attention_mask"],
+                    use_cache=True,
+                    num_beams=model.config.num_beams,
+                    max_length=max_length,
+                )
+                # in case the batch is shorter than max length, the output should be padded
+                generated_tokens = self._pad_tensors_to_max_len(
+                    generated_tokens, max_length, model.config.pad_token_id
+                )
+
+            labels_out = inputs.get("labels")
+            outputs = model(**inputs)
+            logits = outputs[1]
+            loss = self._compute_loss(logits, labels_out, model.config.pad_token_id)
+            loss = loss.mean().item()
+            if self.args.prediction_loss_only:
+                logits = None
+            else:
+                logits = generated_tokens if self.args.predict_with_generate else logits
+
+        if self.args.prediction_loss_only:
+            return (loss, None, None)
+
+        labels_out = labels_out.detach()
+        labels = self._pad_tensors_to_max_len(labels_out, max_length, model.config.pad_token_id)
+        return (loss, logits.detach(), labels)
+
+    def _pad_tensors_to_max_len(self, tensor, max_length, pad_token_id):
+        padded_tensor = pad_token_id * torch.ones(
+            (tensor.shape[0], max_length), dtype=tensor.dtype, device=tensor.device
+        )
+        padded_tensor[:, : tensor.shape[-1]] = tensor
+        return padded_tensor
@@ -1,3 +1,5 @@
+#!/usr/bin/env python
+
 import argparse
 import os
 import sys
@@ -0,0 +1,80 @@
+from collections import defaultdict
+from pathlib import Path
+
+import pandas as pd
+
+from rouge_cli import calculate_rouge_path
+from utils import calculate_rouge
+
+
+PRED = [
+    'Prosecutor: "No videos were used in the crash investigation" German papers say they saw a cell phone video of the final seconds on board Flight 9525. The Germanwings co-pilot says he had a "previous episode of severe depression" German airline confirms it knew of Andreas Lubitz\'s depression years before he took control.',
+    "The Palestinian Authority officially becomes the 123rd member of the International Criminal Court. The formal accession was marked with a ceremony at The Hague, in the Netherlands. The Palestinians signed the ICC's founding Rome Statute in January. Israel and the United States opposed the Palestinians' efforts to join the body.",
+    "Amnesty International releases its annual report on the death penalty. The report catalogs the use of state-sanctioned killing as a punitive measure across the globe. At least 607 people were executed around the world in 2014, compared to 778 in 2013. The U.S. remains one of the worst offenders for imposing capital punishment.",
+]
+
+TGT = [
+    'Marseille prosecutor says "so far no videos were used in the crash investigation" despite media reports . Journalists at Bild and Paris Match are "very confident" the video clip is real, an editor says . Andreas Lubitz had informed his Lufthansa training school of an episode of severe depression, airline says .',
+    "Membership gives the ICC jurisdiction over alleged crimes committed in Palestinian territories since last June . Israel and the United States opposed the move, which could open the door to war crimes investigations against Israelis .",
+    "Amnesty's annual death penalty report catalogs encouraging signs, but setbacks in numbers of those sentenced to death . Organization claims that governments around the world are using the threat of terrorism to advance executions . The number of executions worldwide has gone down by almost 22% compared with 2013, but death sentences up by 28% .",
+]
+
+
+def test_disaggregated_scores_are_determinstic():
+    no_aggregation = calculate_rouge(PRED, TGT, bootstrap_aggregation=False, rouge_keys=["rouge2", "rougeL"])
+    assert isinstance(no_aggregation, defaultdict)
+    no_aggregation_just_r2 = calculate_rouge(PRED, TGT, bootstrap_aggregation=False, rouge_keys=["rouge2"])
+    assert (
+        pd.DataFrame(no_aggregation["rouge2"]).fmeasure.mean()
+        == pd.DataFrame(no_aggregation_just_r2["rouge2"]).fmeasure.mean()
+    )
+
+
+def test_newline_cnn_improvement():
+    k = "rougeLsum"
+    score = calculate_rouge(PRED, TGT, newline_sep=True, rouge_keys=[k])[k]
+    score_no_sep = calculate_rouge(PRED, TGT, newline_sep=False, rouge_keys=[k])[k]
+    assert score > score_no_sep
+
+
+def test_newline_irrelevant_for_other_metrics():
+    k = ["rouge1", "rouge2", "rougeL"]
+    score_sep = calculate_rouge(PRED, TGT, newline_sep=True, rouge_keys=k)
+    score_no_sep = calculate_rouge(PRED, TGT, newline_sep=False, rouge_keys=k)
+    assert score_sep == score_no_sep
+
+
+def test_single_sent_scores_dont_depend_on_newline_sep():
+    pred = [
+        "Her older sister, Margot Frank, died in 1945, a month earlier than previously thought.",
+        'Marseille prosecutor says "so far no videos were used in the crash investigation" despite media reports .',
+    ]
+    tgt = [
+        "Margot Frank, died in 1945, a month earlier than previously thought.",
+        'Prosecutor: "No videos were used in the crash investigation" German papers say they saw a cell phone video of the final seconds on board Flight 9525.',
+    ]
+    assert calculate_rouge(pred, tgt, newline_sep=True) == calculate_rouge(pred, tgt, newline_sep=False)
+
+
+def test_pegasus_newline():
+
+    pred = [
+        """" "a person who has such a video needs to immediately give it to the investigators," prosecutor says .<n> "it is a very disturbing scene," editor-in-chief of bild online tells "erin burnett: outfront" """
+    ]
+    tgt = [
+        """ Marseille prosecutor says "so far no videos were used in the crash investigation" despite media reports . Journalists at Bild and Paris Match are "very confident" the video clip is real, an editor says . Andreas Lubitz had informed his Lufthansa training school of an episode of severe depression, airline says ."""
+    ]
+
+    prev_score = calculate_rouge(pred, tgt, rouge_keys=["rougeLsum"], newline_sep=False)["rougeLsum"]
+    new_score = calculate_rouge(pred, tgt, rouge_keys=["rougeLsum"])["rougeLsum"]
+    assert new_score > prev_score
+
+
+def test_rouge_cli():
+    data_dir = Path("examples/seq2seq/test_data/wmt_en_ro")
+    metrics = calculate_rouge_path(data_dir.joinpath("test.source"), data_dir.joinpath("test.target"))
+    assert isinstance(metrics, dict)
+    metrics_default_dict = calculate_rouge_path(
+        data_dir.joinpath("test.source"), data_dir.joinpath("test.target"), bootstrap_aggregation=False
+    )
+    assert isinstance(metrics_default_dict, defaultdict)
@@ -6,14 +6,13 @@ import numpy as np
 import pytest
 from torch.utils.data import DataLoader

+from pack_dataset import pack_data_dir
+from save_len_file import save_len_file
+from test_seq2seq_examples import ARTICLES, BART_TINY, MARIAN_TINY, MBART_TINY, SUMMARIES, T5_TINY, make_test_data_dir
 from transformers import AutoTokenizer
 from transformers.modeling_bart import shift_tokens_right
 from transformers.testing_utils import slow
-
-from .pack_dataset import pack_data_dir
-from .save_len_file import save_len_file
-from .test_seq2seq_examples import ARTICLES, BART_TINY, MARIAN_TINY, MBART_TINY, SUMMARIES, T5_TINY, make_test_data_dir
-from .utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeq2SeqDataset, Seq2SeqDataset
+from utils import FAIRSEQ_AVAILABLE, DistributedSortishSampler, LegacySeq2SeqDataset, Seq2SeqDataset


 BERT_BASE_CASED = "bert-base-cased"
@@ -0,0 +1,96 @@
+import os
+import sys
+import tempfile
+from unittest.mock import patch
+
+from transformers import BartForConditionalGeneration, MarianMTModel
+from transformers.testing_utils import slow
+
+from .finetune_trainer import main
+from .test_seq2seq_examples import 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
+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)
+
+
+@slow
+def test_finetune_trainer():
+    data_dir = "examples/seq2seq/test_data/wmt_en_ro"
+    output_dir = tempfile.mkdtemp(prefix="marian_output")
+    max_len = "128"
+    num_train_epochs = 4
+    eval_steps = 2
+    argv = [
+        "--model_name_or_path",
+        MARIAN_MODEL,
+        "--data_dir",
+        data_dir,
+        "--output_dir",
+        output_dir,
+        "--overwrite_output_dir",
+        "--n_train",
+        "8",
+        "--n_val",
+        "8",
+        "--max_source_length",
+        max_len,
+        "--max_target_length",
+        max_len,
+        "--val_max_target_length",
+        max_len,
+        "--do_train",
+        "--do_eval",
+        "--do_predict",
+        "--num_train_epochs",
+        str(num_train_epochs),
+        "--per_device_train_batch_size",
+        "4",
+        "--per_device_eval_batch_size",
+        "4",
+        "--learning_rate",
+        "3e-4",
+        "--warmup_steps",
+        "8",
+        "--evaluate_during_training",
+        "--predict_with_generate",
+        "--logging_steps",
+        0,
+        "--save_steps",
+        str(eval_steps),
+        "--eval_steps",
+        str(eval_steps),
+        "--sortish_sampler",
+        "--label_smoothing",
+        "0.1",
+        "--task",
+        "translation",
+    ]
+
+    testargs = ["finetune_trainer.py"] + argv
+    with patch.object(sys, "argv", testargs):
+        main()
+
+    # Check metrics
+    logs = load_json(os.path.join(output_dir, "log_history.json"))
+    eval_metrics = [log for log in logs if "eval_loss" in log.keys()]
+    first_step_stats = eval_metrics[0]
+    last_step_stats = eval_metrics[-1]
+
+    assert first_step_stats["eval_bleu"] < last_step_stats["eval_bleu"]  # model learned nothing
+    assert isinstance(last_step_stats["eval_bleu"], float)
+
+    # test if do_predict saves generations and metrics
+    contents = os.listdir(output_dir)
+    contents = {os.path.basename(p) for p in contents}
+    assert "test_generations.txt" in contents
+    assert "test_results.json" in contents
@@ -14,19 +14,13 @@
 # limitations under the License.

 import io
-import unittest
-
-
-try:
-    from .utils import calculate_bleu
-except ImportError:
-    from utils import calculate_bleu
-
 import json
+import unittest

 from parameterized import parameterized
 from transformers import FSMTForConditionalGeneration, FSMTTokenizer
 from transformers.testing_utils import get_tests_dir, require_torch, slow, torch_device
+from utils import calculate_bleu


 filename = get_tests_dir() + "/test_data/fsmt/fsmt_val_data.json"
@@ -0,0 +1,39 @@
+import tempfile
+import unittest
+
+from make_student import create_student_by_copying_alternating_layers
+from transformers import AutoConfig
+from transformers.file_utils import cached_property
+from transformers.testing_utils import require_torch
+
+
+TINY_BART = "sshleifer/bart-tiny-random"
+TINY_T5 = "patrickvonplaten/t5-tiny-random"
+
+
+@require_torch
+class MakeStudentTester(unittest.TestCase):
+    @cached_property
+    def teacher_config(self):
+        return AutoConfig.from_pretrained(TINY_BART)
+
+    def test_valid_t5(self):
+        student, *_ = create_student_by_copying_alternating_layers(TINY_T5, tempfile.mkdtemp(), e=1, d=1)
+        self.assertEqual(student.config.num_hidden_layers, 1)
+
+    def test_asymmetric_t5(self):
+        student, *_ = create_student_by_copying_alternating_layers(TINY_T5, tempfile.mkdtemp(), e=1, d=None)
+
+    def test_same_decoder_small_encoder(self):
+        student, *_ = create_student_by_copying_alternating_layers(TINY_BART, tempfile.mkdtemp(), e=1, d=None)
+        self.assertEqual(student.config.encoder_layers, 1)
+        self.assertEqual(student.config.decoder_layers, self.teacher_config.encoder_layers)
+
+    def test_small_enc_small_dec(self):
+        student, *_ = create_student_by_copying_alternating_layers(TINY_BART, tempfile.mkdtemp(), e=1, d=1)
+        self.assertEqual(student.config.encoder_layers, 1)
+        self.assertEqual(student.config.decoder_layers, 1)
+
+    def test_raises_assert(self):
+        with self.assertRaises(AssertionError):
+            create_student_by_copying_alternating_layers(TINY_BART, tempfile.mkdtemp(), e=None, d=None)
@@ -20,7 +20,7 @@ from run_eval_search import run_search
 from transformers import AutoConfig, AutoModelForSeq2SeqLM
 from transformers.hf_api import HfApi
 from transformers.testing_utils import CaptureStderr, CaptureStdout, require_multigpu, require_torch_and_cuda, slow
-from utils import label_smoothed_nll_loss, lmap, load_json
+from utils import ROUGE_KEYS, label_smoothed_nll_loss, lmap, load_json


 logging.basicConfig(level=logging.DEBUG)
@@ -145,6 +145,7 @@ class TestSummarizationDistiller(unittest.TestCase):
        assert not failures, f"The following models could not be loaded through AutoConfig: {failures}"

    @require_multigpu
+    @unittest.skip("Broken at the moment")
    def test_multigpu(self):
        updates = dict(
            no_teacher=True,
@@ -364,7 +365,7 @@ def test_run_eval_search(model):
        if "translation" in task:
            expected_strings.append("bleu")
        else:
-            expected_strings.extend(["rouge1", "rouge2", "rougeL"])
+            expected_strings.extend(ROUGE_KEYS)
        for w in expected_strings:
            assert w in cs.out
        for w in un_expected_strings:
@@ -1,21 +1,20 @@
 #!/usr/bin/env bash
 export PYTHONPATH="../":"${PYTHONPATH}"
-export BS=16
-export GAS=2
 python distillation.py \
+  --teacher facebook/bart-large-xsum --data_dir xsum \
+  --student_decoder_layers 6 --student_encoder_layers 12 \
+  --freeze_encoder --freeze_embeds \
  --learning_rate=3e-4 \
  --do_train \
  --do_predict \
-  --fp16 \
-  --val_check_interval 0.1 --n_val 1000 \
-  --teacher facebook/bart-large-xsum --data_dir $XSUM_DIR \
+  --fp16 --fp16_opt_level=O1 \
+  --val_check_interval 0.1 --n_val 1000 --eval_beams 2 --length_penalty=0.5 \
  --max_target_length=60 --val_max_target_length=60 --test_max_target_length=100 \
-  --student_decoder_layers 6 --student_encoder_layers 12 \
-  --freeze_encoder --freeze_embeds \
  --model_name_or_path IGNORED \
-  --alpha_hid=3. --length_penalty=0.5 \
-  --train_batch_size=$BS --eval_batch_size=$BS --gradient_accumulation_steps=$GAS --num_train_epochs=6 \
-  --tokenizer_name facebook/bart-large \
+  --alpha_hid=3. \
+  --train_batch_size=16 --eval_batch_size=16 --gradient_accumulation_steps=2 \
+  --sortish_sampler \
+  --num_train_epochs=6 \
  --warmup_steps 500 \
  --output_dir distilbart_xsum_12_6 \
  "$@"
@@ -18,6 +18,7 @@ from sacrebleu import corpus_bleu
 from torch import nn
 from torch.utils.data import Dataset, Sampler

+from sentence_splitter import add_newline_to_end_of_each_sentence
 from transformers import BartTokenizer
 from transformers.file_utils import cached_property

@@ -378,19 +379,63 @@ def get_git_info():
    return repo_infos


-ROUGE_KEYS = ["rouge1", "rouge2", "rougeL"]
+ROUGE_KEYS = ["rouge1", "rouge2", "rougeL", "rougeLsum"]


-def calculate_rouge(output_lns: List[str], reference_lns: List[str], use_stemmer=True) -> Dict:
-    scorer = rouge_scorer.RougeScorer(ROUGE_KEYS, use_stemmer=use_stemmer)
+def extract_rouge_mid_statistics(dct):
+    new_dict = {}
+    for k1, v1 in dct.items():
+        mid = v1.mid
+        new_dict[k1] = {stat: round(getattr(mid, stat), 4) for stat in ["precision", "recall", "fmeasure"]}
+    return new_dict
+
+
+def calculate_rouge(
+    pred_lns: List[str],
+    tgt_lns: List[str],
+    use_stemmer=True,
+    rouge_keys=ROUGE_KEYS,
+    return_precision_and_recall=False,
+    bootstrap_aggregation=True,
+    newline_sep=True,
+) -> Dict:
+    """Calculate rouge using rouge_scorer package.
+
+    Args:
+        pred_lns: list of summaries generated by model
+        tgt_lns: list of groundtruth summaries (e.g. contents of val.target)
+        use_stemmer:  Bool indicating whether Porter stemmer should be used to
+        strip word suffixes to improve matching.
+        rouge_keys:  which metrics to compute, defaults to rouge1, rouge2, rougeL, rougeLsum
+        return_precision_and_recall: (False) whether to also return precision and recall.
+        bootstrap_aggregation: whether to do the typical bootstrap resampling of scores. Defaults to True, if False
+            this function returns a collections.defaultdict[metric: list of values for each observation for each subscore]``
+        newline_sep:(default=True) whether to add newline between sentences. This is essential for calculation rougeL
+        on multi sentence summaries (CNN/DM dataset).
+
+    Returns:
+         Dict[score: value] if aggregate else defaultdict(list) keyed by rouge_keys
+
+    """
+    scorer = rouge_scorer.RougeScorer(rouge_keys, use_stemmer=use_stemmer)
    aggregator = scoring.BootstrapAggregator()
-
-    for reference_ln, output_ln in zip(reference_lns, output_lns):
-        scores = scorer.score(reference_ln, output_ln)
+    for pred, tgt in zip(tgt_lns, pred_lns):
+        # rougeLsum expects "\n" separated sentences within a summary
+        if newline_sep:
+            pred = add_newline_to_end_of_each_sentence(pred)
+            tgt = add_newline_to_end_of_each_sentence(tgt)
+        scores = scorer.score(pred, tgt)
        aggregator.add_scores(scores)

-    result = aggregator.aggregate()
-    return {k: round(v.mid.fmeasure * 100, 4) for k, v in result.items()}
+    if bootstrap_aggregation:
+        result = aggregator.aggregate()
+        if return_precision_and_recall:
+            return extract_rouge_mid_statistics(result)  # here we return dict
+        else:
+            return {k: round(v.mid.fmeasure * 100, 4) for k, v in result.items()}
+
+    else:
+        return aggregator._scores  # here we return defaultdict(list)


 # Utilities for freezing parameters and checking whether they are frozen
@@ -423,9 +468,6 @@ def assert_not_all_frozen(model):
    assert any(model_grads), f"none of {npars} weights require grad"


-# CLI Parsing utils
-
-
 def parse_numeric_n_bool_cl_kwargs(unparsed_args: List[str]) -> Dict[str, Union[int, float, bool]]:
    """
    Parse an argv list of unspecified command line args to a dict.
@@ -456,3 +498,9 @@ def write_txt_file(ordered_tgt, path):
    for ln in ordered_tgt:
        f.write(ln + "\n")
        f.flush()
+
+
+def chunks(lst, n):
+    """Yield successive n-sized chunks from lst."""
+    for i in range(0, len(lst), n):
+        yield lst[i : i + n]
@@ -0,0 +1,72 @@
+"""
+A simple launcher script for TPU training
+
+Inspired by https://github.com/pytorch/pytorch/blob/master/torch/distributed/launch.py
+
+::
+    >>> python xla_spawn.py --num_cores=NUM_CORES_YOU_HAVE
+               YOUR_TRAINING_SCRIPT.py (--arg1 --arg2 --arg3 and all other
+               arguments of your training script)
+
+"""
+
+
+import importlib
+import sys
+from argparse import REMAINDER, ArgumentParser
+from pathlib import Path
+
+import torch_xla.distributed.xla_multiprocessing as xmp
+
+
+def parse_args():
+    """
+    Helper function parsing the command line options
+    @retval ArgumentParser
+    """
+    parser = ArgumentParser(
+        description=(
+            "PyTorch TPU distributed training launch "
+            "helper utility that will spawn up "
+            "multiple distributed processes"
+        )
+    )
+
+    # Optional arguments for the launch helper
+    parser.add_argument("--num_cores", type=int, default=1, help="Number of TPU cores to use (1 or 8).")
+
+    # positional
+    parser.add_argument(
+        "training_script",
+        type=str,
+        help=(
+            "The full path to the single TPU training "
+            "program/script to be launched in parallel, "
+            "followed by all the arguments for the "
+            "training script"
+        ),
+    )
+
+    # rest from the training program
+    parser.add_argument("training_script_args", nargs=REMAINDER)
+
+    return parser.parse_args()
+
+
+def main():
+    args = parse_args()
+
+    # Import training_script as a module.
+    script_fpath = Path(args.training_script)
+    sys.path.append(str(script_fpath.parent.resolve()))
+    mod_name = script_fpath.stem
+    mod = importlib.import_module(mod_name)
+
+    # Patch sys.argv
+    sys.argv = [args.training_script] + args.training_script_args + ["--tpu_num_cores", str(args.num_cores)]
+
+    xmp.spawn(mod._mp_fn, args=(), nprocs=args.num_cores)
+
+
+if __name__ == "__main__":
+    main()
@@ -0,0 +1,55 @@
+---
+language: fr
+widget:
+- text: "Je m'appelle Hicham et je vis a Fès"
+---
+
+# MagBERT-NER: a state-of-the-art NER model for Moroccan French language (Maghreb)
+
+## Introduction
+
+[MagBERT-NER] is a state-of-the-art NER model for Moroccan French language (Maghreb). The MagBERT-NER model was fine-tuned for NER Task based the language model for French Camembert (based on the RoBERTa architecture).
+
+For further information or requests, please go to [Typica.AI Website](https://typicasoft.io/)
+
+## How to use MagBERT-NER with HuggingFace
+
+##### Load MagBERT-NER and its sub-word tokenizer :
+```python
+from transformers import AutoTokenizer, AutoModelForTokenClassification
+
+tokenizer = AutoTokenizer.from_pretrained("TypicaAI/magbert-ner")
+model = AutoModelForTokenClassification.from_pretrained("TypicaAI/magbert-ner")
+
+
+##### Process text sample (from wikipedia about the current Prime Minister of Morocco) Using NER pipeline  
+
+from transformers import pipeline
+
+nlp = pipeline('ner', model=model, tokenizer=tokenizer, grouped_entities=True)
+nlp("Saad Dine El Otmani, né le 16 janvier 1956 à Inezgane, est un homme d'État marocain, chef du gouvernement du Maroc depuis le 5 avril 2017")
+
+
+#[{'entity_group': 'I-PERSON',
+#  'score': 0.8941445276141167,
+#  'word': 'Saad Dine El Otmani'},
+# {'entity_group': 'B-DATE',
+#  'score': 0.5967703461647034,
+#  'word': '16 janvier 1956'},
+# {'entity_group': 'B-GPE', 'score': 0.7160899192094803, 'word': 'Inezgane'},
+# {'entity_group': 'B-NORP', 'score': 0.7971733212471008, 'word': 'marocain'},
+# {'entity_group': 'B-GPE', 'score': 0.8921478390693665, 'word': 'Maroc'},
+# {'entity_group': 'B-DATE',
+#  'score': 0.5760444005330404,
+#  'word': '5 avril 2017'}]
+
+```
+
+```
+
+
+## Authors 
+
+MagBert-NER was trained and evaluated by Hicham Assoudi, Ph.D.
+
+
@@ -1,22 +1,50 @@
 ## RAG

-This is a "base" version of the RAG-Sequence Model of the the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://arxiv.org/pdf/2005.11401.pdf) 
+This is a non-finetuned version of the RAG-Sequence model of the the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://arxiv.org/pdf/2005.11401.pdf) 
 by Patrick Lewis, Ethan Perez, Aleksandara Piktus et al.

+Rag consits of a *question encoder*, *retriever* and a *generator*. The retriever should be a `RagRetriever` instance. The *question encoder* can be any model that can be loaded with `AutoModel` and the *generator* can be any model that can be loaded with `AutoModelForSeq2SeqLM`. 
+
+This model is a non-finetuned RAG-Sequence model and was created as follows:
+
+```python
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration, AutoTokenizer
+
+model = RagSequenceForGeneration.from_pretrained_question_encoder_generator("facebook/dpr-question_encoder-single-nq-base", "facebook/bart-large")
+
+question_encoder_tokenizer = AutoTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+generator_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
+
+tokenizer = RagTokenizer(question_encoder_tokenizer, generator_tokenizer)
+model.config.use_dummy_dataset = True
+model.config.index_name = "exact"
+retriever = RagRetriever(model.config, question_encoder_tokenizer, generator_tokenizer)
+
+model.save_pretrained("./")
+tokenizer.save_pretrained("./")
+retriever.save_pretrained("./")
+```
+
+Note that the model is *uncased* so that all capital input letters are converted to lower-case.
+
 ## Usage:

+*Note*: the model uses the *dummy* retriever as a default. Better results are obtained by using the full retriever, 
+by setting `config.index_name="legacy"` and `config.use_dummy_dataset=False`.
+The model can be fine-tuned as follows:
+
 ```python
-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration

 tokenizer = RagTokenizer.from_pretrained("facebook/rag-sequence-base")
-retriever = RagRetriever.from_pretrained("facebook/rag-sequence-base", index_name="exact", use_dummy_dataset=True)
-model = RagSequenceForGeneration.from_pretrained("facebook/rag-sequence-base", retriever=retriever)
+retriever = RagRetriever.from_pretrained("facebook/rag-sequence-base")
+model = RagTokenForGeneration.from_pretrained("facebook/rag-sequence-base", retriever=retriever)

-input_ids = tokenizer("What is the largest country in the world?", return_tensors="pt").input_ids
+input_dict = tokenizer.prepare_seq2seq_batch("who holds the record in 100m freestyle", "michael phelps", return_tensors="pt") 

-generated = model.generate(input_ids=input_ids)
-generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
+outputs = model(input_dict["input_ids"], labels=input_dict["labels"])

-# => should give ["Asia ended in 2010 when China overtook Japan to become the world's second largest economy."]
-# Interesting answer. Definitely on topic, but might factual probably not fully correct.
+loss = outputs.loss
+
+# train on loss
 ```
@@ -1,25 +1,36 @@
+---
+language: en
+license: apache-2.0
+datasets:
+- wiki_dpr
+---
 ## RAG

 This is the RAG-Sequence Model of the the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://arxiv.org/pdf/2005.11401.pdf) 
 by Patrick Lewis, Ethan Perez, Aleksandara Piktus et al.

+The model is a *uncased* model, which means that capital letters are simply converted to lower-case letters.
+
+The model consits of a *question_encoder*, *retriever* and a *generator*. The retriever extracts relevant passages from the *wiki_dpr* `train` datasets, which is linked above.
+The question_encoder and retriever are based on `facebook/dpr-question_encoder-single-nq-base` and `facebook/bart-large`, which were jointly finetuned on 
+on the *wiki_dpr* QA dataset in an end-to-end fashion.
+
 ## Usage:

+**Note**: In the usage example below only the *dummy* retriever of *wiki_dpr* is used because the complete *lecagy* index requires over 75 GB of RAM.
+The model can generate answers to any factoid question as follows:
+
 ```python
+from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration 
+ 
+tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq") 
+retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True) 
+model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever) 
+ 
+input_dict = tokenizer.prepare_seq2seq_batch("how many countries are in europe", return_tensors="pt") 

-from transformers import RagTokenizer, RagRetriever, RagSequenceForGeneration
+generated = model.generate(input_ids=input_dict["input_ids"]) 
+print(tokenizer.batch_decode(generated, skip_special_tokens=True)[0]) 

-tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
-retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
-model = RagSequenceForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)
-
-input_dict = tokenizer.prepare_seq2seq_batch("How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="pt")
-outputs = model(input_ids=input_dict["input_ids"], labels=input_dict["labels"])
-
-# outputs.loss should give 76.2978
-
-generated = model.generate(input_ids=input_dict["input_ids"])
-generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
-
-# generated_string should give 270,000,000 -> not quite correct the answer, but it also only uses a dummy index
+# should give 54 => google says either 44 or 51
 ```
@@ -1,21 +1,50 @@
 ## RAG

-This is a "base" version of the RAG-Token Model of the the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://arxiv.org/pdf/2005.11401.pdf) 
+This is a non-finetuned version of the RAG-Token model of the the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://arxiv.org/pdf/2005.11401.pdf) 
 by Patrick Lewis, Ethan Perez, Aleksandara Piktus et al.

+Rag consits of a *question encoder*, *retriever* and a *generator*. The retriever should be a `RagRetriever` instance. The *question encoder* can be any model that can be loaded with `AutoModel` and the *generator* can be any model that can be loaded with `AutoModelForSeq2SeqLM`. 
+
+This model is a non-finetuned RAG-Token model and was created as follows:
+
+```python
+from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration, AutoTokenizer
+
+model = RagTokenForGeneration.from_pretrained_question_encoder_generator("facebook/dpr-question_encoder-single-nq-base", "facebook/bart-large")
+
+question_encoder_tokenizer = AutoTokenizer.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+generator_tokenizer = AutoTokenizer.from_pretrained("facebook/bart-large")
+
+tokenizer = RagTokenizer(question_encoder_tokenizer, generator_tokenizer)
+model.config.use_dummy_dataset = True
+model.config.index_name = "exact"
+retriever = RagRetriever(model.config, question_encoder_tokenizer, generator_tokenizer)
+
+model.save_pretrained("./")
+tokenizer.save_pretrained("./")
+retriever.save_pretrained("./")
+```
+
+Note that the model is *uncased* so that all capital input letters are converted to lower-case.
+
 ## Usage:

+*Note*: the model uses the *dummy* retriever as a default. Better results are obtained by using the full retriever, 
+by setting `config.index_name="legacy"` and `config.use_dummy_dataset=False`.
+The model can be fine-tuned as follows:
+
 ```python
 from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration

 tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-base")
-retriever = RagRetriever.from_pretrained("facebook/rag-token-base", index_name="exact", use_dummy_dataset=True)
+retriever = RagRetriever.from_pretrained("facebook/rag-token-base")
 model = RagTokenForGeneration.from_pretrained("facebook/rag-token-base", retriever=retriever)

-input_ids = tokenizer("What is the largest country in the world?", return_tensors="pt").input_ids
+input_dict = tokenizer.prepare_seq2seq_batch("who holds the record in 100m freestyle", "michael phelps", return_tensors="pt") 

-generated = model.generate(input_ids=input_ids)
-generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
+outputs = model(input_dict["input_ids"], labels=input_dict["labels"])

-# => should give [' russia']. Pretty good answer for just having just a dummy dataset.
+loss = outputs.loss
+
+# train on loss
 ```
@@ -1,25 +1,36 @@
+---
+language: en
+license: apache-2.0
+datasets:
+- wiki_dpr
+---
 ## RAG

 This is the RAG-Token Model of the the paper [Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks](https://arxiv.org/pdf/2005.11401.pdf) 
 by Patrick Lewis, Ethan Perez, Aleksandara Piktus et al.

+The model is a *uncased* model, which means that capital letters are simply converted to lower-case letters.
+
+The model consits of a *question_encoder*, *retriever* and a *generator*. The retriever extracts relevant passages from the *wiki_dpr* `train` datasets, which is linked above.
+The question_encoder and retriever are based on `facebook/dpr-question_encoder-single-nq-base` and `facebook/bart-large`, which were jointly finetuned on 
+on the *wiki_dpr* QA dataset in an end-to-end fashion.
+
 ## Usage:

-```python
+**Note**: In the usage example below only the *dummy* retriever of *wiki_dpr* is used because the complete *lecagy* index requires over 75 GB of RAM.
+The model can generate answers to any factoid question as follows:

+```python
 from transformers import RagTokenizer, RagRetriever, RagTokenForGeneration

 tokenizer = RagTokenizer.from_pretrained("facebook/rag-token-nq")
 retriever = RagRetriever.from_pretrained("facebook/rag-token-nq", index_name="exact", use_dummy_dataset=True)
 model = RagTokenForGeneration.from_pretrained("facebook/rag-token-nq", retriever=retriever)

-input_dict = tokenizer.prepare_seq2seq_batch("How many people live in Paris?", "In Paris, there are 10 million people.", return_tensors="pt")
-outputs = model(input_ids=input_dict["input_ids"], labels=input_dict["labels"])
+input_dict = tokenizer.prepare_seq2seq_batch("who holds the record in 100m freestyle", return_tensors="pt") 

-# outputs.loss should give 76.1230
+generated = model.generate(input_ids=input_dict["input_ids"]) 
+print(tokenizer.batch_decode(generated, skip_special_tokens=True)[0]) 

-generated = model.generate(input_ids=input_dict["input_ids"])
-generated_string = tokenizer.batch_decode(generated, skip_special_tokens=True)
-
-# generated_string should give 270,000 -> not quite correct the answer, but it also only uses a dummy index
+# should give michael phelps => sounds reasonable
 ```
@@ -59,9 +59,21 @@ predictions = torch.round((torch.sign(discriminator_outputs[0]) + 1) / 2)
 el rapido  zorro  marro    ##n   amar  sobre     el  perro   pere ##zoso    0.0    0.0    0.0    0.0    0.0    0.0    1.0    1.0    0.0    0.0    0.0    0.0    0.0[None, None, None, None, None, None, None, None, None, None, None, None, None
 '''
 ```
-
 As you can see there are **1s** in the places where the model detected a fake token. So, it works! 🎉

+
+### Some models fine-tuned on a downstream task 🛠️
+
+[Question Answering](https://huggingface.co/mrm8488/electricidad-base-finetuned-squadv1-es)
+
+[POS](https://huggingface.co/mrm8488/electricidad-base-finetuned-pos)
+
+[NER](https://huggingface.co/mrm8488/electricidad-base-finetuned-ner)
+
+[Paraphrase Identification](https://huggingface.co/mrm8488/RuPERTa-base-finetuned-pawsx-es)
+
+
+
 ## Acknowledgments

 I thank [🤗/transformers team](https://github.com/huggingface/transformers) for allowing me to train the model (specially to [Julien Chaumond](https://twitter.com/julien_c)).
@@ -12,15 +12,22 @@ widget:
  
 ## Model description  
  
-This model performs sentiment analysis on Italian political twitter sentences. It was trained starting from an instance of "bert-base-italian-uncased-xxl" and fine-tuned on an Italian dataset of tweets.
+This model performs sentiment analysis on Italian political twitter sentences. It was trained starting from an instance of "bert-base-italian-uncased-xxl" and fine-tuned on an Italian dataset of tweets. You can try it out at https://www.unideeplearning.com/twitter_sa/ (in italian!)
  
 #### Hands-on  
  
 ```python
 import torch
 from torch import nn 
+from transformers import AutoTokenizer, AutoModelForSequenceClassification

-text = "Giueseppe Rossi è un pessimo politico"
+tokenizer = AutoTokenizer.from_pretrained("unideeplearning/polibert_sa")
+model = AutoModelForSequenceClassification.from_pretrained("unideeplearning/polibert_sa")
+			
+
+
+
+text = "Giuseppe Rossi è un pessimo politico"
 input_ids = tokenizer.encode(text, add_special_tokens=True, return_tensors= 'pt')

 logits, = model(input_ids)
@@ -41,4 +48,6 @@ print(prob.argmax().tolist())
 ## Acknowledgments

 Thanks to the support from: 
-the [Hugging Face](https://huggingface.co/), Unione Professionisti (https://www.unioneprofessionisti.com/)
+the [Hugging Face](https://huggingface.co/), https://www.unioneprofessionisti.com
+
+https://www.unideeplearning.com/
@@ -1,4 +1,4 @@
-#/usr/bin/env bash
+#!/usr/bin/env bash

 # this script acquires data and converts it to fsmt model
 # it covers:
@@ -1,4 +1,4 @@
-#/usr/bin/env bash
+#!/usr/bin/env bash

 # this script acquires data and converts it to fsmt model
 # it covers:
@@ -1,4 +1,4 @@
-#/usr/bin/env bash
+#!/usr/bin/env bash

 # this script acquires data and converts it to fsmt model
 # it covers:
@@ -1,4 +1,4 @@
-#/usr/bin/env bash
+#!/usr/bin/env bash

 # this script evals the following fsmt models
 # it covers:
@@ -1,4 +1,4 @@
-#/usr/bin/env bash
+#!/usr/bin/env bash

 # this script evals the following fsmt models
 # it covers:
@@ -1,4 +1,4 @@
-#/usr/bin/env bash
+#!/usr/bin/env bash

 # this script evals the following fsmt models
 # it covers:
@@ -0,0 +1,103 @@
+
+# this is the process of uploading the updated models to s3. As I can't upload them directly to the correct orgs, this script shows how this is done
+
+1. upload updated models to my account
+
+transformers-cli upload -y wmt19-ru-en
+transformers-cli upload -y wmt19-en-ru
+transformers-cli upload -y wmt19-de-en
+transformers-cli upload -y wmt19-en-de
+transformers-cli upload -y wmt19-de-en-6-6-base
+transformers-cli upload -y wmt19-de-en-6-6-big
+transformers-cli upload -y wmt16-en-de-dist-12-1
+transformers-cli upload -y wmt16-en-de-dist-6-1
+transformers-cli upload -y wmt16-en-de-12-1
+
+
+2. ask someone to move them to:
+
+* to facebook: "wmt19-ru-en", "wmt19-en-ru", "wmt19-en-de", "wmt19-de-en"
+* to allenai: "wmt16-en-de-dist-12-1", "wmt16-en-de-dist-6-1", "wmt16-en-de-12-1", "wmt19-de-en-6-6-base", "wmt19-de-en-6-6-big"
+
+export b="s3://models.huggingface.co/bert"
+stas_to_fb () {
+	src=$1
+	shift
+	aws s3 sync $b/stas/$src $b/facebook/$src $@
+}
+
+stas_to_allenai () {
+	src=$1
+	shift
+	aws s3 sync $b/stas/$src $b/allenai/$src $@
+}
+
+stas_to_fb wmt19-en-ru
+stas_to_fb wmt19-ru-en
+stas_to_fb wmt19-en-de
+stas_to_fb wmt19-de-en
+
+stas_to_allenai wmt16-en-de-dist-12-1
+stas_to_allenai wmt16-en-de-dist-6-1
+stas_to_allenai wmt16-en-de-6-1
+stas_to_allenai wmt16-en-de-12-1
+stas_to_allenai wmt19-de-en-6-6-base
+stas_to_allenai wmt19-de-en-6-6-big
+
+
+3. and then remove all these model files from my account
+
+transformers-cli s3 rm wmt16-en-de-12-1/config.json
+transformers-cli s3 rm wmt16-en-de-12-1/merges.txt
+transformers-cli s3 rm wmt16-en-de-12-1/pytorch_model.bin
+transformers-cli s3 rm wmt16-en-de-12-1/tokenizer_config.json
+transformers-cli s3 rm wmt16-en-de-12-1/vocab-src.json
+transformers-cli s3 rm wmt16-en-de-12-1/vocab-tgt.json
+transformers-cli s3 rm wmt16-en-de-dist-12-1/config.json
+transformers-cli s3 rm wmt16-en-de-dist-12-1/merges.txt
+transformers-cli s3 rm wmt16-en-de-dist-12-1/pytorch_model.bin
+transformers-cli s3 rm wmt16-en-de-dist-12-1/tokenizer_config.json
+transformers-cli s3 rm wmt16-en-de-dist-12-1/vocab-src.json
+transformers-cli s3 rm wmt16-en-de-dist-12-1/vocab-tgt.json
+transformers-cli s3 rm wmt16-en-de-dist-6-1/config.json
+transformers-cli s3 rm wmt16-en-de-dist-6-1/merges.txt
+transformers-cli s3 rm wmt16-en-de-dist-6-1/pytorch_model.bin
+transformers-cli s3 rm wmt16-en-de-dist-6-1/tokenizer_config.json
+transformers-cli s3 rm wmt16-en-de-dist-6-1/vocab-src.json
+transformers-cli s3 rm wmt16-en-de-dist-6-1/vocab-tgt.json
+transformers-cli s3 rm wmt19-de-en-6-6-base/config.json
+transformers-cli s3 rm wmt19-de-en-6-6-base/merges.txt
+transformers-cli s3 rm wmt19-de-en-6-6-base/pytorch_model.bin
+transformers-cli s3 rm wmt19-de-en-6-6-base/tokenizer_config.json
+transformers-cli s3 rm wmt19-de-en-6-6-base/vocab-src.json
+transformers-cli s3 rm wmt19-de-en-6-6-base/vocab-tgt.json
+transformers-cli s3 rm wmt19-de-en-6-6-big/config.json
+transformers-cli s3 rm wmt19-de-en-6-6-big/merges.txt
+transformers-cli s3 rm wmt19-de-en-6-6-big/pytorch_model.bin
+transformers-cli s3 rm wmt19-de-en-6-6-big/tokenizer_config.json
+transformers-cli s3 rm wmt19-de-en-6-6-big/vocab-src.json
+transformers-cli s3 rm wmt19-de-en-6-6-big/vocab-tgt.json
+transformers-cli s3 rm wmt19-de-en/config.json
+transformers-cli s3 rm wmt19-de-en/merges.txt
+transformers-cli s3 rm wmt19-de-en/pytorch_model.bin
+transformers-cli s3 rm wmt19-de-en/tokenizer_config.json
+transformers-cli s3 rm wmt19-de-en/vocab-src.json
+transformers-cli s3 rm wmt19-de-en/vocab-tgt.json
+transformers-cli s3 rm wmt19-en-de/config.json
+transformers-cli s3 rm wmt19-en-de/merges.txt
+transformers-cli s3 rm wmt19-en-de/pytorch_model.bin
+transformers-cli s3 rm wmt19-en-de/tokenizer_config.json
+transformers-cli s3 rm wmt19-en-de/vocab-src.json
+transformers-cli s3 rm wmt19-en-de/vocab-tgt.json
+transformers-cli s3 rm wmt19-en-ru/config.json
+transformers-cli s3 rm wmt19-en-ru/merges.txt
+transformers-cli s3 rm wmt19-en-ru/pytorch_model.bin
+transformers-cli s3 rm wmt19-en-ru/tokenizer_config.json
+transformers-cli s3 rm wmt19-en-ru/vocab-src.json
+transformers-cli s3 rm wmt19-en-ru/vocab-tgt.json
+transformers-cli s3 rm wmt19-ru-en/config.json
+transformers-cli s3 rm wmt19-ru-en/merges.txt
+transformers-cli s3 rm wmt19-ru-en/pytorch_model.bin
+transformers-cli s3 rm wmt19-ru-en/tokenizer_config.json
+transformers-cli s3 rm wmt19-ru-en/vocab-src.json
+transformers-cli s3 rm wmt19-ru-en/vocab-tgt.json
@@ -0,0 +1,6 @@
+#!/usr/bin/env bash
+
+# these scripts need to be run before any changes to FSMT-related code - it should cover all bases
+
+USE_CUDA=0 RUN_SLOW=1 pytest --disable-warnings tests/test_tokenization_fsmt.py tests/test_configuration_auto.py tests/test_modeling_fsmt.py examples/seq2seq/test_fsmt_bleu_score.py
+USE_CUDA=1 RUN_SLOW=1 pytest --disable-warnings tests/test_tokenization_fsmt.py tests/test_configuration_auto.py tests/test_modeling_fsmt.py examples/seq2seq/test_fsmt_bleu_score.py
@@ -5,7 +5,7 @@ To create the package for pypi.

 1. Change the version in __init__.py, setup.py as well as docs/source/conf.py.

-2. Unpin specific versions from setup.py (like isort).
+2. Unpin specific versions from setup.py that use a git install.

 2. Commit these changes with the message: "Release: VERSION"

@@ -98,7 +98,7 @@ extras["dev"] = extras["testing"] + extras["quality"] + extras["ja"] + ["scikit-

 setup(
    name="transformers",
-    version="3.2.0",
+    version="3.3.1",
    author="Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Sam Shleifer, Patrick von Platen, Sylvain Gugger, Google AI Language Team Authors, Open AI team Authors, Facebook AI Authors, Carnegie Mellon University Authors",
    author_email="thomas@huggingface.co",
    description="State-of-the-art Natural Language Processing for TensorFlow 2.0 and PyTorch",
@@ -2,7 +2,7 @@
 # There's no way to ignore "F401 '...' imported but unused" warnings in this
 # module, but to preserve other warnings. So, don't check this module at all.

-__version__ = "3.2.0"
+__version__ = "3.3.1"

 # Work around to update TensorFlow's absl.logging threshold which alters the
 # default Python logging output behavior when present.
@@ -103,6 +103,8 @@ class GPT2Config(PretrainedConfig):
            :class:`~transformers.GPT2DoubleHeadsModel` and :class:`~transformers.TFGPT2DoubleHeadsModel`.

            The dropout ratio to be used after the projection and activation.
+        gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            Whether or not to use gradient checkpointing to save memory at the expense of slower backward pass.

    Example::

@@ -142,6 +144,7 @@ class GPT2Config(PretrainedConfig):
        summary_first_dropout=0.1,
        bos_token_id=50256,
        eos_token_id=50256,
+        gradient_checkpointing=False,
        **kwargs
    ):
        super().__init__(bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)
@@ -164,6 +167,7 @@ class GPT2Config(PretrainedConfig):
        self.summary_activation = summary_activation
        self.summary_first_dropout = summary_first_dropout
        self.summary_proj_to_labels = summary_proj_to_labels
+        self.gradient_checkpointing = gradient_checkpointing

        self.bos_token_id = bos_token_id
        self.eos_token_id = eos_token_id
@@ -67,6 +67,5 @@ class LongformerConfig(RobertaConfig):
    model_type = "longformer"

    def __init__(self, attention_window: Union[List[int], int] = 512, sep_token_id: int = 2, **kwargs):
-        super().__init__(**kwargs)
+        super().__init__(sep_token_id=sep_token_id, **kwargs)
        self.attention_window = attention_window
-        self.sep_token_id = sep_token_id
@@ -57,6 +57,8 @@ class T5Config(PretrainedConfig):
            Size of the intermediate feed forward layer in each :obj:`T5Block`.
        num_layers (:obj:`int`, `optional`, defaults to 6):
            Number of hidden layers in the Transformer encoder.
+        num_decoder_layers (:obj:`int`, `optional`):
+            Number of hidden layers in the Transformer decoder. Will use the same value as :obj:`num_layers` if not set.
        num_heads (:obj:`int`, `optional`, defaults to 8):
            Number of attention heads for each attention layer in
            the Transformer encoder.
@@ -80,6 +82,7 @@ class T5Config(PretrainedConfig):
        d_kv=64,
        d_ff=2048,
        num_layers=6,
+        num_decoder_layers=None,
        num_heads=8,
        relative_attention_num_buckets=32,
        dropout_rate=0.1,
@@ -102,6 +105,9 @@ class T5Config(PretrainedConfig):
        self.d_kv = d_kv
        self.d_ff = d_ff
        self.num_layers = num_layers
+        self.num_decoder_layers = (
+            num_decoder_layers if num_decoder_layers is not None else self.num_layers
+        )  # default = symmetry
        self.num_heads = num_heads
        self.relative_attention_num_buckets = relative_attention_num_buckets
        self.dropout_rate = dropout_rate
@@ -130,6 +130,7 @@ class PretrainedConfig(object):
        - **eos_token_id** (:obj:`int`, `optional`)) -- The id of the `end-of-stream` token.
        - **decoder_start_token_id** (:obj:`int`, `optional`)) -- If an encoder-decoder model starts decoding with
          a different token than `bos`, the id of that token.
+        - **sep_token_id** (:obj:`int`, `optional`)) -- The id of the `separation` token.

    PyTorch specific parameters
        - **torchscript** (:obj:`bool`, `optional`, defaults to :obj:`False`) -- Whether or not the model should be
@@ -195,6 +196,8 @@ class PretrainedConfig(object):
        self.bos_token_id = kwargs.pop("bos_token_id", None)
        self.pad_token_id = kwargs.pop("pad_token_id", None)
        self.eos_token_id = kwargs.pop("eos_token_id", None)
+        self.sep_token_id = kwargs.pop("sep_token_id", None)
+
        self.decoder_start_token_id = kwargs.pop("decoder_start_token_id", None)

        # task specific arguments
@@ -68,8 +68,12 @@ except (ImportError, AssertionError):
 try:
    import datasets  # noqa: F401

-    _datasets_available = True
-    logger.debug(f"Succesfully imported datasets version {datasets.__version__}")
+    # Check we're not importing a "datasets" directory somewhere
+    _datasets_available = hasattr(datasets, "__version__") and hasattr(datasets, "load_dataset")
+    if _datasets_available:
+        logger.debug(f"Succesfully imported datasets version {datasets.__version__}")
+    else:
+        logger.debug("Imported a datasets object but this doesn't seem to be the 🤗 datasets library.")

 except ImportError:
    _datasets_available = False
@@ -587,14 +587,18 @@ class AlbertModel(AlbertPreTrainedModel):
    load_tf_weights = load_tf_weights_in_albert
    base_model_prefix = "albert"

-    def __init__(self, config):
+    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)

        self.config = config
        self.embeddings = AlbertEmbeddings(config)
        self.encoder = AlbertTransformer(config)
-        self.pooler = nn.Linear(config.hidden_size, config.hidden_size)
-        self.pooler_activation = nn.Tanh()
+        if add_pooling_layer:
+            self.pooler = nn.Linear(config.hidden_size, config.hidden_size)
+            self.pooler_activation = nn.Tanh()
+        else:
+            self.pooler = None
+            self.pooler_activation = None

        self.init_weights()

@@ -688,7 +692,7 @@ class AlbertModel(AlbertPreTrainedModel):

        sequence_output = encoder_outputs[0]

-        pooled_output = self.pooler_activation(self.pooler(sequence_output[:, 0]))
+        pooled_output = self.pooler_activation(self.pooler(sequence_output[:, 0])) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]
@@ -859,10 +863,13 @@ class AlbertSOPHead(nn.Module):
    ALBERT_START_DOCSTRING,
 )
 class AlbertForMaskedLM(AlbertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)

-        self.albert = AlbertModel(config)
+        self.albert = AlbertModel(config, add_pooling_layer=False)
        self.predictions = AlbertMLMHead(config)

        self.init_weights()
@@ -1034,11 +1041,14 @@ class AlbertForSequenceClassification(AlbertPreTrainedModel):
    ALBERT_START_DOCSTRING,
 )
 class AlbertForTokenClassification(AlbertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.albert = AlbertModel(config)
+        self.albert = AlbertModel(config, add_pooling_layer=False)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)

@@ -1118,11 +1128,14 @@ class AlbertForTokenClassification(AlbertPreTrainedModel):
    ALBERT_START_DOCSTRING,
 )
 class AlbertForQuestionAnswering(AlbertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.albert = AlbertModel(config)
+        self.albert = AlbertModel(config, add_pooling_layer=False)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
@@ -544,9 +544,8 @@ class AutoModel:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = AutoModel.from_config(config)
        """
-        for config_class, model_class in MODEL_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_MAPPING.keys():
+            return MODEL_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -585,9 +584,10 @@ class AutoModel:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_MAPPING.keys():
+            return MODEL_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -638,9 +638,8 @@ class AutoModelForPreTraining:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = AutoModelForPreTraining.from_config(config)
        """
-        for config_class, model_class in MODEL_FOR_PRETRAINING_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_FOR_PRETRAINING_MAPPING.keys():
+            return MODEL_FOR_PRETRAINING_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -679,9 +678,10 @@ class AutoModelForPreTraining:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_FOR_PRETRAINING_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_FOR_PRETRAINING_MAPPING.keys():
+            return MODEL_FOR_PRETRAINING_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -744,9 +744,8 @@ class AutoModelWithLMHead:
            "`AutoModelForSeq2SeqLM` for encoder-decoder models.",
            FutureWarning,
        )
-        for config_class, model_class in MODEL_WITH_LM_HEAD_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_WITH_LM_HEAD_MAPPING.keys():
+            return MODEL_WITH_LM_HEAD_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -791,9 +790,10 @@ class AutoModelWithLMHead:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_WITH_LM_HEAD_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_WITH_LM_HEAD_MAPPING.keys():
+            return MODEL_WITH_LM_HEAD_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -844,9 +844,8 @@ class AutoModelForCausalLM:
            >>> config = AutoConfig.from_pretrained('gpt2')
            >>> model = AutoModelForCausalLM.from_config(config)
        """
-        for config_class, model_class in MODEL_FOR_CAUSAL_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_FOR_CAUSAL_LM_MAPPING.keys():
+            return MODEL_FOR_CAUSAL_LM_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -885,9 +884,10 @@ class AutoModelForCausalLM:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_FOR_CAUSAL_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_FOR_CAUSAL_LM_MAPPING.keys():
+            return MODEL_FOR_CAUSAL_LM_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -938,9 +938,8 @@ class AutoModelForMaskedLM:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = AutoModelForMaskedLM.from_config(config)
        """
-        for config_class, model_class in MODEL_FOR_MASKED_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_FOR_MASKED_LM_MAPPING.keys():
+            return MODEL_FOR_MASKED_LM_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -979,9 +978,10 @@ class AutoModelForMaskedLM:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_FOR_MASKED_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_FOR_MASKED_LM_MAPPING.keys():
+            return MODEL_FOR_MASKED_LM_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1032,9 +1032,8 @@ class AutoModelForSeq2SeqLM:
            >>> config = AutoConfig.from_pretrained('t5')
            >>> model = AutoModelForSeq2SeqLM.from_config(config)
        """
-        for config_class, model_class in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.keys():
+            return MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1075,9 +1074,10 @@ class AutoModelForSeq2SeqLM:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.keys():
+            return MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1130,9 +1130,8 @@ class AutoModelForSequenceClassification:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = AutoModelForSequenceClassification.from_config(config)
        """
-        for config_class, model_class in MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.keys():
+            return MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1173,9 +1172,10 @@ class AutoModelForSequenceClassification:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.keys():
+            return MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1227,9 +1227,8 @@ class AutoModelForQuestionAnswering:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = AutoModelForQuestionAnswering.from_config(config)
        """
-        for config_class, model_class in MODEL_FOR_QUESTION_ANSWERING_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_FOR_QUESTION_ANSWERING_MAPPING.keys():
+            return MODEL_FOR_QUESTION_ANSWERING_MAPPING[type(config)](config)

        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
@@ -1271,9 +1270,10 @@ class AutoModelForQuestionAnswering:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_FOR_QUESTION_ANSWERING_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_FOR_QUESTION_ANSWERING_MAPPING.keys():
+            return MODEL_FOR_QUESTION_ANSWERING_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )

        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
@@ -1326,9 +1326,8 @@ class AutoModelForTokenClassification:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = AutoModelForTokenClassification.from_config(config)
        """
-        for config_class, model_class in MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.keys():
+            return MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING[type(config)](config)

        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
@@ -1370,9 +1369,10 @@ class AutoModelForTokenClassification:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.keys():
+            return MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )

        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
@@ -1426,9 +1426,8 @@ class AutoModelForMultipleChoice:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = AutoModelForMultipleChoice.from_config(config)
        """
-        for config_class, model_class in MODEL_FOR_MULTIPLE_CHOICE_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in MODEL_FOR_MULTIPLE_CHOICE_MAPPING.keys():
+            return MODEL_FOR_MULTIPLE_CHOICE_MAPPING[type(config)](config)

        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
@@ -1470,9 +1469,10 @@ class AutoModelForMultipleChoice:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in MODEL_FOR_MULTIPLE_CHOICE_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in MODEL_FOR_MULTIPLE_CHOICE_MAPPING.keys():
+            return MODEL_FOR_MULTIPLE_CHOICE_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )

        raise ValueError(
            "Unrecognized configuration class {} for this kind of AutoModel: {}.\n"
@@ -307,7 +307,7 @@ class BartEncoder(nn.Module):
        self.layers = nn.ModuleList([EncoderLayer(config) for _ in range(config.encoder_layers)])
        self.layernorm_embedding = LayerNorm(embed_dim) if config.normalize_embedding else nn.Identity()
        # mbart has one extra layer_norm
-        self.layer_norm = LayerNorm(config.d_model) if config.normalize_before else None
+        self.layer_norm = LayerNorm(config.d_model) if config.add_final_layer_norm else None

    def forward(
        self, input_ids, attention_mask=None, output_attentions=False, output_hidden_states=False, return_dict=False
@@ -551,8 +551,7 @@ class BartDecoder(nn.Module):

        if use_cache:
            input_ids = input_ids[:, -1:]
-            positions = positions[:, -1:]  # happens after we embed them
-            # assert input_ids.ne(self.padding_idx).any()
+            positions = positions[:, -1:]

        x = self.embed_tokens(input_ids) * self.embed_scale
        x += positions
@@ -590,11 +589,12 @@ class BartDecoder(nn.Module):
            if use_cache:
                next_decoder_cache.append(layer_past.copy())

-            if self.layer_norm and (idx == len(self.layers) - 1):  # if config.add_final_layer_norm (mBART)
-                x = self.layer_norm(x)
            if output_attentions:
                all_self_attns += (layer_self_attn,)

+        if self.layer_norm:  # if config.add_final_layer_norm (mBART)
+            x = self.layer_norm(x)
+
        # Convert to standard output format: (seq_len, BS, model_dim) -> (BS, seq_len, model_dim)
        if output_hidden_states:
            all_hidden_states = tuple(hidden_state.transpose(0, 1) for hidden_state in all_hidden_states)
@@ -863,7 +863,7 @@ class BartModel(PretrainedBartModel):
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="facebook/bart-large",
-        output_type=BaseModelOutputWithPast,
+        output_type=Seq2SeqModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
    def forward(
@@ -725,13 +725,14 @@ class BertModel(BertPreTrainedModel):
    :obj:`encoder_hidden_states` is then expected as an input to the forward pass.
    """

-    def __init__(self, config):
+    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = BertEmbeddings(config)
        self.encoder = BertEncoder(config)
-        self.pooler = BertPooler(config)
+
+        self.pooler = BertPooler(config) if add_pooling_layer else None

        self.init_weights()

@@ -840,7 +841,7 @@ class BertModel(BertPreTrainedModel):
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
-        pooled_output = self.pooler(sequence_output)
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]
@@ -966,13 +967,17 @@ class BertForPreTraining(BertPreTrainedModel):
    """Bert Model with a `language modeling` head on top for CLM fine-tuning. """, BERT_START_DOCSTRING
 )
 class BertLMHeadModel(BertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+    authorized_missing_keys = [r"position_ids", r"predictions.decoder.bias"]
+
    def __init__(self, config):
        super().__init__(config)

        if not config.is_decoder:
            logger.warning("If you want to use `BertLMHeadModel` as a standalone, add `is_decoder=True.`")

-        self.bert = BertModel(config)
+        self.bert = BertModel(config, add_pooling_layer=False)
        self.cls = BertOnlyMLMHead(config)

        self.init_weights()
@@ -1081,6 +1086,10 @@ class BertLMHeadModel(BertPreTrainedModel):

@add_start_docstrings("""Bert Model with a `language modeling` head on top. """, BERT_START_DOCSTRING)
 class BertForMaskedLM(BertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+    authorized_missing_keys = [r"position_ids", r"predictions.decoder.bias"]
+
    def __init__(self, config):
        super().__init__(config)

@@ -1090,7 +1099,7 @@ class BertForMaskedLM(BertPreTrainedModel):
                "bi-directional self-attention."
            )

-        self.bert = BertModel(config)
+        self.bert = BertModel(config, add_pooling_layer=False)
        self.cls = BertOnlyMLMHead(config)

        self.init_weights()
@@ -1457,11 +1466,14 @@ class BertForMultipleChoice(BertPreTrainedModel):
    BERT_START_DOCSTRING,
 )
 class BertForTokenClassification(BertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.bert = BertModel(config)
+        self.bert = BertModel(config, add_pooling_layer=False)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

@@ -1543,11 +1555,14 @@ class BertForTokenClassification(BertPreTrainedModel):
    BERT_START_DOCSTRING,
 )
 class BertForQuestionAnswering(BertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.bert = BertModel(config)
+        self.bert = BertModel(config, add_pooling_layer=False)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
@@ -367,7 +367,6 @@ class FunnelAttentionStructure(nn.Module):
        # Stride is applied on the second-to-last dimension.
        stride = (stride, 1)

-        tensor = tensor.float()
        if mode == "mean":
            tensor = F.avg_pool2d(tensor, stride, stride=stride, ceil_mode=True)
        elif mode == "max":
@@ -554,7 +553,7 @@ class FunnelRelMultiheadAttention(nn.Module):
        attn_score = attn_score.float()
        # perform masking
        if attention_mask is not None:
-            attn_score = attn_score - INF * attention_mask[:, None, None].float()
+            attn_score = attn_score - INF * (1 - attention_mask[:, None, None].float())
        # attention probability
        attn_prob = torch.softmax(attn_score, dim=-1, dtype=dtype)
        attn_prob = self.attention_dropout(attn_prob)
@@ -856,7 +855,9 @@ FUNNEL_INPUTS_DOCSTRING = r"""
        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`({0})`, `optional`):
            Mask to avoid performing attention on padding token indices.
            Mask values selected in ``[0, 1]``:
-            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
+
+            - 1 for tokens that are **not masked**,
+            - 0 for tokens that are **maked**.

            `What are attention masks? <../glossary.html#attention-mask>`__
        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`({0})`, `optional`):
@@ -15,7 +15,6 @@
 # limitations under the License.
 """PyTorch OpenAI GPT-2 model."""

-
 import os
 import warnings
 from dataclasses import dataclass
@@ -624,16 +623,35 @@ class GPT2Model(GPT2PreTrainedModel):
            if output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)

-            outputs = block(
-                hidden_states,
-                layer_past=layer_past,
-                attention_mask=attention_mask,
-                head_mask=head_mask[i],
-                encoder_hidden_states=encoder_hidden_states,
-                encoder_attention_mask=encoder_attention_mask,
-                use_cache=use_cache,
-                output_attentions=output_attentions,
-            )
+            if getattr(self.config, "gradient_checkpointing", False):
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        # checkpointing only works with tuple returns, not with lists
+                        return tuple(output for output in module(*inputs, use_cache, output_attentions))
+
+                    return custom_forward
+
+                outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(block),
+                    hidden_states,
+                    layer_past,
+                    attention_mask,
+                    head_mask[i],
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                outputs = block(
+                    hidden_states,
+                    layer_past=layer_past,
+                    attention_mask=attention_mask,
+                    head_mask=head_mask[i],
+                    encoder_hidden_states=encoder_hidden_states,
+                    encoder_attention_mask=encoder_attention_mask,
+                    use_cache=use_cache,
+                    output_attentions=output_attentions,
+                )

            hidden_states, present = outputs[:2]
            if use_cache is True:
@@ -1081,10 +1081,7 @@ class LongformerModel(LongformerPreTrainedModel):

    """

-    config_class = LongformerConfig
-    base_model_prefix = "longformer"
-
-    def __init__(self, config):
+    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

@@ -1100,7 +1097,7 @@ class LongformerModel(LongformerPreTrainedModel):

        self.embeddings = LongformerEmbeddings(config)
        self.encoder = LongformerEncoder(config)
-        self.pooler = LongformerPooler(config)
+        self.pooler = LongformerPooler(config) if add_pooling_layer else None

        self.init_weights()

@@ -1270,7 +1267,7 @@ class LongformerModel(LongformerPreTrainedModel):
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
-        pooled_output = self.pooler(sequence_output)
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        # undo padding
        if padding_len > 0:
@@ -1290,13 +1287,13 @@ class LongformerModel(LongformerPreTrainedModel):

@add_start_docstrings("""Longformer Model with a `language modeling` head on top. """, LONGFORMER_START_DOCSTRING)
 class LongformerForMaskedLM(LongformerPreTrainedModel):
-    config_class = LongformerConfig
-    base_model_prefix = "longformer"
+
+    authorized_unexpected_keys = [r"pooler"]

    def __init__(self, config):
        super().__init__(config)

-        self.longformer = LongformerModel(config)
+        self.longformer = LongformerModel(config, add_pooling_layer=False)
        self.lm_head = LongformerLMHead(config)

        self.init_weights()
@@ -1395,11 +1392,14 @@ class LongformerForMaskedLM(LongformerPreTrainedModel):
    LONGFORMER_START_DOCSTRING,
 )
 class LongformerForSequenceClassification(LongformerPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.longformer = LongformerModel(config)
+        self.longformer = LongformerModel(config, add_pooling_layer=False)
        self.classifier = LongformerClassificationHead(config)

        self.init_weights()
@@ -1500,11 +1500,14 @@ class LongformerClassificationHead(nn.Module):
    LONGFORMER_START_DOCSTRING,
 )
 class LongformerForQuestionAnswering(LongformerPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.longformer = LongformerModel(config)
+        self.longformer = LongformerModel(config, add_pooling_layer=False)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
@@ -1628,11 +1631,14 @@ class LongformerForQuestionAnswering(LongformerPreTrainedModel):
    LONGFORMER_START_DOCSTRING,
 )
 class LongformerForTokenClassification(LongformerPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.longformer = LongformerModel(config)
+        self.longformer = LongformerModel(config, add_pooling_layer=False)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

@@ -676,6 +676,7 @@ class MobileBertPreTrainedModel(PreTrainedModel):
    pretrained_model_archive_map = MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST
    load_tf_weights = load_tf_weights_in_mobilebert
    base_model_prefix = "mobilebert"
+    authorized_missing_keys = [r"position_ids"]

    def _init_weights(self, module):
        """ Initialize the weights """
@@ -813,14 +814,13 @@ class MobileBertModel(MobileBertPreTrainedModel):
    https://arxiv.org/pdf/2004.02984.pdf
    """

-    authorized_missing_keys = [r"position_ids"]
-
-    def __init__(self, config):
+    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config
        self.embeddings = MobileBertEmbeddings(config)
        self.encoder = MobileBertEncoder(config)
-        self.pooler = MobileBertPooler(config)
+
+        self.pooler = MobileBertPooler(config) if add_pooling_layer else None

        self.init_weights()

@@ -919,7 +919,7 @@ class MobileBertModel(MobileBertPreTrainedModel):
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
-        pooled_output = self.pooler(sequence_output)
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]
@@ -1054,9 +1054,12 @@ class MobileBertForPreTraining(MobileBertPreTrainedModel):

@add_start_docstrings("""MobileBert Model with a `language modeling` head on top. """, MOBILEBERT_START_DOCSTRING)
 class MobileBertForMaskedLM(MobileBertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
-        self.mobilebert = MobileBertModel(config)
+        self.mobilebert = MobileBertModel(config, add_pooling_layer=False)
        self.cls = MobileBertOnlyMLMHead(config)
        self.config = config

@@ -1346,11 +1349,14 @@ class MobileBertForSequenceClassification(MobileBertPreTrainedModel):
    MOBILEBERT_START_DOCSTRING,
 )
 class MobileBertForQuestionAnswering(MobileBertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.mobilebert = MobileBertModel(config)
+        self.mobilebert = MobileBertModel(config, add_pooling_layer=False)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
@@ -1532,11 +1538,14 @@ class MobileBertForMultipleChoice(MobileBertPreTrainedModel):
    MOBILEBERT_START_DOCSTRING,
 )
 class MobileBertForTokenClassification(MobileBertPreTrainedModel):
+
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.mobilebert = MobileBertModel(config)
+        self.mobilebert = MobileBertModel(config, add_pooling_layer=False)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

@@ -814,7 +814,8 @@ class RagSequenceForGeneration(RagPreTrainedModel):
    @torch.no_grad()
    def generate(
        self,
-        input_ids,
+        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
        context_input_ids=None,
        do_deduplication=None,  # defaults to True
        num_return_sequences=None,  # defaults to 1
@@ -830,6 +831,14 @@ class RagSequenceForGeneration(RagPreTrainedModel):
            input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
                The sequence used as a prompt for the generation. If :obj:`input_ids` is not passed, then
                :obj:`context_input_ids` has to be provided.
+            attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+                Mask to avoid performing attention on padding token indices.
+                Mask values selected in ``[0, 1]``:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **maked**.
+
+                `What are attention masks? <../glossary.html#attention-mask>`__
            context_input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size * config.n_docs, config.max_combined_length)`, `optional`, returned when `output_retrieved=True`):
                Input IDs post-processed from the retrieved documents and the question encoder input_ids by the
                retriever.
@@ -859,7 +868,7 @@ class RagSequenceForGeneration(RagPreTrainedModel):

        # TODO(patrick) - clean up generate here
        if self.retriever is not None and context_input_ids is None:
-            question_hidden_states = self.question_encoder(input_ids)[0]
+            question_hidden_states = self.question_encoder(input_ids, attention_mask=attention_mask)[0]
            context_input_ids = self.retriever(
                input_ids,
                question_hidden_states.cpu().detach().to(torch.float32).numpy(),
@@ -873,7 +882,7 @@ class RagSequenceForGeneration(RagPreTrainedModel):

        hypos = []
        kwargs["num_beams"] = num_beams
-        kwargs["num_return_sequences"] = num_return_sequences
+        kwargs["num_return_sequences"] = num_beams
        kwargs["attention_mask"] = None

        for index in range(len(input_ids)):
@@ -907,7 +916,8 @@ class RagSequenceForGeneration(RagPreTrainedModel):
        )

        # bos_token_id is None for T5
-        use_bos = self.config.bos_token_id is not None and target[:, 0].eq(self.config.bos_token_id).all()
+        bos_token_id = self.config.bos_token_id or self.config.generator.bos_token_id
+        use_bos = bos_token_id is not None and target[:, 0].eq(bos_token_id).all()

        def _mask_pads(ll, smooth_obj):
            pad_mask = target.eq(self.config.generator.pad_token_id)
@@ -1180,6 +1190,7 @@ class RagTokenForGeneration(RagPreTrainedModel):
    def generate(
        self,
        input_ids: Optional[torch.LongTensor] = None,
+        attention_mask: Optional[torch.LongTensor] = None,
        context_input_ids=None,
        context_attention_mask=None,
        doc_scores=None,
@@ -1205,6 +1216,14 @@ class RagTokenForGeneration(RagPreTrainedModel):
            input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
                The sequence used as a prompt for the generation. If :obj:`input_ids` is not passed, then
                :obj:`context_input_ids` has to be provided.
+            attention_mask (:obj:`torch.Tensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+                Mask to avoid performing attention on padding token indices.
+                Mask values selected in ``[0, 1]``:
+
+                - 1 for tokens that are **not masked**,
+                - 0 for tokens that are **maked**.
+
+                `What are attention masks? <../glossary.html#attention-mask>`__
            context_input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size * config.n_docs, config.max_combined_length)`, `optional`, returned when `output_retrieved=True`):
                Input IDs post-processed from the retrieved documents and the question encoder :obj:`input_ids` by the
                retriever.
@@ -1293,7 +1312,7 @@ class RagTokenForGeneration(RagPreTrainedModel):

        # retrieve docs
        if self.retriever is not None and context_input_ids is None:
-            question_hidden_states = self.question_encoder(input_ids)[0]
+            question_hidden_states = self.question_encoder(input_ids, attention_mask=attention_mask)[0]
            out = self.retriever(
                input_ids,
                question_hidden_states.cpu().detach().to(torch.float32).numpy(),
@@ -460,7 +460,6 @@ class RobertaPreTrainedModel(PreTrainedModel):

    config_class = RobertaConfig
    base_model_prefix = "roberta"
-    authorized_missing_keys = [r"position_ids"]

    # Copied from transformers.modeling_bert.BertPreTrainedModel._init_weights
    def _init_weights(self, module):
@@ -568,14 +567,17 @@ class RobertaModel(RobertaPreTrainedModel):

    """

+    authorized_missing_keys = [r"position_ids"]
+
    # Copied from transformers.modeling_bert.BertModel.__init__ with Bert->Roberta
-    def __init__(self, config):
+    def __init__(self, config, add_pooling_layer=True):
        super().__init__(config)
        self.config = config

        self.embeddings = RobertaEmbeddings(config)
        self.encoder = RobertaEncoder(config)
-        self.pooler = RobertaPooler(config)
+
+        self.pooler = RobertaPooler(config) if add_pooling_layer else None

        self.init_weights()

@@ -683,7 +685,7 @@ class RobertaModel(RobertaPreTrainedModel):
            return_dict=return_dict,
        )
        sequence_output = encoder_outputs[0]
-        pooled_output = self.pooler(sequence_output)
+        pooled_output = self.pooler(sequence_output) if self.pooler is not None else None

        if not return_dict:
            return (sequence_output, pooled_output) + encoder_outputs[1:]
@@ -700,13 +702,16 @@ class RobertaModel(RobertaPreTrainedModel):
    """RoBERTa Model with a `language modeling` head on top for CLM fine-tuning. """, ROBERTA_START_DOCSTRING
 )
 class RobertaForCausalLM(RobertaPreTrainedModel):
+    authorized_missing_keys = [r"position_ids", r"predictions.decoder.bias"]
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)

        if not config.is_decoder:
            logger.warning("If you want to use `RobertaLMHeadModel` as a standalone, add `is_decoder=True.`")

-        self.roberta = RobertaModel(config)
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
        self.lm_head = RobertaLMHead(config)

        self.init_weights()
@@ -816,6 +821,9 @@ class RobertaForCausalLM(RobertaPreTrainedModel):

@add_start_docstrings("""RoBERTa Model with a `language modeling` head on top. """, ROBERTA_START_DOCSTRING)
 class RobertaForMaskedLM(RobertaPreTrainedModel):
+    authorized_missing_keys = [r"position_ids", r"predictions.decoder.bias"]
+    authorized_unexpected_keys = [r"pooler"]
+
    def __init__(self, config):
        super().__init__(config)

@@ -825,7 +833,7 @@ class RobertaForMaskedLM(RobertaPreTrainedModel):
                "bi-directional self-attention."
            )

-        self.roberta = RobertaModel(config)
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
        self.lm_head = RobertaLMHead(config)

        self.init_weights()
@@ -938,11 +946,13 @@ class RobertaLMHead(nn.Module):
    ROBERTA_START_DOCSTRING,
 )
 class RobertaForSequenceClassification(RobertaPreTrainedModel):
+    authorized_missing_keys = [r"position_ids"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.roberta = RobertaModel(config)
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
        self.classifier = RobertaClassificationHead(config)

        self.init_weights()
@@ -1018,6 +1028,8 @@ class RobertaForSequenceClassification(RobertaPreTrainedModel):
    ROBERTA_START_DOCSTRING,
 )
 class RobertaForMultipleChoice(RobertaPreTrainedModel):
+    authorized_missing_keys = [r"position_ids"]
+
    def __init__(self, config):
        super().__init__(config)

@@ -1106,11 +1118,14 @@ class RobertaForMultipleChoice(RobertaPreTrainedModel):
    ROBERTA_START_DOCSTRING,
 )
 class RobertaForTokenClassification(RobertaPreTrainedModel):
+    authorized_unexpected_keys = [r"pooler"]
+    authorized_missing_keys = [r"position_ids"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.roberta = RobertaModel(config)
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)

@@ -1211,11 +1226,14 @@ class RobertaClassificationHead(nn.Module):
    ROBERTA_START_DOCSTRING,
 )
 class RobertaForQuestionAnswering(RobertaPreTrainedModel):
+    authorized_unexpected_keys = [r"pooler"]
+    authorized_missing_keys = [r"position_ids"]
+
    def __init__(self, config):
        super().__init__(config)
        self.num_labels = config.num_labels

-        self.roberta = RobertaModel(config)
+        self.roberta = RobertaModel(config, add_pooling_layer=False)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()
@@ -907,7 +907,7 @@ T5_INPUTS_DOCSTRING = r"""
    T5_START_DOCSTRING,
 )
 class T5Model(T5PreTrainedModel):
-    def __init__(self, config):
+    def __init__(self, config: T5Config):
        super().__init__(config)
        self.shared = nn.Embedding(config.vocab_size, config.d_model)

@@ -919,6 +919,7 @@ class T5Model(T5PreTrainedModel):
        decoder_config = copy.deepcopy(config)
        decoder_config.is_decoder = True
        decoder_config.is_encoder_decoder = False
+        decoder_config.num_layers = config.num_decoder_layers
        self.decoder = T5Stack(decoder_config, self.shared)

        self.init_weights()
@@ -1077,6 +1078,7 @@ class T5ForConditionalGeneration(T5PreTrainedModel):
        decoder_config = copy.deepcopy(config)
        decoder_config.is_decoder = True
        decoder_config.is_encoder_decoder = False
+        decoder_config.num_layers = config.num_decoder_layers
        self.decoder = T5Stack(decoder_config, self.shared)

        self.lm_head = nn.Linear(config.d_model, config.vocab_size, bias=False)
@@ -826,6 +826,9 @@ class TFAlbertSOPHead(tf.keras.layers.Layer):

@add_start_docstrings("""Albert Model with a `language modeling` head on top. """, ALBERT_START_DOCSTRING)
 class TFAlbertForMaskedLM(TFAlbertPreTrainedModel, TFMaskedLanguageModelingLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -991,6 +994,9 @@ class TFAlbertForSequenceClassification(TFAlbertPreTrainedModel, TFSequenceClass
    ALBERT_START_DOCSTRING,
 )
 class TFAlbertForTokenClassification(TFAlbertPreTrainedModel, TFTokenClassificationLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -1073,6 +1079,9 @@ class TFAlbertForTokenClassification(TFAlbertPreTrainedModel, TFTokenClassificat
    ALBERT_START_DOCSTRING,
 )
 class TFAlbertForQuestionAnswering(TFAlbertPreTrainedModel, TFQuestionAnsweringLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -453,9 +453,8 @@ class TFAutoModel(object):
            >>> config = TFAutoConfig.from_pretrained('bert-base-uncased')
            >>> model = TFAutoModel.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_MAPPING.keys():
+            return TF_MODEL_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -494,9 +493,10 @@ class TFAutoModel(object):
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_MAPPING.keys():
+            return TF_MODEL_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -547,9 +547,8 @@ class TFAutoModelForPreTraining(object):
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = TFAutoModelForPreTraining.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_FOR_PRETRAINING_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_FOR_PRETRAINING_MAPPING.keys():
+            return TF_MODEL_FOR_PRETRAINING_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -588,9 +587,10 @@ class TFAutoModelForPreTraining(object):
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_FOR_PRETRAINING_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_FOR_PRETRAINING_MAPPING.keys():
+            return TF_MODEL_FOR_PRETRAINING_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -653,9 +653,8 @@ class TFAutoModelWithLMHead(object):
            "and `TFAutoModelForSeq2SeqLM` for encoder-decoder models.",
            FutureWarning,
        )
-        for config_class, model_class in TF_MODEL_WITH_LM_HEAD_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_WITH_LM_HEAD_MAPPING.keys():
+            return TF_MODEL_WITH_LM_HEAD_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -701,10 +700,10 @@ class TFAutoModelWithLMHead(object):
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_WITH_LM_HEAD_MAPPING.items():
-            # Not using isinstance() here to do not take into account inheritance
-            if config_class == type(config):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_WITH_LM_HEAD_MAPPING.keys():
+            return TF_MODEL_WITH_LM_HEAD_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -755,9 +754,8 @@ class TFAutoModelForCausalLM:
            >>> config = AutoConfig.from_pretrained('gpt2')
            >>> model = TFAutoModelForCausalLM.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_FOR_CAUSAL_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_FOR_CAUSAL_LM_MAPPING.keys():
+            return TF_MODEL_FOR_CAUSAL_LM_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -796,9 +794,10 @@ class TFAutoModelForCausalLM:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_FOR_CAUSAL_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_FOR_CAUSAL_LM_MAPPING.keys():
+            return TF_MODEL_FOR_CAUSAL_LM_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -849,9 +848,8 @@ class TFAutoModelForMaskedLM:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = TFAutoModelForMaskedLM.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_FOR_MASKED_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_FOR_MASKED_LM_MAPPING.keys():
+            return TF_MODEL_FOR_MASKED_LM_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -890,9 +888,10 @@ class TFAutoModelForMaskedLM:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_FOR_MASKED_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_FOR_MASKED_LM_MAPPING.keys():
+            return TF_MODEL_FOR_MASKED_LM_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -943,9 +942,8 @@ class TFAutoModelForSeq2SeqLM:
            >>> config = AutoConfig.from_pretrained('t5')
            >>> model = TFAutoModelForSeq2SeqLM.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.keys():
+            return TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -986,9 +984,10 @@ class TFAutoModelForSeq2SeqLM:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING.keys():
+            return TF_MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1041,9 +1040,8 @@ class TFAutoModelForSequenceClassification(object):
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = TFAutoModelForSequenceClassification.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.keys():
+            return TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1084,9 +1082,10 @@ class TFAutoModelForSequenceClassification(object):
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING.keys():
+            return TF_MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1138,9 +1137,8 @@ class TFAutoModelForQuestionAnswering(object):
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = TFAutoModelForQuestionAnswering.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING.keys():
+            return TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1181,9 +1179,10 @@ class TFAutoModelForQuestionAnswering(object):
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING.keys():
+            return TF_MODEL_FOR_QUESTION_ANSWERING_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1235,9 +1234,8 @@ class TFAutoModelForTokenClassification:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = TFAutoModelForTokenClassification.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.keys():
+            return TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1278,9 +1276,10 @@ class TFAutoModelForTokenClassification:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING.keys():
+            return TF_MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1333,9 +1332,8 @@ class TFAutoModelForMultipleChoice:
            >>> config = AutoConfig.from_pretrained('bert-base-uncased')
            >>> model = TFAutoModelForMultipleChoice.from_config(config)
        """
-        for config_class, model_class in TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class(config)
+        if type(config) in TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING.keys():
+            return TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING[type(config)](config)
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -1376,9 +1374,10 @@ class TFAutoModelForMultipleChoice:
                pretrained_model_name_or_path, return_unused_kwargs=True, **kwargs
            )

-        for config_class, model_class in TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING.items():
-            if isinstance(config, config_class):
-                return model_class.from_pretrained(pretrained_model_name_or_path, *model_args, config=config, **kwargs)
+        if type(config) in TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING.keys():
+            return TF_MODEL_FOR_MULTIPLE_CHOICE_MAPPING[type(config)].from_pretrained(
+                pretrained_model_name_or_path, *model_args, config=config, **kwargs
+            )
        raise ValueError(
            "Unrecognized configuration class {} for this kind of TFAutoModel: {}.\n"
            "Model type should be one of {}.".format(
@@ -853,6 +853,9 @@ class TFBertForPreTraining(TFBertPreTrainedModel):

@add_start_docstrings("""Bert Model with a `language modeling` head on top. """, BERT_START_DOCSTRING)
 class TFBertForMaskedLM(TFBertPreTrainedModel, TFMaskedLanguageModelingLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -935,6 +938,9 @@ class TFBertForMaskedLM(TFBertPreTrainedModel, TFMaskedLanguageModelingLoss):


 class TFBertLMHeadModel(TFBertPreTrainedModel, TFCausalLanguageModelingLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -1279,6 +1285,9 @@ class TFBertForMultipleChoice(TFBertPreTrainedModel, TFMultipleChoiceLoss):
    BERT_START_DOCSTRING,
 )
 class TFBertForTokenClassification(TFBertPreTrainedModel, TFTokenClassificationLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -1359,6 +1368,9 @@ class TFBertForTokenClassification(TFBertPreTrainedModel, TFTokenClassificationL
    BERT_START_DOCSTRING,
 )
 class TFBertForQuestionAnswering(TFBertPreTrainedModel, TFQuestionAnsweringLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -555,7 +555,7 @@ class TFFunnelRelMultiheadAttention(tf.keras.layers.Layer):
            attn_score = tf.cast(attn_score, tf.float32)
        # perform masking
        if attention_mask is not None:
-            attn_score = attn_score - INF * tf.cast(attention_mask[:, None, None], tf.float32)
+            attn_score = attn_score - INF * (1 - tf.cast(attention_mask[:, None, None], tf.float32))
        # attention probability
        attn_prob = tf.nn.softmax(attn_score, axis=-1)
        if dtype != tf.float32:
@@ -1618,6 +1618,9 @@ class TFLongformerModel(TFLongformerPreTrainedModel):
    LONGFORMER_START_DOCSTRING,
 )
 class TFLongformerForMaskedLM(TFLongformerPreTrainedModel, TFMaskedLanguageModelingLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -1700,6 +1703,9 @@ class TFLongformerForMaskedLM(TFLongformerPreTrainedModel, TFMaskedLanguageModel
    LONGFORMER_START_DOCSTRING,
 )
 class TFLongformerForQuestionAnswering(TFLongformerPreTrainedModel, TFQuestionAnsweringLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -1019,6 +1019,9 @@ class TFMobileBertForPreTraining(TFMobileBertPreTrainedModel):

@add_start_docstrings("""MobileBert Model with a `language modeling` head on top. """, MOBILEBERT_START_DOCSTRING)
 class TFMobileBertForMaskedLM(TFMobileBertPreTrainedModel, TFMaskedLanguageModelingLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -1241,6 +1244,9 @@ class TFMobileBertForSequenceClassification(TFMobileBertPreTrainedModel, TFSeque
    MOBILEBERT_START_DOCSTRING,
 )
 class TFMobileBertForQuestionAnswering(TFMobileBertPreTrainedModel, TFQuestionAnsweringLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -1463,6 +1469,9 @@ class TFMobileBertForMultipleChoice(TFMobileBertPreTrainedModel, TFMultipleChoic
    MOBILEBERT_START_DOCSTRING,
 )
 class TFMobileBertForTokenClassification(TFMobileBertPreTrainedModel, TFTokenClassificationLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -160,6 +160,10 @@ def load_pytorch_weights_in_tf2_model(tf_model, pt_state_dict, tf_inputs=None, a
            if allow_missing_keys:
                missing_keys.append(name)
                continue
+            elif tf_model.authorized_missing_keys is not None:
+                # authorized missing keys don't have to be loaded
+                if any(re.search(pat, name) is not None for pat in tf_model.authorized_missing_keys):
+                    continue

            raise AttributeError("{} not found in PyTorch model".format(name))

@@ -194,6 +198,10 @@ def load_pytorch_weights_in_tf2_model(tf_model, pt_state_dict, tf_inputs=None, a

    unexpected_keys = list(all_pytorch_weights)

+    if tf_model.authorized_missing_keys is not None:
+        for pat in tf_model.authorized_missing_keys:
+            missing_keys = [k for k in missing_keys if re.search(pat, k) is None]
+
    if len(unexpected_keys) > 0:
        logger.warning(
            f"Some weights of the PyTorch model were not used when "
@@ -751,6 +751,9 @@ class TFRobertaLMHead(tf.keras.layers.Layer):

@add_start_docstrings("""RoBERTa Model with a `language modeling` head on top. """, ROBERTA_START_DOCSTRING)
 class TFRobertaForMaskedLM(TFRobertaPreTrainedModel, TFMaskedLanguageModelingLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)

@@ -859,6 +862,9 @@ class TFRobertaClassificationHead(tf.keras.layers.Layer):
    ROBERTA_START_DOCSTRING,
 )
 class TFRobertaForSequenceClassification(TFRobertaPreTrainedModel, TFSequenceClassificationLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -1059,6 +1065,9 @@ class TFRobertaForMultipleChoice(TFRobertaPreTrainedModel, TFMultipleChoiceLoss)
    ROBERTA_START_DOCSTRING,
 )
 class TFRobertaForTokenClassification(TFRobertaPreTrainedModel, TFTokenClassificationLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -1140,6 +1149,9 @@ class TFRobertaForTokenClassification(TFRobertaPreTrainedModel, TFTokenClassific
    ROBERTA_START_DOCSTRING,
 )
 class TFRobertaForQuestionAnswering(TFRobertaPreTrainedModel, TFQuestionAnsweringLoss):
+
+    authorized_missing_keys = [r"pooler"]
+
    def __init__(self, config, *inputs, **kwargs):
        super().__init__(config, *inputs, **kwargs)
        self.num_labels = config.num_labels
@@ -16,6 +16,7 @@
 """TF general model utils."""
 import functools
 import os
+import re
 import warnings
 from typing import Dict, List, Optional, Union

@@ -233,6 +234,7 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
    """
    config_class = None
    base_model_prefix = ""
+    authorized_missing_keys = None

    @property
    def dummy_inputs(self) -> Dict[str, tf.Tensor]:
@@ -630,6 +632,10 @@ class TFPreTrainedModel(tf.keras.Model, TFModelUtilsMixin, TFGenerationMixin):
        unexpected_keys = list(hdf5_layer_names - model_layer_names)
        error_msgs = []

+        if cls.authorized_missing_keys is not None:
+            for pat in cls.authorized_missing_keys:
+                missing_keys = [k for k in missing_keys if re.search(pat, k) is None]
+
        if len(unexpected_keys) > 0:
            logger.warning(
                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when "
@@ -398,6 +398,7 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
    config_class = None
    base_model_prefix = ""
    authorized_missing_keys = None
+    authorized_unexpected_keys = None
    keys_to_never_save = None

    @property
@@ -1013,6 +1014,10 @@ class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin):
                for pat in cls.authorized_missing_keys:
                    missing_keys = [k for k in missing_keys if re.search(pat, k) is None]

+            if cls.authorized_unexpected_keys is not None:
+                for pat in cls.authorized_unexpected_keys:
+                    unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]
+
            if len(unexpected_keys) > 0:
                logger.warning(
                    f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when "
Author	SHA1	Message	Date
Sylvain Gugger	1ba08dc221	Release: v3.3.1	2020-09-29 14:17:34 -04:00
Sylvain Gugger	8546dc55c2	Fix Trainer tests in a multiGPU env (#7458 )	2020-09-29 14:06:41 -04:00
Sylvain Gugger	d0fd7154c5	Catch import datasets common errors (#7456 )	2020-09-29 13:42:09 -04:00
Sylvain Gugger	f1220c5fe2	Add a code of conduct (#7433 )	2020-09-29 13:38:47 -04:00
Teven Patrick von Platen Sylvain Gugger	9e9a1fb8c7	Adding gradient checkpointing to GPT2 (#7446 ) * GPT2 gradient checkpointing * find_unused_parameters removed if checkpointing * find_unused_parameters removed if checkpointing * Update src/transformers/configuration_gpt2.py Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> * Added a test for generation with checkpointing * Update src/transformers/configuration_gpt2.py Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com> Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com> Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>	2020-09-29 12:26:26 -04:00
Sylvain Gugger	52e8392b7e	Add automatic best model loading to Trainer (#7431 ) * Add automatic best model loading to Trainer * Some small fixes * Formatting	2020-09-29 10:41:18 -04:00
Sylvain Gugger	1fc4de69ed	Document new features of make fixup (#7434 )	2020-09-29 03:56:57 -04:00
GmailB	205bf0b7ea	Update README.md (#7444 ) Hi, just corrected the example code, add 2 links and fixed some typos	2020-09-29 03:18:01 -04:00
Sam Shleifer	74d8d69bd4	[s2s] consistent output format across eval scripts (#7435 )	2020-09-28 23:20:03 -04:00
Typicasoft	671b278e25	Create README.md (#7436 ) * Create README.md MagBERT-NER : Added widget (Text) * Rename model_cards/README.md to model_cards/TypicaAI/magbert-ner/README.md	2020-09-28 18:25:25 -04:00
Manuel Romero	a1a8ffa512	Update README.md (#7429 ) Add links to models fine-tuned on a downstream task	2020-09-28 13:40:09 -04:00
Stas Bekman	f62f2ffdcc	[makefile] 10x speed up checking/fixing (#7403 ) * [makefile] check/fix only modified since branching files * fix phonies * parametrize dirs * have only one source for dirs to check * look ma, no autoformatters here	2020-09-28 10:45:42 -04:00
Lysandre	16c213820e	Update docs to version v3.3.0	2020-09-28 16:32:00 +02:00
Lysandre	0613f05226	Release: v3.3.0	2020-09-28 16:24:43 +02:00
Sylvain Gugger	ca3fc36de3	Reorganize documentation navbar (#7423 ) * Reorganize documentation navbar * Update css to have clear sections	2020-09-28 16:22:58 +02:00
Lysandre Debutandsgugger	7f4115c099	Pull request template (#7392 ) co-authored-by: sgugger <sylvain.gugger@gmail.com> Co-authored-by: sgugger <sylvain.gugger@gmail.com>	2020-09-28 09:51:49 -04:00
Sylvain Gugger	0611eab5e3	Document RAG again (#7377 ) Do not merge before Monday	2020-09-28 08:31:46 -04:00
Sylvain Gugger	7563d5a3cf	Catch PyTorch warning when saving/loading scheduler (#7401 )	2020-09-28 08:20:10 -04:00
Boris Dayma Julien Chaumond Sylvain Gugger	1749ca317e	docs: fix model sharing file names (#5855 ) * docs: fix model sharing file names * Update docs/source/model_sharing.rst Co-authored-by: Julien Chaumond <chaumond@gmail.com> * docs(model_sharing.rst): fix new line Co-authored-by: Julien Chaumond <chaumond@gmail.com> Co-authored-by: Sylvain Gugger <35901082+sgugger@users.noreply.github.com>	2020-09-28 08:17:30 -04:00
Patrick von Platen	8279471506	correct RAG model cards (#7420 )	2020-09-28 11:08:39 +02:00
Marcin Zabłocki	4083a55ab0	Flos fix (#7384 )	2020-09-28 04:09:26 -04:00
Ola PiktusandYour Name	ae3e84f3ba	[RAG] Clean Rag readme in examples (#7413 ) * Improve README + consolidation script * Reformat README * Reformat README Co-authored-by: Your Name <you@example.com>	2020-09-28 10:06:39 +02:00
Sam Shleifer	748425d47d	[T5] allow config.decoder_layers to control decoder size (#7409 ) * Working assymmetrical T5 * rename decoder_layers -> num_decoder_layers * Fix docstring * Allow creation of asymmetric t5 students	2020-09-28 03:08:04 -04:00
Sam ShleiferandSwetha Mandava	7296fea1d6	[s2s] rougeLSum expects \n between sentences (#7410 ) Co-authored-by: Swetha Mandava <smandava@nvidia.com>	2020-09-27 16:27:19 -04:00
Suraj PatilandSam Shleifer	eab5f59682	[s2s] add create student script (#7290 ) Co-authored-by: Suraj Patil <surajp815@gmail.com> Co-authored-by: Sam Shleifer <sshleifer@gmail.com>	2020-09-27 15:10:46 -04:00
Patrick von Platen	e50a931c11	[Longformer, Bert, Roberta, ...] Fix multi gpu training (#7272 ) * fix multi-gpu * fix longformer * force to delete unnecessary layers * fix notifications * fix warning * fix roberta * fix tests * remove hasattr * fix tests * fix roberta * merge and clean authorized keys	2020-09-25 20:33:21 +02:00
Patrick von Platen	2c8ecdf8a8	fix rag retriever save pretrained (#7399 )	2020-09-25 19:47:12 +02:00
Patrick von Platen	1a14687e6f	Update README.md	2020-09-25 19:43:48 +02:00
Patrick von Platen	3327c2b0f6	Update README.md	2020-09-25 19:43:36 +02:00
Ola PiktusandYour Name	fe326bd5cf	Remove dependency on examples/seq2seq from rag (#7395 ) Co-authored-by: Your Name <you@example.com>	2020-09-25 18:20:49 +02:00
Sylvain Gugger	ad39271ae8	Fix FP16 and attention masks in FunnelTransformer (#7374 ) * Fix #7371 * Fix training * Fix test values * Apply the fix to TF as well	2020-09-25 12:20:39 -04:00
Patrick von Platen	4e5b036bdd	Update README.md	2020-09-25 18:16:46 +02:00
Patrick von Platen	55eccfbb49	Update README.md	2020-09-25 18:16:44 +02:00
Sylvain Gugger	e2e77f02c2	Fix BartModel output documentation (#7390 )	2020-09-25 11:48:13 -04:00
Sylvain Gugger	bbb07830ff	Speedup check_copies script (#7394 )	2020-09-25 11:47:22 -04:00
Stas Bekman	8859c4f841	[code quality] new make target that combines style and quality targets (#7310 ) * [code quality] merge style and quality targets Any reason why we don't run `flake8` in `make style`? I find myself needing to run `make style` and `make quality` all the time, but I need the latter just for the last 2 checks. Since we have no control over the source code why bother with separating checking and fixing - let's just have one target that fixes and then performs the remaining checks, as we know the first two have been done already. This PR suggests to merge the 2 targets into one efficient target. I will edit the docs if this change resonates with the team. * move checks into style, re-use target * better name * add fixup target * document new target	2020-09-25 11:37:40 -04:00
Sam Shleifer	38a1b03f4d	Remove unhelpful bart warning (#7391 )	2020-09-25 11:01:07 -04:00
Patrick von Platen	5ff0d6d7d0	Update README.md	2020-09-25 16:58:29 +02:00
Quentin LhoestandPatrick von Platen	cf1c88e092	[RAG] Fix retrieval offset in RAG's HfIndex and better integration tests (#7372 ) * Fix retrieval offset in RAG's HfIndex * update slow tests * style * fix new test * style * add better tests Co-authored-by: Patrick von Platen <patrick.v.platen@gmail.com>	2020-09-25 16:12:46 +02:00
Patrick von Platen	571c7a11c1	[Rag] Fix wrong usage of `num_beams` and `bos_token_id` in Rag Sequence generation (#7386 ) * fix_rag_sequence * add second bug fix	2020-09-25 14:35:49 +02:00
Suraj Patil	415071b4c2	doc changes (#7385 )	2020-09-25 08:00:36 -04:00
Patrick von Platen	2dd652d757	[RAG] Add missing doc and attention_mask to rag (#7382 ) * add docs * add missing docs and attention_mask in fine-tune	2020-09-25 11:23:55 +02:00
Lysandre Debut	7cdd9da5bf	Check config type using `type` instead of `isinstance` (#7363 ) * Check config type instead of instance Bad merge * Remove for loops * Style	2020-09-25 05:09:09 -04:00
Sam Shleifer	3c6bf8998f	modeling_bart: 3 small cleanups that dont change outputs (#7381 ) * Mbart passing * boom boom * cleaner assert * add assert * Fix tests	2020-09-25 04:24:14 -04:00
Suraj PatilandSam Shleifer	9e68d075a4	Seq2SeqTrainer (#6769 ) Co-authored-by: Sam Shleifer <sshleifer@gmail.com>	2020-09-24 18:46:58 -04:00
Sam Shleifer	d9d0f1140b	[s2s] distributed eval allows num_return_sequences > 1 (#7254 )	2020-09-24 17:30:09 -04:00
Patrick von Platen	0804d077c6	correct attention mask (#7373 )	2020-09-24 23:22:04 +02:00
Stas BekmanandSam Shleifer	a8cbc4269c	[fsmt] build/test scripts (#7257 ) Co-authored-by: Sam Shleifer <sshleifer@gmail.com>	2020-09-24 17:10:26 -04:00
Sylvain Gugger	a8e7982f84	Remove mentions of RAG from the docs (#7376 ) * Remove mentions of RAG from the docs * Deactivate check	2020-09-24 17:07:14 -04:00
Stas Bekman	eadd870b2f	[seq2seq] make it easier to run the scripts (#7274 )	2020-09-24 15:23:48 -04:00
Lysandre Debut	8d3bb781ee	Formatter (#7368 ) * Formatter * Docs	2020-09-24 10:59:21 -04:00