PyTorch initial token classification pipeline.

Signed-off-by: Morgan Funtowicz <funtowiczmo@gmail.com>

.circleci/config.yml

@@ -247,6 +247,7 @@ jobs:
             - run: black --check --line-length 119 --target-version py35 examples templates tests src utils
             - run: isort --check-only examples templates tests src utils
             - run: flake8 examples templates tests src utils
+            - run: python utils/check_copies.py
             - run: python utils/check_repo.py
     check_repository_consistency:
         working_directory: ~/transformers

.circleci/deploy.sh

@@ -48,4 +48,5 @@ deploy_doc "7cb203f" v2.9.1
 deploy_doc "10d7239" v2.10.0 
 deploy_doc "b42586e" v2.11.0
 deploy_doc "7fb8bdf" v3.0.2
-deploy_doc "4b3ee9c" # v3.1.0 Latest stable release
+deploy_doc "4b3ee9c" v3.1.0
+deploy_doc "3ebb1b3" # v3.2.0 Latest stable release

.gitignore

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

Makefile

@@ -6,6 +6,7 @@ quality:
 	black --check --line-length 119 --target-version py35 examples templates tests src utils
 	isort --check-only examples templates tests src utils
 	flake8 examples templates tests src utils
+	python utils/check_copies.py
 	python utils/check_repo.py
 
 # Format source code automatically
@@ -14,6 +15,11 @@ style:
 	black --line-length 119 --target-version py35 examples templates tests src utils
 	isort examples templates tests src utils
 
+# Make marked copies of snippets of codes conform to the original
+
+fix-copies:
+	python utils/check_copies.py --fix_and_overwrite
+
 # Run tests for the library
 
 test:

README.md

@@ -183,8 +183,9 @@ Min, Patrick Lewis, Ledell Wu, Sergey Edunov, Danqi Chen, and Wen-tau Yih.
 24. **[MBart](https://github.com/pytorch/fairseq/tree/master/examples/mbart)** (from Facebook) released with the paper  [Multilingual Denoising Pre-training for Neural Machine Translation](https://arxiv.org/abs/2001.08210) by Yinhan Liu, Jiatao Gu, Naman Goyal, Xian Li, Sergey Edunov, Marjan Ghazvininejad, Mike Lewis, Luke Zettlemoyer.  
 25. **[LXMERT](https://github.com/airsplay/lxmert)** (from UNC Chapel Hill) released with the paper [LXMERT: Learning Cross-Modality Encoder Representations from Transformers for Open-Domain Question Answering](https://arxiv.org/abs/1908.07490) by Hao Tan and Mohit Bansal.
 26. **[Funnel Transformer](https://github.com/laiguokun/Funnel-Transformer)** (from CMU/Google Brain) released with the paper [Funnel-Transformer: Filtering out Sequential Redundancy for Efficient Language Processing](https://arxiv.org/abs/2006.03236) by Zihang Dai, Guokun Lai, Yiming Yang, Quoc V. Le.
-27. **[Other community models](https://huggingface.co/models)**, contributed by the [community](https://huggingface.co/users).
-28. Want to contribute a new model? We have added a **detailed guide and templates** to guide you in the process of adding a new model. You can find them in the [`templates`](./templates) folder of the repository. Be sure to check the [contributing guidelines](./CONTRIBUTING.md) and contact the maintainers or open an issue to collect feedbacks before starting your PR.
+27. **[LayoutLM](https://github.com/microsoft/unilm/tree/master/layoutlm)** (from Microsoft Research Asia) released with the paper [LayoutLM: Pre-training of Text and Layout for Document Image Understanding](https://arxiv.org/abs/1912.13318) by Yiheng Xu, Minghao Li, Lei Cui, Shaohan Huang, Furu Wei, Ming Zhou.
+28. **[Other community models](https://huggingface.co/models)**, contributed by the [community](https://huggingface.co/users).
+29. Want to contribute a new model? We have added a **detailed guide and templates** to guide you in the process of adding a new model. You can find them in the [`templates`](./templates) folder of the repository. Be sure to check the [contributing guidelines](./CONTRIBUTING.md) and contact the maintainers or open an issue to collect feedbacks before starting your PR.
 
 These implementations have been tested on several datasets (see the example scripts) and should match the performances of the original implementations. You can find more details on the performances in the Examples section of the [documentation](https://huggingface.co/transformers/examples.html).
 

docs/source/_static/js/custom.js

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

docs/source/conf.py

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

docs/source/custom_datasets.rst

@@ -324,7 +324,7 @@ which we'll use in a moment:
     id2tag = {id: tag for tag, id in tag2id.items()}
 
 To encode the tokens, we'll use a pre-trained DistilBert tokenizer. We can tell the tokenizer that we're dealing
-with ready-split tokens rather than full sentence strings by passing ``is_pretokenized=True``. We'll also pass
+with ready-split tokens rather than full sentence strings by passing ``is_split_into_words=True``. We'll also pass
 ``padding=True`` and ``truncation=True`` to pad the sequences to be the same length. Lastly, we can tell the model
 to return information about the tokens which are split by the wordpiece tokenization process, which we will need in
 a moment.
@@ -333,8 +333,8 @@ a moment.
 
     from transformers import DistilBertTokenizerFast
     tokenizer = DistilBertTokenizerFast.from_pretrained('distilbert-base-cased')
-    train_encodings = tokenizer(train_texts, is_pretokenized=True, return_offsets_mapping=True, padding=True, truncation=True)
-    val_encodings = tokenizer(val_texts, is_pretokenized=True, return_offsets_mapping=True, padding=True, truncation=True)
+    train_encodings = tokenizer(train_texts, is_split_into_words=True, return_offsets_mapping=True, padding=True, truncation=True)
+    val_encodings = tokenizer(val_texts, is_split_into_words=True, return_offsets_mapping=True, padding=True, truncation=True)
 
 Great, so now our tokens are nicely encoded in the format that they need to be in to feed them into our DistilBert
 model below.

docs/source/index.rst

@@ -137,7 +137,10 @@ conversion utilities for the following models:
 27. `Bert For Sequence Generation <https://tfhub.dev/s?module-type=text-generation&subtype=module,placeholder>`_ (from Google) released with the paper
     `Leveraging Pre-trained Checkpoints for Sequence Generation Tasks
     <https://arxiv.org/abs/1907.12461>`_ by Sascha Rothe, Shashi Narayan, Aliaksei Severyn.
-28. `Other community models <https://huggingface.co/models>`_, contributed by the `community
+28. `LayoutLM <https://github.com/microsoft/unilm/tree/master/layoutlm>`_ (from Microsoft Research Asia) released with the paper
+    `LayoutLM: Pre-training of Text and Layout for Document Image Understanding
+    <https://arxiv.org/abs/1912.13318>`_ by Yiheng Xu, Minghao Li, Lei Cui, Shaohan Huang, Furu Wei, Ming Zhou.
+29. `Other community models <https://huggingface.co/models>`_, contributed by the `community
     <https://huggingface.co/users>`_.
 
 .. toctree::
@@ -227,6 +230,8 @@ conversion utilities for the following models:
     model_doc/funnel
     model_doc/lxmert
     model_doc/bertgeneration
+    model_doc/layoutlm
+    model_doc/rag
     internal/modeling_utils
     internal/tokenization_utils
     internal/pipelines_utils

docs/source/model_doc/layoutlm.rst

@@ -0,0 +1,55 @@
+LayoutLM
+----------------------------------------------------
+
+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
+The LayoutLM model was proposed in `LayoutLM: Pre-training of Text and Layout for Document Image Understanding <https://arxiv.org/abs/1912.13318>`__
+by Yiheng Xu, Minghao Li, Lei Cui, Shaohan Huang, Furu Wei, and Ming Zhou. It's a simple but effective pre-training method 
+of text and layout for document image understanding and information extraction tasks, such as form understanding and receipt understanding.
+
+The abstract from the paper is the following:
+
+*Pre-training techniques have been verified successfully in a variety of NLP tasks in recent years. Despite the widespread use of pre-training models for NLP applications, they almost exclusively focus on text-level manipulation, while neglecting layout and style information that is vital for document image understanding. In this paper, we propose the \textbf{LayoutLM} to jointly model interactions between text and layout information across scanned document images, which is beneficial for a great number of real-world document image understanding tasks such as information extraction from scanned documents. Furthermore, we also leverage image features to incorporate words' visual information into LayoutLM. To the best of our knowledge, this is the first time that text and layout are jointly learned in a single framework for document-level pre-training. It achieves new state-of-the-art results in several downstream tasks, including form understanding (from 70.72 to 79.27), receipt understanding (from 94.02 to 95.24) and document image classification (from 93.07 to 94.42).*
+
+Tips:
+
+- LayoutLM has an extra input called :obj:`bbox`, which is the bounding boxes of the input tokens.
+- The :obj:`bbox` requires the data that on 0-1000 scale, which means you should normalize the bounding box before passing them into model.
+
+The original code can be found `here <https://github.com/microsoft/unilm/tree/master/layoutlm>`_.
+
+
+LayoutLMConfig
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LayoutLMConfig
+    :members:
+
+
+LayoutLMTokenizer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LayoutLMTokenizer
+    :members:
+
+
+LayoutLMModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LayoutLMModel
+    :members:
+
+
+LayoutLMForMaskedLM
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LayoutLMForMaskedLM
+    :members:
+
+
+LayoutLMForTokenClassification
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.LayoutLMForTokenClassification
+    :members:

docs/source/model_doc/rag.rst

@@ -0,0 +1,88 @@
+RAG
+----------------------------------------------------
+
+Overview
+~~~~~~~~~~~~~~~~~~~~~
+
+Retrieval-augmented generation ("RAG") models combine the powers of pretrained dense retrieval (DPR) and Seq2Seq models.
+RAG models retrieve docs, pass them to a seq2seq model, then marginalize to generate outputs.
+The retriever and seq2seq modules are initialized from pretrained models, and fine-tuned jointly, allowing both retrieval and generation to adapt to downstream tasks.
+
+It is based on the paper `Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks <https://arxiv.org/abs/2005.11401>`__ by Patrick Lewis, Ethan Perez, Aleksandara Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, Douwe Kiela.
+
+The abstract from the paper is the following:
+
+*Large pre-trained language models have been shown to store factual knowledge
+in their parameters, and achieve state-of-the-art results when fine-tuned on
+downstream NLP tasks. However, their ability to access and precisely manipulate
+knowledge is still limited, and hence on knowledge-intensive tasks, their
+performance lags behind task-specific architectures. Additionally, providing
+provenance for their decisions and updating their world knowledge remain open
+research problems. Pre-trained models with a differentiable access mechanism to
+explicit nonparametric memory can overcome this issue, but have so far been only
+investigated for extractive downstream tasks. We explore a general-purpose
+fine-tuning recipe for retrieval-augmented generation (RAG) — models which combine
+pre-trained parametric and non-parametric memory for language generation. We
+introduce RAG models where the parametric memory is a pre-trained seq2seq model and
+the non-parametric memory is a dense vector index of Wikipedia, accessed with
+a pre-trained neural retriever. We compare two RAG formulations, one which
+conditions on the same retrieved passages across the whole generated sequence, the
+other can use different passages per token. We fine-tune and evaluate our models
+on a wide range of knowledge-intensive NLP tasks and set the state-of-the-art
+on three open domain QA tasks, outperforming parametric seq2seq models and
+task-specific retrieve-and-extract architectures. For language generation tasks, we
+find that RAG models generate more specific, diverse and factual language than a
+state-of-the-art parametric-only seq2seq baseline.*
+
+
+
+RagConfig
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.RagConfig
+    :members:
+
+
+RagTokenizer
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.RagTokenizer
+    :members:
+
+
+Rag specific outputs
+~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.modeling_rag.RetrievAugLMMarginOutput
+    :members:
+
+.. autoclass:: transformers.modeling_rag.RetrievAugLMOutput
+    :members:
+
+
+RAGRetriever
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.RagRetriever
+    :members:
+
+
+RagModel
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.RagModel
+    :members: forward
+
+
+RagSequenceForGeneration
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.RagSequenceForGeneration
+    :members: forward, generate
+
+
+RagTokenForGeneration
+~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~~
+
+.. autoclass:: transformers.RagTokenForGeneration
+    :members: forward, generate

docs/source/model_summary.rst

@@ -654,7 +654,7 @@ DPR
    <a href="https://huggingface.co/models?filter=dpr">
        <img alt="Models" src="https://img.shields.io/badge/All_model_pages-dpr-blueviolet">
    </a>
-   <a href="model_doc/ctrl.dpr">
+   <a href="model_doc/dpr.html">
        <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-dpr-blueviolet">
    </a>
 
@@ -672,6 +672,27 @@ DPR consists in three models:
 
 DPR's pipeline (not implemented yet) uses a retrieval step to find the top k contexts given a certain question, and then it calls the reader with the question and the retrieved documents to get the answer.
 
+RAG
+----------------------------------------------
+
+.. raw:: html
+
+   <a href="https://huggingface.co/models?filter=rag">
+       <img alt="Models" src="https://img.shields.io/badge/All_model_pages-rag-blueviolet">
+   </a>
+   <a href="model_doc/rag.html">
+       <img alt="Doc" src="https://img.shields.io/badge/Model_documentation-rag-blueviolet">
+   </a>
+
+`Retrieval-Augmented Generation for Knowledge-Intensive NLP Tasks <https://arxiv.org/abs/2005.11401>`_,
+Patrick Lewis, Ethan Perez, Aleksandara Piktus, Fabio Petroni, Vladimir Karpukhin, Naman Goyal, Heinrich Küttler, Mike Lewis, Wen-tau Yih, Tim Rocktäschel, Sebastian Riedel, Douwe Kiela
+
+Retrieval-augmented generation ("RAG") models combine the powers of pretrained dense retrieval (DPR) and Seq2Seq models.
+RAG models retrieve docs, pass them to a seq2seq model, then marginalize to generate outputs.
+The retriever and seq2seq modules are initialized from pretrained models, and fine-tuned jointly, allowing both retrieval and generation to adapt to downstream tasks.
+
+The two models RAG-Token and RAG-Sequence are available for generation.
+
 More technical aspects
 ^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^^
 

docs/source/preprocessing.rst

@@ -290,12 +290,12 @@ predictions in `named entity recognition (NER) <https://en.wikipedia.org/wiki/Na
     if that was the case) but just split into words (which is often the first step in subword tokenization algorithms
     like BPE).
 
-If you want to use pre-tokenized inputs, just set :obj:`is_pretokenized=True` when passing your inputs to the
+If you want to use pre-tokenized inputs, just set :obj:`is_split_into_words=True` when passing your inputs to the
 tokenizer. For instance, we have:
 
 .. code-block::
 
-    >>> encoded_input = tokenizer(["Hello", "I'm", "a", "single", "sentence"], is_pretokenized=True)
+    >>> encoded_input = tokenizer(["Hello", "I'm", "a", "single", "sentence"], is_split_into_words=True)
     >>> print(encoded_input)
     {'input_ids': [101, 8667, 146, 112, 182, 170, 1423, 5650, 102],
      'token_type_ids': [0, 0, 0, 0, 0, 0, 0, 0, 0], 
@@ -312,7 +312,7 @@ like this:
     batch_sentences = [["Hello", "I'm", "a", "single", "sentence"],
                        ["And", "another", "sentence"],
                        ["And", "the", "very", "very", "last", "one"]]
-    encoded_inputs = tokenizer(batch_sentences, is_pretokenized=True)
+    encoded_inputs = tokenizer(batch_sentences, is_split_into_words=True)
 
 or a batch of pair sentences like this:
 
@@ -321,7 +321,7 @@ or a batch of pair sentences like this:
     batch_of_second_sentences = [["I'm", "a", "sentence", "that", "goes", "with", "the", "first", "sentence"],
                                  ["And", "I", "should", "be", "encoded", "with", "the", "second", "sentence"],
                                  ["And", "I", "go", "with", "the", "very", "last", "one"]]
-    encoded_inputs = tokenizer(batch_sentences, batch_of_second_sentences, is_pretokenized=True)
+    encoded_inputs = tokenizer(batch_sentences, batch_of_second_sentences, is_split_into_words=True)
 
 And you can add padding, truncation as well as directly return tensors like before:
 
@@ -330,14 +330,14 @@ And you can add padding, truncation as well as directly return tensors like befo
     ## PYTORCH CODE
     batch = tokenizer(batch_sentences,
                       batch_of_second_sentences,
-                      is_pretokenized=True,
+                      is_split_into_words=True,
                       padding=True,
                       truncation=True,
                       return_tensors="pt")
     ## TENSORFLOW CODE
     batch = tokenizer(batch_sentences,
                       batch_of_second_sentences,
-                      is_pretokenized=True,
+                      is_split_into_words=True,
                       padding=True,
                       truncation=True,
                       return_tensors="tf")

docs/source/pretrained_models.rst

@@ -408,3 +408,11 @@ For a list that includes community-uploaded models, refer to `https://huggingfac
 |                    |                                                            |                                                                                                                                       |
 |                    |                                                            | (see `details <https://github.com/laiguokun/Funnel-Transformer>`__)                                                                   |
 +--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+| LayoutLM           | ``microsoft/layoutlm-base-uncased``                        | | 12 layers, 768-hidden, 12-heads, 113M parameters                                                                                    |
+|                    |                                                            |                                                                                                                                       |
+|                    |                                                            | (see `details <https://github.com/microsoft/unilm/tree/master/layoutlm>`__)                                                           |
++                    +------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+
+|                    | ``microsoft/layoutlm-large-uncased``                       | | 24 layers, 1024-hidden, 16-heads, 343M parameters                                                                                   |
+|                    |                                                            |                                                                                                                                       |
+|                    |                                                            | (see `details <https://github.com/microsoft/unilm/tree/master/layoutlm>`__)                                                           |
++--------------------+------------------------------------------------------------+---------------------------------------------------------------------------------------------------------------------------------------+

examples/lightning_base.py

@@ -119,7 +119,7 @@ class BaseTransformer(pl.LightningModule):
     def get_lr_scheduler(self):
         get_schedule_func = arg_to_scheduler[self.hparams.lr_scheduler]
         scheduler = get_schedule_func(
-            self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.total_steps()
+            self.opt, num_warmup_steps=self.hparams.warmup_steps, num_training_steps=self.total_steps
         )
         scheduler = {"scheduler": scheduler, "interval": "step", "frequency": 1}
         return scheduler
@@ -159,20 +159,19 @@ class BaseTransformer(pl.LightningModule):
     def test_epoch_end(self, outputs):
         return self.validation_end(outputs)
 
+    @property
     def total_steps(self) -> int:
         """The number of total training steps that will be run. Used for lr scheduler purposes."""
         num_devices = max(1, self.hparams.gpus)  # TODO: consider num_tpu_cores
         effective_batch_size = self.hparams.train_batch_size * self.hparams.accumulate_grad_batches * num_devices
-        return (self.dataset_size / effective_batch_size) * self.hparams.max_epochs
+        dataset_size = len(self.train_loader.dataset)
+        return (dataset_size / effective_batch_size) * self.hparams.max_epochs
 
     def setup(self, mode):
-        if mode == "test":
-            self.dataset_size = len(self.test_dataloader().dataset)
-        else:
+        if mode == "fit":
             self.train_loader = self.get_dataloader("train", self.hparams.train_batch_size, shuffle=True)
-            self.dataset_size = len(self.train_loader.dataset)
 
-    def get_dataloader(self, type_path: str, batch_size: int, shuffle: bool = False):
+    def get_dataloader(self, type_path, batch_size, shuffle=False):
         raise NotImplementedError("You must implement this for your task")
 
     def train_dataloader(self):
@@ -367,6 +366,8 @@ def generic_train(
     if args.gpus > 1:
         train_params["distributed_backend"] = "ddp"
 
+    train_params["accumulate_grad_batches"] = args.accumulate_grad_batches
+
     trainer = pl.Trainer.from_argparse_args(
         args,
         weights_summary=None,

examples/longform-qa/eli5_app.py

@@ -1,10 +1,10 @@
 import datasets
-import faiss
 import numpy as np
 import streamlit as st
 import torch
 from elasticsearch import Elasticsearch
 
+import faiss
 import transformers
 from eli5_utils import (
     embed_questions_for_retrieval,

examples/longform-qa/eli5_utils.py

@@ -5,7 +5,6 @@ from random import choice, randint
 from time import time
 
 import datasets  # noqa: F401
-import faiss  # noqa: F401
 import numpy as np
 import pandas as pd
 import torch
@@ -15,6 +14,7 @@ from elasticsearch.helpers import bulk, streaming_bulk  # noqa: F401
 from torch.utils.data import DataLoader, Dataset, RandomSampler, SequentialSampler
 from tqdm import tqdm
 
+import faiss  # noqa: F401
 from transformers import AdamW, AutoModel, AutoModelForSeq2SeqLM, AutoTokenizer, get_linear_schedule_with_warmup
 
 

examples/rag/README.md

@@ -0,0 +1,88 @@
+# Intro
+RAG is a seq2seq model which encapsulates two core components: a question encoder and a generator.
+During a forward pass, we encode the input with the question encoder and pass it
+to the retriever to extract relevant context documents. The documents are then prepended to the input.
+Such contextualized 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.
+
+
+# Finetuning
+Our finetuning logic is based on scripts from [`examples/seq2seq`](https://github.com/huggingface/transformers/tree/master/examples/seq2seq).
+Follow instructions there regarding data preprocessing. A sample finetuning command:
+
+```
+python examples/rag/finetune.py \
+    --data_dir $DATA_DIR \
+    --output_dir $OUTPUT_DIR \
+    --model_name_or_path $MODEL_NAME_OR_PATH \
+    --model_type rag_sequence \
+    --fp16 \
+    --gpus 8
+```
+
+
+# Evaluation
+Apart 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`.
+
+We expect the following formats of the gold data file:
+
+- 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:
+```
+Does He Love You	Does He Love You	Red Sandy Spika dress of Reba McEntire	Greatest Hits Volume Two (Reba McEntire album)	Shoot for the Moon (album)
+```
+
+## Retrieval evaluation
+
+We demonstrate how to evaluate retrieval against DPR evaluation data. You can download respective files from links listed [here](https://github.com/facebookresearch/DPR/blob/master/data/download_data.py#L39-L45).
+
+1. Download and unzip the gold data file. We use the `biencoder-nq-dev` from https://dl.fbaipublicfiles.com/dpr/data/retriever/biencoder-nq-dev.json.gz.
+2. Parse the unziped file using the `parse_dpr_relevance_data.py`
+```
+python examples/rag/parse_dpr_relevance_data.py --src_path path/to/unziped/biencoder-nq-dev.json --evaluation_set path/to/output/biencoder-nq-dev.questions --gold_data_path path/to/output/biencoder-nq-dev.pages
+```
+3. Run evaluation:
+```
+python examples/rag/eval_rag.py \
+    --model_name_or_path $MODEL_NAME_OR_PATH \ # model name or path of the model we're evaluating
+    --model_type rag_sequence \ # RAG model type (rag_token or rag_sequence)
+    --evaluation_set path/to/output/biencoder-nq-dev.questions \ # an input dataset for evaluation
+    --gold_data_path path/to/output/biencoder-nq-dev.pages \ # a dataset containing ground truth answers for samples from the evaluation_set
+    --predictions_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.
+```
+
+
+## End-to-end evaluation
+```
+python examples/rag/eval_rag.py \
+	--model_name_or_path $MODEL_NAME_OR_PATH \
+    --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)
+    --n_docs 5 \ # You can experiment with retrieving different number of documents at evaluation time
+    --print_predictions
+```

examples/rag/callbacks.py

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

examples/rag/distributed_retriever.py

@@ -0,0 +1,135 @@
+import logging
+import os
+from typing import List, Tuple
+
+import numpy as np
+import psutil
+import torch
+import torch.distributed as dist
+
+from transformers import RagRetriever
+
+
+logger = logging.getLogger(__name__)
+
+
+class RagPyTorchDistributedRetriever(RagRetriever):
+    """
+    A distributed retriever built on top of the ``torch.distributed`` communication package. During training all workers
+    initalize their own instance of the retriever, however, only the main worker loads the index into memory. The index is stored
+    in cpu memory. The index will also work well in a non-distributed setup.
+
+    Args:
+        config (:class:`~transformers.RagConfig`):
+            The configuration of the RAG model this Retriever is used with. Contains parameters indicating which ``Index`` to build.
+        question_encoder_tokenizer (:class:`~transformers.PretrainedTokenizer`):
+            The tokenizer that was used to tokenize the question.
+            It is used to decode the question and then use the generator_tokenizer.
+        generator_tokenizer (:class:`~transformers.PretrainedTokenizer`):
+            The tokenizer used for the generator part of the RagModel.
+    """
+
+    _init_retrieval = False
+
+    def __init__(self, config, question_encoder_tokenizer, generator_tokenizer):
+        super().__init__(
+            config, question_encoder_tokenizer=question_encoder_tokenizer, generator_tokenizer=generator_tokenizer
+        )
+
+        self.process_group = None
+
+    def init_retrieval(self, distributed_port: int):
+        """
+        Retriever initalization function, needs to be called from the training process. The function sets some common parameters
+        and environment variables. On top of that, (only) the main process in the process group loads the index into memory.
+
+        Args:
+            distributed_port (:obj:`int`):
+                The port on which the main communication of the training run is carried out. We set the port for retrieval-related
+                communication as ``distributed_port + 1``.
+        """
+
+        logger.info("initializing retrieval")
+
+        # initializing a separate process group for retrievel as the default
+        # nccl backend doesn't support gather/scatter operations while gloo
+        # is too slow to replace nccl for the core gpu communication
+        if dist.is_initialized():
+            logger.info("dist initialized")
+            # needs to be set manually
+            os.environ["GLOO_SOCKET_IFNAME"] = self._infer_socket_ifname()
+            # avoid clash with the NCCL port
+            os.environ["MASTER_PORT"] = str(distributed_port + 1)
+            self.process_group = dist.new_group(ranks=None, backend="gloo")
+
+        # initialize retriever only on the main worker
+        if not dist.is_initialized() or self._is_main():
+            logger.info("dist not initialized / main")
+            self.index.init_index()
+
+        # all processes wait untill the retriever is initialized by the main process
+        if dist.is_initialized():
+            torch.distributed.barrier(group=self.process_group)
+
+    def _is_main(self):
+        return dist.get_rank(group=self.process_group) == 0
+
+    def _scattered(self, scatter_list, target_shape, target_type=torch.float32):
+        target_tensor = torch.empty(target_shape, dtype=target_type)
+        dist.scatter(target_tensor, src=0, scatter_list=scatter_list, group=self.process_group)
+        return target_tensor
+
+    def _infer_socket_ifname(self):
+        addrs = psutil.net_if_addrs()
+        # a hacky way to deal with varying network interface names
+        ifname = next((addr for addr in addrs if addr.startswith("e")), None)
+        return ifname
+
+    def retrieve(self, question_hidden_states: np.ndarray, n_docs: int) -> Tuple[np.ndarray, List[dict]]:
+        """
+        Retrieves documents for specified ``question_hidden_states``. The main process, which has the access to the index stored in memory, gathers queries
+        from all the processes in the main training process group, performs the retrieval and scatters back the results.
+
+        Args:
+            question_hidden_states (:obj:`np.ndarray` of shape :obj:`(batch_size, vector_size)`):
+                A batch of query vectors to retrieve with.
+            n_docs (:obj:`int`):
+                The number of docs retrieved per query.
+
+        Ouput:
+            retrieved_doc_embeds (:obj:`np.ndarray` of shape :obj:`(batch_size, n_docs, dim)`
+                The retrieval embeddings of the retrieved docs per query.
+            doc_ids (:obj:`np.ndarray` of shape :obj:`batch_size, n_docs`)
+                The ids of the documents in the index
+            doc_dicts (:obj:`List[dict]`):
+                The retrieved_doc_embeds examples per query.
+        """
+
+        # single GPU training
+        if not dist.is_initialized():
+            doc_ids, retrieved_doc_embeds = self._main_retrieve(question_hidden_states, n_docs)
+            return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(doc_ids)
+
+        # distributed training
+        world_size = dist.get_world_size(group=self.process_group)
+
+        # gather logic
+        gather_list = None
+        if self._is_main():
+            gather_list = [torch.empty(question_hidden_states.shape, dtype=torch.float32) for _ in range(world_size)]
+        dist.gather(torch.tensor(question_hidden_states), dst=0, gather_list=gather_list, group=self.process_group)
+
+        # scatter logic
+        n_queries = question_hidden_states.shape[0]
+        scatter_ids = []
+        scatter_vectors = []
+        if self._is_main():
+            assert len(gather_list) == world_size
+            ids, vectors = self._main_retrieve(torch.cat(gather_list).numpy(), n_docs)
+            ids, vectors = torch.tensor(ids), torch.tensor(vectors)
+            scatter_ids = self._chunk_tensor(ids, n_queries)
+            scatter_vectors = self._chunk_tensor(vectors, n_queries)
+        doc_ids = self._scattered(scatter_ids, [n_queries, n_docs], target_type=torch.int64)
+        retrieved_doc_embeds = self._scattered(scatter_vectors, [n_queries, n_docs, question_hidden_states.shape[1]])
+
+        return retrieved_doc_embeds.numpy(), doc_ids.numpy(), self.index.get_doc_dicts(doc_ids)

examples/rag/eval_rag.py

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

examples/rag/finetune.py

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

examples/rag/finetune.sh

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

examples/rag/parse_dpr_relevance_data.py

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

examples/rag/requirements.txt

@@ -0,0 +1,4 @@
+faiss-cpu >= 1.6.3
+datasets >= 1.0.1
+psutil >= 5.7.0
+torch >= 1.4.0

examples/rag/test_distributed_retriever.py

@@ -0,0 +1,156 @@
+import json
+import os
+import shutil
+import sys
+import tempfile
+import unittest
+from unittest import TestCase
+from unittest.mock import patch
+
+import numpy as np
+from datasets import Dataset
+
+import faiss
+from transformers.configuration_bart import BartConfig
+from transformers.configuration_dpr import DPRConfig
+from transformers.configuration_rag import RagConfig
+from transformers.file_utils import is_datasets_available, is_faiss_available, is_psutil_available, is_torch_available
+from transformers.tokenization_bart import BartTokenizer
+from transformers.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
+from transformers.tokenization_dpr import DPRQuestionEncoderTokenizer
+from transformers.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
+
+
+sys.path.append(os.path.join(os.getcwd()))  # noqa: E402 # noqa: E402 # isort:skip
+
+from examples.rag.distributed_retriever import RagPyTorchDistributedRetriever  # noqa: E402 # isort:skip
+
+
+def require_distributed_retrieval(test_case):
+    """
+    Decorator marking a test that requires a set of dependencies necessary for pefrorm retrieval with
+    :class:`~transformers.RagRetriever`.
+
+    These tests are skipped when respective libraries are not installed.
+
+    """
+    if not (is_torch_available() and is_datasets_available() and is_faiss_available() and is_psutil_available()):
+        test_case = unittest.skip("test requires PyTorch, Datasets, Faiss, psutil")(test_case)
+    return test_case
+
+
+@require_distributed_retrieval
+class RagRetrieverTest(TestCase):
+    def setUp(self):
+        self.tmpdirname = tempfile.mkdtemp()
+        self.retrieval_vector_size = 8
+
+        # DPR tok
+        vocab_tokens = [
+            "[UNK]",
+            "[CLS]",
+            "[SEP]",
+            "[PAD]",
+            "[MASK]",
+            "want",
+            "##want",
+            "##ed",
+            "wa",
+            "un",
+            "runn",
+            "##ing",
+            ",",
+            "low",
+            "lowest",
+        ]
+        dpr_tokenizer_path = os.path.join(self.tmpdirname, "dpr_tokenizer")
+        os.makedirs(dpr_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(dpr_tokenizer_path, DPR_VOCAB_FILES_NAMES["vocab_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+        # BART tok
+        vocab = [
+            "l",
+            "o",
+            "w",
+            "e",
+            "r",
+            "s",
+            "t",
+            "i",
+            "d",
+            "n",
+            "\u0120",
+            "\u0120l",
+            "\u0120n",
+            "\u0120lo",
+            "\u0120low",
+            "er",
+            "\u0120lowest",
+            "\u0120newer",
+            "\u0120wider",
+            "<unk>",
+        ]
+        vocab_tokens = dict(zip(vocab, range(len(vocab))))
+        merges = ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
+        self.special_tokens_map = {"unk_token": "<unk>"}
+
+        bart_tokenizer_path = os.path.join(self.tmpdirname, "bart_tokenizer")
+        os.makedirs(bart_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(bart_tokenizer_path, BART_VOCAB_FILES_NAMES["vocab_file"])
+        self.merges_file = os.path.join(bart_tokenizer_path, BART_VOCAB_FILES_NAMES["merges_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as fp:
+            fp.write(json.dumps(vocab_tokens) + "\n")
+        with open(self.merges_file, "w", encoding="utf-8") as fp:
+            fp.write("\n".join(merges))
+
+    def get_dpr_tokenizer(self) -> DPRQuestionEncoderTokenizer:
+        return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname, "dpr_tokenizer"))
+
+    def get_bart_tokenizer(self) -> BartTokenizer:
+        return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname, "bart_tokenizer"))
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdirname)
+
+    def get_dummy_pytorch_distributed_retriever(self, init_retrieval, port=12345) -> RagPyTorchDistributedRetriever:
+        dataset = Dataset.from_dict(
+            {
+                "id": ["0", "1"],
+                "text": ["foo", "bar"],
+                "title": ["Foo", "Bar"],
+                "embeddings": [np.ones(self.retrieval_vector_size), 2 * np.ones(self.retrieval_vector_size)],
+            }
+        )
+        dataset.add_faiss_index("embeddings", string_factory="Flat", metric_type=faiss.METRIC_INNER_PRODUCT)
+        config = RagConfig(
+            retrieval_vector_size=self.retrieval_vector_size,
+            question_encoder=DPRConfig().to_dict(),
+            generator=BartConfig().to_dict(),
+        )
+        with patch("transformers.retrieval_rag.load_dataset") as mock_load_dataset:
+            mock_load_dataset.return_value = dataset
+            retriever = RagPyTorchDistributedRetriever(
+                config,
+                question_encoder_tokenizer=self.get_dpr_tokenizer(),
+                generator_tokenizer=self.get_bart_tokenizer(),
+            )
+            if init_retrieval:
+                retriever.init_retrieval(port)
+        return retriever
+
+    def test_pytorch_distributed_retriever_retrieve(self):
+        n_docs = 1
+        retriever = self.get_dummy_pytorch_distributed_retriever(init_retrieval=True)
+        hidden_states = np.array(
+            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+        )
+        retrieved_doc_embeds, doc_ids, doc_dicts = retriever.retrieve(hidden_states, n_docs=n_docs)
+        self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
+        self.assertEqual(len(doc_dicts), 2)
+        self.assertEqual(sorted(doc_dicts[0]), ["embeddings", "id", "text", "title"])
+        self.assertEqual(len(doc_dicts[0]["id"]), n_docs)
+        self.assertEqual(doc_dicts[0]["id"][0], "1")  # max inner product is reached with second doc
+        self.assertEqual(doc_dicts[1]["id"][0], "0")  # max inner product is reached with first doc
+        self.assertListEqual(list(doc_ids), [1, 0])

examples/rag/utils.py

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

examples/requirements.txt

@@ -5,10 +5,10 @@ psutil
 sacrebleu
 rouge-score
 tensorflow_datasets
-pytorch-lightning==0.9.0
+pytorch-lightning==0.8.5
 matplotlib
 git-python==1.0.3
-faiss
+faiss-cpu
 streamlit
 elasticsearch
 pandas

examples/seq2seq/README.md

@@ -1,10 +1,18 @@
 ## Sequence to Sequence
 
 This directory contains examples for finetuning and evaluating transformers on summarization and translation tasks.
-Summarization support is more mature than translation support.
 Please tag @sshleifer with any issues/unexpected behaviors, or send a PR!
 For `bertabs` instructions, see [`bertabs/README.md`](bertabs/README.md).
 
+### Supported Architectures
+
+- `BartForConditionalGeneration` (and anything that inherits from it)
+- `MarianMTModel`
+- `PegasusForConditionalGeneration`
+- `MBartForConditionalGeneration`
+- `FSMTForConditionalGeneration`
+- `T5ForConditionalGeneration`
+    
 
 ## Datasets
 
@@ -227,9 +235,7 @@ python run_eval.py sshleifer/distilbart-cnn-12-6 $DATA_DIR/val.source dbart_val_
     --fp16 \
     --bs 32
 ```
-
-
-### Multi-GPU Evaluation
+### Multi-GPU Evalulation
 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.
@@ -355,17 +361,6 @@ runtime: 13H on V-100 16GB GPU.
 pytest examples/seq2seq/
 ```
 
-### Converting pytorch-lightning checkpoints
-pytorch lightning ``-do_predict`` often fails, after you are done training, the best way to evaluate your model is to convert it.
-
-This should be done for you, with a file called `{save_dir}/best_tfmr`. 
-
-If that file doesn't exist but you have a lightning `.ckpt` file, you can run
-```bash
-python convert_pl_checkpoint_to_hf.py PATH_TO_CKPT  randomly_initialized_hf_model_path save_dir/best_tfmr
-```
-Then either `run_eval` or `run_distributed_eval` with `save_dir/best_tfmr` (see previous sections)
-
 
 ## Experimental Features 
 These features are harder to use and not always useful.
@@ -394,3 +389,4 @@ uses 12,723 batches of length 48 and takes slightly more time 9.5 minutes.
 The feature is still experimental, because:
 + we can make it much more robust if we have memory mapped/preprocessed datasets.
 + The speedup over sortish sampler is not that large at the moment.
+

examples/seq2seq/distillation.py

@@ -13,6 +13,7 @@ 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 transformers.modeling_bart import shift_tokens_right
 from utils import (
@@ -21,6 +22,7 @@ from utils import (
     calculate_bleu,
     freeze_params,
     label_smoothed_nll_loss,
+    pickle_load,
     use_task_specific_params,
 )
 
@@ -410,6 +412,30 @@ def create_module(args):
     return model
 
 
+def evaluate_checkpoint(ckpt_path: Path, dest_dir=None):
+    # TODO(SS): DELETE? Better to convert_pl_ckpt_to_hf and run_eval.py
+    exp_dir = ckpt_path.parent
+    if dest_dir is None:
+        dest_dir = exp_dir
+    clash = list(dest_dir.glob("test_generations*"))
+    if clash:
+        print(f"SKIPPING to avoid overwriting {clash}")
+    ckpt = torch.load(ckpt_path, map_location="cpu")
+    if "hparams" in ckpt:
+        args = argparse.Namespace(**ckpt["hparams"])
+    else:
+        args = argparse.Namespace(**pickle_load(exp_dir / "hparams.pkl"))
+    args.resume_from_checkpoint = str(ckpt_path)
+    args.do_train = False
+    args.output_dir = str(dest_dir)
+    args.n_gpu = 1
+    args.eval_batch_size = 16
+    Path(args.output_dir).mkdir(exist_ok=True)
+    model = create_module(args)
+    trainer: pl.Trainer = generic_train(model, args, early_stopping_callback=False)
+    trainer.test(model)
+
+
 LAYERS_TO_COPY = {
     # maps  num layers in student -> which teacher layers to copy.
     # 12: bart, 16: pegasus, 6: marian/Helsinki-NLP

examples/seq2seq/finetune.py

@@ -36,7 +36,7 @@ from utils import (
 
 
 logger = logging.getLogger(__name__)
-from pytorch_lightning.utilities import rank_zero_only
+
 
 class SummarizationModule(BaseTransformer):
     mode = "summarization"
@@ -61,6 +61,8 @@ class SummarizationModule(BaseTransformer):
         pickle_save(self.hparams, self.hparams_save_path)
         self.step_count = 0
         self.metrics = defaultdict(list)
+        self.model_type = self.config.model_type
+        self.vocab_size = self.config.tgt_vocab_size if self.model_type == "fsmt" else self.config.vocab_size
 
         self.dataset_kwargs: dict = dict(
             data_dir=self.hparams.data_dir,
@@ -106,14 +108,18 @@ class SummarizationModule(BaseTransformer):
 
     def freeze_embeds(self):
         """Freeze token embeddings and positional embeddings for bart, just token embeddings for t5."""
-        try:
-            freeze_params(self.model.model.shared)
+        if self.model_type == "t5":
+            freeze_params(self.model.shared)
+            for d in [self.model.encoder, self.model.decoder]:
+                freeze_params(d.embed_tokens)
+        elif self.model_type == "fsmt":
             for d in [self.model.model.encoder, self.model.model.decoder]:
                 freeze_params(d.embed_positions)
                 freeze_params(d.embed_tokens)
-        except AttributeError:
-            freeze_params(self.model.shared)
-            for d in [self.model.encoder, self.model.decoder]:
+        else:
+            freeze_params(self.model.model.shared)
+            for d in [self.model.model.encoder, self.model.model.decoder]:
+                freeze_params(d.embed_positions)
                 freeze_params(d.embed_tokens)
 
     def forward(self, input_ids, **kwargs):
@@ -140,7 +146,7 @@ class SummarizationModule(BaseTransformer):
             # Same behavior as modeling_bart.py, besides ignoring pad_token_id
             ce_loss_fct = torch.nn.CrossEntropyLoss(ignore_index=pad_token_id)
 
-            assert lm_logits.shape[-1] == self.model.config.vocab_size
+            assert lm_logits.shape[-1] == self.vocab_size
             loss = ce_loss_fct(lm_logits.view(-1, lm_logits.shape[-1]), tgt_ids.view(-1))
         else:
             lprobs = torch.nn.functional.log_softmax(lm_logits, dim=-1)
@@ -183,20 +189,6 @@ class SummarizationModule(BaseTransformer):
         losses.update(generative_metrics)
         all_metrics = {f"{prefix}_avg_{k}": x for k, x in losses.items()}
         all_metrics["step_count"] = self.step_count
-
-        def get_date_str(seconds=True) -> str:
-            """Returns 2019-09-25-10:02:07, for example."""
-            if seconds:
-                return time.strftime('%Y-%m-%d-%H:%M:%S')
-            else:
-                return time.strftime('%Y-%m-%d-%H:%M')
-
-        all_metrics['Time'] = get_date_str(seconds=True)
-        #all_metrics['n_obs'] =
-
-
-        all_metrics['rank'] = getattr(self.train_dataloader(), '_rank', -1.)
-
         self.save_metrics(all_metrics, prefix)  # writes to self.metrics_save_path
         preds = flatten_list([x["preds"] for x in outputs])
         return {

examples/seq2seq/test_seq2seq_examples.py

@@ -13,7 +13,7 @@ import torch
 
 import lightning_base
 from convert_pl_checkpoint_to_hf import convert_pl_to_hf
-from distillation import distill_main
+from distillation import distill_main, evaluate_checkpoint
 from finetune import SummarizationModule, main
 from run_eval import generate_summaries_or_translations, run_generate
 from run_eval_search import run_search
@@ -86,7 +86,7 @@ CHEAP_ARGS = {
     "n_val": -1,
     "n_test": -1,
     "student_encoder_layers": 1,
-    "alpha_loss_encoder": 0.0,
+    "alpha_encoder_loss": 0.0,
     "freeze_encoder": False,
     "auto_scale_batch_size": False,
 }
@@ -103,6 +103,7 @@ T5_TINY = "patrickvonplaten/t5-tiny-random"
 BART_TINY = "sshleifer/bart-tiny-random"
 MBART_TINY = "sshleifer/tiny-mbart"
 MARIAN_TINY = "sshleifer/tiny-marian-en-de"
+FSMT_TINY = "stas/tiny-wmt19-en-de"
 
 
 stream_handler = logging.StreamHandler(sys.stdout)
@@ -177,6 +178,7 @@ class TestSummarizationDistiller(unittest.TestCase):
         generate_summaries_or_translations(examples, out_path, str(model.output_dir / "best_tfmr"))
         self.assertTrue(Path(out_path).exists())
 
+        evaluate_checkpoint(ckpts[0], dest_dir=Path(tempfile.mkdtemp()))
         out_path_new = tempfile.mkdtemp()
         convert_pl_to_hf(ckpts[0], transformer_ckpts[0].parent, out_path_new)
         assert os.path.exists(os.path.join(out_path_new, "pytorch_model.bin"))
@@ -225,6 +227,8 @@ class TestSummarizationDistiller(unittest.TestCase):
         assert len(all_files) > 2
         self.assertEqual(len(transformer_ckpts), 2)
 
+        evaluate_checkpoint(ckpts[0], dest_dir=Path(tempfile.mkdtemp()))
+
     @unittest.skip("T5 distillation is broken at the moment")
     def test_distill_t5(self):
         updates = dict(
@@ -343,14 +347,14 @@ def test_run_eval_search(model):
     task = "translation_en_to_de" if model == T5_TINY else "summarization"
     testargs = f"""
         run_eval_search.py
-        --model_name {model}
-        --data_dir {str(input_file_name)}
-        --save_dir {str(output_file_name)}
+        {model}
+        {str(input_file_name)}
+        {str(output_file_name)}
         --score_path {score_path}
-        --reference_path {reference_path},
+        --reference_path {reference_path}
         --task {task}
-        --search num_beams=1:2 length_penalty=0.9:1.0
         """.split()
+    testargs.extend(["--search", "num_beams=1:2 length_penalty=0.9:1.0"])
 
     with patch.object(sys, "argv", testargs):
         with CaptureStdout() as cs:
@@ -371,11 +375,11 @@ def test_run_eval_search(model):
 
 @pytest.mark.parametrize(
     "model",
-    [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY],
+    [T5_TINY, BART_TINY, MBART_TINY, MARIAN_TINY, FSMT_TINY],
 )
 def test_finetune(model):
     args_d: dict = CHEAP_ARGS.copy()
-    task = "translation" if model in [MBART_TINY, MARIAN_TINY] else "summarization"
+    task = "translation" if model in [MBART_TINY, MARIAN_TINY, FSMT_TINY] else "summarization"
     args_d["label_smoothing"] = 0.1 if task == "translation" else 0
 
     tmp_dir = make_test_data_dir()
@@ -404,7 +408,13 @@ def test_finetune(model):
         lm_head = module.model.lm_head
         assert not lm_head.weight.requires_grad
         assert (lm_head.weight == input_embeds.weight).all().item()
-
+    elif model == FSMT_TINY:
+        fsmt = module.model.model
+        embed_pos = fsmt.decoder.embed_positions
+        assert not embed_pos.weight.requires_grad
+        assert not fsmt.decoder.embed_tokens.weight.requires_grad
+        # check that embeds are not the same
+        assert fsmt.decoder.embed_tokens != fsmt.encoder.embed_tokens
     else:
         bart = module.model.model
         embed_pos = bart.decoder.embed_positions

examples/seq2seq/utils.py

@@ -4,6 +4,7 @@ import linecache
 import math
 import os
 import pickle
+import socket
 from logging import getLogger
 from pathlib import Path
 from typing import Callable, Dict, Iterable, List, Union
@@ -276,12 +277,10 @@ class DistributedSortishSampler(Sampler):
             if not dist.is_available():
                 raise RuntimeError("Requires distributed package to be available")
             num_replicas = dist.get_world_size()
-            self._rank = -1
         if rank is None:
             if not dist.is_available():
                 raise RuntimeError("Requires distributed package to be available")
             rank = dist.get_rank()
-            self._rank = rank
         self.dataset = dataset
         self.num_replicas = num_replicas
         self.rank = rank
@@ -374,6 +373,7 @@ def get_git_info():
         "repo_id": str(repo),
         "repo_sha": str(repo.head.object.hexsha),
         "repo_branch": str(repo.active_branch),
+        "hostname": str(socket.gethostname()),
     }
     return repo_infos
 

examples/text-classification/README.md

@@ -23,6 +23,31 @@ Quick benchmarks from the script (no other modifications):
 Mixed precision (AMP) reduces the training time considerably for the same hardware and hyper-parameters (same batch size was used).
 
 
+## Run generic text classification script in TensorFlow
+
+The script [run_tf_text_classification.py](https://github.com/huggingface/transformers/blob/master/examples/text-classification/run_tf_text_classification.py) allows users to run a text classification on their own CSV files. For now there are few restrictions, the CSV files must have a header corresponding to the column names and not more than three columns: one column for the id, one column for the text and another column for a second piece of text in case of an entailment classification for example.
+
+To use the script, one as to run the following command line:
+```bash
+python run_tf_text_classification.py \
+  --train_file train.csv \ ### training dataset file location (mandatory if running with --do_train option)
+  --dev_file dev.csv \ ### development dataset file location (mandatory if running with --do_eval option)
+  --test_file test.csv \ ### test dataset file location (mandatory if running with --do_predict option)
+  --label_column_id 0 \ ### which column corresponds to the labels
+  --model_name_or_path bert-base-multilingual-uncased \
+  --output_dir model \
+  --num_train_epochs 4 \
+  --per_device_train_batch_size 16 \
+  --per_device_eval_batch_size 32 \
+  --do_train \
+  --do_eval \
+  --do_predict \
+  --logging_steps 10 \
+  --evaluate_during_training \
+  --save_steps 10 \
+  --overwrite_output_dir \
+  --max_seq_length 128
+```
 
 # Run PyTorch version
 

examples/text-classification/run_tf_text_classification.py

@@ -0,0 +1,283 @@
+# coding=utf-8
+""" Fine-tuning the library models for sequence classification."""
+
+
+import logging
+import os
+from dataclasses import dataclass, field
+from typing import Dict, Optional
+
+import datasets
+import numpy as np
+import tensorflow as tf
+
+from transformers import (
+    AutoConfig,
+    AutoTokenizer,
+    EvalPrediction,
+    HfArgumentParser,
+    PreTrainedTokenizer,
+    TFAutoModelForSequenceClassification,
+    TFTrainer,
+    TFTrainingArguments,
+)
+
+
+def get_tfds(
+    train_file: str,
+    eval_file: str,
+    test_file: str,
+    tokenizer: PreTrainedTokenizer,
+    label_column_id: int,
+    max_seq_length: Optional[int] = None,
+):
+    files = {}
+
+    if train_file is not None:
+        files[datasets.Split.TRAIN] = [train_file]
+    if eval_file is not None:
+        files[datasets.Split.VALIDATION] = [eval_file]
+    if test_file is not None:
+        files[datasets.Split.TEST] = [test_file]
+
+    ds = datasets.load_dataset("csv", data_files=files)
+    features_name = list(ds[list(files.keys())[0]].features.keys())
+    label_name = features_name.pop(label_column_id)
+    label_list = list(set(ds[list(files.keys())[0]][label_name]))
+    label2id = {label: i for i, label in enumerate(label_list)}
+    input_names = ["input_ids"] + tokenizer.model_input_names
+    transformed_ds = {}
+
+    if len(features_name) == 1:
+        for k in files.keys():
+            transformed_ds[k] = ds[k].map(
+                lambda example: tokenizer.batch_encode_plus(
+                    example[features_name[0]], truncation=True, max_length=max_seq_length, padding="max_length"
+                ),
+                batched=True,
+            )
+    elif len(features_name) == 2:
+        for k in files.keys():
+            transformed_ds[k] = ds[k].map(
+                lambda example: tokenizer.batch_encode_plus(
+                    (example[features_name[0]], features_name[1]),
+                    truncation=True,
+                    max_length=max_seq_length,
+                    padding="max_length",
+                ),
+                batched=True,
+            )
+
+    def gen_train():
+        for ex in transformed_ds[datasets.Split.TRAIN]:
+            d = {k: v for k, v in ex.items() if k in input_names}
+            label = label2id[ex[label_name]]
+            yield (d, label)
+
+    def gen_val():
+        for ex in transformed_ds[datasets.Split.VALIDATION]:
+            d = {k: v for k, v in ex.items() if k in input_names}
+            label = label2id[ex[label_name]]
+            yield (d, label)
+
+    def gen_test():
+        for ex in transformed_ds[datasets.Split.TEST]:
+            d = {k: v for k, v in ex.items() if k in input_names}
+            label = label2id[ex[label_name]]
+            yield (d, label)
+
+    train_ds = (
+        tf.data.Dataset.from_generator(
+            gen_train,
+            ({k: tf.int32 for k in input_names}, tf.int64),
+            ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])),
+        )
+        if datasets.Split.TRAIN in transformed_ds
+        else None
+    )
+
+    val_ds = (
+        tf.data.Dataset.from_generator(
+            gen_val,
+            ({k: tf.int32 for k in input_names}, tf.int64),
+            ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])),
+        )
+        if datasets.Split.VALIDATION in transformed_ds
+        else None
+    )
+
+    test_ds = (
+        tf.data.Dataset.from_generator(
+            gen_test,
+            ({k: tf.int32 for k in input_names}, tf.int64),
+            ({k: tf.TensorShape([None]) for k in input_names}, tf.TensorShape([])),
+        )
+        if datasets.Split.TEST in transformed_ds
+        else None
+    )
+
+    return train_ds, val_ds, test_ds, label2id
+
+
+logger = logging.getLogger(__name__)
+
+
+@dataclass
+class DataTrainingArguments:
+    """
+    Arguments pertaining to what data we are going to input our model for training and eval.
+
+    Using `HfArgumentParser` we can turn this class
+    into argparse arguments to be able to specify them on
+    the command line.
+    """
+
+    label_column_id: int = field(metadata={"help": "Which column contains the label"})
+    train_file: str = field(default=None, metadata={"help": "The path of the training file"})
+    dev_file: Optional[str] = field(default=None, metadata={"help": "The path of the development file"})
+    test_file: Optional[str] = field(default=None, metadata={"help": "The path of the test file"})
+    max_seq_length: int = field(
+        default=128,
+        metadata={
+            "help": "The maximum total input sequence length after tokenization. Sequences longer "
+            "than this will be truncated, sequences shorter will be padded."
+        },
+    )
+    overwrite_cache: bool = field(
+        default=False, metadata={"help": "Overwrite the cached training and evaluation sets"}
+    )
+
+
+@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"}
+    )
+    use_fast: bool = field(default=False, metadata={"help": "Set this flag to use fast tokenization."})
+    # If you want to tweak more attributes on your tokenizer, you should do it in a distinct script,
+    # or just modify its tokenizer_config.json.
+    cache_dir: Optional[str] = field(
+        default=None, metadata={"help": "Where do you want to store the pretrained models downloaded from s3"}
+    )
+
+
+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, TFTrainingArguments))
+    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,
+    )
+    logger.info(
+        "n_replicas: %s, distributed training: %s, 16-bits training: %s",
+        training_args.n_replicas,
+        bool(training_args.n_replicas > 1),
+        training_args.fp16,
+    )
+    logger.info("Training/evaluation parameters %s", training_args)
+
+    # Load pretrained model and tokenizer
+    #
+    # Distributed training:
+    # The .from_pretrained methods guarantee that only one local process can concurrently
+    # download model & vocab.
+
+    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,
+    )
+
+    train_dataset, eval_dataset, test_ds, label2id = get_tfds(
+        train_file=data_args.train_file,
+        eval_file=data_args.dev_file,
+        test_file=data_args.test_file,
+        tokenizer=tokenizer,
+        label_column_id=data_args.label_column_id,
+        max_seq_length=data_args.max_seq_length,
+    )
+
+    config = AutoConfig.from_pretrained(
+        model_args.config_name if model_args.config_name else model_args.model_name_or_path,
+        num_labels=len(label2id),
+        label2id=label2id,
+        id2label={id: label for label, id in label2id.items()},
+        finetuning_task="text-classification",
+        cache_dir=model_args.cache_dir,
+    )
+
+    with training_args.strategy.scope():
+        model = TFAutoModelForSequenceClassification.from_pretrained(
+            model_args.model_name_or_path,
+            from_pt=bool(".bin" in model_args.model_name_or_path),
+            config=config,
+            cache_dir=model_args.cache_dir,
+        )
+
+    def compute_metrics(p: EvalPrediction) -> Dict:
+        preds = np.argmax(p.predictions, axis=1)
+
+        return {"acc": (preds == p.label_ids).mean()}
+
+    # Initialize our Trainer
+    trainer = TFTrainer(
+        model=model,
+        args=training_args,
+        train_dataset=train_dataset,
+        eval_dataset=eval_dataset,
+        compute_metrics=compute_metrics,
+    )
+
+    # Training
+    if training_args.do_train:
+        trainer.train()
+        trainer.save_model()
+        tokenizer.save_pretrained(training_args.output_dir)
+
+    # Evaluation
+    results = {}
+    if training_args.do_eval:
+        logger.info("*** Evaluate ***")
+
+        result = trainer.evaluate()
+        output_eval_file = os.path.join(training_args.output_dir, "eval_results.txt")
+
+        with open(output_eval_file, "w") as writer:
+            logger.info("***** Eval results *****")
+
+            for key, value in result.items():
+                logger.info("  %s = %s", key, value)
+                writer.write("%s = %s\n" % (key, value))
+
+            results.update(result)
+
+    return results
+
+
+if __name__ == "__main__":
+    main()

model_cards/etalab-ia/camembert-base-squadFR-fquad-piaf/README.md

@@ -45,7 +45,7 @@ python run_squad.py \
 ```
 ### SQuAD-FR Evaluation
 ```shell
-{"f1": 59.54, "exact_match": 80.61}
+{"f1": 80.61, "exact_match": 59.54}
 ```
 
 ## Usage

model_cards/ixa-ehu/ixambert-base-cased/README.md

@@ -0,0 +1,31 @@
+---
+language: 
+- en
+- es
+- eu
+---
+
+# IXAmBERT base cased
+
+This is a multilingual language pretrained for English, Spanish and Basque. The training corpora is composed by the English, Spanish and Basque Wikipedias, together with Basque crawled news articles from online newspapers. The model has been successfully used to transfer knowledge from English to Basque in a conversational QA system, as reported in the paper [Conversational Question Answering in Low Resource Scenarios: A Dataset and Case Study for Basque](http://www.lrec-conf.org/proceedings/lrec2020/pdf/2020.lrec-1.55.pdf). In the paper, IXAmBERT performed better than mBERT when transferring knowledge from English to Basque, as shown in the following Table:
+
+| Model              | Zero-shot | Transfer learning |
+|--------------------|-----------|-------------------|
+| Baseline           |      28.7 |              28.7 |
+| mBERT              |      31.5 |              37.4 |
+| IXAmBERT           |      38.9 |          **41.2** |
+| mBERT + history    |      33.3 |              28.7 |
+| IXAmBERT + history |  **40.7** |              40.0 |
+
+This Table shows the results on a Basque CQA dataset. *Zero-shot* means that the model is fine-tuned using using QuaC, an English CQA dataset. In the *Transfer Learning* setting the model is first fine-tuned on QuaC, and then on a Basque CQA dataset. 
+
+If using this model, please cite the following paper:
+```
+@inproceedings{otegi2020conversational,
+  title={Conversational Question Answering in Low Resource Scenarios: A Dataset and Case Study for Basque},
+  author={Otegi, Arantxa and Agirre, Aitor and Campos, Jon Ander and Soroa, Aitor and Agirre, Eneko},
+  booktitle={Proceedings of The 12th Language Resources and Evaluation Conference},
+  pages={436--442},
+  year={2020}
+}
+```

model_cards/pdelobelle/robbert-v2-dutch-base/README.md

@@ -0,0 +1,164 @@
+---
+language: "nl"
+thumbnail: "https://github.com/iPieter/RobBERT/raw/master/res/robbert_logo.png"
+tags:
+- Dutch
+- RoBERTa
+- RobBERT
+license: "MIT"
+datasets:
+- oscar
+- Shuffled Dutch section of the OSCAR corpus (https://oscar-corpus.com/)
+---
+
+# RobBERT
+
+## Model description
+
+[RobBERT v2](https://github.com/iPieter/RobBERT) is a Dutch state-of-the-art [RoBERTa](https://arxiv.org/abs/1907.11692)-based language model.
+
+More detailled information can be found in the [RobBERT paper](https://arxiv.org/abs/2001.06286).
+
+## How to use
+
+```python
+from transformers import RobertaTokenizer, RobertaForSequenceClassification
+tokenizer = RobertaTokenizer.from_pretrained("pdelobelle/robbert-v2-dutch-base")
+model = RobertaForSequenceClassification.from_pretrained("pdelobelle/robbert-v2-dutch-base")
+```
+
+## Performance Evaluation Results
+
+All experiments are described in more detail in our [paper](https://arxiv.org/abs/2001.06286).
+
+### Sentiment analysis
+Predicting whether a review is positive or negative using the [Dutch Book Reviews Dataset](https://github.com/benjaminvdb/110kDBRD).
+
+|   Model           | Accuracy [%]             |
+|-------------------|--------------------------|
+| ULMFiT            | 93.8                     |
+| BERTje            | 93.0                     |
+| RobBERT v2        | **95.1**                 |
+
+### Die/Dat (coreference resolution)
+
+We measured how well the models are able to do coreference resolution by predicting whether "die" or "dat" should be filled into a sentence.
+For this, we used the [EuroParl corpus](https://www.statmt.org/europarl/).
+
+#### Finetuning on whole dataset
+
+|   Model           | Accuracy [%]             |  F1 [%]    |
+|-------------------|--------------------------|--------------|
+| [Baseline](https://arxiv.org/abs/2001.02943) (LSTM)   |                          | 75.03        |
+| mBERT             | 98.285                   | 98.033       |
+| BERTje            | 98.268                   | 98.014       |
+| RobBERT v2        | **99.232**               | **99.121**   |
+
+#### Finetuning on 10K examples
+
+We also measured the performance using only 10K training examples.
+This experiment clearly illustrates that RobBERT outperforms other models when there is little data available.
+
+|   Model           | Accuracy [%]             |  F1 [%]      |
+|-------------------|--------------------------|--------------|
+| mBERT             | 92.157                   | 90.898       |
+| BERTje            | 93.096                   | 91.279       |
+| RobBERT v2        | **97.816**               | **97.514**   |
+
+#### Using zero-shot word masking task
+
+Since BERT models are pre-trained using the word masking task, we can use this to predict whether "die" or "dat" is more likely.
+This experiment shows that RobBERT has internalised more information about Dutch than other models.
+
+|   Model           | Accuracy [%]             |
+|-------------------|--------------------------|
+| ZeroR             | 66.70                    |
+| mBERT             | 90.21                    |
+| BERTje            | 94.94                    |
+| RobBERT v2        | **98.75**                |
+
+### Part-of-Speech Tagging.
+
+Using the [Lassy UD dataset](https://universaldependencies.org/treebanks/nl_lassysmall/index.html).
+
+
+|   Model           | Accuracy [%]             |
+|-------------------|--------------------------|
+| Frog              | 91.7                     |
+| mBERT             | **96.5**                 |
+| BERTje            | 96.3                     |
+| RobBERT v2        | 96.4                     |
+
+Interestingly, we found that when dealing with **small data sets**, RobBERT v2 **significantly outperforms** other models.
+
+<p align="center"> 
+    <img src="https://github.com/iPieter/RobBERT/blob/master/res/robbert_pos_accuracy.png" alt="RobBERT's performance on smaller datasets">
+ </p>
+
+### Named Entity Recognition
+
+Using the [CoNLL 2002 evaluation script](https://www.clips.uantwerpen.be/conll2002/ner/).
+
+
+|   Model           | Accuracy [%]             |
+|-------------------|--------------------------|
+| Frog              | 57.31                    |
+| mBERT             | **90.94**                |
+| BERT-NL           | 89.7                     |
+| BERTje            | 88.3                     |
+| RobBERT v2        | 89.08                    |
+
+
+## Training procedure
+
+We pre-trained RobBERT using the RoBERTa training regime.
+We pre-trained our model on the Dutch section of the [OSCAR corpus](https://oscar-corpus.com/), a large multilingual corpus which was obtained by language classification in the Common Crawl corpus.
+This Dutch corpus is 39GB large, with 6.6 billion words spread over 126 million lines of text, where each line could contain multiple sentences, thus using more data than concurrently developed Dutch BERT models.
+
+
+RobBERT shares its architecture with [RoBERTa's base model](https://github.com/pytorch/fairseq/tree/master/examples/roberta), which itself is a replication and improvement over BERT.
+Like BERT, it's architecture consists of 12 self-attention layers with 12 heads with 117M trainable parameters.
+One difference with the original BERT model is due to the different pre-training task specified by RoBERTa, using only the MLM task and not the NSP task.
+During pre-training, it thus only predicts which words are masked in certain positions of given sentences.
+The training process uses the Adam optimizer with polynomial decay of the learning rate l_r=10^-6 and a ramp-up period of 1000 iterations, with hyperparameters beta_1=0.9
+and RoBERTa's default beta_2=0.98.
+Additionally, a weight decay of 0.1 and a small dropout of 0.1 helps prevent the model from overfitting. 
+
+
+RobBERT was trained on a computing cluster with 4 Nvidia P100 GPUs per node, where the number of nodes was dynamically adjusted while keeping a fixed batch size of 8192 sentences.
+At most 20 nodes were used (i.e. 80 GPUs), and the median was 5 nodes.
+By using gradient accumulation, the batch size could be set independently of the number of GPUs available, in order to maximally utilize the cluster.
+Using the [Fairseq library](https://github.com/pytorch/fairseq/tree/master/examples/roberta), the model trained for two epochs, which equals over 16k batches in total, which took about three days on the computing cluster.
+In between training jobs on the computing cluster, 2 Nvidia 1080 Ti's also covered some parameter updates for RobBERT v2.
+
+
+## Limitations and bias
+
+In the [RobBERT paper](https://arxiv.org/abs/2001.06286), we also investigated potential sources of bias in RobBERT.
+
+We found that the zeroshot model estimates the probability of *hij* (he) to be higher than *zij* (she) for most occupations in bleached template sentences, regardless of their actual job gender ratio in reality.
+
+<p align="center"> 
+    <img src="https://github.com/iPieter/RobBERT/blob/master/res/gender_diff.png" alt="RobBERT's performance on smaller datasets">
+ </p>
+
+By augmenting the DBRB Dutch Book sentiment analysis dataset with the stated gender of the author of the review, we found that highly positive reviews written by women were generally more accurately detected by RobBERT as being positive than those written by men.
+
+<p align="center"> 
+    <img src="https://github.com/iPieter/RobBERT/blob/master/res/dbrd.png" alt="RobBERT's performance on smaller datasets">
+ </p>
+
+
+
+## BibTeX entry and citation info
+
+```bibtex
+@misc{delobelle2020robbert,
+    title={RobBERT: a Dutch RoBERTa-based Language Model},
+    author={Pieter Delobelle and Thomas Winters and Bettina Berendt},
+    year={2020},
+    eprint={2001.06286},
+    archivePrefix={arXiv},
+    primaryClass={cs.CL}
+}
+```

setup.cfg

@@ -10,7 +10,7 @@ known_third_party =
     datasets
     elasticsearch
     fairseq
-    faiss
+    faiss-cpu
     fastprogress
     fire
     fugashi

setup.py

@@ -89,7 +89,8 @@ extras["onnxruntime"] = ["onnxruntime>=1.4.0", "onnxruntime-tools>=1.4.2"]
 extras["serving"] = ["pydantic", "uvicorn", "fastapi", "starlette"]
 extras["all"] = extras["serving"] + ["tensorflow", "torch"]
 
-extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator", "psutil", "parameterized"]
+extras["retrieval"] = ["faiss-cpu", "datasets"]
+extras["testing"] = ["pytest", "pytest-xdist", "timeout-decorator", "parameterized", "psutil"] + extras["retrieval"]
 # sphinx-rtd-theme==0.5.0 introduced big changes in the style.
 extras["docs"] = ["recommonmark", "sphinx", "sphinx-markdown-tables", "sphinx-rtd-theme==0.4.3", "sphinx-copybutton"]
 extras["quality"] = ["black >= 20.8b1", "isort >= 5", "flake8 >= 3.8.3"]
@@ -97,7 +98,7 @@ extras["dev"] = extras["testing"] + extras["quality"] + extras["ja"] + ["scikit-
 
 setup(
     name="transformers",
-    version="3.1.0",
+    version="3.2.0",
     author="Thomas Wolf, Lysandre Debut, Victor Sanh, Julien Chaumond, Sam Shleifer, Patrick von Platen, Sylvain Gugger, Google AI Language Team Authors, Open AI team Authors, Facebook AI Authors, Carnegie Mellon University Authors",
     author_email="thomas@huggingface.co",
     description="State-of-the-art Natural Language Processing for TensorFlow 2.0 and PyTorch",

src/transformers/__init__.py

@@ -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.1.0"
+__version__ = "3.2.0"
 
 # Work around to update TensorFlow's absl.logging threshold which alters the
 # default Python logging output behavior when present.
@@ -33,6 +33,7 @@ from .configuration_flaubert import FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, Flau
 from .configuration_fsmt import FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP, FSMTConfig
 from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
 from .configuration_gpt2 import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPT2Config
+from .configuration_layoutlm import LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMConfig
 from .configuration_longformer import LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, LongformerConfig
 from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
 from .configuration_marian import MarianConfig
@@ -41,6 +42,7 @@ from .configuration_mmbt import MMBTConfig
 from .configuration_mobilebert import MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, MobileBertConfig
 from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
 from .configuration_pegasus import PegasusConfig
+from .configuration_rag import RagConfig
 from .configuration_reformer import REFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, ReformerConfig
 from .configuration_retribert import RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RetriBertConfig
 from .configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig
@@ -85,6 +87,7 @@ from .file_utils import (
     cached_path,
     is_apex_available,
     is_datasets_available,
+    is_faiss_available,
     is_psutil_available,
     is_py3nvml_available,
     is_tf_available,
@@ -117,27 +120,30 @@ from .modeling_tf_pytorch_utils import (
 )
 
 # Pipelines
-from .pipelines import (
-    Conversation,
-    ConversationalPipeline,
-    CsvPipelineDataFormat,
-    FeatureExtractionPipeline,
-    FillMaskPipeline,
-    JsonPipelineDataFormat,
-    NerPipeline,
-    PipedPipelineDataFormat,
-    Pipeline,
-    PipelineDataFormat,
-    QuestionAnsweringPipeline,
-    SummarizationPipeline,
-    Text2TextGenerationPipeline,
-    TextClassificationPipeline,
-    TextGenerationPipeline,
-    TokenClassificationPipeline,
-    TranslationPipeline,
-    ZeroShotClassificationPipeline,
-    pipeline,
-)
+# from .pipelines import (
+#     Conversation,
+#     ConversationalPipeline,
+#     CsvPipelineDataFormat,
+#     FeatureExtractionPipeline,
+#     FillMaskPipeline,
+#     JsonPipelineDataFormat,
+#     NerPipeline,
+#     PipedPipelineDataFormat,
+#     Pipeline,
+#     PipelineDataFormat,
+#     QuestionAnsweringPipeline,
+#     SummarizationPipeline,
+#     Text2TextGenerationPipeline,
+#     TextClassificationPipeline,
+#     TextGenerationPipeline,
+#     TokenClassificationPipeline,
+#     TranslationPipeline,
+#     ZeroShotClassificationPipeline,
+#     pipeline,
+# )
+
+# Retriever
+from .retrieval_rag import RagRetriever
 
 # Tokenizers
 from .tokenization_albert import AlbertTokenizer
@@ -163,6 +169,7 @@ from .tokenization_flaubert import FlaubertTokenizer
 from .tokenization_fsmt import FSMTTokenizer
 from .tokenization_funnel import FunnelTokenizer, FunnelTokenizerFast
 from .tokenization_gpt2 import GPT2Tokenizer, GPT2TokenizerFast
+from .tokenization_layoutlm import LayoutLMTokenizer, LayoutLMTokenizerFast
 from .tokenization_longformer import LongformerTokenizer, LongformerTokenizerFast
 from .tokenization_lxmert import LxmertTokenizer, LxmertTokenizerFast
 from .tokenization_mbart import MBartTokenizer
@@ -170,6 +177,7 @@ from .tokenization_mobilebert import MobileBertTokenizer, MobileBertTokenizerFas
 from .tokenization_openai import OpenAIGPTTokenizer, OpenAIGPTTokenizerFast
 from .tokenization_pegasus import PegasusTokenizer
 from .tokenization_phobert import PhobertTokenizer
+from .tokenization_rag import RagTokenizer
 from .tokenization_reformer import ReformerTokenizer
 from .tokenization_retribert import RetriBertTokenizer, RetriBertTokenizerFast
 from .tokenization_roberta import RobertaTokenizer, RobertaTokenizerFast
@@ -363,6 +371,12 @@ if is_torch_available():
         GPT2PreTrainedModel,
         load_tf_weights_in_gpt2,
     )
+    from .modeling_layoutlm import (
+        LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST,
+        LayoutLMForMaskedLM,
+        LayoutLMForTokenClassification,
+        LayoutLMModel,
+    )
     from .modeling_longformer import (
         LONGFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
         LongformerForMaskedLM,
@@ -408,6 +422,7 @@ if is_torch_available():
         load_tf_weights_in_openai_gpt,
     )
     from .modeling_pegasus import PegasusForConditionalGeneration
+    from .modeling_rag import RagModel, RagSequenceForGeneration, RagTokenForGeneration
     from .modeling_reformer import (
         REFORMER_PRETRAINED_MODEL_ARCHIVE_LIST,
         ReformerAttention,

src/transformers/activations.py

@@ -55,3 +55,7 @@ def get_activation(activation_string):
         return ACT2FN[activation_string]
     else:
         raise KeyError("function {} not found in ACT2FN mapping {}".format(activation_string, list(ACT2FN.keys())))
+
+
+def mish(x):
+    return x * torch.tanh(torch.nn.functional.softplus(x))

src/transformers/configuration_auto.py

@@ -24,12 +24,14 @@ from .configuration_bert_generation import BertGenerationConfig
 from .configuration_camembert import CAMEMBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, CamembertConfig
 from .configuration_ctrl import CTRL_PRETRAINED_CONFIG_ARCHIVE_MAP, CTRLConfig
 from .configuration_distilbert import DISTILBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, DistilBertConfig
+from .configuration_dpr import DPR_PRETRAINED_CONFIG_ARCHIVE_MAP, DPRConfig
 from .configuration_electra import ELECTRA_PRETRAINED_CONFIG_ARCHIVE_MAP, ElectraConfig
 from .configuration_encoder_decoder import EncoderDecoderConfig
 from .configuration_flaubert import FLAUBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, FlaubertConfig
 from .configuration_fsmt import FSMT_PRETRAINED_CONFIG_ARCHIVE_MAP, FSMTConfig
 from .configuration_funnel import FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP, FunnelConfig
 from .configuration_gpt2 import GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP, GPT2Config
+from .configuration_layoutlm import LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP, LayoutLMConfig
 from .configuration_longformer import LONGFORMER_PRETRAINED_CONFIG_ARCHIVE_MAP, LongformerConfig
 from .configuration_lxmert import LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP, LxmertConfig
 from .configuration_marian import MarianConfig
@@ -37,6 +39,7 @@ from .configuration_mbart import MBART_PRETRAINED_CONFIG_ARCHIVE_MAP, MBartConfi
 from .configuration_mobilebert import MobileBertConfig
 from .configuration_openai import OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP, OpenAIGPTConfig
 from .configuration_pegasus import PegasusConfig
+from .configuration_rag import RagConfig
 from .configuration_reformer import ReformerConfig
 from .configuration_retribert import RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP, RetriBertConfig
 from .configuration_roberta import ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP, RobertaConfig
@@ -73,6 +76,8 @@ ALL_PRETRAINED_CONFIG_ARCHIVE_MAP = dict(
         RETRIBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
         FUNNEL_PRETRAINED_CONFIG_ARCHIVE_MAP,
         LXMERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
+        LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP,
+        DPR_PRETRAINED_CONFIG_ARCHIVE_MAP,
     ]
     for key, value, in pretrained_map.items()
 )
@@ -108,6 +113,9 @@ CONFIG_MAPPING = OrderedDict(
         ("encoder-decoder", EncoderDecoderConfig),
         ("funnel", FunnelConfig),
         ("lxmert", LxmertConfig),
+        ("dpr", DPRConfig),
+        ("layoutlm", LayoutLMConfig),
+        ("rag", RagConfig),
     ]
 )
 
@@ -141,6 +149,9 @@ MODEL_NAMES_MAPPING = OrderedDict(
         ("encoder-decoder", "Encoder decoder"),
         ("funnel", "Funnel Transformer"),
         ("lxmert", "LXMERT"),
+        ("layoutlm", "LayoutLM"),
+        ("dpr", "DPR"),
+        ("rag", "RAG"),
     ]
 )
 

src/transformers/configuration_dpr.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """ DPR model configuration """
 
-from .configuration_bert import BertConfig
+from .configuration_utils import PretrainedConfig
 from .utils import logging
 
 
@@ -27,7 +27,7 @@ DPR_PRETRAINED_CONFIG_ARCHIVE_MAP = {
 }
 
 
-class DPRConfig(BertConfig):
+class DPRConfig(PretrainedConfig):
     r"""
     :class:`~transformers.DPRConfig` is the configuration class to store the configuration of a
     `DPRModel`.
@@ -36,12 +36,73 @@ class DPRConfig(BertConfig):
     It is used to instantiate the components of the DPR model.
 
     Args:
-        projection_dim (:obj:`int`, optional, defaults to 0):
+        vocab_size (:obj:`int`, `optional`, defaults to 30522):
+            Vocabulary size of the DPR model. Defines the different tokens that
+            can be represented by the `inputs_ids` passed to the forward method of :class:`~transformers.BertModel`.
+        hidden_size (:obj:`int`, `optional`, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_hidden_layers (:obj:`int`, `optional`, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (:obj:`int`, `optional`, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (:obj:`int`, `optional`, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        hidden_act (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler.
+            If string, :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
+            The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
+            The dropout ratio for the attention probabilities.
+        max_position_embeddings (:obj:`int`, `optional`, defaults to 512):
+            The maximum sequence length that this model might ever be used with.
+            Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
+        type_vocab_size (:obj:`int`, `optional`, defaults to 2):
+            The vocabulary size of the `token_type_ids` passed into :class:`~transformers.BertModel`.
+        initializer_range (:obj:`float`, `optional`, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_eps (:obj:`float`, `optional`, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If True, use gradient checkpointing to save memory at the expense of slower backward pass.
+        projection_dim (:obj:`int`, `optional`, defaults to 0):
             Dimension of the projection for the context and question encoders.
             If it is set to zero (default), then no projection is done.
     """
     model_type = "dpr"
 
-    def __init__(self, projection_dim: int = 0, **kwargs):  # projection of the encoders, 0 for no projection
-        super().__init__(**kwargs)
+    def __init__(
+        self,
+        vocab_size=30522,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=512,
+        type_vocab_size=2,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        pad_token_id=0,
+        gradient_checkpointing=False,
+        projection_dim: int = 0,
+        **kwargs
+    ):
+        super().__init__(pad_token_id=pad_token_id, **kwargs)
+
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.hidden_act = hidden_act
+        self.intermediate_size = intermediate_size
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.initializer_range = initializer_range
+        self.layer_norm_eps = layer_norm_eps
+        self.gradient_checkpointing = gradient_checkpointing
         self.projection_dim = projection_dim

src/transformers/configuration_layoutlm.py

@@ -0,0 +1,128 @@
+# coding=utf-8
+# Copyright 2010, The Microsoft Research Asia LayoutLM Team authors
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" LayoutLM model configuration """
+
+
+from .configuration_bert import BertConfig
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+LAYOUTLM_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+    "layoutlm-base-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/microsoft/layoutlm-base-uncased/config.json",
+    "layoutlm-large-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/microsoft/layoutlm-large-uncased/config.json",
+}
+
+
+class LayoutLMConfig(BertConfig):
+    r"""
+    This is the configuration class to store the configuration of a :class:`~transformers.LayoutLMModel`.
+    It is used to instantiate a LayoutLM model according to the specified arguments, defining the model
+    architecture. Instantiating a configuration with the defaults will yield a similar configuration to that of
+    the LayoutLM `layoutlm-base-uncased <https://huggingface.co/microsoft/layoutlm-base-uncased>`__ architecture.
+
+    Configuration objects inherit from :class:`~transformers.BertConfig` and can be used
+    to control the model outputs. Read the documentation from :class:`~transformers.BertConfig`
+    for more information.
+
+
+    Args:
+        vocab_size (:obj:`int`, `optional`, defaults to 30522):
+            Vocabulary size of the LayoutLM model. Defines the different tokens that
+            can be represented by the `inputs_ids` passed to the forward method of :class:`~transformers.LayoutLMModel`.
+        hidden_size (:obj:`int`, `optional`, defaults to 768):
+            Dimensionality of the encoder layers and the pooler layer.
+        num_hidden_layers (:obj:`int`, `optional`, defaults to 12):
+            Number of hidden layers in the Transformer encoder.
+        num_attention_heads (:obj:`int`, `optional`, defaults to 12):
+            Number of attention heads for each attention layer in the Transformer encoder.
+        intermediate_size (:obj:`int`, `optional`, defaults to 3072):
+            Dimensionality of the "intermediate" (i.e., feed-forward) layer in the Transformer encoder.
+        hidden_act (:obj:`str` or :obj:`function`, `optional`, defaults to :obj:`"gelu"`):
+            The non-linear activation function (function or string) in the encoder and pooler.
+            If string, :obj:`"gelu"`, :obj:`"relu"`, :obj:`"swish"` and :obj:`"gelu_new"` are supported.
+        hidden_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
+            The dropout probabilitiy for all fully connected layers in the embeddings, encoder, and pooler.
+        attention_probs_dropout_prob (:obj:`float`, `optional`, defaults to 0.1):
+            The dropout ratio for the attention probabilities.
+        max_position_embeddings (:obj:`int`, `optional`, defaults to 512):
+            The maximum sequence length that this model might ever be used with.
+            Typically set this to something large just in case (e.g., 512 or 1024 or 2048).
+        type_vocab_size (:obj:`int`, `optional`, defaults to 2):
+            The vocabulary size of the :obj:`token_type_ids` passed into :class:`~transformers.LayoutLMModel`.
+        initializer_range (:obj:`float`, `optional`, defaults to 0.02):
+            The standard deviation of the truncated_normal_initializer for initializing all weight matrices.
+        layer_norm_eps (:obj:`float`, `optional`, defaults to 1e-12):
+            The epsilon used by the layer normalization layers.
+        gradient_checkpointing (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If True, use gradient checkpointing to save memory at the expense of slower backward pass.
+        max_2d_position_embeddings (:obj:`int`, `optional`, defaults to 1024):
+            The maximum value that the 2D position embedding might ever used.
+            Typically set this to something large just in case (e.g., 1024).
+
+    Examples::
+
+        >>> from transformers import LayoutLMModel, LayoutLMConfig
+
+        >>> # Initializing a LayoutLM configuration
+        >>> configuration = LayoutLMConfig()
+
+        >>> # Initializing a model from the configuration
+        >>> model = LayoutLMModel(configuration)
+
+        >>> # Accessing the model configuration
+        >>> configuration = model.config
+
+    """
+    model_type = "layoutlm"
+
+    def __init__(
+        self,
+        vocab_size=30522,
+        hidden_size=768,
+        num_hidden_layers=12,
+        num_attention_heads=12,
+        intermediate_size=3072,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=512,
+        type_vocab_size=2,
+        initializer_range=0.02,
+        layer_norm_eps=1e-12,
+        pad_token_id=0,
+        gradient_checkpointing=False,
+        max_2d_position_embeddings=1024,
+        **kwargs
+    ):
+        super().__init__(
+            vocab_size=vocab_size,
+            hidden_size=hidden_size,
+            num_hidden_layers=num_hidden_layers,
+            num_attention_heads=num_attention_heads,
+            intermediate_size=intermediate_size,
+            hidden_act=hidden_act,
+            hidden_dropout_prob=hidden_dropout_prob,
+            attention_probs_dropout_prob=attention_probs_dropout_prob,
+            max_position_embeddings=max_position_embeddings,
+            type_vocab_size=type_vocab_size,
+            initializer_range=initializer_range,
+            layer_norm_eps=layer_norm_eps,
+            pad_token_id=pad_token_id,
+            gradient_checkpointing=gradient_checkpointing,
+            **kwargs,
+        )
+        self.max_2d_position_embeddings = max_2d_position_embeddings

src/transformers/configuration_rag.py

@@ -0,0 +1,176 @@
+# coding=utf-8
+# Copyright 2020, The RAG Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" RAG model configuration """
+
+import copy
+
+from .configuration_utils import PretrainedConfig
+from .file_utils import add_start_docstrings
+
+
+RAG_CONFIG_DOC = r"""
+    :class:`~transformers.RagConfig` stores the configuration of a `RagModel`.
+    Configuration objects inherit from  :class:`~transformers.PretrainedConfig` and can be used
+    to control the model outputs. Read the documentation from  :class:`~transformers.PretrainedConfig`
+    for more information.
+
+    Args:
+        title_sep (:obj:`str`, `optional`, defaults to  ``" / "``):
+            Separator inserted between the title and the text of the retrieved document when calling :class:`~transformers.RagRetriever`.
+        doc_sep (:obj:`str`, `optional`, defaults to  ``" // "``):
+            Separator inserted between the the text of the retrieved document and the original input when calliang :class:`~transformers.RagRetriever`.
+        n_docs (:obj:`int`, `optional`, defaults to 5):
+            Number of documents to retrieve.
+        max_combined_length (:obj:`int`, `optional`, defaults to 300):
+            Max length of contextualized input returned by :meth:`~transformers.RagRetriever.__call__`.
+        retrieval_vector_size (:obj:`int`, `optional`, defaults to 768):
+            Dimensionality of the document embeddings indexed by :class:`~transformers.RagRetriever`.
+        retrieval_batch_size (:obj:`int`, `optional`, defaults to 8):
+            Retrieval batch size, defined as the number of queries issues concurrently to the faiss index excapsulated :class:`~transformers.RagRetriever`.
+        dataset (:obj:`str`, `optional`, defaults to :obj:`"wiki_dpr"`):
+            A datatset identifier of the indexed dataset on HuggingFace AWS bucket (list all available datasets and ids using :obj:`datasets.list_datasets()`).
+        dataset_split (:obj:`str`, `optional`, defaults to :obj:`train`)
+            Which split of the ``dataset`` to load.
+        index_name (:obj:`str`, `optional`, defaults to :obj:`compressed`)
+            The index_name of the index associated with the :obj:`dataset`. One can choose between :obj:`legacy`, :obj:`exact` and :obj:`compressed`.
+        index_path (:obj:`str`, `optional`)
+            The path to the serialized faiss index on disk.
+        passages_path: (:obj:`str`, `optional`):
+            A path to text passages compatible with the faiss index. Required if using :class:`~transformers.retrieval_rag.LegacyIndex`
+        use_dummy_dataset (:obj:`bool`, `optional`, defaults to ``False``)
+            Whether to load a "dummy" variant of the dataset specified by :obj:`dataset`.
+        label_smoothing (:obj:`float`, `optional`, defaults to 0.0):
+            Only relevant if ``return_loss`` is set to :obj:`True`. Controls the ``epsilon`` parameter value for label smoothing in the loss calculation.
+            If set to ``0.0``, no label smoothing is performed.
+        do_marginalize (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If :obj:`True`, the logits are marginalized over all documents
+            by making use of ``torch.nn.functional.log_softmax``.
+        reduce_loss (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If :obj:`True`, the NLL loss is reduced using the ``torch.Tensor.sum`` operation.
+        do_deduplication (:obj:`bool`, `optional`, defaults to :obj:`True`):
+            Controls whether we want to deduplicate the generations from different context documents for a given input.
+            Has to be set to :obj:`False` if used while training with distributed backend.
+        exclude_bos_score (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If :obj:`True`, the score of the BOS token is disregarded when computing
+            the loss.
+        output_retrieved(:obj:`bool`, `optional`, defaults to :obj:`False`):
+            If set to ``True``, :obj:`retrieved_doc_embeds`, :obj:`retrieved_doc_ids`, :obj:`context_input_ids` and :obj:`context_attention_mask` are returned. See returned tensors for more detail.
+"""
+
+
+@add_start_docstrings(RAG_CONFIG_DOC)
+class RagConfig(PretrainedConfig):
+    model_type = "rag"
+
+    def __init__(
+        self,
+        vocab_size=None,
+        is_encoder_decoder=True,
+        prefix=None,
+        bos_token_id=None,
+        pad_token_id=None,
+        eos_token_id=None,
+        decoder_start_token_id=None,
+        title_sep=" / ",
+        doc_sep=" // ",
+        n_docs=5,
+        max_combined_length=300,
+        retrieval_vector_size=768,
+        retrieval_batch_size=8,
+        dataset="wiki_dpr",
+        dataset_split="train",
+        index_name="compressed",
+        index_path=None,
+        passages_path=None,
+        use_dummy_dataset=False,
+        reduce_loss=False,
+        label_smoothing=0.0,
+        do_deduplication=True,
+        exclude_bos_score=False,
+        do_marginalize=False,
+        output_retrieved=False,
+        **kwargs
+    ):
+        super().__init__(
+            bos_token_id=bos_token_id,
+            pad_token_id=pad_token_id,
+            eos_token_id=eos_token_id,
+            decoder_start_token_id=decoder_start_token_id,
+            is_encoder_decoder=is_encoder_decoder,
+            prefix=prefix,
+            vocab_size=vocab_size,
+            **kwargs,
+        )
+        assert (
+            "question_encoder" in kwargs and "generator" in kwargs
+        ), "Config has to be initialized with question_encoder and generator config"
+        question_encoder_config = kwargs.pop("question_encoder")
+        question_encoder_model_type = question_encoder_config.pop("model_type")
+        decoder_config = kwargs.pop("generator")
+        decoder_model_type = decoder_config.pop("model_type")
+
+        from .configuration_auto import AutoConfig
+
+        self.question_encoder = AutoConfig.for_model(question_encoder_model_type, **question_encoder_config)
+        self.generator = AutoConfig.for_model(decoder_model_type, **decoder_config)
+
+        self.reduce_loss = reduce_loss
+        self.label_smoothing = label_smoothing
+        self.exclude_bos_score = exclude_bos_score
+        self.do_marginalize = do_marginalize
+
+        self.title_sep = title_sep
+        self.doc_sep = doc_sep
+        self.n_docs = n_docs
+        self.max_combined_length = max_combined_length
+
+        self.dataset = dataset
+        self.dataset_split = dataset_split
+        self.index_name = index_name
+
+        self.retrieval_vector_size = retrieval_vector_size
+        self.retrieval_batch_size = retrieval_batch_size
+        self.passages_path = passages_path
+        self.index_path = index_path
+        self.use_dummy_dataset = use_dummy_dataset
+
+        self.output_retrieved = output_retrieved
+
+        self.do_deduplication = do_deduplication
+
+    @classmethod
+    def from_question_encoder_generator_configs(
+        cls, question_encoder_config: PretrainedConfig, generator_config: PretrainedConfig, **kwargs
+    ) -> PretrainedConfig:
+        r"""
+        Instantiate a :class:`~transformers.EncoderDecoderConfig` (or a derived class) from a pre-trained encoder model configuration and decoder model configuration.
+
+        Returns:
+            :class:`EncoderDecoderConfig`: An instance of a configuration object
+        """
+        return cls(question_encoder=question_encoder_config.to_dict(), generator=generator_config.to_dict(), **kwargs)
+
+    def to_dict(self):
+        """
+        Serializes this instance to a Python dictionary. Override the default :meth:`~transformers.PretrainedConfig.to_dict`.
+
+        Returns:
+            :obj:`Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+        """
+        output = copy.deepcopy(self.__dict__)
+        output["question_encoder"] = self.question_encoder.to_dict()
+        output["generator"] = self.generator.to_dict()
+        output["model_type"] = self.__class__.model_type
+        return output

src/transformers/data/data_collator.py

@@ -1,4 +1,3 @@
-import random
 from dataclasses import dataclass
 from typing import Any, Callable, Dict, List, NewType, Optional, Tuple, Union
 
@@ -402,8 +401,8 @@ class DataCollatorForPermutationLanguageModeling:
 @dataclass
 class DataCollatorForNextSentencePrediction:
     """
-    Data collator used for language modeling.
-    - collates batches of tensors, honoring their tokenizer's pad_token
+    Data collator used for next sentence prediction.
+    - collates examples which contains pre-generated negative examples
     - preprocesses batches for masked language modeling
     """
 
@@ -414,21 +413,30 @@ class DataCollatorForNextSentencePrediction:
     nsp_probability: float = 0.5
     mlm_probability: float = 0.15
 
-    def __call__(self, examples: List[Union[List[List[int]], Dict[str, torch.Tensor]]]) -> Dict[str, torch.Tensor]:
-        if isinstance(examples[0], (dict, BatchEncoding)):
-            examples = [e["input_ids"] for e in examples]
+    def __call__(self, examples: List[Dict[str, torch.Tensor]]) -> Dict[str, torch.Tensor]:
+        """
+        The input should contain negative examples, :class:`~transformers.DataCollatorForNextSentencePrediction` will not generate any negative examples.
+        Args:
+            examples (:obj:`List[Dict]`): Each dictionary should have the following keys:
+                  - ``tokens_a``: A sequence of tokens, which should appear before ``tokens_b`` in the text.
+                  - ``tokens_b``: A sequence of tokens, which should appear after ``tokens_a`` in the text.
+                  - ``is_random_next``: 1 if this pair is generated randomly, else 0.
+        """
+
+        tokens_a = [e["tokens_a"] for e in examples]
+        tokens_b = [e["tokens_b"] for e in examples]
+        nsp_labels = [1 if e["is_random_next"] else 0 for e in examples]
 
         input_ids = []
         segment_ids = []
         attention_masks = []
-        nsp_labels = []
 
-        for i, doc in enumerate(examples):
-            input_id, segment_id, attention_mask, label = self.create_examples_from_document(doc, i, examples)
-            input_ids.extend(input_id)
-            segment_ids.extend(segment_id)
-            attention_masks.extend(attention_mask)
-            nsp_labels.extend(label)
+        assert len(tokens_a) == len(tokens_b)
+        for i in range(len(tokens_a)):
+            input_id, attention_mask, segment_id = self.create_features_from_example(tokens_a[i], tokens_b[i])
+            input_ids.append(input_id)
+            segment_ids.append(segment_id)
+            attention_masks.append(attention_mask)
         if self.mlm:
             input_ids, mlm_labels = self.mask_tokens(self._tensorize_batch(input_ids))
         else:
@@ -438,6 +446,7 @@ class DataCollatorForNextSentencePrediction:
             "input_ids": input_ids,
             "attention_mask": self._tensorize_batch(attention_masks),
             "token_type_ids": self._tensorize_batch(segment_ids),
+            "masked_lm_labels": mlm_labels if self.mlm else None,
             "next_sentence_label": torch.tensor(nsp_labels),
         }
         if self.mlm:
@@ -457,111 +466,34 @@ class DataCollatorForNextSentencePrediction:
                 )
             return pad_sequence(examples, batch_first=True, padding_value=self.tokenizer.pad_token_id)
 
-    def create_examples_from_document(
-        self, document: List[List[int]], doc_index: int, examples: List[List[List[int]]]
-    ):
+    def create_features_from_example(self, tokens_a, tokens_b):
         """Creates examples for a single document."""
 
         max_num_tokens = self.block_size - self.tokenizer.num_special_tokens_to_add(pair=True)
 
-        # We *usually* want to fill up the entire sequence since we are padding
-        # to `block_size` anyways, so short sequences are generally wasted
-        # computation. However, we *sometimes*
-        # (i.e., short_seq_prob == 0.1 == 10% of the time) want to use shorter
-        # sequences to minimize the mismatch between pre-training and fine-tuning.
-        # The `target_seq_length` is just a rough target however, whereas
-        # `block_size` is a hard limit.
-        target_seq_length = max_num_tokens
-        if random.random() < self.short_seq_probability:
-            target_seq_length = random.randint(2, max_num_tokens)
+        tokens_a, tokens_b, _ = self.tokenizer.truncate_sequences(
+            tokens_a,
+            tokens_b,
+            num_tokens_to_remove=len(tokens_a) + len(tokens_b) - max_num_tokens,
+            truncation_strategy="longest_first",
+        )
 
-        current_chunk = []  # a buffer stored current working segments
-        current_length = 0
-        i = 0
-        input_ids = []
-        segment_ids = []
-        attention_masks = []
-        labels = []
-        while i < len(document):
-            segment = document[i]
-            current_chunk.append(segment)
-            current_length += len(segment)
-            if i == len(document) - 1 or current_length >= target_seq_length:
-                if current_chunk:
-                    # `a_end` is how many segments from `current_chunk` go into the `A`
-                    # (first) sentence.
-                    a_end = 1
-                    if len(current_chunk) >= 2:
-                        a_end = random.randint(1, len(current_chunk) - 1)
+        input_id = self.tokenizer.build_inputs_with_special_tokens(tokens_a, tokens_b)
+        attention_mask = [1] * len(input_id)
+        segment_id = self.tokenizer.create_token_type_ids_from_sequences(tokens_a, tokens_b)
+        assert len(input_id) <= self.block_size
 
-                    tokens_a = []
-                    for j in range(a_end):
-                        tokens_a.extend(current_chunk[j])
+        # pad
+        while len(input_id) < self.block_size:
+            input_id.append(0)
+            attention_mask.append(0)
+            segment_id.append(0)
 
-                    tokens_b = []
+        input_id = torch.tensor(input_id)
+        attention_mask = torch.tensor(attention_mask)
+        segment_id = torch.tensor(segment_id)
 
-                    if len(current_chunk) == 1 or random.random() < self.nsp_probability:
-                        is_random_next = True
-                        target_b_length = target_seq_length - len(tokens_a)
-
-                        # This should rarely go for more than one iteration for large
-                        # corpora. However, just to be careful, we try to make sure that
-                        # the random document is not the same as the document
-                        # we're processing. Also check to make sure that the random document
-                        # is not empty.
-                        for _ in range(10):
-                            random_document_index = random.randint(0, len(examples) - 1)
-                            if random_document_index != doc_index and len(examples[random_document_index]) > 0:
-                                break
-
-                        random_document = examples[random_document_index]
-                        random_start = random.randint(0, len(random_document) - 1)
-                        for j in range(random_start, len(random_document)):
-                            tokens_b.extend(random_document[j])
-                            if len(tokens_b) >= target_b_length:
-                                break
-                        # We didn't actually use these segments so we "put them back" so
-                        # they don't go to waste.
-                        num_unused_segments = len(current_chunk) - a_end
-                        i -= num_unused_segments
-                    # Actual next
-                    else:
-                        is_random_next = False
-                        for j in range(a_end, len(current_chunk)):
-                            tokens_b.extend(current_chunk[j])
-
-                    assert len(tokens_a) >= 1
-                    assert len(tokens_b) >= 1
-
-                    tokens_a, tokens_b, _ = self.tokenizer.truncate_sequences(
-                        tokens_a,
-                        tokens_b,
-                        num_tokens_to_remove=len(tokens_a) + len(tokens_b) - max_num_tokens,
-                        truncation_strategy="longest_first",
-                    )
-
-                    input_id = self.tokenizer.build_inputs_with_special_tokens(tokens_a, tokens_b)
-                    attention_mask = [1] * len(input_id)
-                    segment_id = self.tokenizer.create_token_type_ids_from_sequences(tokens_a, tokens_b)
-                    assert len(input_id) <= self.block_size
-
-                    # pad
-                    while len(input_id) < self.block_size:
-                        input_id.append(0)
-                        attention_mask.append(0)
-                        segment_id.append(0)
-
-                    input_ids.append(torch.tensor(input_id))
-                    segment_ids.append(torch.tensor(segment_id))
-                    attention_masks.append(torch.tensor(attention_mask))
-                    labels.append(torch.tensor(1 if is_random_next else 0))
-
-                current_chunk = []
-                current_length = 0
-
-            i += 1
-
-        return input_ids, segment_ids, attention_masks, labels
+        return input_id, attention_mask, segment_id
 
     def mask_tokens(self, inputs: torch.Tensor) -> Tuple[torch.Tensor, torch.Tensor]:
         """

src/transformers/data/datasets/language_modeling.py

@@ -2,7 +2,7 @@ import os
 import pickle
 import random
 import time
-from typing import Dict, Optional
+from typing import Dict, List, Optional
 
 import torch
 from torch.utils.data.dataset import Dataset
@@ -267,10 +267,14 @@ class TextDatasetForNextSentencePrediction(Dataset):
         file_path: str,
         block_size: int,
         overwrite_cache=False,
+        short_seq_probability=0.1,
+        nsp_probability=0.5,
     ):
         assert os.path.isfile(file_path), f"Input file path {file_path} not found"
 
-        block_size = block_size - tokenizer.num_special_tokens_to_add(pair=True)
+        self.block_size = block_size - tokenizer.num_special_tokens_to_add(pair=True)
+        self.short_seq_probability = short_seq_probability
+        self.nsp_probability = nsp_probability
 
         directory, filename = os.path.split(file_path)
         cached_features_file = os.path.join(
@@ -283,7 +287,6 @@ class TextDatasetForNextSentencePrediction(Dataset):
         )
 
         self.tokenizer = tokenizer
-        self.examples = []
 
         # Make sure only the first process in distributed training processes the dataset,
         # and the others will use the cache.
@@ -313,7 +316,7 @@ class TextDatasetForNextSentencePrediction(Dataset):
             else:
                 logger.info(f"Creating features from dataset file at {directory}")
 
-                self.examples = [[]]
+                self.documents = [[]]
                 with open(file_path, encoding="utf-8") as f:
                     while True:
                         line = f.readline()
@@ -322,12 +325,17 @@ class TextDatasetForNextSentencePrediction(Dataset):
                         line = line.strip()
 
                         # Empty lines are used as document delimiters
-                        if not line and len(self.examples[-1]) != 0:
-                            self.examples.append([])
+                        if not line and len(self.documents[-1]) != 0:
+                            self.documents.append([])
                         tokens = tokenizer.tokenize(line)
                         tokens = tokenizer.convert_tokens_to_ids(tokens)
                         if tokens:
-                            self.examples[-1].append(tokens)
+                            self.documents[-1].append(tokens)
+
+                logger.info(f"Creating examples from {len(self.documents)} documents.")
+                self.examples = []
+                for doc_index, document in enumerate(self.documents):
+                    self.create_examples_from_document(document, doc_index)
 
                 start = time.time()
                 with open(cached_features_file, "wb") as handle:
@@ -336,6 +344,85 @@ class TextDatasetForNextSentencePrediction(Dataset):
                     "Saving features into cached file %s [took %.3f s]", cached_features_file, time.time() - start
                 )
 
+    def create_examples_from_document(self, document: List[List[int]], doc_index: int):
+        """Creates examples for a single document."""
+
+        max_num_tokens = self.block_size - self.tokenizer.num_special_tokens_to_add(pair=True)
+
+        # We *usually* want to fill up the entire sequence since we are padding
+        # to `block_size` anyways, so short sequences are generally wasted
+        # computation. However, we *sometimes*
+        # (i.e., short_seq_prob == 0.1 == 10% of the time) want to use shorter
+        # sequences to minimize the mismatch between pre-training and fine-tuning.
+        # The `target_seq_length` is just a rough target however, whereas
+        # `block_size` is a hard limit.
+        target_seq_length = max_num_tokens
+        if random.random() < self.short_seq_probability:
+            target_seq_length = random.randint(2, max_num_tokens)
+
+        current_chunk = []  # a buffer stored current working segments
+        current_length = 0
+        i = 0
+
+        while i < len(document):
+            segment = document[i]
+            current_chunk.append(segment)
+            current_length += len(segment)
+            if i == len(document) - 1 or current_length >= target_seq_length:
+                if current_chunk:
+                    # `a_end` is how many segments from `current_chunk` go into the `A`
+                    # (first) sentence.
+                    a_end = 1
+                    if len(current_chunk) >= 2:
+                        a_end = random.randint(1, len(current_chunk) - 1)
+
+                    tokens_a = []
+                    for j in range(a_end):
+                        tokens_a.extend(current_chunk[j])
+
+                    tokens_b = []
+
+                    if len(current_chunk) == 1 or random.random() < self.nsp_probability:
+                        is_random_next = True
+                        target_b_length = target_seq_length - len(tokens_a)
+
+                        # This should rarely go for more than one iteration for large
+                        # corpora. However, just to be careful, we try to make sure that
+                        # the random document is not the same as the document
+                        # we're processing.
+                        for _ in range(10):
+                            random_document_index = random.randint(0, len(self.documents) - 1)
+                            if random_document_index != doc_index:
+                                break
+
+                        random_document = self.documents[random_document_index]
+                        random_start = random.randint(0, len(random_document) - 1)
+                        for j in range(random_start, len(random_document)):
+                            tokens_b.extend(random_document[j])
+                            if len(tokens_b) >= target_b_length:
+                                break
+                        # We didn't actually use these segments so we "put them back" so
+                        # they don't go to waste.
+                        num_unused_segments = len(current_chunk) - a_end
+                        i -= num_unused_segments
+                    # Actual next
+                    else:
+                        is_random_next = False
+                        for j in range(a_end, len(current_chunk)):
+                            tokens_b.extend(current_chunk[j])
+
+                    assert len(tokens_a) >= 1
+                    assert len(tokens_b) >= 1
+
+                    self.examples.append(
+                        {"tokens_a": tokens_a, "tokens_b": tokens_b, "is_random_next": is_random_next}
+                    )
+
+                current_chunk = []
+                current_length = 0
+
+            i += 1
+
     def __len__(self):
         return len(self.examples)
 

src/transformers/file_utils.py

@@ -69,6 +69,7 @@ try:
     import datasets  # noqa: F401
 
     _datasets_available = True
+    logger.debug(f"Succesfully imported datasets version {datasets.__version__}")
 
 except ImportError:
     _datasets_available = False
@@ -119,6 +120,16 @@ try:
 except ImportError:
     _has_apex = False
 
+
+try:
+    import faiss  # noqa: F401
+
+    _faiss_available = True
+    logger.debug(f"Succesfully imported faiss version {faiss.__version__}")
+except ImportError:
+    _faiss_available = False
+
+
 default_cache_path = os.path.join(torch_cache_home, "transformers")
 
 
@@ -175,6 +186,10 @@ def is_apex_available():
     return _has_apex
 
 
+def is_faiss_available():
+    return _faiss_available
+
+
 def add_start_docstrings(*docstr):
     def docstring_decorator(fn):
         fn.__doc__ = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")

src/transformers/modeling_auto.py

@@ -27,12 +27,14 @@ from .configuration_auto import (
     CamembertConfig,
     CTRLConfig,
     DistilBertConfig,
+    DPRConfig,
     ElectraConfig,
     EncoderDecoderConfig,
     FlaubertConfig,
     FSMTConfig,
     FunnelConfig,
     GPT2Config,
+    LayoutLMConfig,
     LongformerConfig,
     LxmertConfig,
     MBartConfig,
@@ -96,6 +98,7 @@ from .modeling_distilbert import (
     DistilBertForTokenClassification,
     DistilBertModel,
 )
+from .modeling_dpr import DPRQuestionEncoder
 from .modeling_electra import (
     ElectraForMaskedLM,
     ElectraForMultipleChoice,
@@ -124,6 +127,7 @@ from .modeling_funnel import (
     FunnelModel,
 )
 from .modeling_gpt2 import GPT2LMHeadModel, GPT2Model
+from .modeling_layoutlm import LayoutLMForMaskedLM, LayoutLMForTokenClassification, LayoutLMModel
 from .modeling_longformer import (
     LongformerForMaskedLM,
     LongformerForMultipleChoice,
@@ -146,6 +150,11 @@ from .modeling_mobilebert import (
 )
 from .modeling_openai import OpenAIGPTLMHeadModel, OpenAIGPTModel
 from .modeling_pegasus import PegasusForConditionalGeneration
+from .modeling_rag import (  # noqa: F401 - need to import all RagModels to be in globals() function
+    RagModel,
+    RagSequenceForGeneration,
+    RagTokenForGeneration,
+)
 from .modeling_reformer import (
     ReformerForMaskedLM,
     ReformerForQuestionAnswering,
@@ -206,6 +215,7 @@ MODEL_MAPPING = OrderedDict(
         (BartConfig, BartModel),
         (LongformerConfig, LongformerModel),
         (RobertaConfig, RobertaModel),
+        (LayoutLMConfig, LayoutLMModel),
         (BertConfig, BertModel),
         (OpenAIGPTConfig, OpenAIGPTModel),
         (GPT2Config, GPT2Model),
@@ -221,11 +231,13 @@ MODEL_MAPPING = OrderedDict(
         (FunnelConfig, FunnelModel),
         (LxmertConfig, LxmertModel),
         (BertGenerationConfig, BertGenerationEncoder),
+        (DPRConfig, DPRQuestionEncoder),
     ]
 )
 
 MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(
     [
+        (LayoutLMConfig, LayoutLMForMaskedLM),
         (RetriBertConfig, RetriBertModel),
         (T5Config, T5ForConditionalGeneration),
         (DistilBertConfig, DistilBertForMaskedLM),
@@ -252,12 +264,14 @@ MODEL_FOR_PRETRAINING_MAPPING = OrderedDict(
 
 MODEL_WITH_LM_HEAD_MAPPING = OrderedDict(
     [
+        (LayoutLMConfig, LayoutLMForMaskedLM),
         (T5Config, T5ForConditionalGeneration),
         (DistilBertConfig, DistilBertForMaskedLM),
         (AlbertConfig, AlbertForMaskedLM),
         (CamembertConfig, CamembertForMaskedLM),
         (XLMRobertaConfig, XLMRobertaForMaskedLM),
         (MarianConfig, MarianMTModel),
+        (FSMTConfig, FSMTForConditionalGeneration),
         (BartConfig, BartForConditionalGeneration),
         (LongformerConfig, LongformerForMaskedLM),
         (RobertaConfig, RobertaForMaskedLM),
@@ -299,6 +313,7 @@ MODEL_FOR_CAUSAL_LM_MAPPING = OrderedDict(
 
 MODEL_FOR_MASKED_LM_MAPPING = OrderedDict(
     [
+        (LayoutLMConfig, LayoutLMForMaskedLM),
         (DistilBertConfig, DistilBertForMaskedLM),
         (AlbertConfig, AlbertForMaskedLM),
         (BartConfig, BartForConditionalGeneration),
@@ -369,6 +384,7 @@ MODEL_FOR_QUESTION_ANSWERING_MAPPING = OrderedDict(
 
 MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING = OrderedDict(
     [
+        (LayoutLMConfig, LayoutLMForTokenClassification),
         (DistilBertConfig, DistilBertForTokenClassification),
         (CamembertConfig, CamembertForTokenClassification),
         (FlaubertConfig, FlaubertForTokenClassification),
@@ -404,7 +420,6 @@ MODEL_FOR_MULTIPLE_CHOICE_MAPPING = OrderedDict(
     ]
 )
 
-
 AUTO_MODEL_PRETRAINED_DOCSTRING = r"""
 
         The model class to instantiate is selected based on the :obj:`model_type` property of the config object

src/transformers/modeling_dpr.py

@@ -272,6 +272,7 @@ class DPRPretrainedContextEncoder(PreTrainedModel):
     config_class = DPRConfig
     load_tf_weights = None
     base_model_prefix = "ctx_encoder"
+    authorized_missing_keys = [r"position_ids"]
 
     def init_weights(self):
         self.ctx_encoder.init_weights()
@@ -285,6 +286,7 @@ class DPRPretrainedQuestionEncoder(PreTrainedModel):
     config_class = DPRConfig
     load_tf_weights = None
     base_model_prefix = "question_encoder"
+    authorized_missing_keys = [r"position_ids"]
 
     def init_weights(self):
         self.question_encoder.init_weights()
@@ -298,6 +300,7 @@ class DPRPretrainedReader(PreTrainedModel):
     config_class = DPRConfig
     load_tf_weights = None
     base_model_prefix = "span_predictor"
+    authorized_missing_keys = [r"position_ids"]
 
     def init_weights(self):
         self.span_predictor.encoder.init_weights()

src/transformers/modeling_fsmt.py

@@ -1150,13 +1150,14 @@ class SinusoidalPositionalEmbedding(nn.Embedding):
     def __init__(self, num_positions, embedding_dim, padding_idx):
         self.make_weight(num_positions, embedding_dim, padding_idx)
 
-    def make_weight(self, num_positions, embedding_dim, padding_idx, device=None):
+    def make_weight(self, num_positions, embedding_dim, padding_idx):
         weight = self.get_embedding(num_positions, embedding_dim, padding_idx)
-        if device is not None:
-            weight = weight.to(device)
         if not hasattr(self, "weight"):
+            # in ___init__
             super().__init__(num_positions, embedding_dim, padding_idx, _weight=weight)
         else:
+            # in forward
+            weight = weight.to(self.weight.device)
             self.weight = nn.Parameter(weight)
         self.weight.detach_()
         self.weight.requires_grad = False
@@ -1204,6 +1205,6 @@ class SinusoidalPositionalEmbedding(nn.Embedding):
         max_pos = self.padding_idx + 1 + seq_len
         if max_pos > self.weight.size(0):
             # expand embeddings if needed
-            self.make_weight(max_pos, self.embedding_dim, self.padding_idx, device=input.device)
+            self.make_weight(max_pos, self.embedding_dim, self.padding_idx)
         positions = self.make_positions(input, self.padding_idx)
         return super().forward(positions)

src/transformers/modeling_layoutlm.py

@@ -0,0 +1,889 @@
+# coding=utf-8
+# Copyright 2018 The Microsoft Research Asia LayoutLM Team Authors and the HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" PyTorch LayoutLM model. """
+
+
+import math
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss
+
+from .activations import ACT2FN
+from .configuration_layoutlm import LayoutLMConfig
+from .file_utils import add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_callable
+from .modeling_outputs import BaseModelOutput, BaseModelOutputWithPooling, MaskedLMOutput, TokenClassifierOutput
+from .modeling_utils import (
+    PreTrainedModel,
+    apply_chunking_to_forward,
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+_CONFIG_FOR_DOC = "LayoutLMConfig"
+_TOKENIZER_FOR_DOC = "LayoutLMTokenizer"
+
+LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST = [
+    "layoutlm-base-uncased",
+    "layoutlm-large-uncased",
+]
+
+
+LayoutLMLayerNorm = torch.nn.LayerNorm
+
+
+class LayoutLMEmbeddings(nn.Module):
+    """Construct the embeddings from word, position and token_type embeddings."""
+
+    def __init__(self, config):
+        super(LayoutLMEmbeddings, self).__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.x_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size)
+        self.y_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size)
+        self.h_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size)
+        self.w_position_embeddings = nn.Embedding(config.max_2d_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+        self.LayerNorm = LayoutLMLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+
+    def forward(
+        self,
+        input_ids=None,
+        bbox=None,
+        token_type_ids=None,
+        position_ids=None,
+        inputs_embeds=None,
+    ):
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, :seq_length]
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+
+        words_embeddings = inputs_embeds
+        position_embeddings = self.position_embeddings(position_ids)
+        left_position_embeddings = self.x_position_embeddings(bbox[:, :, 0])
+        upper_position_embeddings = self.y_position_embeddings(bbox[:, :, 1])
+        right_position_embeddings = self.x_position_embeddings(bbox[:, :, 2])
+        lower_position_embeddings = self.y_position_embeddings(bbox[:, :, 3])
+        h_position_embeddings = self.h_position_embeddings(bbox[:, :, 3] - bbox[:, :, 1])
+        w_position_embeddings = self.w_position_embeddings(bbox[:, :, 2] - bbox[:, :, 0])
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = (
+            words_embeddings
+            + position_embeddings
+            + left_position_embeddings
+            + upper_position_embeddings
+            + right_position_embeddings
+            + lower_position_embeddings
+            + h_position_embeddings
+            + w_position_embeddings
+            + token_type_embeddings
+        )
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
+
+
+# Copied from transformers.modeling_bert.BertSelfAttention with Bert->LayoutLM
+class LayoutLMSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+    ):
+        mixed_query_layer = self.query(hidden_states)
+
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        if encoder_hidden_states is not None:
+            mixed_key_layer = self.key(encoder_hidden_states)
+            mixed_value_layer = self.value(encoder_hidden_states)
+            attention_mask = encoder_attention_mask
+        else:
+            mixed_key_layer = self.key(hidden_states)
+            mixed_value_layer = self.value(hidden_states)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+        key_layer = self.transpose_for_scores(mixed_key_layer)
+        value_layer = self.transpose_for_scores(mixed_value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in LayoutLMModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs = attention_probs * head_mask
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+        return outputs
+
+
+# Copied from transformers.modeling_bert.BertSelfOutput with Bert->LayoutLM
+class LayoutLMSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = LayoutLMLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.modeling_bert.BertAttention with Bert->LayoutLM
+class LayoutLMAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self = LayoutLMSelfAttention(config)
+        self.output = LayoutLMSelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+    ):
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+# Copied from transformers.modeling_bert.BertIntermediate
+class LayoutLMIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.modeling_bert.BertOutput with Bert->LayoutLM
+class LayoutLMOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = LayoutLMLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.modeling_bert.BertLayer with Bert->LayoutLM
+class LayoutLMLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = LayoutLMAttention(config)
+        self.is_decoder = config.is_decoder
+        self.add_cross_attention = config.add_cross_attention
+        if self.add_cross_attention:
+            assert self.is_decoder, f"{self} should be used as a decoder model if cross attention is added"
+            self.crossattention = LayoutLMAttention(config)
+        self.intermediate = LayoutLMIntermediate(config)
+        self.output = LayoutLMOutput(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+    ):
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+        )
+        attention_output = self_attention_outputs[0]
+        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights
+
+        if self.is_decoder and encoder_hidden_states is not None:
+            assert hasattr(
+                self, "crossattention"
+            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"
+            cross_attention_outputs = self.crossattention(
+                attention_output,
+                attention_mask,
+                head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions,
+            )
+            attention_output = cross_attention_outputs[0]
+            outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights
+
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+        )
+        outputs = (layer_output,) + outputs
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+# Copied from transformers.modeling_bert.BertEncoder with Bert->LayoutLM
+class LayoutLMEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([LayoutLMLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=False,
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+
+            if getattr(self.config, "gradient_checkpointing", False):
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    output_attentions,
+                )
+            hidden_states = layer_outputs[0]
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
+        )
+
+
+# Copied from transformers.modeling_bert.BertPooler
+class LayoutLMPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+# Copied from transformers.modeling_bert.BertPredictionHeadTransform with Bert->LayoutLM
+class LayoutLMPredictionHeadTransform(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        if isinstance(config.hidden_act, str):
+            self.transform_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.transform_act_fn = config.hidden_act
+        self.LayerNorm = LayoutLMLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.transform_act_fn(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.modeling_bert.BertLMPredictionHead with Bert->LayoutLM
+class LayoutLMLMPredictionHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.transform = LayoutLMPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.modeling_bert.BertOnlyMLMHead with Bert->LayoutLM
+class LayoutLMOnlyMLMHead(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.predictions = LayoutLMLMPredictionHead(config)
+
+    def forward(self, sequence_output):
+        prediction_scores = self.predictions(sequence_output)
+        return prediction_scores
+
+
+class LayoutLMPreTrainedModel(PreTrainedModel):
+    """An abstract class to handle weights initialization and
+    a simple interface for downloading and loading pretrained models.
+    """
+
+    config_class = LayoutLMConfig
+    base_model_prefix = "layoutlm"
+    authorized_missing_keys = [r"position_ids"]
+
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, LayoutLMLayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
+LAYOUTLM_START_DOCSTRING = r"""    The LayoutLM model was proposed in
+    `LayoutLM: Pre-training of Text and Layout for Document Image Understanding
+    <https://arxiv.org/abs/1912.13318>`__ by....
+
+    This model is a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`_ sub-class.
+    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general
+    usage and behavior.
+
+    Parameters:
+        config (:class:`~transformers.LayoutLMConfig`): Model configuration class with all the parameters of the model.
+            Initializing with a config file does not load the weights associated with the model, only the configuration.
+            Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights.
+"""
+
+LAYOUTLM_INPUTS_DOCSTRING = r"""
+    Inputs:
+        input_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`):
+            Indices of input sequence tokens in the vocabulary.
+
+            Indices can be obtained using :class:`transformers.LayoutLMTokenizer`.
+            See :func:`transformers.PreTrainedTokenizer.encode` and
+            :func:`transformers.PreTrainedTokenizer.__call__` for details.
+
+            `What are input IDs? <../glossary.html#input-ids>`__
+        bbox (:obj:`torch.LongTensor` of shape :obj:`{0}`, `optional`):
+            Bounding Boxes of each input sequence tokens.
+            Selected in the range ``[0, config.max_2d_position_embeddings - 1]``.
+
+            `What are bboxes? <../glossary.html#position-ids>`_
+        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`{0}`, `optional`):
+            Mask to avoid performing attention on padding token indices.
+            Mask values selected in ``[0, 1]``:
+            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
+
+            `What are attention masks? <../glossary.html#attention-mask>`__
+        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`, `optional`):
+            Segment token indices to indicate first and second portions of the inputs.
+            Indices are selected in ``[0, 1]``: ``0`` corresponds to a `sentence A` token, ``1``
+            corresponds to a `sentence B` token
+
+            `What are token type IDs? <../glossary.html#token-type-ids>`_
+        position_ids (:obj:`torch.LongTensor` of shape :obj:`{0}`, `optional`):
+            Indices of positions of each input sequence tokens in the position embeddings.
+            Selected in the range ``[0, config.max_position_embeddings - 1]``.
+
+            `What are position IDs? <../glossary.html#position-ids>`_
+        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`):
+            Mask to nullify selected heads of the self-attention modules.
+            Mask values selected in ``[0, 1]``:
+            :obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
+        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
+            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+            than the model's internal embedding lookup matrix.
+        output_attentions (:obj:`bool`, `optional`):
+            If set to ``True``, the attentions tensors of all attention layers are returned. See ``attentions`` under returned tensors for more detail.
+        output_hidden_states (:obj:`bool`, `optional`):
+            If set to ``True``, the hidden states of all layers are returned. See ``hidden_states`` under returned tensors for more detail.
+        return_dict (:obj:`bool`, `optional`):
+            If set to ``True``, the model will return a :class:`~transformers.file_utils.ModelOutput` instead of a
+            plain tuple.
+"""
+
+
+@add_start_docstrings(
+    "The bare LayoutLM Model transformer outputting raw hidden-states without any specific head on top.",
+    LAYOUTLM_START_DOCSTRING,
+)
+class LayoutLMModel(LayoutLMPreTrainedModel):
+
+    config_class = LayoutLMConfig
+    pretrained_model_archive_map = LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST
+    base_model_prefix = "layoutlm"
+
+    def __init__(self, config):
+        super(LayoutLMModel, self).__init__(config)
+        self.config = config
+
+        self.embeddings = LayoutLMEmbeddings(config)
+        self.encoder = LayoutLMEncoder(config)
+        self.pooler = LayoutLMPooler(config)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.embeddings.word_embeddings
+
+    def set_input_embeddings(self, value):
+        self.embeddings.word_embeddings = value
+
+    def _prune_heads(self, heads_to_prune):
+        """Prunes heads of the model.
+        heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
+        See base class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    @add_start_docstrings_to_callable(LAYOUTLM_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="layoutlm-base-uncased",
+        output_type=BaseModelOutputWithPooling,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids=None,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        """
+        input_ids (torch.LongTensor of shape (batch_size, sequence_length)):
+            Indices of input sequence tokens in the vocabulary.
+        attention_mask (torch.FloatTensor of shape (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 MASKED tokens.
+        token_type_ids (torch.LongTensor of shape (batch_size, sequence_length), optional):
+            Segment token indices to indicate first and second portions of the inputs.
+            Indices are selected in [0, 1]: 0 corresponds to a sentence A token, 1 corresponds to a sentence B token
+        position_ids (torch.LongTensor of shape (batch_size, sequence_length), optional):
+            Indices of positions of each input sequence tokens in the position embeddings.
+            Selected in the range [0, config.max_position_embeddings - 1].
+        head_mask (torch.FloatTensor of shape (num_heads,) or (num_layers, num_heads), optional):
+            Mask to nullify selected heads of the self-attention modules.
+            Mask values selected in [0, 1]: 1 indicates the head is not masked, 0 indicates the head is masked.
+        inputs_embeds (torch.FloatTensor of shape (batch_size, sequence_length, hidden_size), optional):
+            Optionally, instead of passing input_ids you can choose to directly pass an embedded representation.
+            This is useful if you want more control over how to convert input_ids indices into associated vectors than the model’s internal embedding lookup matrix.
+        output_attentions (bool, optional):
+            If set to True, the attentions tensors of all attention layers are returned.
+        output_hidden_states (bool, optional):
+            If set to True, the hidden states of all layers are returned.
+        return_dict (bool, optional):
+            If set to True, the model will return a ModelOutput instead of a plain tuple.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        if bbox is None:
+            bbox = torch.zeros(tuple(list(input_shape) + [4]), dtype=torch.long, device=device)
+
+        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
+
+        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)
+        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+
+        if head_mask is not None:
+            if head_mask.dim() == 1:
+                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+                head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1)
+            elif head_mask.dim() == 2:
+                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
+            head_mask = head_mask.to(dtype=next(self.parameters()).dtype)
+        else:
+            head_mask = [None] * self.config.num_hidden_layers
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids,
+            bbox=bbox,
+            position_ids=position_ids,
+            token_type_ids=token_type_ids,
+            inputs_embeds=inputs_embeds,
+        )
+        encoder_outputs = self.encoder(
+            embedding_output,
+            extended_attention_mask,
+            head_mask=head_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output)
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
+
+@add_start_docstrings("""LayoutLM Model with a `language modeling` head on top. """, LAYOUTLM_START_DOCSTRING)
+class LayoutLMForMaskedLM(LayoutLMPreTrainedModel):
+    config_class = LayoutLMConfig
+    pretrained_model_archive_map = LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST
+    base_model_prefix = "layoutlm"
+
+    def __init__(self, config):
+        super().__init__(config)
+
+        self.layoutlm = LayoutLMModel(config)
+        self.cls = LayoutLMOnlyMLMHead(config)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.layoutlm.embeddings.word_embeddings
+
+    def get_output_embeddings(self):
+        return self.cls.predictions.decoder
+
+    @add_start_docstrings_to_callable(LAYOUTLM_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="layoutlm-base-uncased",
+        output_type=MaskedLMOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids=None,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.layoutlm(
+            input_ids,
+            bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+        prediction_scores = self.cls(sequence_output)
+
+        masked_lm_loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+            masked_lm_loss = loss_fct(
+                prediction_scores.view(-1, self.config.vocab_size),
+                labels.view(-1),
+            )
+
+        if not return_dict:
+            output = (prediction_scores,) + outputs[2:]
+            return ((masked_lm_loss,) + output) if masked_lm_loss is not None else output
+
+        return MaskedLMOutput(
+            loss=masked_lm_loss,
+            logits=prediction_scores,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )
+
+
+@add_start_docstrings(
+    """LayoutLM Model with a token classification head on top (a linear layer on top of
+    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
+    LAYOUTLM_START_DOCSTRING,
+)
+class LayoutLMForTokenClassification(LayoutLMPreTrainedModel):
+    config_class = LayoutLMConfig
+    pretrained_model_archive_map = LAYOUTLM_PRETRAINED_MODEL_ARCHIVE_LIST
+    base_model_prefix = "layoutlm"
+
+    def __init__(self, config):
+        super().__init__(config)
+        self.num_labels = config.num_labels
+        self.layoutlm = LayoutLMModel(config)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+        self.init_weights()
+
+    def get_input_embeddings(self):
+        return self.layoutlm.embeddings.word_embeddings
+
+    @add_start_docstrings_to_callable(LAYOUTLM_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="layoutlm-base-uncased",
+        output_type=TokenClassifierOutput,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    def forward(
+        self,
+        input_ids=None,
+        bbox=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        labels=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        outputs = self.layoutlm(
+            input_ids=input_ids,
+            bbox=bbox,
+            attention_mask=attention_mask,
+            token_type_ids=token_type_ids,
+            position_ids=position_ids,
+            head_mask=head_mask,
+            inputs_embeds=inputs_embeds,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+
+        sequence_output = outputs[0]
+
+        sequence_output = self.dropout(sequence_output)
+        logits = self.classifier(sequence_output)
+
+        loss = None
+        if labels is not None:
+            loss_fct = CrossEntropyLoss()
+
+            if attention_mask is not None:
+                active_loss = attention_mask.view(-1) == 1
+                active_logits = logits.view(-1, self.num_labels)[active_loss]
+                active_labels = labels.view(-1)[active_loss]
+                loss = loss_fct(active_logits, active_labels)
+            else:
+                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+
+        if not return_dict:
+            output = (logits,) + outputs[2:]
+            return ((loss,) + output) if loss is not None else output
+
+        return TokenClassifierOutput(
+            loss=loss,
+            logits=logits,
+            hidden_states=outputs.hidden_states,
+            attentions=outputs.attentions,
+        )

src/transformers/modeling_roberta.py

@@ -15,7 +15,7 @@
 # limitations under the License.
 """PyTorch RoBERTa model. """
 
-
+import math
 import warnings
 
 import torch
@@ -29,8 +29,10 @@ from .file_utils import (
     add_start_docstrings_to_callable,
     replace_return_docstrings,
 )
-from .modeling_bert import BertEmbeddings, BertLayerNorm, BertModel, BertPreTrainedModel, gelu
+from .modeling_bert import ACT2FN, gelu
 from .modeling_outputs import (
+    BaseModelOutput,
+    BaseModelOutputWithPooling,
     CausalLMOutput,
     MaskedLMOutput,
     MultipleChoiceModelOutput,
@@ -38,6 +40,12 @@ from .modeling_outputs import (
     SequenceClassifierOutput,
     TokenClassifierOutput,
 )
+from .modeling_utils import (
+    PreTrainedModel,
+    apply_chunking_to_forward,
+    find_pruneable_heads_and_indices,
+    prune_linear_layer,
+)
 from .utils import logging
 
 
@@ -57,15 +65,31 @@ ROBERTA_PRETRAINED_MODEL_ARCHIVE_LIST = [
 ]
 
 
-class RobertaEmbeddings(BertEmbeddings):
+RobertaLayerNorm = torch.nn.LayerNorm
+
+
+class RobertaEmbeddings(nn.Module):
     """
     Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
     """
 
+    # Copied from transformers.modeling_bert.BertEmbeddings.__init__ with Bert->Roberta
     def __init__(self, config):
-        super().__init__(config)
+        super().__init__()
+        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=config.pad_token_id)
+        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
+        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
+
+        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
+        # any TensorFlow checkpoint file
+        self.LayerNorm = RobertaLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+        # position_ids (1, len position emb) is contiguous in memory and exported when serialized
+        self.register_buffer("position_ids", torch.arange(config.max_position_embeddings).expand((1, -1)))
+
+        # End copy
         self.padding_idx = config.pad_token_id
-        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=self.padding_idx)
         self.position_embeddings = nn.Embedding(
             config.max_position_embeddings, config.hidden_size, padding_idx=self.padding_idx
         )
@@ -78,9 +102,29 @@ class RobertaEmbeddings(BertEmbeddings):
             else:
                 position_ids = self.create_position_ids_from_inputs_embeds(inputs_embeds)
 
-        return super().forward(
-            input_ids, token_type_ids=token_type_ids, position_ids=position_ids, inputs_embeds=inputs_embeds
-        )
+        # Copied from transformers.modeling_bert.BertEmbeddings.forward
+        if input_ids is not None:
+            input_shape = input_ids.size()
+        else:
+            input_shape = inputs_embeds.size()[:-1]
+
+        seq_length = input_shape[1]
+
+        if position_ids is None:
+            position_ids = self.position_ids[:, :seq_length]
+
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=self.position_ids.device)
+
+        if inputs_embeds is None:
+            inputs_embeds = self.word_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+        token_type_embeddings = self.token_type_embeddings(token_type_ids)
+
+        embeddings = inputs_embeds + position_embeddings + token_type_embeddings
+        embeddings = self.LayerNorm(embeddings)
+        embeddings = self.dropout(embeddings)
+        return embeddings
 
     def create_position_ids_from_inputs_embeds(self, inputs_embeds):
         """We are provided embeddings directly. We cannot infer which are padded so just generate
@@ -98,6 +142,343 @@ class RobertaEmbeddings(BertEmbeddings):
         return position_ids.unsqueeze(0).expand(input_shape)
 
 
+# Copied from transformers.modeling_bert.BertSelfAttention with Bert->Roberta
+class RobertaSelfAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        if config.hidden_size % config.num_attention_heads != 0 and not hasattr(config, "embedding_size"):
+            raise ValueError(
+                "The hidden size (%d) is not a multiple of the number of attention "
+                "heads (%d)" % (config.hidden_size, config.num_attention_heads)
+            )
+
+        self.num_attention_heads = config.num_attention_heads
+        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
+        self.all_head_size = self.num_attention_heads * self.attention_head_size
+
+        self.query = nn.Linear(config.hidden_size, self.all_head_size)
+        self.key = nn.Linear(config.hidden_size, self.all_head_size)
+        self.value = nn.Linear(config.hidden_size, self.all_head_size)
+
+        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
+
+    def transpose_for_scores(self, x):
+        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
+        x = x.view(*new_x_shape)
+        return x.permute(0, 2, 1, 3)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+    ):
+        mixed_query_layer = self.query(hidden_states)
+
+        # If this is instantiated as a cross-attention module, the keys
+        # and values come from an encoder; the attention mask needs to be
+        # such that the encoder's padding tokens are not attended to.
+        if encoder_hidden_states is not None:
+            mixed_key_layer = self.key(encoder_hidden_states)
+            mixed_value_layer = self.value(encoder_hidden_states)
+            attention_mask = encoder_attention_mask
+        else:
+            mixed_key_layer = self.key(hidden_states)
+            mixed_value_layer = self.value(hidden_states)
+
+        query_layer = self.transpose_for_scores(mixed_query_layer)
+        key_layer = self.transpose_for_scores(mixed_key_layer)
+        value_layer = self.transpose_for_scores(mixed_value_layer)
+
+        # Take the dot product between "query" and "key" to get the raw attention scores.
+        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
+        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
+        if attention_mask is not None:
+            # Apply the attention mask is (precomputed for all layers in RobertaModel forward() function)
+            attention_scores = attention_scores + attention_mask
+
+        # Normalize the attention scores to probabilities.
+        attention_probs = nn.Softmax(dim=-1)(attention_scores)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        attention_probs = self.dropout(attention_probs)
+
+        # Mask heads if we want to
+        if head_mask is not None:
+            attention_probs = attention_probs * head_mask
+
+        context_layer = torch.matmul(attention_probs, value_layer)
+
+        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
+        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        outputs = (context_layer, attention_probs) if output_attentions else (context_layer,)
+        return outputs
+
+
+# Copied from transformers.modeling_bert.BertSelfOutput with Bert->Roberta
+class RobertaSelfOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.LayerNorm = RobertaLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.modeling_bert.BertAttention with Bert->Roberta
+class RobertaAttention(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.self = RobertaSelfAttention(config)
+        self.output = RobertaSelfOutput(config)
+        self.pruned_heads = set()
+
+    def prune_heads(self, heads):
+        if len(heads) == 0:
+            return
+        heads, index = find_pruneable_heads_and_indices(
+            heads, self.self.num_attention_heads, self.self.attention_head_size, self.pruned_heads
+        )
+
+        # Prune linear layers
+        self.self.query = prune_linear_layer(self.self.query, index)
+        self.self.key = prune_linear_layer(self.self.key, index)
+        self.self.value = prune_linear_layer(self.self.value, index)
+        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
+
+        # Update hyper params and store pruned heads
+        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
+        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads
+        self.pruned_heads = self.pruned_heads.union(heads)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+    ):
+        self_outputs = self.self(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            encoder_hidden_states,
+            encoder_attention_mask,
+            output_attentions,
+        )
+        attention_output = self.output(self_outputs[0], hidden_states)
+        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+        return outputs
+
+
+# Copied from transformers.modeling_bert.BertIntermediate
+class RobertaIntermediate(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
+        if isinstance(config.hidden_act, str):
+            self.intermediate_act_fn = ACT2FN[config.hidden_act]
+        else:
+            self.intermediate_act_fn = config.hidden_act
+
+    def forward(self, hidden_states):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.intermediate_act_fn(hidden_states)
+        return hidden_states
+
+
+# Copied from transformers.modeling_bert.BertOutput with Bert->Roberta
+class RobertaOutput(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+        self.LayerNorm = RobertaLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+    def forward(self, hidden_states, input_tensor):
+        hidden_states = self.dense(hidden_states)
+        hidden_states = self.dropout(hidden_states)
+        hidden_states = self.LayerNorm(hidden_states + input_tensor)
+        return hidden_states
+
+
+# Copied from transformers.modeling_bert.BertLayer with Bert->Roberta
+class RobertaLayer(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.chunk_size_feed_forward = config.chunk_size_feed_forward
+        self.seq_len_dim = 1
+        self.attention = RobertaAttention(config)
+        self.is_decoder = config.is_decoder
+        self.add_cross_attention = config.add_cross_attention
+        if self.add_cross_attention:
+            assert self.is_decoder, f"{self} should be used as a decoder model if cross attention is added"
+            self.crossattention = RobertaAttention(config)
+        self.intermediate = RobertaIntermediate(config)
+        self.output = RobertaOutput(config)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+    ):
+        self_attention_outputs = self.attention(
+            hidden_states,
+            attention_mask,
+            head_mask,
+            output_attentions=output_attentions,
+        )
+        attention_output = self_attention_outputs[0]
+        outputs = self_attention_outputs[1:]  # add self attentions if we output attention weights
+
+        if self.is_decoder and encoder_hidden_states is not None:
+            assert hasattr(
+                self, "crossattention"
+            ), f"If `encoder_hidden_states` are passed, {self} has to be instantiated with cross-attention layers by setting `config.add_cross_attention=True`"
+            cross_attention_outputs = self.crossattention(
+                attention_output,
+                attention_mask,
+                head_mask,
+                encoder_hidden_states,
+                encoder_attention_mask,
+                output_attentions,
+            )
+            attention_output = cross_attention_outputs[0]
+            outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights
+
+        layer_output = apply_chunking_to_forward(
+            self.feed_forward_chunk, self.chunk_size_feed_forward, self.seq_len_dim, attention_output
+        )
+        outputs = (layer_output,) + outputs
+        return outputs
+
+    def feed_forward_chunk(self, attention_output):
+        intermediate_output = self.intermediate(attention_output)
+        layer_output = self.output(intermediate_output, attention_output)
+        return layer_output
+
+
+# Copied from transformers.modeling_bert.BertEncoder with Bert->Roberta
+class RobertaEncoder(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.config = config
+        self.layer = nn.ModuleList([RobertaLayer(config) for _ in range(config.num_hidden_layers)])
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask=None,
+        head_mask=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=False,
+        output_hidden_states=False,
+        return_dict=False,
+    ):
+        all_hidden_states = () if output_hidden_states else None
+        all_attentions = () if output_attentions else None
+        for i, layer_module in enumerate(self.layer):
+            if output_hidden_states:
+                all_hidden_states = all_hidden_states + (hidden_states,)
+
+            layer_head_mask = head_mask[i] if head_mask is not None else None
+
+            if getattr(self.config, "gradient_checkpointing", False):
+
+                def create_custom_forward(module):
+                    def custom_forward(*inputs):
+                        return module(*inputs, output_attentions)
+
+                    return custom_forward
+
+                layer_outputs = torch.utils.checkpoint.checkpoint(
+                    create_custom_forward(layer_module),
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                )
+            else:
+                layer_outputs = layer_module(
+                    hidden_states,
+                    attention_mask,
+                    layer_head_mask,
+                    encoder_hidden_states,
+                    encoder_attention_mask,
+                    output_attentions,
+                )
+            hidden_states = layer_outputs[0]
+            if output_attentions:
+                all_attentions = all_attentions + (layer_outputs[1],)
+
+        if output_hidden_states:
+            all_hidden_states = all_hidden_states + (hidden_states,)
+
+        if not return_dict:
+            return tuple(v for v in [hidden_states, all_hidden_states, all_attentions] if v is not None)
+        return BaseModelOutput(
+            last_hidden_state=hidden_states, hidden_states=all_hidden_states, attentions=all_attentions
+        )
+
+
+# Copied from transformers.modeling_bert.BertPooler
+class RobertaPooler(nn.Module):
+    def __init__(self, config):
+        super().__init__()
+        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+        self.activation = nn.Tanh()
+
+    def forward(self, hidden_states):
+        # We "pool" the model by simply taking the hidden state corresponding
+        # to the first token.
+        first_token_tensor = hidden_states[:, 0]
+        pooled_output = self.dense(first_token_tensor)
+        pooled_output = self.activation(pooled_output)
+        return pooled_output
+
+
+class RobertaPreTrainedModel(PreTrainedModel):
+    """An abstract class to handle weights initialization and
+    a simple interface for downloading and loading pretrained models.
+    """
+
+    config_class = RobertaConfig
+    base_model_prefix = "roberta"
+    authorized_missing_keys = [r"position_ids"]
+
+    # Copied from transformers.modeling_bert.BertPreTrainedModel._init_weights with Bert->Roberta
+    def _init_weights(self, module):
+        """ Initialize the weights """
+        if isinstance(module, (nn.Linear, nn.Embedding)):
+            # Slightly different from the TF version which uses truncated_normal for initialization
+            # cf https://github.com/pytorch/pytorch/pull/5617
+            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
+        elif isinstance(module, RobertaLayerNorm):
+            module.bias.data.zero_()
+            module.weight.data.fill_(1.0)
+        if isinstance(module, nn.Linear) and module.bias is not None:
+            module.bias.data.zero_()
+
+
 ROBERTA_START_DOCSTRING = r"""
 
     This model is a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`_ sub-class.
@@ -159,19 +540,34 @@ ROBERTA_INPUTS_DOCSTRING = r"""
     "The bare RoBERTa Model transformer outputting raw hidden-states without any specific head on top.",
     ROBERTA_START_DOCSTRING,
 )
-class RobertaModel(BertModel):
-    """
-    This class overrides :class:`~transformers.BertModel`. Please check the
-    superclass for the appropriate documentation alongside usage examples.
+class RobertaModel(RobertaPreTrainedModel):
     """
 
-    config_class = RobertaConfig
-    base_model_prefix = "roberta"
+    The model can behave as an encoder (with only self-attention) as well
+    as a decoder, in which case a layer of cross-attention is added between
+    the self-attention layers, following the architecture described in `Attention is all you need`_ by Ashish Vaswani,
+    Noam Shazeer, Niki Parmar, Jakob Uszkoreit, Llion Jones, Aidan N. Gomez, Lukasz Kaiser and Illia Polosukhin.
 
+    To behave as an decoder the model needs to be initialized with the
+    :obj:`is_decoder` argument of the configuration set to :obj:`True`.
+    To be used in a Seq2Seq model, the model needs to initialized with both :obj:`is_decoder`
+    argument and :obj:`add_cross_attention` set to :obj:`True`; an
+    :obj:`encoder_hidden_states` is then expected as an input to the forward pass.
+
+    .. _`Attention is all you need`:
+        https://arxiv.org/abs/1706.03762
+
+    """
+
+    # Copied from transformers.modeling_bert.BertModel.__init__ with Bert->Roberta
     def __init__(self, config):
         super().__init__(config)
+        self.config = config
 
         self.embeddings = RobertaEmbeddings(config)
+        self.encoder = RobertaEncoder(config)
+        self.pooler = RobertaPooler(config)
+
         self.init_weights()
 
     def get_input_embeddings(self):
@@ -180,14 +576,121 @@ class RobertaModel(BertModel):
     def set_input_embeddings(self, value):
         self.embeddings.word_embeddings = value
 
+    def _prune_heads(self, heads_to_prune):
+        """Prunes heads of the model.
+        heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
+        See base class PreTrainedModel
+        """
+        for layer, heads in heads_to_prune.items():
+            self.encoder.layer[layer].attention.prune_heads(heads)
+
+    @add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING.format("(batch_size, sequence_length)"))
+    @add_code_sample_docstrings(
+        tokenizer_class=_TOKENIZER_FOR_DOC,
+        checkpoint="roberta-base",
+        output_type=BaseModelOutputWithPooling,
+        config_class=_CONFIG_FOR_DOC,
+    )
+    # Copied from transformers.modeling_bert.BertModel.forward
+    def forward(
+        self,
+        input_ids=None,
+        attention_mask=None,
+        token_type_ids=None,
+        position_ids=None,
+        head_mask=None,
+        inputs_embeds=None,
+        encoder_hidden_states=None,
+        encoder_attention_mask=None,
+        output_attentions=None,
+        output_hidden_states=None,
+        return_dict=None,
+    ):
+        r"""
+        encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`):
+            Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
+            if the model is configured as a decoder.
+        encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`):
+            Mask to avoid performing attention on the padding token indices of the encoder input. This mask
+            is used in the cross-attention if the model is configured as a decoder.
+            Mask values selected in ``[0, 1]``:
+            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
+        """
+        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
+        output_hidden_states = (
+            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
+        )
+        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
+
+        if input_ids is not None and inputs_embeds is not None:
+            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
+        elif input_ids is not None:
+            input_shape = input_ids.size()
+        elif inputs_embeds is not None:
+            input_shape = inputs_embeds.size()[:-1]
+        else:
+            raise ValueError("You have to specify either input_ids or inputs_embeds")
+
+        device = input_ids.device if input_ids is not None else inputs_embeds.device
+
+        if attention_mask is None:
+            attention_mask = torch.ones(input_shape, device=device)
+        if token_type_ids is None:
+            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
+
+        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
+        # ourselves in which case we just need to make it broadcastable to all heads.
+        extended_attention_mask: torch.Tensor = self.get_extended_attention_mask(attention_mask, input_shape, device)
+
+        # If a 2D or 3D attention mask is provided for the cross-attention
+        # we need to make broadcastable to [batch_size, num_heads, seq_length, seq_length]
+        if self.config.is_decoder and encoder_hidden_states is not None:
+            encoder_batch_size, encoder_sequence_length, _ = encoder_hidden_states.size()
+            encoder_hidden_shape = (encoder_batch_size, encoder_sequence_length)
+            if encoder_attention_mask is None:
+                encoder_attention_mask = torch.ones(encoder_hidden_shape, device=device)
+            encoder_extended_attention_mask = self.invert_attention_mask(encoder_attention_mask)
+        else:
+            encoder_extended_attention_mask = None
+
+        # Prepare head mask if needed
+        # 1.0 in head_mask indicate we keep the head
+        # attention_probs has shape bsz x n_heads x N x N
+        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
+        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
+
+        embedding_output = self.embeddings(
+            input_ids=input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
+        )
+        encoder_outputs = self.encoder(
+            embedding_output,
+            attention_mask=extended_attention_mask,
+            head_mask=head_mask,
+            encoder_hidden_states=encoder_hidden_states,
+            encoder_attention_mask=encoder_extended_attention_mask,
+            output_attentions=output_attentions,
+            output_hidden_states=output_hidden_states,
+            return_dict=return_dict,
+        )
+        sequence_output = encoder_outputs[0]
+        pooled_output = self.pooler(sequence_output)
+
+        if not return_dict:
+            return (sequence_output, pooled_output) + encoder_outputs[1:]
+
+        return BaseModelOutputWithPooling(
+            last_hidden_state=sequence_output,
+            pooler_output=pooled_output,
+            hidden_states=encoder_outputs.hidden_states,
+            attentions=encoder_outputs.attentions,
+        )
+
 
 @add_start_docstrings(
     """RoBERTa Model with a `language modeling` head on top for CLM fine-tuning. """, ROBERTA_START_DOCSTRING
 )
-class RobertaForCausalLM(BertPreTrainedModel):
-    config_class = RobertaConfig
-    base_model_prefix = "roberta"
-
+class RobertaForCausalLM(RobertaPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
@@ -300,11 +803,7 @@ class RobertaForCausalLM(BertPreTrainedModel):
 
 
 @add_start_docstrings("""RoBERTa Model with a `language modeling` head on top. """, ROBERTA_START_DOCSTRING)
-class RobertaForMaskedLM(BertPreTrainedModel):
-    config_class = RobertaConfig
-    base_model_prefix = "roberta"
-    authorized_missing_keys = [r"position_ids", r"lm_head\.decoder\.bias"]
-
+class RobertaForMaskedLM(RobertaPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
@@ -402,7 +901,7 @@ class RobertaLMHead(nn.Module):
     def __init__(self, config):
         super().__init__()
         self.dense = nn.Linear(config.hidden_size, config.hidden_size)
-        self.layer_norm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+        self.layer_norm = RobertaLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
 
         self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
         self.bias = nn.Parameter(torch.zeros(config.vocab_size))
@@ -426,10 +925,7 @@ class RobertaLMHead(nn.Module):
     on top of the pooled output) e.g. for GLUE tasks. """,
     ROBERTA_START_DOCSTRING,
 )
-class RobertaForSequenceClassification(BertPreTrainedModel):
-    config_class = RobertaConfig
-    base_model_prefix = "roberta"
-
+class RobertaForSequenceClassification(RobertaPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
@@ -509,10 +1005,7 @@ class RobertaForSequenceClassification(BertPreTrainedModel):
     the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
     ROBERTA_START_DOCSTRING,
 )
-class RobertaForMultipleChoice(BertPreTrainedModel):
-    config_class = RobertaConfig
-    base_model_prefix = "roberta"
-
+class RobertaForMultipleChoice(RobertaPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
 
@@ -600,10 +1093,7 @@ class RobertaForMultipleChoice(BertPreTrainedModel):
     the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
     ROBERTA_START_DOCSTRING,
 )
-class RobertaForTokenClassification(BertPreTrainedModel):
-    config_class = RobertaConfig
-    base_model_prefix = "roberta"
-
+class RobertaForTokenClassification(RobertaPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels
@@ -708,10 +1198,7 @@ class RobertaClassificationHead(nn.Module):
     the hidden-states output to compute `span start logits` and `span end logits`). """,
     ROBERTA_START_DOCSTRING,
 )
-class RobertaForQuestionAnswering(BertPreTrainedModel):
-    config_class = RobertaConfig
-    base_model_prefix = "roberta"
-
+class RobertaForQuestionAnswering(RobertaPreTrainedModel):
     def __init__(self, config):
         super().__init__(config)
         self.num_labels = config.num_labels

src/transformers/modeling_tf_utils.py

@@ -85,20 +85,20 @@ def keras_serializable(cls):
 
     @functools.wraps(initializer)
     def wrapped_init(self, *args, **kwargs):
-        transformers_config = kwargs.pop("transformers_config", None)
-        config = args[0] if args and isinstance(args[0], PretrainedConfig) else kwargs.get("config", None)
-        if config is not None and transformers_config is not None:
-            raise ValueError("Must pass either `config` or `transformers_config`, not both")
-        elif config is not None:
-            # normal layer construction, call with unchanged args (config is already in there)
-            initializer(self, *args, **kwargs)
-        elif transformers_config is not None:
-            # Keras deserialization, convert dict to config
-            config = config_class.from_dict(transformers_config)
+        config = args[0] if args and isinstance(args[0], PretrainedConfig) else kwargs.pop("config", None)
+
+        if isinstance(config, dict):
+            config = config_class.from_dict(config)
             initializer(self, config, *args, **kwargs)
+        elif isinstance(config, PretrainedConfig):
+            if len(args) > 0:
+                initializer(self, *args, **kwargs)
+            else:
+                initializer(self, config, *args, **kwargs)
         else:
-            raise ValueError("Must pass either `config` (PretrainedConfig) or `transformers_config` (dict)")
-        self._transformers_config = config
+            raise ValueError("Must pass either `config` (PretrainedConfig) or `config` (dict)")
+
+        self._config = config
         self._kwargs = kwargs
 
     cls.__init__ = wrapped_init
@@ -109,7 +109,7 @@ def keras_serializable(cls):
 
         def get_config(self):
             cfg = super(cls, self).get_config()
-            cfg["transformers_config"] = self._transformers_config.to_dict()
+            cfg["config"] = self._config.to_dict()
             cfg.update(self._kwargs)
             return cfg
 

src/transformers/pipelines/__init__.py

@@ -0,0 +1,13 @@
+from pathlib import Path
+from typing import Any, Union
+
+from transformers import PreTrainedTokenizer
+
+from .base import PipelineConfigType, PipelineInputType, PipelineIntermediateType, MaybeBatch, ModelType, \
+    PipelineOutputType, Pipeline, PipelineConfig
+from .configs import TokenClassificationConfig
+from .outputs import TokenClassificationOutput
+
+
+def pipeline(task: str, model: Union[str, Path, Any], tokenizer: Union[str, Path, PreTrainedTokenizer]):
+    pass

src/transformers/pipelines/base.py

@@ -0,0 +1,168 @@
+from abc import ABC, abstractmethod
+from dataclasses import dataclass, field
+from typing import Any, ClassVar, Dict, Generic, Mapping, Optional, TypeVar, Union, List
+
+from transformers import BatchEncoding, PreTrainedTokenizer
+
+
+# Syntactic sugar to indicate it can take multiple inputs of type T at once.
+T = TypeVar("T")
+MaybeBatch = Union[T, List[T]]
+
+
+@dataclass
+class PipelineConfig:
+    tokenizer_kwargs: Dict[str, Any] = field(default_factory=dict)
+    model_kwargs: Dict[str, Any] = field(default_factory=dict)
+
+    def __post_init__(self):
+        self.model_kwargs.setdefault("return_dict", True)
+
+
+# Define all the generic types for a Pipeline
+PipelineInputType = TypeVar("PipelineInputType")  # Pipeline input type
+PipelineOutputType = TypeVar("PipelineOutputType")  # Pipeline output type
+PipelineIntermediateType = TypeVar("PipelineIntermediateType")  # Model output type (i.e. after forward())
+PipelineConfigType = TypeVar("PipelineConfigType", bound=PipelineConfig)  # Pipeline configuration type
+ModelType = TypeVar("ModelType")  # Pipeline model type
+
+
+class Pipeline(ABC, Generic[PipelineConfigType, ModelType, PipelineInputType, PipelineIntermediateType, PipelineOutputType]):
+    """"""
+
+    __slots__ = ["_tokenizer", "_model", "_configs"]
+
+    # TODO: Check if the $task variable is required and won't introduce dependency on it
+    task: ClassVar[str]
+    default_config: ClassVar[PipelineConfigType]
+
+    def __init__(self, tokenizer: PreTrainedTokenizer, model: ModelType):
+        self._tokenizer = tokenizer
+        self._model = model
+
+        self._configs: Dict[str, PipelineConfigType] = dict()
+        self._configs["default"] = self.default_config
+
+    @property
+    def tokenizer(self) -> PreTrainedTokenizer:
+        """
+        Tokenizer associated with this pipeline
+        Returns:
+            PreTrainedTokenizer instance
+        """
+        return self._tokenizer
+
+    @property
+    def model(self) -> ModelType:
+        """
+        Model associated with this pipeline
+        Returns:
+            Pretrained model instance
+        """
+        raise self._model
+
+    @property
+    def configs(self) -> Mapping[str, PipelineConfigType]:
+        """
+        Return all registered configuration's mapping on the pipeline
+        Returns:
+            Mapping[str, ConfigType]
+
+        """
+        return self._configs
+
+    @property
+    def default_config(self) -> PipelineConfigType:
+        return self.default_config
+
+    def get_config(self, name: str) -> Optional[PipelineConfigType]:
+        """
+        Attempt to retrieve a configuration from its registered name.
+        If no configuration matches the requested name, None is returned.
+        Args:
+            name (:obj:`str`): Name of the configuration to look for
+
+        Returns:
+            Instance of ConfigType if the configuration associated with the requested name is present in the registry
+            None if the requested named cannot be found in the registry
+        """
+        return self._configs.get(name, None)
+
+    def register_config(self, name: str, config: PipelineConfigType) -> PipelineConfigType:
+        """
+        Register a configuration with the specified name uniquely identifying a configuration.
+        Args:
+            name (:obj:`str`): Name used to reference the configuration
+            config (:obj:`ConfigType`): The configuration instance to associate to the specified name
+        Returns:
+            The provided ConfigType instance
+        """
+        self._configs[name] = config
+        return config
+
+    def delete_config(self, name: str) -> Optional[PipelineConfigType]:
+        """
+        Attempt to delete a configuration if registered on the pipeline.
+        The ConfigType instance associated with the provided name is returned if present on the pipeline,
+        None is returned otherwise.
+        Args:
+            name (:obj:`str`): Name identifying the configuration to delete
+
+        Returns:
+            ConfigType instance if the name identifier was found on the pipeline.
+            None if the name identifier is unknown on the pipeline.
+        """
+        return self._configs.pop(name, None)
+
+    @abstractmethod
+    def __call__(
+            self, inputs: MaybeBatch[PipelineInputType], config: Optional[Union[str, PipelineConfigType]], **kwargs
+    ) -> MaybeBatch[PipelineIntermediateType]:
+
+        """
+        Args:
+            inputs (:obj:`InputType`):
+            config (:obj:`ConfigType`)
+            **kwargs:
+
+        Returns:
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def preprocess(self, inputs: MaybeBatch[PipelineInputType], config: PipelineConfigType) -> BatchEncoding:
+        """
+        Process the raw inputs to generate a BatchEncoding representation through the use of a tokenizer.
+        Args:
+            inputs (:obj:`InputType`):
+            config (:obj:`ConfigType`)
+
+        Returns:
+        """
+        # TODO: preprocess while calling batch_encode_plus might support truncation and thus all the remaining
+        #       steps [forward(), postprocess()] should supports inference and input reconstruction.
+        raise NotImplementedError()
+
+    @abstractmethod
+    def forward(self, encodings: BatchEncoding, config: PipelineConfigType) -> PipelineIntermediateType:
+        """
+
+        Args:
+            encodings (:obj:`BatchEncoding`):
+            config (:obj:`ConfigType`)
+        Returns:
+        """
+        raise NotImplementedError()
+
+    @abstractmethod
+    def postprocess(self, encodings: BatchEncoding, model_output: PipelineIntermediateType, config: PipelineConfigType) -> MaybeBatch[PipelineOutputType]:
+        """
+
+        Args:
+            encodings:
+            model_output (:obj:`PipelineIntermediateType`):
+            config (:obj:`ConfigType`)
+
+        Returns:
+        """
+        raise NotImplementedError()

src/transformers/pipelines/configs.py

@@ -0,0 +1,11 @@
+from dataclasses import dataclass, field
+from typing import List
+
+from transformers.pipelines import PipelineConfig
+
+
+@dataclass
+class TokenClassificationConfig(PipelineConfig):
+    group_entities: bool = False
+    ignore_labels: List[int] = field(default_factory=list)
+

src/transformers/pipelines/outputs.py

@@ -0,0 +1,13 @@
+from dataclasses import dataclass
+from typing import Union, List, NamedTuple
+
+from transformers import BatchEncoding
+
+Entity = NamedTuple("Entity", [("token", str), ("index", int), ("label", str), ("score", float)])
+
+
+@dataclass
+class TokenClassificationOutput:
+    input: Union[str, List[str]]
+    encodings: BatchEncoding
+    entities: List[Entity]

src/transformers/pipelines/torch/__init__.py

@@ -0,0 +1,8 @@
+from .pipeline_utils import PreTrainedPipeline
+from .token_classification import (
+    TokenClassificationConfig,
+    TokenClassificationInput,
+    TokenClassificationOutput,
+    TokenClassificationPipeline,
+    TokenClassifierOutput,
+)

src/transformers/pipelines/torch/pipeline_utils.py

@@ -0,0 +1,55 @@
+from abc import ABC
+from typing import Optional, Union
+
+import torch
+
+from transformers import PreTrainedModel, BatchEncoding, TensorType
+from transformers.pipelines import MaybeBatch, Pipeline
+from transformers.pipelines.base import PipelineConfigType, PipelineInputType, PipelineOutputType, \
+    PipelineIntermediateType
+
+
+class PreTrainedPipeline(Pipeline[PipelineConfigType, PreTrainedModel, PipelineInputType, PipelineIntermediateType, PipelineOutputType], ABC):
+
+    @property
+    def device(self) -> torch.device:
+        return self._model.device
+
+    def cpu(self) -> "PreTrainedPipeline":
+        self._model.cpu()
+        return self
+
+    def gpu(self) -> "PreTrainedPipeline":
+        self._model.gpu()
+        return self
+
+    def preprocess(self, inputs: MaybeBatch[PipelineInputType], config: PipelineConfigType) -> BatchEncoding:
+        return self._tokenizer(
+            inputs,
+            return_tensors=TensorType.PYTORCH,
+            **(config.tokenizer_kwargs or {}),
+        )
+
+    def __call__(self, inputs: MaybeBatch[PipelineInputType], config: Optional[Union[str, PipelineConfigType]], **kwargs) -> MaybeBatch[PipelineIntermediateType]:
+        # Retrieve the appropriate config if identifier is provided or None
+        if isinstance(config, str):
+            config = self.get_config(config)
+        elif not config:
+            config = self.default_config
+
+        # Preprocess the input
+        encodings = self.preprocess(inputs, config)
+
+        # Forward the encoding through the model
+        model_output = self.forward(encodings, config)
+
+        # Apply any postprocessing steps required
+        return self.postprocess(encodings, model_output, config)
+
+    def forward(self, encodings, config) -> PipelineIntermediateType:
+        return self._model(**encodings, **config.model_kwargs)
+
+
+
+
+

src/transformers/pipelines/torch/token_classification.py

@@ -0,0 +1,48 @@
+from transformers import BatchEncoding
+
+from transformers.modeling_outputs import TokenClassifierOutput
+from transformers.pipelines import MaybeBatch, TokenClassificationConfig, TokenClassificationOutput
+from transformers.pipelines.outputs import Entity
+from transformers.pipelines.torch import PreTrainedPipeline
+
+
+# Token Classification
+
+TokenClassificationInput = str
+
+
+class TokenClassificationPipeline(
+    PreTrainedPipeline[
+        TokenClassificationConfig, TokenClassificationInput, TokenClassifierOutput, TokenClassificationOutput
+    ]
+):
+    task = "token-classification"
+    default_config = TokenClassificationConfig()
+
+    def postprocess(self, encodings: BatchEncoding, model_output: TokenClassifierOutput, config: TokenClassificationConfig) -> MaybeBatch[TokenClassificationOutput]:
+        # Retrieve labels
+        labels = self._model.config.id2label
+
+        # Convert to probabilities
+        probs = model_output.logits.softmax(dim=-1)
+
+        # Retrieve the argmax for each token over sequence axis
+        batch_entities_labels = probs.argmax(dim=-1)
+
+        batch_entities = []
+        for batch_idx, entities_labels in enumerate(batch_entities_labels):
+            entities = []
+            for token_idx, label_idx in enumerate(entities_labels):
+                label_idx = label_idx.item()
+                if label_idx not in config.ignore_labels:
+                    entities.append(Entity(
+                        token=self._tokenizer.convert_ids_to_tokens([encodings["input_ids"][batch_idx][token_idx]])[0],
+                        index=token_idx,
+                        label=labels[label_idx],
+                        score=probs[batch_idx, token_idx, label_idx].item(),
+                    ))
+            batch_entities.append(entities)
+
+        return batch_entities
+
+

src/transformers/retrieval_rag.py

@@ -0,0 +1,470 @@
+# coding=utf-8
+# Copyright 2020, The RAG Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""RAG Retriever model implementation."""
+
+import os
+import pickle
+import time
+from typing import Iterable, List, Optional, Tuple
+
+import numpy as np
+
+from .configuration_rag import RagConfig
+from .file_utils import cached_path, is_datasets_available, is_faiss_available, is_remote_url
+from .tokenization_rag import RagTokenizer
+from .tokenization_utils_base import BatchEncoding
+from .utils import logging
+
+
+if is_datasets_available() and is_faiss_available():
+    from datasets import load_dataset
+
+    import faiss
+
+
+logger = logging.get_logger(__name__)
+
+
+LEGACY_INDEX_PATH = "https://storage.googleapis.com/huggingface-nlp/datasets/wiki_dpr/"
+
+
+class Index:
+    """
+    A base class for the Indices encapsulated by the :class:`~transformers.RagRetriever`.
+    """
+
+    def get_doc_dicts(self, doc_ids: np.ndarray) -> List[dict]:
+        """
+        Returns a list of dictionaries, containing titles and text of the retrieved documents.
+
+        Args:
+            doc_ids (:obj:`np.ndarray` of shape :obj:`(batch_size, n_docs)`):
+                A tensor of document indices.
+        """
+        raise NotImplementedError
+
+    def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
+        """
+        For each query in the batch, retrieves ``n_docs`` documents.
+
+        Args:
+            question_hidden_states (:obj:`np.ndarray` of shape :obj:`(batch_size, vector_size):
+                An array of query vectors.
+            n_docs (:obj:`int`):
+                The number of docs retrieved per query.
+
+        Returns:
+            :obj:`np.ndarray` of shape :obj:`(batch_size, n_docs)`: A tensor of indices of retrieved documents.
+            :obj:`np.ndarray` of shape :obj:`(batch_size, vector_size)`: A tensor of vector representations of retrieved documents.
+        """
+        raise NotImplementedError
+
+    def is_initialized(self):
+        """
+        Returns :obj:`True` if index is already initialized.
+        """
+        raise NotImplementedError
+
+    def init_index(self):
+        """
+        A function responsible for loading the index into memory. Should be called only once per training run of a RAG model.
+        E.g. if the model is trained on multiple GPUs in a distributed setup, only one of the workers will load the index.
+        """
+        raise NotImplementedError
+
+
+class LegacyIndex:
+    """
+    An index which can be deserialized from the files built using https://github.com/facebookresearch/DPR.
+    We use default faiss index parameters as specified in that repository.
+
+    Args:
+        vector_size (:obj:`int`):
+            The dimension of indexed vectors.
+        index_path (:obj:`str`):
+            A path to a `directory` containing index files compatible with
+            :class:`~transformers.retrieval_rag.LegacyIndex`
+    """
+
+    INDEX_FILENAME = "hf_bert_base.hnswSQ8_correct_phi_128.c_index"
+    PASSAGE_FILENAME = "psgs_w100.tsv.pkl"
+
+    def __init__(self, vector_size, index_path):
+        self.index_id_to_db_id = []
+        self.index_path = index_path
+        self.passages = self._load_passages()
+        self.vector_size = vector_size
+        self.index = None
+        self._index_initialize = False
+
+    def _resolve_path(self, index_path, filename):
+        assert os.path.isdir(index_path) or is_remote_url(index_path), "Please specify a valid ``index_path``."
+        archive_file = os.path.join(index_path, filename)
+        try:
+            # Load from URL or cache if already cached
+            resolved_archive_file = cached_path(archive_file)
+            if resolved_archive_file is None:
+                raise EnvironmentError
+        except EnvironmentError:
+            msg = (
+                f"Can't load '{archive_file}'. Make sure that:\n\n"
+                f"- '{index_path}' is a correct remote path to a directory containing a file named {filename}"
+                f"- or '{index_path}' is the correct path to a directory containing a file named {filename}.\n\n"
+            )
+            raise EnvironmentError(msg)
+        if resolved_archive_file == archive_file:
+            logger.info("loading file {}".format(archive_file))
+        else:
+            logger.info("loading file {} from cache at {}".format(archive_file, resolved_archive_file))
+        return resolved_archive_file
+
+    def _load_passages(self):
+        logger.info("Loading passages from {}".format(self.index_path))
+        passages_path = self._resolve_path(self.index_path, self.PASSAGE_FILENAME)
+        with open(passages_path, "rb") as passages_file:
+            passages = pickle.load(passages_file)
+        return passages
+
+    def _deserialize_index(self):
+        logger.info("Loading index from {}".format(self.index_path))
+        resolved_index_path = self._resolve_path(self.index_path, self.INDEX_FILENAME + ".index.dpr")
+        self.index = faiss.read_index(resolved_index_path)
+        resolved_meta_path = self._resolve_path(self.index_path, self.INDEX_FILENAME + ".index_meta.dpr")
+        with open(resolved_meta_path, "rb") as metadata_file:
+            self.index_id_to_db_id = pickle.load(metadata_file)
+        assert (
+            len(self.index_id_to_db_id) == self.index.ntotal
+        ), "Deserialized index_id_to_db_id should match faiss index size"
+
+    def is_initialized(self):
+        return self._index_initialize
+
+    def init_index(self):
+        index = faiss.IndexHNSWFlat(self.vector_size + 1, 512)
+        index.hnsw.efSearch = 128
+        index.hnsw.efConstruction = 200
+        self.index = index
+        self._deserialize_index()
+        self._index_initialize = True
+
+    def get_doc_dicts(self, doc_ids: np.array):
+        doc_list = []
+        for doc_ids_i in doc_ids:
+            ids = [str(int(doc_id)) for doc_id in doc_ids_i]
+            docs = [self.passages[doc_id] for doc_id in ids]
+            doc_list.append(docs)
+        doc_dicts = []
+        for docs in doc_list:
+            doc_dict = {}
+            doc_dict["title"] = [doc[1] for doc in docs]
+            doc_dict["text"] = [doc[0] for doc in docs]
+            doc_dicts.append(doc_dict)
+        return doc_dicts
+
+    def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
+        aux_dim = np.zeros(len(question_hidden_states), dtype="float32").reshape(-1, 1)
+        query_nhsw_vectors = np.hstack((question_hidden_states, aux_dim))
+        _, docs_ids = self.index.search(query_nhsw_vectors, n_docs)
+        vectors = [[self.index.reconstruct(int(doc_id))[:-1] for doc_id in doc_ids] for doc_ids in docs_ids]
+        ids = [[int(self.index_id_to_db_id[doc_id]) for doc_id in doc_ids] for doc_ids in docs_ids]
+        return np.array(ids), np.array(vectors)
+
+
+class HFIndex:
+    """
+    A wrapper around an instance of :class:`~datasets.Datasets`. If ``index_path`` is set to ``None``,
+    we load the pre-computed index available with the :class:`~datasets.arrow_dataset.Dataset`, otherwise, we load the index from the indicated path on disk.
+
+    Args:
+        dataset (:obj:`str`, optional, defaults to ``wiki_dpr``):
+            A datatset identifier of the indexed dataset on HuggingFace AWS bucket (list all available datasets and ids with ``datasets.list_datasets()``).
+        dataset_split (:obj:`str`, optional, defaults to ``train``)
+            Which split of the ``dataset`` to load.
+        index_name (:obj:`str`, optional, defaults to ``train``)
+            The index_name of the index associated with the ``dataset``. The index loaded from ``index_path`` will be saved under this name.
+        index_path (:obj:`str`, optional, defaults to ``None``)
+            The path to the serialized faiss index on disk.
+    """
+
+    def __init__(
+        self,
+        dataset_name: str,
+        dataset_split: str,
+        index_name: str,
+        index_path: Optional[str] = None,
+        use_dummy_dataset=False,
+    ):
+        super().__init__()
+        self.dataset_name = dataset_name
+        self.dataset_split = dataset_split
+        self.index_name = index_name
+        self.index_path = index_path
+        self.use_dummy_dataset = use_dummy_dataset
+        self._index_initialize = False
+
+        logger.info("Loading passages from {}".format(self.dataset_name))
+        self.dataset = load_dataset(
+            self.dataset_name, with_index=False, split=self.dataset_split, dummy=self.use_dummy_dataset
+        )
+
+    def is_initialized(self):
+        return self._index_initialize
+
+    def init_index(self):
+        if self.index_path is not None:
+            logger.info("Loading index from {}".format(self.index_path))
+            self.index.load_faiss_index(index_name=self.index_name, file=self.index_path)
+        else:
+            logger.info("Loading index from {}".format(self.dataset_name + " with index name " + self.index_name))
+            self.dataset = load_dataset(
+                self.dataset_name,
+                with_embeddings=True,
+                with_index=True,
+                split=self.dataset_split,
+                index_name=self.index_name,
+                dummy=self.use_dummy_dataset,
+            )
+        self._index_initialize = True
+
+    def get_doc_dicts(self, doc_ids: np.ndarray) -> List[dict]:
+        return [self.dataset[doc_ids[i].tolist()] for i in range(doc_ids.shape[0])]
+
+    def get_top_docs(self, question_hidden_states: np.ndarray, n_docs=5) -> Tuple[np.ndarray, np.ndarray]:
+        _, docs = self.dataset.get_nearest_examples_batch("embeddings", question_hidden_states, n_docs)
+        ids = [[int(i) for i in doc["id"]] for doc in docs]
+        vectors = [doc["embeddings"] for doc in docs]
+        return np.array(ids), np.array(vectors)  # shapes (batch_size, n_docs) and (batch_size, n_docs, d)
+
+
+class RagRetriever:
+    """
+    Retriever used to get documents from vector queries.
+    It retrieves the documents embeddings as well as the documents contents, and it formats them to be used with a RagModel.
+
+    Args:
+        config (:class:`~transformers.RagConfig`):
+            The configuration of the RAG model this Retriever is used with. Contains parameters indicating which ``Index`` to build.
+        question_encoder_tokenizer (:class:`~transformers.PreTrainedTokenizer`):
+            The tokenizer that was used to tokenize the question.
+            It is used to decode the question and then use the generator_tokenizer.
+        generator_tokenizer (:class:`~transformers.PreTrainedTokenizer`):
+            The tokenizer used for the generator part of the RagModel.
+    """
+
+    _init_retrieval = True
+
+    def __init__(self, config, question_encoder_tokenizer, generator_tokenizer):
+        super().__init__()
+        self.index = (
+            LegacyIndex(
+                config.retrieval_vector_size,
+                config.index_path or LEGACY_INDEX_PATH,
+            )
+            if config.index_name == "legacy"
+            else HFIndex(
+                config.dataset, config.dataset_split, config.index_name, config.index_path, config.use_dummy_dataset
+            )
+        )
+        self.generator_tokenizer = generator_tokenizer
+        self.question_encoder_tokenizer = question_encoder_tokenizer
+
+        self.n_docs = config.n_docs
+        self.batch_size = config.retrieval_batch_size
+
+        self.config = config
+        if self._init_retrieval:
+            self.init_retrieval()
+
+    @classmethod
+    def from_pretrained(cls, retriever_name_or_path, **kwargs):
+        config = RagConfig.from_pretrained(retriever_name_or_path, **kwargs)
+        rag_tokenizer = RagTokenizer.from_pretrained(retriever_name_or_path, config=config)
+        question_encoder_tokenizer = rag_tokenizer.question_encoder
+        generator_tokenizer = rag_tokenizer.generator
+        return cls(
+            config, question_encoder_tokenizer=question_encoder_tokenizer, generator_tokenizer=generator_tokenizer
+        )
+
+    def save_pretrained(self, save_directory):
+        self.config.save_pretrained(save_directory)
+        rag_tokenizer = RagTokenizer(
+            question_encoder_tokenizer=self.question_encoder_tokenizer,
+            generator_tokenizer=self.generator_tokenizer,
+        )
+        rag_tokenizer.save_pretrained(save_directory)
+
+    def init_retrieval(self):
+        """
+        Retriever initalization function. It loads the index into memory.
+        """
+
+        logger.info("initializing retrieval")
+        self.index.init_index()
+
+    def postprocess_docs(self, docs, input_strings, prefix, n_docs, return_tensors=None):
+        r"""
+        Postprocessing retrieved ``docs`` and combining them with ``input_strings``.
+
+        Args:
+            doc_scores (:obj:`np.ndarray` of shape :obj:`(batch_size, n_docs)`):
+                Retrieval scores of respective docs - passed for logging.
+            docs  (:obj:`dict`):
+                Retrieved documents.
+            input_strings (:obj:`str`):
+                Input strings decoded by ``preprocess_query``.
+            prefix (:obj:`str`):
+                Prefix added at the beginning of each input, typically used with T5-based models.
+
+        Return:
+            :obj:`tuple(tensors)`:
+                a tuple consisting of two elements: contextualized ``input_ids`` and a compatible ``attention_mask``.
+        """
+
+        def cat_input_and_doc(doc_title, doc_text, input_string, prefix):
+            # TODO(Patrick): if we train more RAG models, I want to put the input first to take advantage of effortless truncation
+            # TODO(piktus): better handling of truncation
+            if doc_title.startswith('"'):
+                doc_title = doc_title[1:]
+            if doc_title.endswith('"'):
+                doc_title = doc_title[:-1]
+            if prefix is None:
+                prefix = ""
+            out = (prefix + doc_title + self.config.title_sep + doc_text + self.config.doc_sep + input_string).replace(
+                "  ", " "
+            )
+            return out
+
+        rag_input_strings = [
+            cat_input_and_doc(
+                docs[i]["title"][j],
+                docs[i]["text"][j],
+                input_strings[i],
+                prefix,
+            )
+            for i in range(len(docs))
+            for j in range(n_docs)
+        ]
+
+        contextualized_inputs = self.generator_tokenizer.batch_encode_plus(
+            rag_input_strings,
+            max_length=self.config.max_combined_length,
+            return_tensors=return_tensors,
+            padding="max_length",
+            truncation=True,
+        )
+
+        return contextualized_inputs["input_ids"], contextualized_inputs["attention_mask"]
+
+    def _chunk_tensor(self, t: Iterable, chunk_size: int) -> List[Iterable]:
+        return [t[i : i + chunk_size] for i in range(0, len(t), chunk_size)]
+
+    def _main_retrieve(self, question_hidden_states: np.ndarray, n_docs: int) -> Tuple[np.ndarray, np.ndarray]:
+        question_hidden_states_batched = self._chunk_tensor(question_hidden_states, self.batch_size)
+        ids_batched = []
+        vectors_batched = []
+        for question_hidden_states in question_hidden_states_batched:
+            start_time = time.time()
+            ids, vectors = self.index.get_top_docs(question_hidden_states, n_docs)
+            logger.debug(
+                "index search time: {} sec, batch size {}".format(
+                    time.time() - start_time, question_hidden_states.shape
+                )
+            )
+            ids_batched.extend(ids)
+            vectors_batched.extend(vectors)
+        return np.array(ids_batched), np.array(
+            vectors_batched
+        )  # shapes (batch_size, n_docs) and (batch_size, n_docs, d)
+
+    def retrieve(self, question_hidden_states: np.ndarray, n_docs: int) -> Tuple[np.ndarray, List[dict]]:
+        """
+        Retrieves documents for specified ``question_hidden_states``.
+
+        Args:
+            question_hidden_states (:obj:`np.ndarray` of shape :obj:`(batch_size, vector_size)`):
+                A batch of query vectors to retrieve with.
+            n_docs (:obj:`int`):
+                The number of docs retrieved per query.
+
+        Return:
+            retrieved_doc_embeds (:obj:`np.ndarray` of shape :obj:`(batch_size, n_docs, dim)`
+                The retrieval embeddings of the retrieved docs per query.
+            doc_ids (:obj:`np.ndarray` of shape :obj:`batch_size, n_docs`)
+                The ids of the documents in the index
+            doc_dicts (:obj:`List[dict]`):
+                The retrieved_doc_embeds examples per query.
+        """
+
+        doc_ids, retrieved_doc_embeds = self._main_retrieve(question_hidden_states, n_docs)
+        return retrieved_doc_embeds, doc_ids, self.index.get_doc_dicts(doc_ids)
+
+    def __call__(
+        self,
+        question_input_ids: List[List[int]],
+        question_hidden_states: np.ndarray,
+        prefix=None,
+        n_docs=None,
+        return_tensors=None,
+    ) -> BatchEncoding:
+        """
+        Retrieves documents for specified :obj:`question_hidden_states`.
+
+        Args:
+            question_input_ids: (:obj:`List[List[int]]`) batch of input ids
+            question_hidden_states (:obj:`np.ndarray` of shape :obj:`(batch_size, vector_size)`:
+                A batch of query vectors to retrieve with.
+            prefix: (:obj:`str`, `optional`):
+                The prefix used by the generator's tokenizer.
+            n_docs (:obj:`int`, `optional`):
+                The number of docs retrieved per query.
+            return_tensors (:obj:`str` or :class:`~transformers.tokenization_utils_base.TensorType`, `optional`, defaults to "pt"):
+                If set, will return tensors instead of list of python integers. Acceptable values are:
+
+                * :obj:`'tf'`: Return TensorFlow :obj:`tf.constant` objects.
+                * :obj:`'pt'`: Return PyTorch :obj:`torch.Tensor` objects.
+                * :obj:`'np'`: Return Numpy :obj:`np.ndarray` objects.
+
+        Output:
+            :class:`~transformers.BatchEncoding`: A :class:`~transformers.BatchEncoding` with the following fields:
+
+            - **context_input_ids** -- List of token ids to be fed to a model.
+
+              `What are input IDs? <../glossary.html#input-ids>`__
+            - **context_attention_mask** -- List of indices specifying which tokens should be attended to by the model (when
+              :obj:`return_attention_mask=True` or if `"attention_mask"` is in :obj:`self.model_input_names`).
+
+              `What are attention masks? <../glossary.html#attention-mask>`__
+            - **retrieved_doc_embeds** -- List of embeddings of the retrieved documents
+            - **doc_ids** -- List of ids of the retrieved documents
+        """
+
+        n_docs = n_docs if n_docs is not None else self.n_docs
+        prefix = prefix if prefix is not None else self.config.generator.prefix
+        retrieved_doc_embeds, doc_ids, docs = self.retrieve(question_hidden_states, n_docs)
+
+        input_strings = self.question_encoder_tokenizer.batch_decode(question_input_ids, skip_special_tokens=True)
+        context_input_ids, context_attention_mask = self.postprocess_docs(
+            docs, input_strings, prefix, n_docs, return_tensors=return_tensors
+        )
+
+        return BatchEncoding(
+            {
+                "context_input_ids": context_input_ids,
+                "context_attention_mask": context_attention_mask,
+                "retrieved_doc_embeds": retrieved_doc_embeds,
+                "doc_ids": doc_ids,
+            },
+            tensor_type=return_tensors,
+        )

src/transformers/testing_utils.py

@@ -10,7 +10,7 @@ from distutils.util import strtobool
 from io import StringIO
 from pathlib import Path
 
-from .file_utils import _tf_available, _torch_available, _torch_tpu_available
+from .file_utils import _datasets_available, _faiss_available, _tf_available, _torch_available, _torch_tpu_available
 
 
 SMALL_MODEL_IDENTIFIER = "julien-c/bert-xsmall-dummy"
@@ -161,6 +161,21 @@ def require_torch_and_cuda(test_case):
         return test_case
 
 
+def require_datasets(test_case):
+    """Decorator marking a test that requires datasets."""
+
+    if not _datasets_available:
+        test_case = unittest.skip("test requires Datasets")(test_case)
+    return test_case
+
+
+def require_faiss(test_case):
+    """Decorator marking a test that requires faiss."""
+    if not _faiss_available:
+        test_case = unittest.skip("test requires Faiss")(test_case)
+    return test_case
+
+
 def get_tests_dir():
     """
     returns the full path to the `tests` dir, so that the tests can be invoked from anywhere

src/transformers/tokenization_auto.py

@@ -26,12 +26,14 @@ from .configuration_auto import (
     CamembertConfig,
     CTRLConfig,
     DistilBertConfig,
+    DPRConfig,
     ElectraConfig,
     EncoderDecoderConfig,
     FlaubertConfig,
     FSMTConfig,
     FunnelConfig,
     GPT2Config,
+    LayoutLMConfig,
     LongformerConfig,
     LxmertConfig,
     MarianConfig,
@@ -39,6 +41,7 @@ from .configuration_auto import (
     MobileBertConfig,
     OpenAIGPTConfig,
     PegasusConfig,
+    RagConfig,
     ReformerConfig,
     RetriBertConfig,
     RobertaConfig,
@@ -59,11 +62,13 @@ from .tokenization_bertweet import BertweetTokenizer
 from .tokenization_camembert import CamembertTokenizer
 from .tokenization_ctrl import CTRLTokenizer
 from .tokenization_distilbert import DistilBertTokenizer, DistilBertTokenizerFast
+from .tokenization_dpr import DPRQuestionEncoderTokenizer, DPRQuestionEncoderTokenizerFast
 from .tokenization_electra import ElectraTokenizer, ElectraTokenizerFast
 from .tokenization_flaubert import FlaubertTokenizer
 from .tokenization_fsmt import FSMTTokenizer
 from .tokenization_funnel import FunnelTokenizer, FunnelTokenizerFast
 from .tokenization_gpt2 import GPT2Tokenizer, GPT2TokenizerFast
+from .tokenization_layoutlm import LayoutLMTokenizer, LayoutLMTokenizerFast
 from .tokenization_longformer import LongformerTokenizer, LongformerTokenizerFast
 from .tokenization_lxmert import LxmertTokenizer, LxmertTokenizerFast
 from .tokenization_marian import MarianTokenizer
@@ -72,6 +77,7 @@ from .tokenization_mobilebert import MobileBertTokenizer, MobileBertTokenizerFas
 from .tokenization_openai import OpenAIGPTTokenizer, OpenAIGPTTokenizerFast
 from .tokenization_pegasus import PegasusTokenizer
 from .tokenization_phobert import PhobertTokenizer
+from .tokenization_rag import RagTokenizer
 from .tokenization_reformer import ReformerTokenizer
 from .tokenization_retribert import RetriBertTokenizer, RetriBertTokenizerFast
 from .tokenization_roberta import RobertaTokenizer, RobertaTokenizerFast
@@ -107,6 +113,8 @@ TOKENIZER_MAPPING = OrderedDict(
         (ElectraConfig, (ElectraTokenizer, ElectraTokenizerFast)),
         (FunnelConfig, (FunnelTokenizer, FunnelTokenizerFast)),
         (LxmertConfig, (LxmertTokenizer, LxmertTokenizerFast)),
+        (LayoutLMConfig, (LayoutLMTokenizer, LayoutLMTokenizerFast)),
+        (DPRConfig, (DPRQuestionEncoderTokenizer, DPRQuestionEncoderTokenizerFast)),
         (BertConfig, (BertTokenizer, BertTokenizerFast)),
         (OpenAIGPTConfig, (OpenAIGPTTokenizer, OpenAIGPTTokenizerFast)),
         (GPT2Config, (GPT2Tokenizer, GPT2TokenizerFast)),
@@ -117,6 +125,8 @@ TOKENIZER_MAPPING = OrderedDict(
         (CTRLConfig, (CTRLTokenizer, None)),
         (FSMTConfig, (FSMTTokenizer, None)),
         (BertGenerationConfig, (BertGenerationTokenizer, None)),
+        (LayoutLMConfig, (LayoutLMTokenizer, None)),
+        (RagConfig, (RagTokenizer, None)),
     ]
 )
 

src/transformers/tokenization_gpt2.py

@@ -17,6 +17,7 @@
 
 import json
 import os
+import warnings
 from functools import lru_cache
 
 import regex as re
@@ -121,7 +122,7 @@ class GPT2Tokenizer(PreTrainedTokenizer):
 
     .. note::
 
-        When used with ``is_pretokenized=True``, this tokenizer will add a space before each word (even the first one).
+        When used with ``is_split_into_words=True``, this tokenizer will add a space before each word (even the first one).
 
     This tokenizer inherits from :class:`~transformers.PreTrainedTokenizer` which contains most of the methods. Users
     should refer to the superclass for more information regarding methods.
@@ -288,9 +289,16 @@ class GPT2Tokenizer(PreTrainedTokenizer):
 
         return vocab_file, merge_file
 
-    def prepare_for_tokenization(self, text, is_pretokenized=False, **kwargs):
+    def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            is_split_into_words = kwargs.pop("is_pretokenized")
+
         add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space)
-        if is_pretokenized or add_prefix_space:
+        if is_split_into_words or add_prefix_space:
             text = " " + text
         return (text, kwargs)
 
@@ -317,7 +325,7 @@ class GPT2TokenizerFast(PreTrainedTokenizerFast):
 
     .. note::
 
-        When used with ``is_pretokenized=True``, this tokenizer needs to be instantiated with
+        When used with ``is_split_into_words=True``, this tokenizer needs to be instantiated with
         ``add_prefix_space=True``.
 
     This tokenizer inherits from :class:`~transformers.PreTrainedTokenizer` which contains most of the methods. Users
@@ -377,9 +385,15 @@ class GPT2TokenizerFast(PreTrainedTokenizerFast):
         self.add_prefix_space = add_prefix_space
 
     def _batch_encode_plus(self, *args, **kwargs) -> BatchEncoding:
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            is_split_into_words = kwargs.pop("is_pretokenized")
 
-        is_pretokenized = kwargs.get("is_pretokenized", False)
-        assert self.add_prefix_space or not is_pretokenized, (
+        is_split_into_words = kwargs.get("is_split_into_words", False)
+        assert self.add_prefix_space or not is_split_into_words, (
             f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
             "to use it with pretokenized inputs."
         )
@@ -387,9 +401,15 @@ class GPT2TokenizerFast(PreTrainedTokenizerFast):
         return super()._batch_encode_plus(*args, **kwargs)
 
     def _encode_plus(self, *args, **kwargs) -> BatchEncoding:
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            is_split_into_words = kwargs.pop("is_pretokenized")
 
-        is_pretokenized = kwargs.get("is_pretokenized", False)
-        assert self.add_prefix_space or not is_pretokenized, (
+        is_split_into_words = kwargs.get("is_split_into_words", False)
+        assert self.add_prefix_space or not is_split_into_words, (
             f"You need to instantiate {self.__class__.__name__} with add_prefix_space=True "
             "to use it with pretokenized inputs."
         )

src/transformers/tokenization_layoutlm.py

@@ -0,0 +1,78 @@
+# coding=utf-8
+# Copyright 2018 The Microsoft Research Asia LayoutLM Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Tokenization class for model LayoutLM."""
+
+
+from .tokenization_bert import BertTokenizer, BertTokenizerFast
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+VOCAB_FILES_NAMES = {"vocab_file": "vocab.txt"}
+
+PRETRAINED_VOCAB_FILES_MAP = {
+    "vocab_file": {
+        "microsoft/layoutlm-base-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-vocab.txt",
+        "microsoft/layoutlm-large-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-vocab.txt",
+    }
+}
+
+
+PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES = {
+    "microsoft/layoutlm-base-uncased": 512,
+    "microsoft/layoutlm-large-uncased": 512,
+}
+
+
+PRETRAINED_INIT_CONFIGURATION = {
+    "microsoft/layoutlm-base-uncased": {"do_lower_case": True},
+    "microsoft/layoutlm-large-uncased": {"do_lower_case": True},
+}
+
+
+class LayoutLMTokenizer(BertTokenizer):
+    r"""
+    Constructs a LayoutLM tokenizer.
+
+    :class:`~transformers.LayoutLMTokenizer is identical to :class:`~transformers.BertTokenizer` and runs end-to-end
+    tokenization: punctuation splitting + wordpiece.
+
+    Refer to superclass :class:`~transformers.BertTokenizer` for usage examples and documentation concerning
+    parameters.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    pretrained_init_configuration = PRETRAINED_INIT_CONFIGURATION
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+
+
+class LayoutLMTokenizerFast(BertTokenizerFast):
+    r"""
+    Constructs a  "Fast" LayoutLMTokenizer.
+
+    :class:`~transformers.LayoutLMTokenizerFast` is identical to :class:`~transformers.BertTokenizerFast` and runs end-to-end
+    tokenization: punctuation splitting + wordpiece.
+
+    Refer to superclass :class:`~transformers.BertTokenizerFast` for usage examples and documentation concerning
+    parameters.
+    """
+
+    vocab_files_names = VOCAB_FILES_NAMES
+    pretrained_vocab_files_map = PRETRAINED_VOCAB_FILES_MAP
+    max_model_input_sizes = PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
+    pretrained_init_configuration = PRETRAINED_INIT_CONFIGURATION
+    model_input_names = ["attention_mask"]

src/transformers/tokenization_rag.py

@@ -0,0 +1,160 @@
+# coding=utf-8
+# Copyright 2020, The RAG Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""Tokenization classes for RAG."""
+import os
+from typing import List, Optional
+
+from .configuration_rag import RagConfig
+from .tokenization_utils_base import BatchEncoding
+from .utils import logging
+
+
+logger = logging.get_logger(__name__)
+
+
+class RagTokenizer:
+    def __init__(self, question_encoder, generator):
+        self.question_encoder = question_encoder
+        self.generator = generator
+
+    def save_pretrained(self, save_directory):
+        if os.path.isfile(save_directory):
+            raise ValueError("Provided path ({}) should be a directory, not a file".format(save_directory))
+        os.makedirs(save_directory, exist_ok=True)
+        question_encoder_path = os.path.join(save_directory, "question_encoder_tokenizer")
+        generator_path = os.path.join(save_directory, "generator_tokenizer")
+        self.question_encoder.save_pretrained(question_encoder_path)
+        self.generator.save_pretrained(generator_path)
+
+    @classmethod
+    def from_pretrained(cls, pretrained_model_name_or_path, **kwargs):
+        # dynamically import AutoTokenizer
+        from .tokenization_auto import AutoTokenizer
+
+        config = kwargs.pop("config", None)
+
+        if config is None:
+            config = RagConfig.from_pretrained(pretrained_model_name_or_path)
+
+        question_encoder_path = os.path.join(pretrained_model_name_or_path, "question_encoder_tokenizer")
+        generator_path = os.path.join(pretrained_model_name_or_path, "generator_tokenizer")
+        question_encoder = AutoTokenizer.from_pretrained(question_encoder_path, config=config.question_encoder)
+        generator = AutoTokenizer.from_pretrained(generator_path, config=config.generator)
+        return cls(question_encoder=question_encoder, generator=generator)
+
+    def __call__(self, *args, **kwargs):
+        return self.question_encoder(*args, **kwargs)
+
+    def batch_decode(self, *args, **kwargs):
+        return self.generator.batch_decode(*args, **kwargs)
+
+    def prepare_seq2seq_batch(
+        self,
+        src_texts: List[str],
+        tgt_texts: Optional[List[str]] = None,
+        max_length: Optional[int] = None,
+        max_target_length: Optional[int] = None,
+        padding: str = "longest",
+        return_tensors: str = "np",
+        truncation=True,
+        **kwargs,
+    ) -> BatchEncoding:
+        r"""
+
+        Prepare a batch that can be passed directly to an instance of :class:`~transformers.RagModel`.
+
+        Args:
+            src_texts: (:obj:`List[str]`):
+                List of documents to summarize or source language texts.
+            tgt_texts: (:obj:`List[str]`, `optional`):
+                List of summaries or target language texts.
+            max_length (:obj:`int`, `optional`):
+                Controls the maximum length for encoder inputs (documents to summarize or source language texts).
+                If left unset or set to :obj:`None`, this will use the predefined model maximum length if a maximum
+                length is required by one of the truncation/padding parameters. If the model has no specific maximum
+                input length (like XLNet) truncation/padding to a maximum length will be deactivated.
+            max_target_length (:obj:`int`, `optional`):
+                Controls the maximum length of decoder inputs (target language texts or summaries).
+                If left unset or set to :obj:`None`, this will use the max_length value.
+            padding (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.PaddingStrategy`, `optional`, defaults to :obj:`False`):
+                Activates and controls padding. Accepts the following values:
+
+                * :obj:`True` or :obj:`'longest'`: Pad to the longest sequence in the batch (or no padding if only a
+                  single sequence if provided).
+                * :obj:`'max_length'`: Pad to a maximum length specified with the argument :obj:`max_length` or to the
+                  maximum acceptable input length for the model if that argument is not provided.
+                * :obj:`False` or :obj:`'do_not_pad'` (default): No padding (i.e., can output a batch with sequences of
+                  different lengths).
+            return_tensors (:obj:`str` or :class:`~transformers.tokenization_utils_base.TensorType`, `optional`, defaults to "pt"):
+                If set, will return tensors instead of list of python integers. Acceptable values are:
+
+                * :obj:`'tf'`: Return TensorFlow :obj:`tf.constant` objects.
+                * :obj:`'pt'`: Return PyTorch :obj:`torch.Tensor` objects.
+                * :obj:`'np'`: Return Numpy :obj:`np.ndarray` objects.
+            truncation (:obj:`bool`, :obj:`str` or :class:`~transformers.tokenization_utils_base.TruncationStrategy`, `optional`, defaults to :obj:`True`):
+                Activates and controls truncation. Accepts the following values:
+
+                * :obj:`True` or :obj:`'longest_first'`: Truncate to a maximum length specified with the argument
+                  :obj:`max_length` or to the maximum acceptable input length for the model if that argument is not
+                  provided. This will truncate token by token, removing a token from the longest sequence in the pair
+                  if a pair of sequences (or a batch of pairs) is provided.
+                * :obj:`'only_first'`: Truncate to a maximum length specified with the argument :obj:`max_length` or to
+                  the maximum acceptable input length for the model if that argument is not provided. This will only
+                  truncate the first sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
+                * :obj:`'only_second'`: Truncate to a maximum length specified with the argument :obj:`max_length` or
+                  to the maximum acceptable input length for the model if that argument is not provided. This will only
+                  truncate the second sequence of a pair if a pair of sequences (or a batch of pairs) is provided.
+                * :obj:`False` or :obj:`'do_not_truncate'` (default): No truncation (i.e., can output batch with
+                  sequence lengths greater than the model maximum admissible input size).
+            **kwargs:
+                Additional keyword arguments passed along to :obj:`self.__call__`.
+
+        Returns:
+            :class:`~transformers.BatchEncoding`: A :class:`~transformers.BatchEncoding` with the following fields:
+
+            - **input_ids** -- List of token ids to be fed to the encoder.
+            - **attention_mask** -- List of indices specifying which tokens should be attended to by the model.
+            - **labels** -- List of token ids for tgt_texts
+
+            The full set of keys ``[input_ids, attention_mask, labels]``,
+            will only be returned if tgt_texts is passed. Otherwise, input_ids, attention_mask will be the only keys.
+        """
+        if max_length is None:
+            max_length = self.question_encoder.model_max_length
+        model_inputs: BatchEncoding = self.question_encoder(
+            src_texts,
+            add_special_tokens=True,
+            return_tensors=return_tensors,
+            max_length=max_length,
+            padding=padding,
+            truncation=truncation,
+            **kwargs,
+        )
+        if tgt_texts is None:
+            return model_inputs
+        # Process tgt_texts
+        if max_target_length is None:
+            max_target_length = self.generator.model_max_length
+        labels = self.generator(
+            tgt_texts,
+            add_special_tokens=True,
+            return_tensors=return_tensors,
+            padding=padding,
+            max_length=max_target_length,
+            truncation=truncation,
+            **kwargs,
+        )["input_ids"]
+        model_inputs["labels"] = labels
+        return model_inputs

src/transformers/tokenization_roberta.py

@@ -14,7 +14,7 @@
 # limitations under the License.
 """Tokenization classes for RoBERTa."""
 
-
+import warnings
 from typing import List, Optional
 
 from tokenizers.processors import RobertaProcessing
@@ -81,7 +81,7 @@ class RobertaTokenizer(GPT2Tokenizer):
 
     .. note::
 
-        When used with ``is_pretokenized=True``, this tokenizer will add a space before each word (even the first one).
+        When used with ``is_split_into_words=True``, this tokenizer will add a space before each word (even the first one).
 
     This tokenizer inherits from :class:`~transformers.PreTrainedTokenizer` which contains most of the methods. Users
     should refer to the superclass for more information regarding methods.
@@ -251,9 +251,16 @@ class RobertaTokenizer(GPT2Tokenizer):
             return len(cls + token_ids_0 + sep) * [0]
         return len(cls + token_ids_0 + sep + sep + token_ids_1 + sep) * [0]
 
-    def prepare_for_tokenization(self, text, is_pretokenized=False, **kwargs):
+    def prepare_for_tokenization(self, text, is_split_into_words=False, **kwargs):
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            is_split_into_words = kwargs.pop("is_pretokenized")
+
         add_prefix_space = kwargs.pop("add_prefix_space", self.add_prefix_space)
-        if (is_pretokenized or add_prefix_space) and (len(text) > 0 and not text[0].isspace()):
+        if (is_split_into_words or add_prefix_space) and (len(text) > 0 and not text[0].isspace()):
             text = " " + text
         return (text, kwargs)
 
@@ -280,7 +287,7 @@ class RobertaTokenizerFast(GPT2TokenizerFast):
 
     .. note::
 
-        When used with ``is_pretokenized=True``, this tokenizer needs to be instantiated with
+        When used with ``is_split_into_words=True``, this tokenizer needs to be instantiated with
         ``add_prefix_space=True``.
 
     This tokenizer inherits from :class:`~transformers.PreTrainedTokenizerFast` which contains most of the methods. Users

src/transformers/tokenization_utils.py

@@ -19,6 +19,7 @@
 import itertools
 import re
 import unicodedata
+import warnings
 from typing import Any, Dict, List, Optional, Tuple, Union, overload
 
 from .file_utils import add_end_docstrings
@@ -250,6 +251,12 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
         Returns:
             :obj:`List[str]`: The list of tokens.
         """
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            kwargs["is_split_into_words"] = kwargs.pop("is_pretokenized")
         # Simple mapping string => AddedToken for special tokens with specific tokenization behaviors
         all_special_tokens_extended = dict(
             (str(t), t) for t in self.all_special_tokens_extended if isinstance(t, AddedToken)
@@ -402,7 +409,7 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
         truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[Union[str, TensorType]] = None,
         return_token_type_ids: Optional[bool] = None,
@@ -419,17 +426,19 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
                 tokens = self.tokenize(text, **kwargs)
                 return self.convert_tokens_to_ids(tokens)
             elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], str):
-                if is_pretokenized:
-                    tokens = list(itertools.chain(*(self.tokenize(t, is_pretokenized=True, **kwargs) for t in text)))
+                if is_split_into_words:
+                    tokens = list(
+                        itertools.chain(*(self.tokenize(t, is_split_into_words=True, **kwargs) for t in text))
+                    )
                     return self.convert_tokens_to_ids(tokens)
                 else:
                     return self.convert_tokens_to_ids(text)
             elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], int):
                 return text
             else:
-                if is_pretokenized:
+                if is_split_into_words:
                     raise ValueError(
-                        f"Input {text} is not valid. Should be a string or a list/tuple of strings when `is_pretokenized=True`."
+                        f"Input {text} is not valid. Should be a string or a list/tuple of strings when `is_split_into_words=True`."
                     )
                 else:
                     raise ValueError(
@@ -445,6 +454,13 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
                 "https://github.com/huggingface/transformers/pull/2674"
             )
 
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            is_split_into_words = kwargs.pop("is_pretokenized")
+
         first_ids = get_input_ids(text)
         second_ids = get_input_ids(text_pair) if text_pair is not None else None
 
@@ -482,7 +498,7 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
         truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[Union[str, TensorType]] = None,
         return_token_type_ids: Optional[bool] = None,
@@ -499,8 +515,10 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
                 tokens = self.tokenize(text, **kwargs)
                 return self.convert_tokens_to_ids(tokens)
             elif isinstance(text, (list, tuple)) and len(text) > 0 and isinstance(text[0], str):
-                if is_pretokenized:
-                    tokens = list(itertools.chain(*(self.tokenize(t, is_pretokenized=True, **kwargs) for t in text)))
+                if is_split_into_words:
+                    tokens = list(
+                        itertools.chain(*(self.tokenize(t, is_split_into_words=True, **kwargs) for t in text))
+                    )
                     return self.convert_tokens_to_ids(tokens)
                 else:
                     return self.convert_tokens_to_ids(text)
@@ -518,11 +536,18 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
                 "transformers.PreTrainedTokenizerFast."
             )
 
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            is_split_into_words = kwargs.pop("is_pretokenized")
+
         input_ids = []
         for ids_or_pair_ids in batch_text_or_text_pairs:
             if not isinstance(ids_or_pair_ids, (list, tuple)):
                 ids, pair_ids = ids_or_pair_ids, None
-            elif is_pretokenized and not isinstance(ids_or_pair_ids[0], (list, tuple)):
+            elif is_split_into_words and not isinstance(ids_or_pair_ids[0], (list, tuple)):
                 ids, pair_ids = ids_or_pair_ids, None
             else:
                 ids, pair_ids = ids_or_pair_ids
@@ -616,7 +641,7 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
         return batch_outputs
 
     def prepare_for_tokenization(
-        self, text: str, is_pretokenized: bool = False, **kwargs
+        self, text: str, is_split_into_words: bool = False, **kwargs
     ) -> Tuple[str, Dict[str, Any]]:
         """
         Performs any necessary transformations before tokenization.
@@ -627,7 +652,7 @@ class PreTrainedTokenizer(PreTrainedTokenizerBase):
         Args:
             test (:obj:`str`):
                 The text to prepare.
-            is_pretokenized (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            is_split_into_words (:obj:`bool`, `optional`, defaults to :obj:`False`):
                 Whether or not the text has been pretokenized.
             kwargs:
                 Keyword arguments to use for the tokenization.

src/transformers/tokenization_utils_base.py

@@ -1088,7 +1088,7 @@ ENCODE_KWARGS_DOCSTRING = r"""
                 :obj:`return_overflowing_tokens=True` will contain some tokens from the end of the truncated sequence
                 returned to provide some overlap between truncated and overflowing sequences. The value of this
                 argument defines the number of overlapping tokens.
-            is_pretokenized (:obj:`bool`, `optional`, defaults to :obj:`False`):
+            is_split_into_words (:obj:`bool`, `optional`, defaults to :obj:`False`):
                 Whether or not the input is already pre-tokenized (e.g., split into words), in which case the tokenizer
                 will skip the pre-tokenization step. This is useful for NER or token classification.
             pad_to_multiple_of (:obj:`int`, `optional`):
@@ -1863,7 +1863,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
         truncation: Union[bool, str, TruncationStrategy] = False,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[Union[str, TensorType]] = None,
         return_token_type_ids: Optional[bool] = None,
@@ -1884,12 +1884,12 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
                 The sequence or batch of sequences to be encoded.
                 Each sequence can be a string or a list of strings (pretokenized string).
                 If the sequences are provided as list of strings (pretokenized), you must set
-                :obj:`is_pretokenized=True` (to lift the ambiguity with a batch of sequences).
+                :obj:`is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
             text_pair (:obj:`str`, :obj:`List[str]`, :obj:`List[List[str]]`):
                 The sequence or batch of sequences to be encoded.
                 Each sequence can be a string or a list of strings (pretokenized string).
                 If the sequences are provided as list of strings (pretokenized), you must set
-                :obj:`is_pretokenized=True` (to lift the ambiguity with a batch of sequences).
+                :obj:`is_split_into_words=True` (to lift the ambiguity with a batch of sequences).
         """
         # Input type checking for clearer error
         assert isinstance(text, str) or (
@@ -1928,8 +1928,10 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
         )
 
         is_batched = bool(
-            (not is_pretokenized and isinstance(text, (list, tuple)))
-            or (is_pretokenized and isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple)))
+            (not is_split_into_words and isinstance(text, (list, tuple)))
+            or (
+                is_split_into_words and isinstance(text, (list, tuple)) and text and isinstance(text[0], (list, tuple))
+            )
         )
 
         if is_batched:
@@ -1941,7 +1943,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
                 truncation=truncation,
                 max_length=max_length,
                 stride=stride,
-                is_pretokenized=is_pretokenized,
+                is_split_into_words=is_split_into_words,
                 pad_to_multiple_of=pad_to_multiple_of,
                 return_tensors=return_tensors,
                 return_token_type_ids=return_token_type_ids,
@@ -1962,7 +1964,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
                 truncation=truncation,
                 max_length=max_length,
                 stride=stride,
-                is_pretokenized=is_pretokenized,
+                is_split_into_words=is_split_into_words,
                 pad_to_multiple_of=pad_to_multiple_of,
                 return_tensors=return_tensors,
                 return_token_type_ids=return_token_type_ids,
@@ -1985,7 +1987,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
         truncation: Union[bool, str, TruncationStrategy] = False,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[Union[str, TensorType]] = None,
         return_token_type_ids: Optional[bool] = None,
@@ -2032,7 +2034,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
             truncation_strategy=truncation_strategy,
             max_length=max_length,
             stride=stride,
-            is_pretokenized=is_pretokenized,
+            is_split_into_words=is_split_into_words,
             pad_to_multiple_of=pad_to_multiple_of,
             return_tensors=return_tensors,
             return_token_type_ids=return_token_type_ids,
@@ -2054,7 +2056,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
         truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[Union[str, TensorType]] = None,
         return_token_type_ids: Optional[bool] = None,
@@ -2084,7 +2086,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
         truncation: Union[bool, str, TruncationStrategy] = False,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[Union[str, TensorType]] = None,
         return_token_type_ids: Optional[bool] = None,
@@ -2126,7 +2128,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
             truncation_strategy=truncation_strategy,
             max_length=max_length,
             stride=stride,
-            is_pretokenized=is_pretokenized,
+            is_split_into_words=is_split_into_words,
             pad_to_multiple_of=pad_to_multiple_of,
             return_tensors=return_tensors,
             return_token_type_ids=return_token_type_ids,
@@ -2154,7 +2156,7 @@ class PreTrainedTokenizerBase(SpecialTokensMixin):
         truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[Union[str, TensorType]] = None,
         return_token_type_ids: Optional[bool] = None,

src/transformers/tokenization_utils_fast.py

@@ -17,6 +17,7 @@
 """
 
 import os
+import warnings
 from collections import defaultdict
 from typing import Any, Dict, List, Optional, Tuple, Union
 
@@ -328,7 +329,7 @@ class PreTrainedTokenizerFast(PreTrainedTokenizerBase):
         truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[str] = None,
         return_token_type_ids: Optional[bool] = None,
@@ -346,6 +347,13 @@ class PreTrainedTokenizerFast(PreTrainedTokenizerBase):
                 "batch_text_or_text_pairs has to be a list (got {})".format(type(batch_text_or_text_pairs))
             )
 
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            is_split_into_words = kwargs.pop("is_pretokenized")
+
         if kwargs:
             raise ValueError(f"Keyword arguments {kwargs} not recognized.")
 
@@ -365,19 +373,21 @@ class PreTrainedTokenizerFast(PreTrainedTokenizerBase):
                 encodings = self._tokenizer.encode(
                     *batch_text_or_text_pairs[0],
                     add_special_tokens=add_special_tokens,
-                    is_pretokenized=is_pretokenized,
+                    is_pretokenized=is_split_into_words,
                 )
             else:
                 # We got a single sequence
                 encodings = self._tokenizer.encode(
                     batch_text_or_text_pairs[0],
                     add_special_tokens=add_special_tokens,
-                    is_pretokenized=is_pretokenized,
+                    is_pretokenized=is_split_into_words,
                 )
             encodings = [encodings]
         else:
             encodings = self._tokenizer.encode_batch(
-                batch_text_or_text_pairs, add_special_tokens=add_special_tokens, is_pretokenized=is_pretokenized
+                batch_text_or_text_pairs,
+                add_special_tokens=add_special_tokens,
+                is_pretokenized=is_split_into_words,
             )
 
         # Convert encoding to dict
@@ -423,7 +433,7 @@ class PreTrainedTokenizerFast(PreTrainedTokenizerBase):
         truncation_strategy: TruncationStrategy = TruncationStrategy.DO_NOT_TRUNCATE,
         max_length: Optional[int] = None,
         stride: int = 0,
-        is_pretokenized: bool = False,
+        is_split_into_words: bool = False,
         pad_to_multiple_of: Optional[int] = None,
         return_tensors: Optional[bool] = None,
         return_token_type_ids: Optional[bool] = None,
@@ -435,11 +445,17 @@ class PreTrainedTokenizerFast(PreTrainedTokenizerBase):
         verbose: bool = True,
         **kwargs
     ) -> BatchEncoding:
+        if "is_pretokenized" in kwargs:
+            warnings.warn(
+                "`is_pretokenized` is deprecated and will be removed in a future version, use `is_split_into_words` instead.",
+                FutureWarning,
+            )
+            is_split_into_words = kwargs.pop("is_pretokenized")
 
         batched_input = [(text, text_pair)] if text_pair else [text]
         batched_output = self._batch_encode_plus(
             batched_input,
-            is_pretokenized=is_pretokenized,
+            is_split_into_words=is_split_into_words,
             add_special_tokens=add_special_tokens,
             padding_strategy=padding_strategy,
             truncation_strategy=truncation_strategy,

src/transformers/trainer.py

@@ -39,6 +39,7 @@ from .trainer_utils import (
     PREFIX_CHECKPOINT_DIR,
     BestRun,
     EvalPrediction,
+    EvaluationStrategy,
     HPSearchBackend,
     PredictionOutput,
     TrainOutput,
@@ -351,6 +352,7 @@ class Trainer:
             sampler=train_sampler,
             collate_fn=self.data_collator,
             drop_last=self.args.dataloader_drop_last,
+            num_workers=self.args.dataloader_num_workers,
         )
 
     def _get_eval_sampler(self, eval_dataset: Dataset) -> Optional[torch.utils.data.sampler.Sampler]:
@@ -390,6 +392,7 @@ class Trainer:
             batch_size=self.args.eval_batch_size,
             collate_fn=self.data_collator,
             drop_last=self.args.dataloader_drop_last,
+            num_workers=self.args.dataloader_num_workers,
         )
 
     def get_test_dataloader(self, test_dataset: Dataset) -> DataLoader:
@@ -782,7 +785,10 @@ class Trainer:
 
                         self.log(logs)
 
-                    if self.args.evaluate_during_training and self.global_step % self.args.eval_steps == 0:
+                    if (
+                        self.args.evaluation_strategy == EvaluationStrategy.STEPS
+                        and self.global_step % self.args.eval_steps == 0
+                    ):
                         metrics = self.evaluate()
                         self._report_to_hp_search(trial, epoch, metrics)
 
@@ -820,6 +826,9 @@ class Trainer:
                             torch.save(self.lr_scheduler.state_dict(), os.path.join(output_dir, "scheduler.pt"))
 
                 epoch_pbar.update(1)
+                if self.args.evaluation_strategy == EvaluationStrategy.EPOCH:
+                    metrics = self.evaluate()
+                    self._report_to_hp_search(trial, epoch, metrics)
                 if self.args.max_steps > 0 and self.global_step >= self.args.max_steps:
                     break
             epoch_pbar.close()
@@ -1334,9 +1343,9 @@ class Trainer:
         elif is_torch_tpu_available():
             # tpu-comment: Get all predictions and labels from all worker shards of eval dataset
             if preds is not None:
-                preds = nested_xla_mesh_reduce("eval_preds", preds)
+                preds = nested_xla_mesh_reduce(preds, "eval_preds")
             if label_ids is not None:
-                label_ids = nested_xla_mesh_reduce("eval_label_ids", label_ids, torch.cat)
+                label_ids = nested_xla_mesh_reduce(label_ids, "eval_label_ids")
             if eval_losses is not None:
                 eval_losses = xm.mesh_reduce("eval_losses", torch.tensor(eval_losses), torch.cat).tolist()
 

src/transformers/trainer_tf.py

@@ -531,10 +531,6 @@ class TFTrainer:
 
             tf.summary.experimental.set_step(self.global_step)
 
-            if self.args.fp16:
-                policy = tf.keras.mixed_precision.experimental.Policy("mixed_float16")
-                tf.keras.mixed_precision.experimental.set_policy(policy)
-
             with self.tb_writer.as_default():
                 tf.summary.text("args", self.args.to_json_string())
 

src/transformers/trainer_utils.py

@@ -60,6 +60,12 @@ class TrainOutput(NamedTuple):
 PREFIX_CHECKPOINT_DIR = "checkpoint"
 
 
+class EvaluationStrategy(ExplicitEnum):
+    NO = "no"
+    STEPS = "steps"
+    EPOCH = "epoch"
+
+
 class BestRun(NamedTuple):
     """
     The best run found by an hyperparameter search (see :class:`~transformers.Trainer.hyperparameter_search`).

src/transformers/training_args.py

@@ -1,10 +1,13 @@
 import dataclasses
 import json
 import os
+import warnings
 from dataclasses import dataclass, field
+from enum import Enum
 from typing import Any, Dict, List, Optional, Tuple
 
 from .file_utils import cached_property, is_torch_available, is_torch_tpu_available, torch_required
+from .trainer_utils import EvaluationStrategy
 from .utils import logging
 
 
@@ -50,8 +53,13 @@ class TrainingArguments:
             Whether to run evaluation on the dev set or not.
         do_predict (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether to run predictions on the test set or not.
-        evaluate_during_training (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Whether to run evaluation during training at each logging step or not.
+        evaluation_strategy(:obj:`str` or :class:`~transformers.trainer_utils.EvaluationStrategy`, `optional`, defaults to :obj:`"no"`):
+            The evaluation strategy to adopt during training. Possible values are:
+
+                * :obj:`"no"`: No evaluation is done during training.
+                * :obj:`"steps"`: Evaluation is done (and logged) every :obj:`eval_steps`.
+                * :obj:`"epoch"`: Evaluation is done at the end of each epoch.
+
         prediction_loss_only (:obj:`bool`, `optional`, defaults to `False`):
             When performing evaluation and predictions, only returns the loss.
         per_device_train_batch_size (:obj:`int`, `optional`, defaults to 8):
@@ -111,8 +119,11 @@ class TrainingArguments:
         dataloader_drop_last (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size)
             or not.
-        eval_steps (:obj:`int`, `optional`, defaults to 1000):
-            Number of update steps between two evaluations.
+        eval_steps (:obj:`int`, `optional`):
+            Number of update steps between two evaluations if :obj:`evaluation_strategy="steps"`. Will default to the
+            same value as :obj:`logging_steps` if not set.
+        dataloader_num_workers (:obj:`int`, `optional`, defaults to 0):
+            Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in the main process.
         past_index (:obj:`int`, `optional`, defaults to -1):
             Some models like :doc:`TransformerXL <../model_doc/transformerxl>` or :doc`XLNet <../model_doc/xlnet>` can
             make use of the past hidden states for their predictions. If this argument is set to a positive int, the
@@ -153,7 +164,11 @@ class TrainingArguments:
     do_eval: bool = field(default=False, metadata={"help": "Whether to run eval on the dev set."})
     do_predict: bool = field(default=False, metadata={"help": "Whether to run predictions on the test set."})
     evaluate_during_training: bool = field(
-        default=False,
+        default=None,
+        metadata={"help": "Run evaluation during training at each logging step."},
+    )
+    evaluation_strategy: EvaluationStrategy = field(
+        default="no",
         metadata={"help": "Run evaluation during training at each logging step."},
     )
     prediction_loss_only: bool = field(
@@ -245,7 +260,13 @@ class TrainingArguments:
     dataloader_drop_last: bool = field(
         default=False, metadata={"help": "Drop the last incomplete batch if it is not divisible by the batch size."}
     )
-    eval_steps: int = field(default=1000, metadata={"help": "Run an evaluation every X steps."})
+    eval_steps: int = field(default=None, metadata={"help": "Run an evaluation every X steps."})
+    dataloader_num_workers: int = field(
+        default=0,
+        metadata={
+            "help": "Number of subprocesses to use for data loading (PyTorch only). 0 means that the data will be loaded in the main process."
+        },
+    )
 
     past_index: int = field(
         default=-1,
@@ -269,6 +290,19 @@ class TrainingArguments:
     def __post_init__(self):
         if self.disable_tqdm is None:
             self.disable_tqdm = logger.getEffectiveLevel() > logging.WARN
+        if self.evaluate_during_training is not None:
+            self.evaluation_strategy = (
+                EvaluationStrategy.STEPS if self.evaluate_during_training else EvaluationStrategy.NO
+            )
+            warnings.warn(
+                "The `evaluate_during_training` argument is deprecated in favor of `evaluation_strategy` (which has more options)",
+                FutureWarning,
+            )
+        else:
+            self.evaluation_strategy = EvaluationStrategy(self.evaluation_strategy)
+
+        if self.eval_steps is None:
+            self.eval_steps = self.logging_steps
 
     @property
     def train_batch_size(self) -> int:
@@ -347,17 +381,27 @@ class TrainingArguments:
         """
         return self._setup_devices[1]
 
+    def to_dict(self):
+        """
+        Serializes this instance while replace `Enum` by their values (for JSON serialization support).
+        """
+        d = dataclasses.asdict(self)
+        for k, v in d.items():
+            if isinstance(v, Enum):
+                d[k] = v.value
+        return d
+
     def to_json_string(self):
         """
         Serializes this instance to a JSON string.
         """
-        return json.dumps(dataclasses.asdict(self), indent=2)
+        return json.dumps(self.to_dict(), indent=2)
 
     def to_sanitized_dict(self) -> Dict[str, Any]:
         """
         Sanitized serialization to use with TensorBoard’s hparams
         """
-        d = dataclasses.asdict(self)
+        d = self.to_dict()
         d = {**d, **{"train_batch_size": self.train_batch_size, "eval_batch_size": self.eval_batch_size}}
 
         valid_types = [bool, int, float, str]

src/transformers/training_args_tf.py

@@ -88,7 +88,7 @@ class TFTrainingArguments(TrainingArguments):
         tpu_num_cores (:obj:`int`, `optional`):
             When training on TPU, the mumber of TPU cores (automatically passed by launcher script).
         debug (:obj:`bool`, `optional`, defaults to :obj:`False`):
-            Wheter to activate the trace to record computation graphs and profiling information or not.
+            Whether to activate the trace to record computation graphs and profiling information or not.
         dataloader_drop_last (:obj:`bool`, `optional`, defaults to :obj:`False`):
             Whether to drop the last incomplete batch (if the length of the dataset is not divisible by the batch size)
             or not.
@@ -103,6 +103,8 @@ class TFTrainingArguments(TrainingArguments):
             The name of the TPU the process is running on.
         run_name (:obj:`str`, `optional`):
             A descriptor for the run. Notably used for wandb logging.
+        xla (:obj:`bool`, `optional`):
+            Whether to activate the XLA compilation or not.
     """
 
     tpu_name: str = field(
@@ -110,12 +112,23 @@ class TFTrainingArguments(TrainingArguments):
         metadata={"help": "Name of TPU"},
     )
 
+    xla: bool = field(default=False, metadata={"help": "Whether to activate the XLA compilation or not"})
+
     @cached_property
     @tf_required
     def _setup_strategy(self) -> Tuple["tf.distribute.Strategy", int]:
         logger.info("Tensorflow: setting up strategy")
+
+        if self.args.xla:
+            tf.config.optimizer.set_jit(True)
+
         gpus = tf.config.list_physical_devices("GPU")
 
+        # Set to float16 at first
+        if self.fp16:
+            policy = tf.keras.mixed_precision.experimental.Policy("mixed_float16")
+            tf.keras.mixed_precision.experimental.set_policy(policy)
+
         if self.no_cuda:
             strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
         else:
@@ -128,10 +141,16 @@ class TFTrainingArguments(TrainingArguments):
                 tpu = None
 
             if tpu:
+                # Set to bfloat16 in case of TPU
+                if self.fp16:
+                    policy = tf.keras.mixed_precision.experimental.Policy("mixed_bfloat16")
+                    tf.keras.mixed_precision.experimental.set_policy(policy)
+
                 tf.config.experimental_connect_to_cluster(tpu)
                 tf.tpu.experimental.initialize_tpu_system(tpu)
 
                 strategy = tf.distribute.experimental.TPUStrategy(tpu)
+
             elif len(gpus) == 0:
                 strategy = tf.distribute.OneDeviceStrategy(device="/cpu:0")
             elif len(gpus) == 1:

tests/test_data_collator.py

@@ -1,7 +1,7 @@
 import unittest
 
 from transformers import AutoTokenizer, is_torch_available
-from transformers.testing_utils import require_torch
+from transformers.testing_utils import require_torch, slow
 
 
 if is_torch_available():
@@ -69,6 +69,7 @@ class DataCollatorIntegrationTest(unittest.TestCase):
         self.assertTrue("labels" not in batch)
         self.assertEqual(batch["inputs"].shape, torch.Size([8, 6]))
 
+    @slow
     def test_default_classification(self):
         MODEL_ID = "bert-base-cased-finetuned-mrpc"
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -80,6 +81,7 @@ class DataCollatorIntegrationTest(unittest.TestCase):
         batch = data_collator(dataset.features)
         self.assertEqual(batch["labels"].dtype, torch.long)
 
+    @slow
     def test_default_regression(self):
         MODEL_ID = "distilroberta-base"
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -91,6 +93,7 @@ class DataCollatorIntegrationTest(unittest.TestCase):
         batch = data_collator(dataset.features)
         self.assertEqual(batch["labels"].dtype, torch.float)
 
+    @slow
     def test_lm_tokenizer_without_padding(self):
         tokenizer = AutoTokenizer.from_pretrained("gpt2")
         data_collator = DataCollatorForLanguageModeling(tokenizer, mlm=False)
@@ -109,6 +112,7 @@ class DataCollatorIntegrationTest(unittest.TestCase):
         self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
         self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
 
+    @slow
     def test_lm_tokenizer_with_padding(self):
         tokenizer = AutoTokenizer.from_pretrained("distilroberta-base")
         data_collator = DataCollatorForLanguageModeling(tokenizer)
@@ -128,6 +132,7 @@ class DataCollatorIntegrationTest(unittest.TestCase):
         self.assertEqual(batch["input_ids"].shape, torch.Size((2, 512)))
         self.assertEqual(batch["labels"].shape, torch.Size((2, 512)))
 
+    @slow
     def test_plm(self):
         tokenizer = AutoTokenizer.from_pretrained("xlnet-base-cased")
         data_collator = DataCollatorForPermutationLanguageModeling(tokenizer)
@@ -156,6 +161,7 @@ class DataCollatorIntegrationTest(unittest.TestCase):
             # Expect error due to odd sequence length
             data_collator(example)
 
+    @slow
     def test_nsp(self):
         tokenizer = AutoTokenizer.from_pretrained("bert-base-cased")
         data_collator = DataCollatorForNextSentencePrediction(tokenizer)
@@ -172,6 +178,7 @@ class DataCollatorIntegrationTest(unittest.TestCase):
         self.assertEqual(batch["masked_lm_labels"].shape, torch.Size((total_samples, 512)))
         self.assertEqual(batch["next_sentence_label"].shape, torch.Size((total_samples,)))
 
+    @slow
     def test_sop(self):
         tokenizer = AutoTokenizer.from_pretrained("albert-base-v2")
         data_collator = DataCollatorForSOP(tokenizer)

tests/test_modeling_layoutlm.py

@@ -0,0 +1,239 @@
+# coding=utf-8
+# Copyright 2018 The Microsoft Research Asia LayoutLM Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import unittest
+
+from transformers import is_torch_available
+from transformers.file_utils import cached_property
+from transformers.testing_utils import require_torch, require_torch_and_cuda, slow, torch_device
+
+from .test_configuration_common import ConfigTester
+from .test_modeling_common import ModelTesterMixin, ids_tensor
+
+
+if is_torch_available():
+    from transformers import LayoutLMConfig, LayoutLMForMaskedLM, LayoutLMForTokenClassification, LayoutLMModel
+
+
+class LayoutLMModelTester:
+    """You can also import this e.g from .test_modeling_bart import BartModelTester """
+
+    def __init__(
+        self,
+        parent,
+        batch_size=13,
+        seq_length=7,
+        is_training=True,
+        use_input_mask=True,
+        use_token_type_ids=True,
+        use_labels=True,
+        vocab_size=99,
+        hidden_size=32,
+        num_hidden_layers=5,
+        num_attention_heads=4,
+        intermediate_size=37,
+        hidden_act="gelu",
+        hidden_dropout_prob=0.1,
+        attention_probs_dropout_prob=0.1,
+        max_position_embeddings=512,
+        type_vocab_size=16,
+        type_sequence_label_size=2,
+        initializer_range=0.02,
+        num_labels=3,
+        num_choices=4,
+        scope=None,
+        range_bbox=1000,
+    ):
+        self.parent = parent
+        self.batch_size = batch_size
+        self.seq_length = seq_length
+        self.is_training = is_training
+        self.use_input_mask = use_input_mask
+        self.use_token_type_ids = use_token_type_ids
+        self.use_labels = use_labels
+        self.vocab_size = vocab_size
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.intermediate_size = intermediate_size
+        self.hidden_act = hidden_act
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.attention_probs_dropout_prob = attention_probs_dropout_prob
+        self.max_position_embeddings = max_position_embeddings
+        self.type_vocab_size = type_vocab_size
+        self.type_sequence_label_size = type_sequence_label_size
+        self.initializer_range = initializer_range
+        self.num_labels = num_labels
+        self.num_choices = num_choices
+        self.scope = scope
+        self.range_bbox = range_bbox
+
+    def prepare_config_and_inputs(self):
+        input_ids = ids_tensor([self.batch_size, self.seq_length], self.vocab_size)
+
+        bbox = ids_tensor([self.batch_size, self.seq_length, 4], self.range_bbox)
+        # Ensure that bbox is legal
+        for i in range(bbox.shape[0]):
+            for j in range(bbox.shape[1]):
+                if bbox[i, j, 3] < bbox[i, j, 1]:
+                    t = bbox[i, j, 3]
+                    bbox[i, j, 3] = bbox[i, j, 1]
+                    bbox[i, j, 1] = t
+                if bbox[i, j, 2] < bbox[i, j, 0]:
+                    t = bbox[i, j, 2]
+                    bbox[i, j, 2] = bbox[i, j, 0]
+                    bbox[i, j, 0] = t
+
+        input_mask = None
+        if self.use_input_mask:
+            input_mask = ids_tensor([self.batch_size, self.seq_length], vocab_size=2)
+
+        token_type_ids = None
+        if self.use_token_type_ids:
+            token_type_ids = ids_tensor([self.batch_size, self.seq_length], self.type_vocab_size)
+
+        sequence_labels = None
+        token_labels = None
+        choice_labels = None
+        if self.use_labels:
+            sequence_labels = ids_tensor([self.batch_size], self.type_sequence_label_size)
+            token_labels = ids_tensor([self.batch_size, self.seq_length], self.num_labels)
+            choice_labels = ids_tensor([self.batch_size], self.num_choices)
+
+        config = LayoutLMConfig(
+            vocab_size=self.vocab_size,
+            hidden_size=self.hidden_size,
+            num_hidden_layers=self.num_hidden_layers,
+            num_attention_heads=self.num_attention_heads,
+            intermediate_size=self.intermediate_size,
+            hidden_act=self.hidden_act,
+            hidden_dropout_prob=self.hidden_dropout_prob,
+            attention_probs_dropout_prob=self.attention_probs_dropout_prob,
+            max_position_embeddings=self.max_position_embeddings,
+            type_vocab_size=self.type_vocab_size,
+            initializer_range=self.initializer_range,
+            return_dict=True,
+        )
+
+        return config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
+
+    def create_and_check_model(
+        self, config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
+    ):
+        model = LayoutLMModel(config=config)
+        model.to(torch_device)
+        model.eval()
+        result = model(input_ids, bbox, attention_mask=input_mask, token_type_ids=token_type_ids)
+        result = model(input_ids, bbox, token_type_ids=token_type_ids)
+        result = model(input_ids, bbox)
+        self.parent.assertEqual(result.last_hidden_state.shape, (self.batch_size, self.seq_length, self.hidden_size))
+        self.parent.assertEqual(result.pooler_output.shape, (self.batch_size, self.hidden_size))
+
+    def create_and_check_for_masked_lm(
+        self, config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
+    ):
+        model = LayoutLMForMaskedLM(config=config)
+        model.to(torch_device)
+        model.eval()
+        result = model(input_ids, bbox, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.vocab_size))
+
+    def create_and_check_for_token_classification(
+        self, config, input_ids, bbox, token_type_ids, input_mask, sequence_labels, token_labels, choice_labels
+    ):
+        config.num_labels = self.num_labels
+        model = LayoutLMForTokenClassification(config=config)
+        model.to(torch_device)
+        model.eval()
+        result = model(input_ids, bbox, attention_mask=input_mask, token_type_ids=token_type_ids, labels=token_labels)
+        self.parent.assertEqual(result.logits.shape, (self.batch_size, self.seq_length, self.num_labels))
+
+    def prepare_config_and_inputs_for_common(self):
+        config_and_inputs = self.prepare_config_and_inputs()
+        (
+            config,
+            input_ids,
+            bbox,
+            token_type_ids,
+            input_mask,
+            sequence_labels,
+            token_labels,
+            choice_labels,
+        ) = config_and_inputs
+        inputs_dict = {
+            "input_ids": input_ids,
+            "bbox": bbox,
+            "token_type_ids": token_type_ids,
+            "attention_mask": input_mask,
+        }
+        return config, inputs_dict
+
+
+@require_torch
+class LayoutLMModelTest(ModelTesterMixin, unittest.TestCase):
+
+    all_model_classes = (
+        (LayoutLMModel, LayoutLMForMaskedLM, LayoutLMForTokenClassification) if is_torch_available() else ()
+    )
+
+    def setUp(self):
+        self.model_tester = LayoutLMModelTester(self)
+        self.config_tester = ConfigTester(self, config_class=LayoutLMConfig, hidden_size=37)
+
+    def test_config(self):
+        self.config_tester.run_common_tests()
+
+    def test_model(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_model(*config_and_inputs)
+
+    def test_for_masked_lm(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_for_masked_lm(*config_and_inputs)
+
+    def test_for_token_classification(self):
+        config_and_inputs = self.model_tester.prepare_config_and_inputs()
+        self.model_tester.create_and_check_for_token_classification(*config_and_inputs)
+
+    @cached_property
+    def big_model(self):
+        """Cached property means this code will only be executed once."""
+        checkpoint_path = "microsoft/layoutlm-large-uncased"
+        model = LayoutLMForMaskedLM.from_pretrained(checkpoint_path).to(
+            torch_device
+        )  # test whether AutoModel can determine your model_class from checkpoint name
+        if torch_device == "cuda":
+            model.half()
+
+    # optional: do more testing! This will save you time later!
+    @slow
+    def test_that_LayoutLM_can_be_used_in_a_pipeline(self):
+        """We can use self.big_model here without calling __init__ again."""
+        pass
+
+    def test_LayoutLM_loss_doesnt_change_if_you_add_padding(self):
+        pass
+
+    def test_LayoutLM_bad_args(self):
+        pass
+
+    def test_LayoutLM_backward_pass_reduces_loss(self):
+        """Test loss/gradients same as reference implementation, for example."""
+        pass
+
+    @require_torch_and_cuda
+    def test_large_inputs_in_fp16_dont_cause_overflow(self):
+        pass

tests/test_modeling_rag.py

@@ -0,0 +1,885 @@
+# coding=utf-8
+# Copyright 2020, The RAG Authors and The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import json
+import os
+import shutil
+import tempfile
+import unittest
+from unittest.mock import patch
+
+import numpy as np
+
+from transformers.file_utils import cached_property, is_datasets_available, is_faiss_available, is_torch_available
+from transformers.testing_utils import require_torch, slow, torch_device
+from transformers.tokenization_bart import BartTokenizer
+from transformers.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
+from transformers.tokenization_dpr import DPRQuestionEncoderTokenizer
+from transformers.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
+from transformers.tokenization_t5 import T5Tokenizer
+
+from .test_modeling_bart import ModelTester as BartModelTester
+from .test_modeling_dpr import DPRModelTester
+from .test_modeling_t5 import T5ModelTester
+
+
+TOLERANCE = 1e-3
+
+T5_SAMPLE_VOCAB = os.path.join(os.path.dirname(os.path.abspath(__file__)), "fixtures/test_sentencepiece.model")
+
+if is_torch_available() and is_datasets_available() and is_faiss_available():
+    import torch
+    from datasets import Dataset
+
+    import faiss
+    from transformers import (
+        AutoConfig,
+        AutoModel,
+        AutoModelForSeq2SeqLM,
+        RagConfig,
+        RagModel,
+        RagRetriever,
+        RagSequenceForGeneration,
+        RagTokenForGeneration,
+    )
+    from transformers.modeling_outputs import BaseModelOutput
+
+
+def _assert_tensors_equal(a, b, atol=1e-12, prefix=""):
+    """If tensors not close, or a and b arent both tensors, raise a nice Assertion error."""
+    if a is None and b is None:
+        return True
+    try:
+        if torch.allclose(a, b, atol=atol):
+            return True
+        raise
+    except Exception:
+        msg = "{} != {}".format(a, b)
+        if prefix:
+            msg = prefix + ": " + msg
+        raise AssertionError(msg)
+
+
+def require_retrieval(test_case):
+    """
+    Decorator marking a test that requires a set of dependencies necessary for pefrorm retrieval with
+    :class:`~transformers.RagRetriever`.
+
+    These tests are skipped when respective libraries are not installed.
+
+    """
+    if not (is_torch_available() and is_datasets_available() and is_faiss_available()):
+        test_case = unittest.skip("test requires PyTorch")(test_case)
+    return test_case
+
+
+@require_torch
+@require_retrieval
+class RagTestMixin:
+
+    all_model_classes = (
+        (RagModel, RagTokenForGeneration, RagSequenceForGeneration)
+        if is_torch_available() and is_datasets_available() and is_faiss_available()
+        else ()
+    )
+
+    retrieval_vector_size = 32
+    n_docs = 2
+    max_combined_length = 16
+
+    def setUp(self):
+        self.tmpdirname = tempfile.mkdtemp()
+
+        # DPR tok
+        vocab_tokens = [
+            "[UNK]",
+            "[CLS]",
+            "[SEP]",
+            "[PAD]",
+            "[MASK]",
+            "want",
+            "##want",
+            "##ed",
+            "wa",
+            "un",
+            "runn",
+            "##ing",
+            ",",
+            "low",
+            "lowest",
+        ]
+        dpr_tokenizer_path = os.path.join(self.tmpdirname, "dpr_tokenizer")
+        os.makedirs(dpr_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(dpr_tokenizer_path, DPR_VOCAB_FILES_NAMES["vocab_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+        # BART tok
+        vocab = [
+            "l",
+            "o",
+            "w",
+            "e",
+            "r",
+            "s",
+            "t",
+            "i",
+            "d",
+            "n",
+            "\u0120",
+            "\u0120l",
+            "\u0120n",
+            "\u0120lo",
+            "\u0120low",
+            "er",
+            "\u0120lowest",
+            "\u0120newer",
+            "\u0120wider",
+            "<unk>",
+        ]
+        vocab_tokens = dict(zip(vocab, range(len(vocab))))
+        merges = ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
+        self.special_tokens_map = {"unk_token": "<unk>"}
+
+        bart_tokenizer_path = os.path.join(self.tmpdirname, "bart_tokenizer")
+        os.makedirs(bart_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(bart_tokenizer_path, BART_VOCAB_FILES_NAMES["vocab_file"])
+        self.merges_file = os.path.join(bart_tokenizer_path, BART_VOCAB_FILES_NAMES["merges_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as fp:
+            fp.write(json.dumps(vocab_tokens) + "\n")
+        with open(self.merges_file, "w", encoding="utf-8") as fp:
+            fp.write("\n".join(merges))
+
+        t5_tokenizer = T5Tokenizer(T5_SAMPLE_VOCAB)
+        t5_tokenizer_path = os.path.join(self.tmpdirname, "t5_tokenizer")
+        t5_tokenizer.save_pretrained(t5_tokenizer_path)
+
+    @cached_property
+    def dpr_tokenizer(self) -> DPRQuestionEncoderTokenizer:
+        return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname, "dpr_tokenizer"))
+
+    @cached_property
+    def bart_tokenizer(self) -> BartTokenizer:
+        return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname, "bart_tokenizer"))
+
+    @cached_property
+    def t5_tokenizer(self) -> BartTokenizer:
+        return T5Tokenizer.from_pretrained(os.path.join(self.tmpdirname, "t5_tokenizer"))
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdirname)
+
+    def get_retriever(self, config):
+        dataset = Dataset.from_dict(
+            {
+                "id": ["0", "1"],
+                "text": ["foo", "bar"],
+                "title": ["Foo", "Bar"],
+                "embeddings": [np.ones(self.retrieval_vector_size), 2 * np.ones(self.retrieval_vector_size)],
+            }
+        )
+        dataset.add_faiss_index("embeddings", string_factory="Flat", metric_type=faiss.METRIC_INNER_PRODUCT)
+        tokenizer = self.bart_tokenizer if config.generator.model_type == "bart" else self.t5_tokenizer
+        with patch("transformers.retrieval_rag.load_dataset") as mock_load_dataset:
+            mock_load_dataset.return_value = dataset
+            retriever = RagRetriever(
+                config,
+                question_encoder_tokenizer=self.dpr_tokenizer,
+                generator_tokenizer=tokenizer,
+            )
+        return retriever
+
+    def check_model_with_retriever(
+        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
+    ):
+        self.assertIsNotNone(config.question_encoder)
+        self.assertIsNotNone(config.generator)
+
+        for model_class in self.all_model_classes:
+            model = model_class(config, retriever=self.get_retriever(config)).to(torch_device)
+            model.eval()
+
+            self.assertTrue(model.config.is_encoder_decoder)
+
+            outputs = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                decoder_input_ids=decoder_input_ids,
+                decoder_attention_mask=decoder_attention_mask,
+            )
+
+            # logits
+            self.assertEqual(
+                outputs.logits.shape,
+                (self.n_docs * decoder_input_ids.shape[0], decoder_input_ids.shape[1], config.generator.vocab_size),
+            )
+            # generator encoder last hidden states
+            self.assertEqual(
+                outputs.generator_enc_last_hidden_state.shape,
+                (self.n_docs * decoder_input_ids.shape[0], self.max_combined_length, config.generator.hidden_size),
+            )
+            # doc scores
+            self.assertEqual(outputs.doc_scores.shape, (input_ids.shape[0], self.n_docs))
+
+    def check_model_generate(
+        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
+    ):
+        self.assertIsNotNone(config.question_encoder)
+        self.assertIsNotNone(config.generator)
+
+        for model_class in self.all_model_classes[1:]:
+            model = model_class(config, retriever=self.get_retriever(config)).to(torch_device)
+            model.eval()
+
+            self.assertTrue(model.config.is_encoder_decoder)
+
+            outputs = model.generate(
+                input_ids=input_ids,
+                num_beams=2,
+                num_return_sequences=2,
+                decoder_start_token_id=config.generator.eos_token_id,
+            )
+
+            self.assertIsNotNone(outputs)
+
+    def check_model_without_retriever(
+        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
+    ):
+        self.assertIsNotNone(config.question_encoder)
+        self.assertIsNotNone(config.generator)
+
+        retriever = self.get_retriever(config)
+
+        for model_class in self.all_model_classes:
+            model = model_class(config).to(torch_device)
+            model.eval()
+            self.assertTrue(model.config.is_encoder_decoder)
+
+            question_hidden_states = model.question_encoder(input_ids, attention_mask=attention_mask)[0]
+
+            out = retriever(
+                input_ids,
+                question_hidden_states.cpu().detach().to(torch.float32).numpy(),
+                prefix=config.generator.prefix,
+                return_tensors="pt",
+            )
+
+            context_input_ids, context_attention_mask, retrieved_doc_embeds = (
+                out["context_input_ids"],
+                out["context_attention_mask"],
+                out["retrieved_doc_embeds"],
+            )
+
+            # cast
+            retrieved_doc_embeds = retrieved_doc_embeds.to(question_hidden_states)
+            context_input_ids = context_input_ids.to(input_ids)
+            context_attention_mask = context_attention_mask.to(input_ids)
+
+            # compute doc_scores
+            doc_scores = torch.bmm(question_hidden_states.unsqueeze(1), retrieved_doc_embeds.transpose(1, 2)).squeeze(
+                1
+            )
+
+            outputs = model(
+                context_input_ids=context_input_ids,
+                context_attention_mask=context_attention_mask,
+                doc_scores=doc_scores,
+                decoder_input_ids=decoder_input_ids,
+                decoder_attention_mask=decoder_attention_mask,
+            )
+
+            # logits
+            self.assertEqual(
+                outputs.logits.shape,
+                (self.n_docs * decoder_input_ids.shape[0], decoder_input_ids.shape[1], config.generator.vocab_size),
+            )
+            # generator encoder last hidden states
+            self.assertEqual(
+                outputs.generator_enc_last_hidden_state.shape,
+                (self.n_docs * decoder_input_ids.shape[0], self.max_combined_length, config.generator.hidden_size),
+            )
+            # doc scores
+            self.assertEqual(outputs.doc_scores.shape, (input_ids.shape[0], self.n_docs))
+
+    def check_model_with_encoder_outputs(
+        self, config, input_ids, attention_mask, decoder_input_ids, decoder_attention_mask, **kwargs
+    ):
+        self.assertIsNotNone(config.question_encoder)
+        self.assertIsNotNone(config.generator)
+
+        for model_class in self.all_model_classes:
+            model = model_class(config, retriever=self.get_retriever(config)).to(torch_device)
+            model.eval()
+
+            self.assertTrue(model.config.is_encoder_decoder)
+
+            outputs = model(
+                input_ids=input_ids,
+                attention_mask=attention_mask,
+                decoder_input_ids=decoder_input_ids,
+                decoder_attention_mask=decoder_attention_mask,
+            )
+
+            encoder_outputs = BaseModelOutput(outputs.generator_enc_last_hidden_state)
+
+            # run only generator
+            outputs = model(
+                encoder_outputs=encoder_outputs,
+                doc_scores=outputs.doc_scores,
+                decoder_input_ids=decoder_input_ids,
+                decoder_attention_mask=decoder_attention_mask,
+            )
+
+            # logits
+            self.assertEqual(
+                outputs.logits.shape,
+                (self.n_docs * decoder_input_ids.shape[0], decoder_input_ids.shape[1], config.generator.vocab_size),
+            )
+            # generator encoder last hidden states
+            self.assertEqual(
+                outputs.generator_enc_last_hidden_state.shape,
+                (self.n_docs * decoder_input_ids.shape[0], self.max_combined_length, config.generator.hidden_size),
+            )
+            # doc scores
+            self.assertEqual(outputs.doc_scores.shape, (input_ids.shape[0], self.n_docs))
+
+    def test_model_with_retriever(self):
+        inputs_dict = self.config_and_inputs
+        self.check_model_with_retriever(**inputs_dict)
+
+    def test_model_without_retriever(self):
+        inputs_dict = self.config_and_inputs
+        self.check_model_without_retriever(**inputs_dict)
+
+    def test_model_with_encoder_outputs(self):
+        inputs_dict = self.config_and_inputs
+        self.check_model_with_encoder_outputs(**inputs_dict)
+
+    def test_model_generate(self):
+        inputs_dict = self.config_and_inputs
+        self.check_model_generate(**inputs_dict)
+
+
+@require_torch
+@require_retrieval
+class RagDPRBartTest(RagTestMixin, unittest.TestCase):
+    @cached_property
+    def config_and_inputs(self):
+        question_encoder_tester = DPRModelTester(self)
+        dpr_config_and_inputs = question_encoder_tester.prepare_config_and_inputs()
+        generator_tester = BartModelTester(self)
+        bart_config_and_inputs = generator_tester.prepare_config_and_inputs_for_common()
+
+        (question_encoder_config, input_ids, _, input_mask, _, _, _) = dpr_config_and_inputs
+        (generator_config, bart_inputs_dict) = bart_config_and_inputs
+        decoder_input_ids, decoder_attention_mask = bart_inputs_dict["input_ids"], bart_inputs_dict["attention_mask"]
+
+        config = RagConfig.from_question_encoder_generator_configs(
+            question_encoder_config,
+            generator_config,
+            n_docs=self.n_docs,
+            retrieval_vector_size=self.retrieval_vector_size,
+            max_combined_length=self.max_combined_length,
+            use_cache=False,
+        )
+
+        return {
+            "config": config,
+            "input_ids": input_ids,
+            "attention_mask": input_mask,
+            "decoder_input_ids": decoder_input_ids,
+            "decoder_attention_mask": decoder_attention_mask,
+        }
+
+
+@require_torch
+@require_retrieval
+class RagDPRT5Test(RagTestMixin, unittest.TestCase):
+    @cached_property
+    def config_and_inputs(self):
+        question_encoder_tester = DPRModelTester(self)
+        dpr_config_and_inputs = question_encoder_tester.prepare_config_and_inputs()
+        generator_tester = T5ModelTester(self, vocab_size=1100, n_positions=30)
+        t5_config_and_inputs = generator_tester.prepare_config_and_inputs()
+
+        (question_encoder_config, input_ids, _, input_mask, _, _, _) = dpr_config_and_inputs
+        #        import ipdb; ipdb.set_trace()
+        (generator_config, _, decoder_input_ids, _, decoder_attention_mask, _) = t5_config_and_inputs
+        config = RagConfig.from_question_encoder_generator_configs(
+            question_encoder_config,
+            generator_config,
+            n_docs=self.n_docs,
+            retrieval_vector_size=self.retrieval_vector_size,
+            max_combined_length=self.max_combined_length,
+            use_cache=False,
+        )
+
+        return {
+            "config": config,
+            "input_ids": input_ids,
+            "attention_mask": input_mask,
+            "decoder_input_ids": decoder_input_ids,
+            "decoder_attention_mask": decoder_attention_mask,
+        }
+
+
+@require_torch
+@require_retrieval
+class RagModelIntegrationTests(unittest.TestCase):
+    @cached_property
+    def sequence_model(self):
+        return (
+            RagSequenceForGeneration.from_pretrained_question_encoder_generator(
+                "facebook/dpr-question_encoder-single-nq-base", "facebook/bart-large-cnn"
+            )
+            .to(torch_device)
+            .eval()
+        )
+
+    @cached_property
+    def token_model(self):
+        return (
+            RagTokenForGeneration.from_pretrained_question_encoder_generator(
+                "facebook/dpr-question_encoder-single-nq-base", "facebook/bart-large-cnn"
+            )
+            .to(torch_device)
+            .eval()
+        )
+
+    def get_rag_config(self):
+        question_encoder_config = AutoConfig.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+        generator_config = AutoConfig.from_pretrained("facebook/bart-large-cnn")
+        return RagConfig.from_question_encoder_generator_configs(
+            question_encoder_config,
+            generator_config,
+            bos_token_id=0,
+            decoder_start_token_id=2,
+            eos_token_id=2,
+            is_encoder_decoder=True,
+            pad_token_id=1,
+            vocab_size=50264,
+            title_sep=" / ",
+            doc_sep=" // ",
+            n_docs=5,
+            max_combined_length=300,
+            dataset="wiki_dpr",
+            dataset_split="train",
+            index_name="exact",
+            index_path=None,
+            use_dummy_dataset=True,
+            retrieval_vector_size=768,
+            retrieval_batch_size=8,
+        )
+
+    @slow
+    def test_rag_sequence_inference(self):
+        rag_config = self.get_rag_config()
+        rag_decoder_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        rag_question_encoder_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base"
+        )
+        rag_retriever = RagRetriever(
+            rag_config,
+            question_encoder_tokenizer=rag_question_encoder_tokenizer,
+            generator_tokenizer=rag_decoder_tokenizer,
+        )
+
+        rag_sequence = self.sequence_model
+        rag_sequence.set_retriever(rag_retriever)
+
+        input_ids = rag_question_encoder_tokenizer(
+            "who sings does he love me with reba", return_tensors="pt"
+        ).input_ids
+        decoder_input_ids = rag_decoder_tokenizer("Linda Davis", return_tensors="pt").input_ids
+
+        input_ids = input_ids.to(torch_device)
+        decoder_input_ids = decoder_input_ids.to(torch_device)
+
+        with torch.no_grad():
+            output = rag_sequence(
+                input_ids,
+                labels=decoder_input_ids,
+            )
+
+        expected_shape = torch.Size([5, 5, 50264])
+        self.assertEqual(output.logits.shape, expected_shape)
+
+        expected_doc_scores = torch.tensor([[75.0286, 74.4998, 74.0804, 74.0306, 73.9504]]).to(torch_device)
+        _assert_tensors_equal(expected_doc_scores, output.doc_scores, atol=TOLERANCE)
+
+        expected_loss = torch.tensor([38.7446]).to(torch_device)
+        _assert_tensors_equal(expected_loss, output.loss, atol=TOLERANCE)
+
+    @slow
+    def test_rag_token_inference(self):
+        rag_config = self.get_rag_config()
+        rag_decoder_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        rag_question_encoder_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base"
+        )
+        rag_retriever = RagRetriever(
+            rag_config,
+            question_encoder_tokenizer=rag_question_encoder_tokenizer,
+            generator_tokenizer=rag_decoder_tokenizer,
+        )
+
+        rag_token = self.token_model
+        rag_token.set_retriever(rag_retriever)
+
+        input_ids = rag_question_encoder_tokenizer(
+            "who sings does he love me with reba", return_tensors="pt"
+        ).input_ids
+        decoder_input_ids = rag_decoder_tokenizer("Linda Davis", return_tensors="pt").input_ids
+
+        input_ids = input_ids.to(torch_device)
+        decoder_input_ids = decoder_input_ids.to(torch_device)
+
+        with torch.no_grad():
+            output = rag_token(
+                input_ids,
+                labels=decoder_input_ids,
+            )
+
+        expected_shape = torch.Size([5, 5, 50264])
+        self.assertEqual(output.logits.shape, expected_shape)
+
+        expected_doc_scores = torch.tensor([[75.0286, 74.4998, 74.0804, 74.0306, 73.9504]]).to(torch_device)
+        _assert_tensors_equal(expected_doc_scores, output.doc_scores, atol=TOLERANCE)
+
+        expected_loss = torch.tensor([38.7045]).to(torch_device)
+        _assert_tensors_equal(expected_loss, output.loss, atol=TOLERANCE)
+
+    @slow
+    def test_rag_token_generate_beam(self):
+        rag_config = self.get_rag_config()
+        rag_decoder_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        rag_question_encoder_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base"
+        )
+        rag_retriever = RagRetriever(
+            rag_config,
+            question_encoder_tokenizer=rag_question_encoder_tokenizer,
+            generator_tokenizer=rag_decoder_tokenizer,
+        )
+
+        rag_token = self.token_model
+        rag_token.set_retriever(rag_retriever)
+
+        input_ids = rag_question_encoder_tokenizer(
+            "who sings does he love me with reba", return_tensors="pt"
+        ).input_ids
+
+        input_ids = input_ids.to(torch_device)
+
+        output_ids = rag_token.generate(
+            input_ids,
+            decoder_start_token_id=rag_token.generator.config.decoder_start_token_id,
+            num_beams=2,
+            num_return_sequences=2,
+        )
+        # sequence generate test
+        output_text_1 = rag_decoder_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        output_text_2 = rag_decoder_tokenizer.decode(output_ids[1], skip_special_tokens=True)
+
+        # Expected outputs as given by model at integration time.
+        EXPECTED_OUTPUT_TEXT_1 = "The songwriting credits are credited to ABBA"
+        EXPECTED_OUTPUT_TEXT_2 = 'The songwriting credits are credited to "B'
+
+        self.assertEqual(output_text_1, EXPECTED_OUTPUT_TEXT_1)
+        self.assertEqual(output_text_2, EXPECTED_OUTPUT_TEXT_2)
+
+    @slow
+    def test_rag_token_generate_batch(self):
+        rag_config = self.get_rag_config()
+        rag_decoder_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        rag_question_encoder_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base"
+        )
+        rag_retriever = RagRetriever(
+            rag_config,
+            question_encoder_tokenizer=rag_question_encoder_tokenizer,
+            generator_tokenizer=rag_decoder_tokenizer,
+        )
+
+        rag_token = self.token_model
+        rag_token.set_retriever(rag_retriever)
+
+        questions = [
+            "who sings does he love me with reba",
+            "how many pages is invisible man by ralph ellison",
+        ]
+        input_ids = rag_question_encoder_tokenizer.batch_encode_plus(
+            questions,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+        ).input_ids
+
+        input_ids = input_ids.to(torch_device)
+
+        output_ids = rag_token.generate(
+            input_ids,
+            decoder_start_token_id=rag_token.generator.config.decoder_start_token_id,
+            num_beams=4,
+            num_return_sequences=1,
+            max_length=10,
+        )
+
+        # sequence generate test
+        output_text_1 = rag_decoder_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        output_text_2 = rag_decoder_tokenizer.decode(output_ids[1], skip_special_tokens=True)
+
+        # Expected outputs as given by model at integration time.
+        EXPECTED_OUTPUT_TEXT_1 = '"People Need Love" is the'
+        EXPECTED_OUTPUT_TEXT_2 = '"How many pages is invisible man'
+
+        self.assertEqual(output_text_1, EXPECTED_OUTPUT_TEXT_1)
+        self.assertEqual(output_text_2, EXPECTED_OUTPUT_TEXT_2)
+
+    @slow
+    def test_rag_sequence_generate_batch(self):
+        # IMPORTAN: This test fails on GPU, but is fine on CPU -> beam search is very sensible
+        rag_config = self.get_rag_config()
+        rag_decoder_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        rag_question_encoder_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base"
+        )
+        rag_retriever = RagRetriever(
+            rag_config,
+            question_encoder_tokenizer=rag_question_encoder_tokenizer,
+            generator_tokenizer=rag_decoder_tokenizer,
+        )
+
+        rag_sequence = self.sequence_model
+        rag_sequence.set_retriever(rag_retriever)
+
+        questions = [
+            "who sings does he love me with reba",
+            "how many pages is invisible man by ralph ellison",
+        ]
+        input_ids = rag_question_encoder_tokenizer.batch_encode_plus(
+            questions,
+            return_tensors="pt",
+            padding=True,
+            truncation=True,
+        ).input_ids
+
+        input_ids = input_ids.to(torch_device)
+
+        output_ids = rag_sequence.generate(
+            input_ids,
+            decoder_start_token_id=rag_sequence.generator.config.decoder_start_token_id,
+            num_beams=4,
+            num_return_sequences=1,
+            max_length=10,
+        )
+
+        # sequence generate test
+        output_text_1 = rag_decoder_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        output_text_2 = rag_decoder_tokenizer.decode(output_ids[1], skip_special_tokens=True)
+
+        # Expected outputs as given by model at integration time.
+        EXPECTED_OUTPUT_TEXT_1 = '"I Know Him So Well"'
+        EXPECTED_OUTPUT_TEXT_2 = '"Howl" chronicles the'
+
+        self.assertEqual(output_text_1, EXPECTED_OUTPUT_TEXT_1)
+        self.assertEqual(output_text_2, EXPECTED_OUTPUT_TEXT_2)
+
+    @slow
+    def test_rag_sequence_generate_beam(self):
+        rag_config = self.get_rag_config()
+        rag_decoder_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        rag_question_encoder_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base"
+        )
+        rag_retriever = RagRetriever(
+            rag_config,
+            question_encoder_tokenizer=rag_question_encoder_tokenizer,
+            generator_tokenizer=rag_decoder_tokenizer,
+        )
+
+        rag_token = self.sequence_model
+        rag_token.set_retriever(rag_retriever)
+
+        input_ids = rag_question_encoder_tokenizer(
+            "who sings does he love me with reba", return_tensors="pt"
+        ).input_ids
+
+        input_ids = input_ids.to(torch_device)
+
+        output_ids = rag_token.generate(
+            input_ids,
+            decoder_start_token_id=rag_token.generator.config.decoder_start_token_id,
+            num_beams=2,
+            num_return_sequences=2,
+        )
+        # sequence generate test
+        output_text_1 = rag_decoder_tokenizer.decode(output_ids[0], skip_special_tokens=True)
+        output_text_2 = rag_decoder_tokenizer.decode(output_ids[1], skip_special_tokens=True)
+
+        # Expected outputs as given by model at integration time.
+        EXPECTED_OUTPUT_TEXT_1 = """ ABBA / small label like Playboy Records did not have the distribution resources to meet the demand for the single from retailers and radio programmers. The foursome decided to record their first album together in late 1972, and sessions began on 26 September 1972. The women shared lead vocals on "Nina, Pretty Ballerina" that day."""
+        EXPECTED_OUTPUT_TEXT_2 = """ ABBA / small label like Playboy Records did not have the distribution resources to meet the demand for the single from retailers and radio programmers. The foursome decided to record their first album together in late 1972, and sessions began on 26 September 1972. The women shared lead vocals on "Nina, Pretty Ballerina" (a top ten hit in Austria)"""
+
+        self.assertEqual(output_text_1, EXPECTED_OUTPUT_TEXT_1)
+        self.assertEqual(output_text_2, EXPECTED_OUTPUT_TEXT_2)
+
+
+@require_torch
+@require_retrieval
+class RagModelSaveLoadTests(unittest.TestCase):
+    def get_rag_config(self):
+        question_encoder_config = AutoConfig.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+        generator_config = AutoConfig.from_pretrained("facebook/bart-large-cnn")
+        return RagConfig.from_question_encoder_generator_configs(
+            question_encoder_config,
+            generator_config,
+            bos_token_id=0,
+            decoder_start_token_id=2,
+            eos_token_id=2,
+            is_encoder_decoder=True,
+            pad_token_id=1,
+            vocab_size=50264,
+            title_sep=" / ",
+            doc_sep=" // ",
+            n_docs=5,
+            max_combined_length=300,
+            dataset="wiki_dpr",
+            dataset_split="train",
+            index_name="exact",
+            index_path=None,
+            use_dummy_dataset=True,
+            retrieval_vector_size=768,
+            retrieval_batch_size=8,
+        )
+
+    @slow
+    def test_rag_sequence_from_pretrained(self):
+        rag_config = self.get_rag_config()
+        rag_decoder_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        rag_question_encoder_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base"
+        )
+        rag_retriever = RagRetriever(
+            rag_config,
+            question_encoder_tokenizer=rag_question_encoder_tokenizer,
+            generator_tokenizer=rag_decoder_tokenizer,
+        )
+
+        input_ids = rag_question_encoder_tokenizer(
+            "who sings does he love me with reba", return_tensors="pt"
+        ).input_ids
+        decoder_input_ids = rag_decoder_tokenizer("Linda Davis", return_tensors="pt").input_ids
+
+        input_ids = input_ids.to(torch_device)
+        decoder_input_ids = decoder_input_ids.to(torch_device)
+
+        with tempfile.TemporaryDirectory() as tmp_dirname:
+            rag_sequence = RagSequenceForGeneration.from_pretrained_question_encoder_generator(
+                "facebook/dpr-question_encoder-single-nq-base",
+                "facebook/bart-large-cnn",
+                retriever=rag_retriever,
+                config=rag_config,
+            ).to(torch_device)
+            # check that the from pretrained methods work
+            rag_sequence.save_pretrained(tmp_dirname)
+            rag_sequence.from_pretrained(tmp_dirname, retriever=rag_retriever)
+            rag_sequence.to(torch_device)
+
+            with torch.no_grad():
+                output = rag_sequence(
+                    input_ids,
+                    labels=decoder_input_ids,
+                )
+
+            loss_pretrained = output.loss
+            del rag_sequence
+
+        question_encoder = AutoModel.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+        generator = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
+        rag_sequence = RagSequenceForGeneration(
+            config=rag_config, question_encoder=question_encoder, generator=generator, retriever=rag_retriever
+        )
+        rag_sequence.to(torch_device)
+
+        with torch.no_grad():
+            output = rag_sequence(
+                input_ids,
+                labels=decoder_input_ids,
+            )
+
+        loss_init = output.loss
+
+        self.assertAlmostEqual(loss_pretrained.item(), loss_init.item(), places=4)
+
+    @slow
+    def test_rag_token_from_pretrained(self):
+        rag_config = self.get_rag_config()
+        rag_decoder_tokenizer = BartTokenizer.from_pretrained("facebook/bart-large-cnn")
+        rag_question_encoder_tokenizer = DPRQuestionEncoderTokenizer.from_pretrained(
+            "facebook/dpr-question_encoder-single-nq-base"
+        )
+        rag_retriever = RagRetriever(
+            rag_config,
+            question_encoder_tokenizer=rag_question_encoder_tokenizer,
+            generator_tokenizer=rag_decoder_tokenizer,
+        )
+
+        input_ids = rag_question_encoder_tokenizer(
+            "who sings does he love me with reba", return_tensors="pt"
+        ).input_ids
+        decoder_input_ids = rag_decoder_tokenizer("Linda Davis", return_tensors="pt").input_ids
+
+        input_ids = input_ids.to(torch_device)
+        decoder_input_ids = decoder_input_ids.to(torch_device)
+
+        with tempfile.TemporaryDirectory() as tmp_dirname:
+            rag_token = RagTokenForGeneration.from_pretrained_question_encoder_generator(
+                "facebook/dpr-question_encoder-single-nq-base",
+                "facebook/bart-large-cnn",
+                retriever=rag_retriever,
+                config=rag_config,
+            ).to(torch_device)
+            # check that the from pretrained methods work
+            rag_token.save_pretrained(tmp_dirname)
+            rag_token.from_pretrained(tmp_dirname, retriever=rag_retriever)
+            rag_token.to(torch_device)
+
+            with torch.no_grad():
+                output = rag_token(
+                    input_ids,
+                    labels=decoder_input_ids,
+                )
+
+            loss_pretrained = output.loss
+            del rag_token
+
+        question_encoder = AutoModel.from_pretrained("facebook/dpr-question_encoder-single-nq-base")
+        generator = AutoModelForSeq2SeqLM.from_pretrained("facebook/bart-large-cnn")
+        rag_token = RagTokenForGeneration(
+            config=rag_config, question_encoder=question_encoder, generator=generator, retriever=rag_retriever
+        )
+        rag_token.to(torch_device)
+
+        with torch.no_grad():
+            output = rag_token(
+                input_ids,
+                labels=decoder_input_ids,
+            )
+
+        loss_init = output.loss
+
+        self.assertAlmostEqual(loss_pretrained.item(), loss_init.item(), places=4)

tests/test_modeling_t5.py

@@ -35,28 +35,53 @@ if is_torch_available():
 
 
 class T5ModelTester:
-    def __init__(self, parent):
+    def __init__(
+        self,
+        parent,
+        vocab_size=99,
+        n_positions=14,
+        batch_size=13,
+        encoder_seq_length=7,
+        decoder_seq_length=9,
+        # For common tests
+        seq_length=7,
+        is_training=True,
+        use_attention_mask=True,
+        use_labels=True,
+        hidden_size=32,
+        num_hidden_layers=5,
+        num_attention_heads=4,
+        d_ff=37,
+        relative_attention_num_buckets=8,
+        dropout_rate=0.1,
+        initializer_factor=0.002,
+        eos_token_id=1,
+        pad_token_id=0,
+        decoder_start_token_id=0,
+        scope=None,
+    ):
+
         self.parent = parent
-        self.batch_size = 13
-        self.encoder_seq_length = 7
-        self.decoder_seq_length = 9
+        self.batch_size = batch_size
+        self.encoder_seq_length = encoder_seq_length
+        self.decoder_seq_length = decoder_seq_length
         # For common tests
         self.seq_length = self.decoder_seq_length
-        self.is_training = True
-        self.use_attention_mask = True
-        self.use_labels = True
-        self.vocab_size = 99
-        self.n_positions = 14
-        self.hidden_size = 32
-        self.num_hidden_layers = 5
-        self.num_attention_heads = 4
-        self.d_ff = 37
-        self.relative_attention_num_buckets = 8
-        self.dropout_rate = 0.1
-        self.initializer_factor = 0.002
-        self.eos_token_id = 1
-        self.pad_token_id = 0
-        self.decoder_start_token_id = 0
+        self.is_training = is_training
+        self.use_attention_mask = use_attention_mask
+        self.use_labels = use_labels
+        self.vocab_size = vocab_size
+        self.n_positions = n_positions
+        self.hidden_size = hidden_size
+        self.num_hidden_layers = num_hidden_layers
+        self.num_attention_heads = num_attention_heads
+        self.d_ff = d_ff
+        self.relative_attention_num_buckets = relative_attention_num_buckets
+        self.dropout_rate = dropout_rate
+        self.initializer_factor = initializer_factor
+        self.eos_token_id = eos_token_id
+        self.pad_token_id = pad_token_id
+        self.decoder_start_token_id = decoder_start_token_id
         self.scope = None
 
     def prepare_config_and_inputs(self):

tests/test_modeling_tf_common.py

@@ -354,6 +354,69 @@ class TFModelTesterMixin:
             max_diff = np.amax(np.abs(tfo - pto))
             self.assertLessEqual(max_diff, 4e-2)
 
+    def test_train_pipeline_custom_model(self):
+        config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
+        tf_main_layer_classes = set(
+            module_member
+            for model_class in self.all_model_classes
+            for module in (import_module(model_class.__module__),)
+            for module_member_name in dir(module)
+            if module_member_name.endswith("MainLayer")
+            for module_member in (getattr(module, module_member_name),)
+            if isinstance(module_member, type)
+            and tf.keras.layers.Layer in module_member.__bases__
+            and getattr(module_member, "_keras_serializable", False)
+        )
+
+        for main_layer_class in tf_main_layer_classes:
+            # T5MainLayer needs an embed_tokens parameter when called without the inputs_embeds parameter
+            if "T5" in main_layer_class.__name__:
+                # Take the same values than in TFT5ModelTester for this shared layer
+                shared = TFSharedEmbeddings(self.model_tester.vocab_size, self.model_tester.hidden_size, name="shared")
+                config.use_cache = False
+                main_layer = main_layer_class(config, embed_tokens=shared)
+                del inputs_dict["use_cache"]
+            else:
+                main_layer = main_layer_class(config)
+
+            symbolic_inputs = {
+                name: tf.keras.Input(tensor.shape[1:], dtype=tensor.dtype) for name, tensor in inputs_dict.items()
+            }
+
+            if hasattr(self.model_tester, "num_labels"):
+                num_labels = self.model_tester.num_labels
+            else:
+                num_labels = 2
+
+            X = tf.data.Dataset.from_tensor_slices(
+                (inputs_dict, np.random.randint(0, num_labels, (self.model_tester.batch_size, 1)))
+            ).batch(1)
+
+            hidden_states = main_layer(symbolic_inputs)[0]
+            outputs = tf.keras.layers.Dense(num_labels, activation="softmax", name="outputs")(hidden_states)
+            model = tf.keras.models.Model(inputs=symbolic_inputs, outputs=[outputs])
+
+            model.compile(loss="binary_crossentropy", optimizer="adam", metrics=["acc"])
+            model.fit(X, epochs=1)
+
+            with tempfile.TemporaryDirectory() as tmpdirname:
+                filepath = os.path.join(tmpdirname, "keras_model.h5")
+                model.save(filepath)
+                if "T5" in main_layer_class.__name__:
+                    model = tf.keras.models.load_model(
+                        filepath,
+                        custom_objects={
+                            main_layer_class.__name__: main_layer_class,
+                            "TFSharedEmbeddings": TFSharedEmbeddings,
+                        },
+                    )
+                else:
+                    model = tf.keras.models.load_model(
+                        filepath, custom_objects={main_layer_class.__name__: main_layer_class}
+                    )
+                assert isinstance(model, tf.keras.Model)
+                model(inputs_dict)
+
     def test_compile_tf_model(self):
         config, inputs_dict = self.model_tester.prepare_config_and_inputs_for_common()
 

tests/test_modeling_tf_funnel.py

@@ -327,7 +327,7 @@ class TFFunnelModelTester:
 
 
 @require_tf
-class FunnelModelTest(TFModelTesterMixin, unittest.TestCase):
+class TFFunnelModelTest(TFModelTesterMixin, unittest.TestCase):
     all_model_classes = (
         (
             TFFunnelModel,

tests/test_pipelines.py

@@ -31,7 +31,7 @@ TF_TRANSLATION_FINETUNED_MODELS = [("patrickvonplaten/t5-tiny-random", "translat
 TEXT2TEXT_FINETUNED_MODELS = ["patrickvonplaten/t5-tiny-random"]
 TF_TEXT2TEXT_FINETUNED_MODELS = ["patrickvonplaten/t5-tiny-random"]
 
-DIALOGUE_FINETUNED_MODELS = ["microsoft/DialoGPT-medium"]
+DIALOGUE_FINETUNED_MODELS = ["microsoft/DialoGPT-medium"]  # @slow
 
 expected_fill_mask_result = [
     [
@@ -350,8 +350,8 @@ class MonoColumnInputTestCase(unittest.TestCase):
                 nlp, VALID_INPUTS, mandatory_keys, invalid_inputs=invalid_inputs, **SUMMARIZATION_KWARGS
             )
 
-    @slow
     @require_torch
+    @slow
     def test_integration_torch_summarization(self):
         nlp = pipeline(task="summarization", device=DEFAULT_DEVICE_NUM)
         cnn_article = ' (CNN)The Palestinian Authority officially became the 123rd member of the International Criminal Court on Wednesday, a step that gives the court jurisdiction over alleged crimes in Palestinian territories. The formal accession was marked with a ceremony at The Hague, in the Netherlands, where the court is based. The Palestinians signed the ICC\'s founding Rome Statute in January, when they also accepted its jurisdiction over alleged crimes committed "in the occupied Palestinian territory, including East Jerusalem, since June 13, 2014." Later that month, the ICC opened a preliminary examination into the situation in Palestinian territories, paving the way for possible war crimes investigations against Israelis. As members of the court, Palestinians may be subject to counter-charges as well. Israel and the United States, neither of which is an ICC member, opposed the Palestinians\' efforts to join the body. But Palestinian Foreign Minister Riad al-Malki, speaking at Wednesday\'s ceremony, said it was a move toward greater justice. "As Palestine formally becomes a State Party to the Rome Statute today, the world is also a step closer to ending a long era of impunity and injustice," he said, according to an ICC news release. "Indeed, today brings us closer to our shared goals of justice and peace." Judge Kuniko Ozaki, a vice president of the ICC, said acceding to the treaty was just the first step for the Palestinians. "As the Rome Statute today enters into force for the State of Palestine, Palestine acquires all the rights as well as responsibilities that come with being a State Party to the Statute. These are substantive commitments, which cannot be taken lightly," she said. Rights group Human Rights Watch welcomed the development. "Governments seeking to penalize Palestine for joining the ICC should immediately end their pressure, and countries that support universal acceptance of the court\'s treaty should speak out to welcome its membership," said Balkees Jarrah, international justice counsel for the group. "What\'s objectionable is the attempts to undermine international justice, not Palestine\'s decision to join a treaty to which over 100 countries around the world are members." In January, when the preliminary ICC examination was opened, Israeli Prime Minister Benjamin Netanyahu described it as an outrage, saying the court was overstepping its boundaries. The United States also said it "strongly" disagreed with the court\'s decision. "As we have said repeatedly, we do not believe that Palestine is a state and therefore we do not believe that it is eligible to join the ICC," the State Department said in a statement. It urged the warring sides to resolve their differences through direct negotiations. "We will continue to oppose actions against Israel at the ICC as counterproductive to the cause of peace," it said. But the ICC begs to differ with the definition of a state for its purposes and refers to the territories as "Palestine." While a preliminary examination is not a formal investigation, it allows the court to review evidence and determine whether to investigate suspects on both sides. Prosecutor Fatou Bensouda said her office would "conduct its analysis in full independence and impartiality." The war between Israel and Hamas militants in Gaza last summer left more than 2,000 people dead. The inquiry will include alleged war crimes committed since June. The International Criminal Court was set up in 2002 to prosecute genocide, crimes against humanity and war crimes. CNN\'s Vasco Cotovio, Kareem Khadder and Faith Karimi contributed to this report.'
@@ -359,8 +359,8 @@ class MonoColumnInputTestCase(unittest.TestCase):
         result = nlp(cnn_article)
         self.assertEqual(result[0]["summary_text"], expected_cnn_summary)
 
-    @slow
     @require_tf
+    @slow
     def test_tf_summarization(self):
         invalid_inputs = [4, "<mask>"]
         mandatory_keys = ["summary_text"]
@@ -433,8 +433,8 @@ class MonoColumnInputTestCase(unittest.TestCase):
             self._test_mono_column_pipeline(nlp, VALID_INPUTS, {})
         self._test_mono_column_pipeline(nlp, VALID_INPUTS, {}, prefix="This is ")
 
-    @slow
     @require_torch
+    @slow
     def test_integration_torch_conversation(self):
         # When
         nlp = pipeline(task="conversational", device=DEFAULT_DEVICE_NUM)
@@ -465,8 +465,8 @@ class MonoColumnInputTestCase(unittest.TestCase):
         self.assertEqual(result.past_user_inputs[1], "Why do you recommend it?")
         self.assertEqual(result.generated_responses[1], "It's a good book.")
 
-    @slow
     @require_torch
+    @slow
     def test_integration_torch_conversation_truncated_history(self):
         # When
         nlp = pipeline(task="conversational", min_length_for_response=24, device=DEFAULT_DEVICE_NUM)
@@ -614,14 +614,14 @@ class ZeroShotClassificationPipelineTests(unittest.TestCase):
             nlp = pipeline(task="zero-shot-classification", model=model_name, tokenizer=model_name, framework="tf")
             self._test_zero_shot_pipeline(nlp)
 
-    @slow
     @require_torch
+    @slow
     def test_torch_zero_shot_outputs(self):
         nlp = pipeline(task="zero-shot-classification", model="roberta-large-mnli")
         self._test_zero_shot_pipeline_outputs(nlp)
 
-    @slow
     @require_tf
+    @slow
     def test_tf_zero_shot_outputs(self):
         nlp = pipeline(task="zero-shot-classification", model="roberta-large-mnli", framework="tf")
         self._test_zero_shot_pipeline_outputs(nlp)
@@ -647,12 +647,14 @@ class DialoguePipelineTests(unittest.TestCase):
         self.assertRaises(Exception, nlp, invalid_inputs)
 
     @require_torch
+    @slow
     def test_torch_conversation(self):
         for model_name in DIALOGUE_FINETUNED_MODELS:
             nlp = pipeline(task="conversational", model=model_name, tokenizer=model_name)
             self._test_conversation_pipeline(nlp)
 
     @require_tf
+    @slow
     def test_tf_conversation(self):
         for model_name in DIALOGUE_FINETUNED_MODELS:
             nlp = pipeline(task="conversational", model=model_name, tokenizer=model_name, framework="tf")
@@ -806,16 +808,16 @@ class NerPipelineTests(unittest.TestCase):
 class PipelineCommonTests(unittest.TestCase):
     pipelines = SUPPORTED_TASKS.keys()
 
-    @slow
     @require_tf
+    @slow
     def test_tf_defaults(self):
         # Test that pipelines can be correctly loaded without any argument
         for task in self.pipelines:
             with self.subTest(msg="Testing TF defaults with TF and {}".format(task)):
                 pipeline(task, framework="tf")
 
-    @slow
     @require_torch
+    @slow
     def test_pt_defaults(self):
         # Test that pipelines can be correctly loaded without any argument
         for task in self.pipelines:

tests/test_retrieval_rag.py

@@ -0,0 +1,223 @@
+import json
+import os
+import pickle
+import shutil
+import tempfile
+from unittest import TestCase
+from unittest.mock import patch
+
+import numpy as np
+from datasets import Dataset
+
+import faiss
+from transformers.configuration_bart import BartConfig
+from transformers.configuration_dpr import DPRConfig
+from transformers.configuration_rag import RagConfig
+from transformers.retrieval_rag import RagRetriever
+from transformers.testing_utils import require_datasets, require_faiss, require_torch
+from transformers.tokenization_bart import BartTokenizer
+from transformers.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
+from transformers.tokenization_dpr import DPRQuestionEncoderTokenizer
+from transformers.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
+
+
+@require_faiss
+@require_datasets
+class RagRetrieverTest(TestCase):
+    def setUp(self):
+        self.tmpdirname = tempfile.mkdtemp()
+        self.retrieval_vector_size = 8
+
+        # DPR tok
+        vocab_tokens = [
+            "[UNK]",
+            "[CLS]",
+            "[SEP]",
+            "[PAD]",
+            "[MASK]",
+            "want",
+            "##want",
+            "##ed",
+            "wa",
+            "un",
+            "runn",
+            "##ing",
+            ",",
+            "low",
+            "lowest",
+        ]
+        dpr_tokenizer_path = os.path.join(self.tmpdirname, "dpr_tokenizer")
+        os.makedirs(dpr_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(dpr_tokenizer_path, DPR_VOCAB_FILES_NAMES["vocab_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+        # BART tok
+        vocab = [
+            "l",
+            "o",
+            "w",
+            "e",
+            "r",
+            "s",
+            "t",
+            "i",
+            "d",
+            "n",
+            "\u0120",
+            "\u0120l",
+            "\u0120n",
+            "\u0120lo",
+            "\u0120low",
+            "er",
+            "\u0120lowest",
+            "\u0120newer",
+            "\u0120wider",
+            "<unk>",
+        ]
+        vocab_tokens = dict(zip(vocab, range(len(vocab))))
+        merges = ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
+        self.special_tokens_map = {"unk_token": "<unk>"}
+
+        bart_tokenizer_path = os.path.join(self.tmpdirname, "bart_tokenizer")
+        os.makedirs(bart_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(bart_tokenizer_path, BART_VOCAB_FILES_NAMES["vocab_file"])
+        self.merges_file = os.path.join(bart_tokenizer_path, BART_VOCAB_FILES_NAMES["merges_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as fp:
+            fp.write(json.dumps(vocab_tokens) + "\n")
+        with open(self.merges_file, "w", encoding="utf-8") as fp:
+            fp.write("\n".join(merges))
+
+    def get_dpr_tokenizer(self) -> DPRQuestionEncoderTokenizer:
+        return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname, "dpr_tokenizer"))
+
+    def get_bart_tokenizer(self) -> BartTokenizer:
+        return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname, "bart_tokenizer"))
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdirname)
+
+    def get_dummy_hf_index_retriever(self):
+        dataset = Dataset.from_dict(
+            {
+                "id": ["0", "1"],
+                "text": ["foo", "bar"],
+                "title": ["Foo", "Bar"],
+                "embeddings": [np.ones(self.retrieval_vector_size), 2 * np.ones(self.retrieval_vector_size)],
+            }
+        )
+        dataset.add_faiss_index("embeddings", string_factory="Flat", metric_type=faiss.METRIC_INNER_PRODUCT)
+        config = RagConfig(
+            retrieval_vector_size=self.retrieval_vector_size,
+            question_encoder=DPRConfig().to_dict(),
+            generator=BartConfig().to_dict(),
+        )
+        with patch("transformers.retrieval_rag.load_dataset") as mock_load_dataset:
+            mock_load_dataset.return_value = dataset
+            retriever = RagRetriever(
+                config,
+                question_encoder_tokenizer=self.get_dpr_tokenizer(),
+                generator_tokenizer=self.get_bart_tokenizer(),
+            )
+        return retriever
+
+    def get_dummy_legacy_index_retriever(self):
+        dataset = Dataset.from_dict(
+            {
+                "id": ["0", "1"],
+                "text": ["foo", "bar"],
+                "title": ["Foo", "Bar"],
+                "embeddings": [np.ones(self.retrieval_vector_size + 1), 2 * np.ones(self.retrieval_vector_size + 1)],
+            }
+        )
+        dataset.add_faiss_index("embeddings", string_factory="Flat", metric_type=faiss.METRIC_INNER_PRODUCT)
+
+        index_file_name = os.path.join(self.tmpdirname, "hf_bert_base.hnswSQ8_correct_phi_128.c_index")
+        dataset.save_faiss_index("embeddings", index_file_name + ".index.dpr")
+        pickle.dump(dataset["id"], open(index_file_name + ".index_meta.dpr", "wb"))
+
+        passages_file_name = os.path.join(self.tmpdirname, "psgs_w100.tsv.pkl")
+        passages = {sample["id"]: [sample["text"], sample["title"]] for sample in dataset}
+        pickle.dump(passages, open(passages_file_name, "wb"))
+
+        config = RagConfig(
+            retrieval_vector_size=self.retrieval_vector_size,
+            question_encoder=DPRConfig().to_dict(),
+            generator=BartConfig().to_dict(),
+            index_name="legacy",
+            index_path=self.tmpdirname,
+            passages_path=self.tmpdirname,
+        )
+        retriever = RagRetriever(
+            config, question_encoder_tokenizer=self.get_dpr_tokenizer(), generator_tokenizer=self.get_bart_tokenizer()
+        )
+        return retriever
+
+    def test_hf_index_retriever_retrieve(self):
+        n_docs = 1
+        retriever = self.get_dummy_hf_index_retriever()
+        hidden_states = np.array(
+            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+        )
+        retrieved_doc_embeds, doc_ids, doc_dicts = retriever.retrieve(hidden_states, n_docs=n_docs)
+        self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
+        self.assertEqual(len(doc_dicts), 2)
+        self.assertEqual(sorted(doc_dicts[0]), ["embeddings", "id", "text", "title"])
+        self.assertEqual(len(doc_dicts[0]["id"]), n_docs)
+        self.assertEqual(doc_dicts[0]["id"][0], "1")  # max inner product is reached with second doc
+        self.assertEqual(doc_dicts[1]["id"][0], "0")  # max inner product is reached with first doc
+        self.assertListEqual(list(doc_ids), [1, 0])
+
+    def test_legacy_index_retriever_retrieve(self):
+        n_docs = 1
+        retriever = self.get_dummy_legacy_index_retriever()
+        hidden_states = np.array(
+            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+        )
+        retrieved_doc_embeds, doc_ids, doc_dicts = retriever.retrieve(hidden_states, n_docs=n_docs)
+        self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
+        self.assertEqual(len(doc_dicts), 2)
+        self.assertEqual(sorted(doc_dicts[0]), ["text", "title"])
+        self.assertEqual(len(doc_dicts[0]["text"]), n_docs)
+        self.assertEqual(doc_dicts[0]["text"][0], "bar")  # max inner product is reached with second doc
+        self.assertEqual(doc_dicts[1]["text"][0], "foo")  # max inner product is reached with first doc
+        self.assertListEqual(list(doc_ids), [1, 0])
+
+    @require_torch
+    def test_hf_index_retriever_call(self):
+        import torch
+
+        n_docs = 1
+        retriever = self.get_dummy_hf_index_retriever()
+        question_input_ids = [[5, 7], [10, 11]]
+        hidden_states = np.array(
+            [np.ones(self.retrieval_vector_size), -np.ones(self.retrieval_vector_size)], dtype=np.float32
+        )
+        out = retriever(question_input_ids, hidden_states, prefix=retriever.config.generator.prefix, n_docs=n_docs)
+        context_input_ids, context_attention_mask, retrieved_doc_embeds = (
+            out["context_input_ids"],
+            out["context_attention_mask"],
+            out["retrieved_doc_embeds"],
+        )
+        self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
+        self.assertIsInstance(context_input_ids, list)
+        self.assertIsInstance(context_attention_mask, list)
+        self.assertIsInstance(retrieved_doc_embeds, np.ndarray)
+
+        out = retriever(
+            question_input_ids,
+            hidden_states,
+            prefix=retriever.config.generator.prefix,
+            n_docs=n_docs,
+            return_tensors="pt",
+        )
+        context_input_ids, context_attention_mask, retrieved_doc_embeds, doc_ids = (  # noqa: F841
+            out["context_input_ids"],
+            out["context_attention_mask"],
+            out["retrieved_doc_embeds"],
+            out["doc_ids"],
+        )
+        self.assertEqual(retrieved_doc_embeds.shape, (2, n_docs, self.retrieval_vector_size))
+        self.assertIsInstance(context_input_ids, torch.Tensor)
+        self.assertIsInstance(context_attention_mask, torch.Tensor)
+        self.assertIsInstance(retrieved_doc_embeds, torch.Tensor)

tests/test_tokenization_common.py

@@ -743,7 +743,7 @@ class TokenizerTesterMixin:
     #             formatted_input = tokenizer.encode(sequence, add_special_tokens=True, add_prefix_space=False)
 
     #             self.assertEqual(
-    #                 tokenizer.encode(tokens, is_pretokenized=True, add_special_tokens=True), formatted_input
+    #                 tokenizer.encode(tokens, is_split_into_words=True, add_special_tokens=True), formatted_input
     #             )
     #             # This is not supported with the Rust tokenizers
     #             # self.assertEqual(tokenizer.encode(input_ids, add_special_tokens=True), formatted_input)
@@ -1250,20 +1250,20 @@ class TokenizerTesterMixin:
                 # sequence_no_prefix_space = sequence.strip()
 
                 # Test encode for pretokenized inputs
-                output = tokenizer.encode(token_sequence, is_pretokenized=True, add_special_tokens=False)
+                output = tokenizer.encode(token_sequence, is_split_into_words=True, add_special_tokens=False)
                 output_sequence = tokenizer.encode(sequence, add_special_tokens=False)
                 self.assertEqual(output, output_sequence)
 
-                output = tokenizer.encode(token_sequence, is_pretokenized=True, add_special_tokens=True)
+                output = tokenizer.encode(token_sequence, is_split_into_words=True, add_special_tokens=True)
                 output_sequence = tokenizer.encode(sequence, add_special_tokens=True)
                 self.assertEqual(output, output_sequence)
 
                 # Test encode_plus for pretokenized inputs
-                output = tokenizer.encode_plus(token_sequence, is_pretokenized=True, add_special_tokens=False)
+                output = tokenizer.encode_plus(token_sequence, is_split_into_words=True, add_special_tokens=False)
                 output_sequence = tokenizer.encode_plus(sequence, add_special_tokens=False)
                 for key in output.keys():
                     self.assertEqual(output[key], output_sequence[key])
-                output = tokenizer.encode_plus(token_sequence, is_pretokenized=True, add_special_tokens=True)
+                output = tokenizer.encode_plus(token_sequence, is_split_into_words=True, add_special_tokens=True)
                 output_sequence = tokenizer.encode_plus(sequence, add_special_tokens=True)
                 for key in output.keys():
                     self.assertEqual(output[key], output_sequence[key])
@@ -1274,7 +1274,7 @@ class TokenizerTesterMixin:
                 sequence_batch_cleaned_up_spaces = [" " + " ".join(s) for s in token_sequence_batch]
 
                 output = tokenizer.batch_encode_plus(
-                    token_sequence_batch, is_pretokenized=True, add_special_tokens=False
+                    token_sequence_batch, is_split_into_words=True, add_special_tokens=False
                 )
                 output_sequence = tokenizer.batch_encode_plus(
                     sequence_batch_cleaned_up_spaces, add_special_tokens=False
@@ -1282,7 +1282,7 @@ class TokenizerTesterMixin:
                 for key in output.keys():
                     self.assertEqual(output[key], output_sequence[key])
                 output = tokenizer.batch_encode_plus(
-                    token_sequence_batch, is_pretokenized=True, add_special_tokens=True
+                    token_sequence_batch, is_split_into_words=True, add_special_tokens=True
                 )
                 output_sequence = tokenizer.batch_encode_plus(
                     sequence_batch_cleaned_up_spaces, add_special_tokens=True
@@ -1292,25 +1292,25 @@ class TokenizerTesterMixin:
 
                 # Test encode for pretokenized inputs pairs
                 output = tokenizer.encode(
-                    token_sequence, token_sequence, is_pretokenized=True, add_special_tokens=False
+                    token_sequence, token_sequence, is_split_into_words=True, add_special_tokens=False
                 )
                 output_sequence = tokenizer.encode(sequence, sequence, add_special_tokens=False)
                 self.assertEqual(output, output_sequence)
                 output = tokenizer.encode(
-                    token_sequence, token_sequence, is_pretokenized=True, add_special_tokens=True
+                    token_sequence, token_sequence, is_split_into_words=True, add_special_tokens=True
                 )
                 output_sequence = tokenizer.encode(sequence, sequence, add_special_tokens=True)
                 self.assertEqual(output, output_sequence)
 
                 # Test encode_plus for pretokenized inputs pairs
                 output = tokenizer.encode_plus(
-                    token_sequence, token_sequence, is_pretokenized=True, add_special_tokens=False
+                    token_sequence, token_sequence, is_split_into_words=True, add_special_tokens=False
                 )
                 output_sequence = tokenizer.encode_plus(sequence, sequence, add_special_tokens=False)
                 for key in output.keys():
                     self.assertEqual(output[key], output_sequence[key])
                 output = tokenizer.encode_plus(
-                    token_sequence, token_sequence, is_pretokenized=True, add_special_tokens=True
+                    token_sequence, token_sequence, is_split_into_words=True, add_special_tokens=True
                 )
                 output_sequence = tokenizer.encode_plus(sequence, sequence, add_special_tokens=True)
                 for key in output.keys():
@@ -1326,7 +1326,7 @@ class TokenizerTesterMixin:
                 ]
 
                 output = tokenizer.batch_encode_plus(
-                    token_sequence_pair_batch, is_pretokenized=True, add_special_tokens=False
+                    token_sequence_pair_batch, is_split_into_words=True, add_special_tokens=False
                 )
                 output_sequence = tokenizer.batch_encode_plus(
                     sequence_pair_batch_cleaned_up_spaces, add_special_tokens=False
@@ -1334,7 +1334,7 @@ class TokenizerTesterMixin:
                 for key in output.keys():
                     self.assertEqual(output[key], output_sequence[key])
                 output = tokenizer.batch_encode_plus(
-                    token_sequence_pair_batch, is_pretokenized=True, add_special_tokens=True
+                    token_sequence_pair_batch, is_split_into_words=True, add_special_tokens=True
                 )
                 output_sequence = tokenizer.batch_encode_plus(
                     sequence_pair_batch_cleaned_up_spaces, add_special_tokens=True

tests/test_tokenization_fast.py

@@ -340,12 +340,12 @@ class CommonFastTokenizerTest(unittest.TestCase):
         pretokenized_input_pair = "This is a sample pair".split()
 
         # Test encode for pretokenized inputs
-        output_r = tokenizer_r.encode(pretokenized_input_simple, is_pretokenized=True)
-        output_p = tokenizer_p.encode(pretokenized_input_simple, is_pretokenized=True)
+        output_r = tokenizer_r.encode(pretokenized_input_simple, is_split_into_words=True)
+        output_p = tokenizer_p.encode(pretokenized_input_simple, is_split_into_words=True)
         self.assertEqual(output_p, output_r)
 
         kwargs = {
-            "is_pretokenized": True,
+            "is_split_into_words": True,
             "return_token_type_ids": True,
             "return_attention_mask": True,
             "return_overflowing_tokens": False,
@@ -353,7 +353,7 @@ class CommonFastTokenizerTest(unittest.TestCase):
             "return_offsets_mapping": False,  # Not implemented in python tokenizers
         }
         batch_kwargs = {
-            "is_pretokenized": True,
+            "is_split_into_words": True,
             "return_token_type_ids": True,
             "return_attention_mask": True,  # we have an 's' here
             "return_overflowing_tokens": False,
@@ -374,8 +374,8 @@ class CommonFastTokenizerTest(unittest.TestCase):
             self.assertEqual(output_p[key], output_r[key])
 
         # Test encode for pretokenized inputs pairs
-        output_r = tokenizer_r.encode(pretokenized_input_simple, pretokenized_input_pair, is_pretokenized=True)
-        output_p = tokenizer_p.encode(pretokenized_input_simple, pretokenized_input_pair, is_pretokenized=True)
+        output_r = tokenizer_r.encode(pretokenized_input_simple, pretokenized_input_pair, is_split_into_words=True)
+        output_p = tokenizer_p.encode(pretokenized_input_simple, pretokenized_input_pair, is_split_into_words=True)
         self.assertEqual(output_p, output_r)
 
         # Test encode_plus for pretokenized inputs

tests/test_tokenization_layoutlm.py

@@ -0,0 +1,68 @@
+# coding=utf-8
+# Copyright 2018 The Microsoft Research Asia LayoutLM Team Authors.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+
+import os
+import unittest
+
+from transformers.tokenization_layoutlm import VOCAB_FILES_NAMES, LayoutLMTokenizer
+
+from .test_tokenization_common import TokenizerTesterMixin
+
+
+class LayoutLMTokenizationTest(TokenizerTesterMixin, unittest.TestCase):
+
+    tokenizer_class = LayoutLMTokenizer
+
+    def setUp(self):
+        super().setUp()
+
+        vocab_tokens = [
+            "[UNK]",
+            "[CLS]",
+            "[SEP]",
+            "want",
+            "##want",
+            "##ed",
+            "wa",
+            "un",
+            "runn",
+            "##ing",
+            ",",
+            "low",
+            "lowest",
+        ]
+        self.vocab_file = os.path.join(self.tmpdirname, VOCAB_FILES_NAMES["vocab_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+    def get_tokenizer(self, **kwargs):
+        return LayoutLMTokenizer.from_pretrained(self.tmpdirname, **kwargs)
+
+    def get_input_output_texts(self, tokenizer):
+        input_text = "UNwant\u00E9d,running"
+        output_text = "unwanted, running"
+        return input_text, output_text
+
+    def test_full_tokenizer(self):
+        tokenizer = self.tokenizer_class(self.vocab_file)
+
+        tokens = tokenizer.tokenize("UNwant\u00E9d,running")
+        self.assertListEqual(tokens, ["un", "##want", "##ed", ",", "runn", "##ing"])
+        self.assertListEqual(tokenizer.convert_tokens_to_ids(tokens), [7, 4, 5, 10, 8, 9])
+
+    def test_special_tokens_as_you_expect(self):
+        """If you are training a seq2seq model that expects a decoder_prefix token make sure it is prepended to decoder_input_ids """
+        pass

tests/test_tokenization_rag.py

@@ -0,0 +1,110 @@
+import json
+import os
+import shutil
+import tempfile
+from unittest import TestCase
+
+from transformers.configuration_bart import BartConfig
+from transformers.configuration_dpr import DPRConfig
+from transformers.file_utils import is_datasets_available, is_faiss_available, is_torch_available
+from transformers.testing_utils import require_datasets, require_faiss, require_torch
+from transformers.tokenization_bart import BartTokenizer
+from transformers.tokenization_bert import VOCAB_FILES_NAMES as DPR_VOCAB_FILES_NAMES
+from transformers.tokenization_dpr import DPRQuestionEncoderTokenizer
+from transformers.tokenization_roberta import VOCAB_FILES_NAMES as BART_VOCAB_FILES_NAMES
+
+
+if is_torch_available() and is_datasets_available() and is_faiss_available():
+    from transformers.configuration_rag import RagConfig
+    from transformers.tokenization_rag import RagTokenizer
+
+
+@require_faiss
+@require_datasets
+@require_torch
+class RagTokenizerTest(TestCase):
+    def setUp(self):
+        self.tmpdirname = tempfile.mkdtemp()
+        self.retrieval_vector_size = 8
+
+        # DPR tok
+        vocab_tokens = [
+            "[UNK]",
+            "[CLS]",
+            "[SEP]",
+            "[PAD]",
+            "[MASK]",
+            "want",
+            "##want",
+            "##ed",
+            "wa",
+            "un",
+            "runn",
+            "##ing",
+            ",",
+            "low",
+            "lowest",
+        ]
+        dpr_tokenizer_path = os.path.join(self.tmpdirname, "dpr_tokenizer")
+        os.makedirs(dpr_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(dpr_tokenizer_path, DPR_VOCAB_FILES_NAMES["vocab_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as vocab_writer:
+            vocab_writer.write("".join([x + "\n" for x in vocab_tokens]))
+
+        # BART tok
+        vocab = [
+            "l",
+            "o",
+            "w",
+            "e",
+            "r",
+            "s",
+            "t",
+            "i",
+            "d",
+            "n",
+            "\u0120",
+            "\u0120l",
+            "\u0120n",
+            "\u0120lo",
+            "\u0120low",
+            "er",
+            "\u0120lowest",
+            "\u0120newer",
+            "\u0120wider",
+            "<unk>",
+        ]
+        vocab_tokens = dict(zip(vocab, range(len(vocab))))
+        merges = ["#version: 0.2", "\u0120 l", "\u0120l o", "\u0120lo w", "e r", ""]
+        self.special_tokens_map = {"unk_token": "<unk>"}
+
+        bart_tokenizer_path = os.path.join(self.tmpdirname, "bart_tokenizer")
+        os.makedirs(bart_tokenizer_path, exist_ok=True)
+        self.vocab_file = os.path.join(bart_tokenizer_path, BART_VOCAB_FILES_NAMES["vocab_file"])
+        self.merges_file = os.path.join(bart_tokenizer_path, BART_VOCAB_FILES_NAMES["merges_file"])
+        with open(self.vocab_file, "w", encoding="utf-8") as fp:
+            fp.write(json.dumps(vocab_tokens) + "\n")
+        with open(self.merges_file, "w", encoding="utf-8") as fp:
+            fp.write("\n".join(merges))
+
+    def get_dpr_tokenizer(self) -> DPRQuestionEncoderTokenizer:
+        return DPRQuestionEncoderTokenizer.from_pretrained(os.path.join(self.tmpdirname, "dpr_tokenizer"))
+
+    def get_bart_tokenizer(self) -> BartTokenizer:
+        return BartTokenizer.from_pretrained(os.path.join(self.tmpdirname, "bart_tokenizer"))
+
+    def tearDown(self):
+        shutil.rmtree(self.tmpdirname)
+
+    def test_save_load_pretrained_with_saved_config(self):
+
+        save_dir = os.path.join(self.tmpdirname, "rag_tokenizer")
+        rag_config = RagConfig(question_encoder=DPRConfig().to_dict(), generator=BartConfig().to_dict())
+        rag_tokenizer = RagTokenizer(question_encoder=self.get_dpr_tokenizer(), generator=self.get_bart_tokenizer())
+        rag_config.save_pretrained(save_dir)
+        rag_tokenizer.save_pretrained(save_dir)
+        new_rag_tokenizer = RagTokenizer.from_pretrained(save_dir, config=rag_config)
+        self.assertIsInstance(new_rag_tokenizer.question_encoder, DPRQuestionEncoderTokenizer)
+        self.assertEqual(new_rag_tokenizer.question_encoder.vocab, rag_tokenizer.question_encoder.vocab)
+        self.assertIsInstance(new_rag_tokenizer.generator, BartTokenizer)
+        self.assertEqual(new_rag_tokenizer.generator.encoder, rag_tokenizer.generator.encoder)

tests/test_trainer.py

@@ -4,7 +4,7 @@ import datasets
 import numpy as np
 
 from transformers import AutoTokenizer, TrainingArguments, is_torch_available
-from transformers.testing_utils import get_tests_dir, require_torch
+from transformers.testing_utils import get_tests_dir, require_torch, slow
 
 
 if is_torch_available():
@@ -287,6 +287,7 @@ class TrainerIntegrationTest(unittest.TestCase):
         trainer.train()
         self.check_trained_model(trainer.model, alternate_seed=True)
 
+    @slow
     def test_trainer_eval_mrpc(self):
         MODEL_ID = "bert-base-cased-finetuned-mrpc"
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)
@@ -301,6 +302,7 @@ class TrainerIntegrationTest(unittest.TestCase):
         result = trainer.evaluate()
         self.assertLess(result["eval_loss"], 0.2)
 
+    @slow
     def test_trainer_eval_lm(self):
         MODEL_ID = "distilroberta-base"
         tokenizer = AutoTokenizer.from_pretrained(MODEL_ID)

tests/test_utils_check_copies.py

@@ -0,0 +1,104 @@
+import os
+import re
+import shutil
+import sys
+import tempfile
+import unittest
+
+
+git_repo_path = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
+sys.path.append(os.path.join(git_repo_path, "utils"))
+
+import check_copies  # noqa: E402
+
+
+# This is the reference code that will be used in the tests.
+# If BertLMPredictionHead is changed in modeling_bert.py, this code needs to be manually updated.
+REFERENCE_CODE = """    def __init__(self, config):
+        super().__init__()
+        self.transform = BertPredictionHeadTransform(config)
+
+        # The output weights are the same as the input embeddings, but there is
+        # an output-only bias for each token.
+        self.decoder = nn.Linear(config.hidden_size, config.vocab_size, bias=False)
+
+        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
+
+        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
+        self.decoder.bias = self.bias
+
+    def forward(self, hidden_states):
+        hidden_states = self.transform(hidden_states)
+        hidden_states = self.decoder(hidden_states)
+        return hidden_states
+"""
+
+
+class CopyCheckTester(unittest.TestCase):
+    def setUp(self):
+        self.transformer_dir = tempfile.mkdtemp()
+        check_copies.TRANSFORMER_PATH = self.transformer_dir
+        shutil.copy(
+            os.path.join(git_repo_path, "src/transformers/modeling_bert.py"),
+            os.path.join(self.transformer_dir, "modeling_bert.py"),
+        )
+
+    def tearDown(self):
+        check_copies.TRANSFORMER_PATH = "src/transformers"
+        shutil.rmtree(self.transformer_dir)
+
+    def check_copy_consistency(self, comment, class_name, class_code, overwrite_result=None):
+        code = comment + f"\nclass {class_name}(nn.Module):\n" + class_code
+        if overwrite_result is not None:
+            expected = comment + f"\nclass {class_name}(nn.Module):\n" + overwrite_result
+        fname = os.path.join(self.transformer_dir, "new_code.py")
+        with open(fname, "w") as f:
+            f.write(code)
+        if overwrite_result is None:
+            self.assertTrue(check_copies.is_copy_consistent(fname))
+        else:
+            check_copies.is_copy_consistent(f.name, overwrite=True)
+            with open(fname, "r") as f:
+                self.assertTrue(f.read(), expected)
+
+    def test_find_code_in_transformers(self):
+        code = check_copies.find_code_in_transformers("modeling_bert.BertLMPredictionHead")
+        self.assertEqual(code, REFERENCE_CODE)
+
+    def test_is_copy_consistent(self):
+        # Base copy consistency
+        self.check_copy_consistency(
+            "# Copied from transformers.modeling_bert.BertLMPredictionHead",
+            "BertLMPredictionHead",
+            REFERENCE_CODE + "\n",
+        )
+
+        # With no empty line at the end
+        self.check_copy_consistency(
+            "# Copied from transformers.modeling_bert.BertLMPredictionHead",
+            "BertLMPredictionHead",
+            REFERENCE_CODE,
+        )
+
+        # Copy consistency with rename
+        self.check_copy_consistency(
+            "# Copied from transformers.modeling_bert.BertLMPredictionHead with Bert->TestModel",
+            "TestModelLMPredictionHead",
+            re.sub("Bert", "TestModel", REFERENCE_CODE),
+        )
+
+        # Copy consistency with a really long name
+        long_class_name = "TestModelWithAReallyLongNameBecauseSomePeopleLikeThatForSomeReasonIReallyDontUnderstand"
+        self.check_copy_consistency(
+            f"# Copied from transformers.modeling_bert.BertLMPredictionHead with Bert->{long_class_name}",
+            f"{long_class_name}LMPredictionHead",
+            re.sub("Bert", long_class_name, REFERENCE_CODE),
+        )
+
+        # Copy consistency with overwrite
+        self.check_copy_consistency(
+            "# Copied from transformers.modeling_bert.BertLMPredictionHead with Bert->TestModel",
+            "TestModelLMPredictionHead",
+            REFERENCE_CODE,
+            overwrite_result=re.sub("Bert", "TestModel", REFERENCE_CODE),
+        )

utils/check_copies.py

@@ -0,0 +1,181 @@
+# coding=utf-8
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
+import argparse
+import glob
+import os
+import re
+import tempfile
+
+
+# All paths are set with the intent you should run this script from the root of the repo with the command
+# python utils/check_copies.py
+TRANSFORMERS_PATH = "src/transformers"
+
+
+def find_code_in_transformers(object_name):
+    """ Find and return the code source code of `object_name`."""
+    parts = object_name.split(".")
+    i = 0
+
+    # First let's find the module where our object lives.
+    module = parts[i]
+    while i < len(parts) and not os.path.isfile(os.path.join(TRANSFORMERS_PATH, f"{module}.py")):
+        i += 1
+        module = os.path.join(module, parts[i])
+    if i >= len(parts):
+        raise ValueError(
+            f"`object_name` should begin with the name of a module of transformers but got {object_name}."
+        )
+
+    with open(os.path.join(TRANSFORMERS_PATH, f"{module}.py"), "r", encoding="utf-8") as f:
+        lines = f.readlines()
+
+    # Now let's find the class / func in the code!
+    indent = ""
+    line_index = 0
+    for name in parts[i + 1 :]:
+        while line_index < len(lines) and re.search(f"^{indent}(class|def)\s+{name}", lines[line_index]) is None:
+            line_index += 1
+        indent += "    "
+        line_index += 1
+
+    if line_index >= len(lines):
+        raise ValueError(f" {object_name} does not match any function or class in {module}.")
+
+    # We found the beginning of the class / func, now let's find the end (when the indent diminishes).
+    start_index = line_index
+    while line_index < len(lines) and (lines[line_index].startswith(indent) or len(lines[line_index]) <= 1):
+        line_index += 1
+    # Clean up empty lines at the end (if any).
+    while len(lines[line_index - 1]) <= 1:
+        line_index -= 1
+
+    code_lines = lines[start_index:line_index]
+    return "".join(code_lines)
+
+
+_re_copy_warning = re.compile(r"^(\s*)#\s*Copied from\s+transformers\.(\S+\.\S+)\s*($|\S.*$)")
+_re_replace_pattern = re.compile(r"with\s+(\S+)->(\S+)(?:\s|$)")
+
+
+def blackify(code):
+    """
+    Applies the black part of our `make style` command to `code`.
+    """
+    has_indent = code.startswith("    ")
+    if has_indent:
+        code = f"class Bla:\n{code}"
+    with tempfile.TemporaryDirectory() as d:
+        fname = os.path.join(d, "tmp.py")
+        with open(fname, "w", encoding="utf-8") as f:
+            f.write(code)
+        os.system(f"black -q --line-length 119 --target-version py35 {fname}")
+        with open(fname, "r", encoding="utf-8") as f:
+            result = f.read()
+            return result[len("class Bla:\n") :] if has_indent else result
+
+
+def is_copy_consistent(filename, overwrite=False):
+    """
+    Check if the code commented as a copy in `filename` matches the original.
+
+    Return the differences or overwrites the content depending on `overwrite`.
+    """
+    with open(filename, "r", encoding="utf-8") as f:
+        lines = f.readlines()
+    found_diff = False
+    line_index = 0
+    # Not a foor loop cause `lines` is going to change (if `overwrite=True`).
+    while line_index < len(lines):
+        search = _re_copy_warning.search(lines[line_index])
+        if search is None:
+            line_index += 1
+            continue
+
+        # There is some copied code here, let's retrieve the original.
+        indent, object_name, replace_pattern = search.groups()
+        theoretical_code = find_code_in_transformers(object_name)
+        theoretical_indent = re.search(r"^(\s*)\S", theoretical_code).groups()[0]
+
+        start_index = line_index + 1 if indent == theoretical_indent else line_index + 2
+        indent = theoretical_indent
+        line_index = start_index
+
+        # Loop to check the observed code, stop when indentation diminishes or if we see a End copy comment.
+        should_continue = True
+        while line_index < len(lines) and should_continue:
+            line_index += 1
+            if line_index >= len(lines):
+                break
+            line = lines[line_index]
+            should_continue = (len(line) <= 1 or line.startswith(indent)) and re.search(
+                f"^{indent}# End copy", line
+            ) is None
+        # Clean up empty lines at the end (if any).
+        while len(lines[line_index - 1]) <= 1:
+            line_index -= 1
+
+        observed_code_lines = lines[start_index:line_index]
+        observed_code = "".join(observed_code_lines)
+
+        # Before comparing, use the `replace_pattern` on the original code.
+        if len(replace_pattern) > 0:
+            search_patterns = _re_replace_pattern.search(replace_pattern)
+            if search_patterns is not None:
+                obj1, obj2 = search_patterns.groups()
+                theoretical_code = re.sub(obj1, obj2, theoretical_code)
+
+        # Blackify each version before comparing them.
+        observed_code = blackify(observed_code)
+        theoretical_code = blackify(theoretical_code)
+
+        # Test for a diff and act accordingly.
+        if observed_code != theoretical_code:
+            found_diff = True
+            if overwrite:
+                lines = lines[:start_index] + [theoretical_code] + lines[line_index:]
+                line_index = start_index + 1
+
+    if overwrite and found_diff:
+        # Warn the user a file has been modified.
+        print(f"Detected changes, rewriting {filename}.")
+        with open(filename, "w", encoding="utf-8") as f:
+            f.writelines(lines)
+    return not found_diff
+
+
+def check_copies(overwrite: bool = False):
+    all_files = glob.glob(os.path.join(TRANSFORMERS_PATH, "**/*.py"), recursive=True)
+    diffs = []
+    for filename in all_files:
+        consistent = is_copy_consistent(filename, overwrite)
+        if not consistent:
+            diffs.append(filename)
+    if not overwrite and len(diffs) > 0:
+        diff = "\n".join(diffs)
+        raise Exception(
+            "Found copy inconsistencies in the following files:\n"
+            + diff
+            + "\nRun `make fix-copies` or `python utils/check_copies --fix_and_overwrite` to fix them."
+        )
+
+
+if __name__ == "__main__":
+    parser = argparse.ArgumentParser()
+    parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
+    args = parser.parse_args()
+
+    check_copies(args.fix_and_overwrite)

utils/check_repo.py

@@ -1,3 +1,18 @@
+# coding=utf-8
+# Copyright 2020 The HuggingFace Inc. team.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+
 import importlib
 import inspect
 import os