Upload folder using huggingface_hub

config.json

@@ -3,6 +3,10 @@
     "EmpathyModel"
   ],
   "attn_heads": 1,
+  "auto_map": {
+    "AutoConfig": "modeling_empathy.EmpathyModelConfig",
+    "AutoModel": "modeling_empathy.EmpathyModel"
+  },
   "empathy_num_labels": 3,
   "empathy_type": "ER",
   "hidden_dropout_prob": 0.2,

modeling_empathy.py

@@ -0,0 +1,59 @@
+"""Custom modeling file for Empathy Classification Model"""
+import torch
+from transformers import PreTrainedModel, PretrainedConfig
+from .models.models import BiEncoderAttentionWithRationaleClassification
+
+
+class EmpathyModelConfig(PretrainedConfig):
+    """Configuration for Empathy Model"""
+    model_type = "empathy_classifier"
+    
+    def __init__(
+        self,
+        hidden_dropout_prob=0.2,
+        rationale_num_labels=2,
+        empathy_num_labels=3,
+        hidden_size=768,
+        attn_heads=1,
+        max_length=64,
+        **kwargs
+    ):
+        super().__init__(**kwargs)
+        self.hidden_dropout_prob = hidden_dropout_prob
+        self.rationale_num_labels = rationale_num_labels
+        self.empathy_num_labels = empathy_num_labels
+        self.hidden_size = hidden_size
+        self.attn_heads = attn_heads
+        self.max_length = max_length
+
+
+class EmpathyModel(PreTrainedModel):
+    """HuggingFace wrapper for Empathy Model"""
+    config_class = EmpathyModelConfig
+    
+    def __init__(self, config):
+        super().__init__(config)
+        self.config = config
+        self.model = BiEncoderAttentionWithRationaleClassification(
+            hidden_dropout_prob=config.hidden_dropout_prob,
+            rationale_num_labels=config.rationale_num_labels,
+            empathy_num_labels=config.empathy_num_labels,
+            hidden_size=config.hidden_size,
+            attn_heads=config.attn_heads
+        )
+    
+    def forward(
+        self,
+        input_ids_SP=None,
+        input_ids_RP=None,
+        attention_mask_SP=None,
+        attention_mask_RP=None,
+        **kwargs
+    ):
+        return self.model(
+            input_ids_SP=input_ids_SP,
+            input_ids_RP=input_ids_RP,
+            attention_mask_SP=attention_mask_SP,
+            attention_mask_RP=attention_mask_RP,
+            **kwargs
+        )

models/__init__.py

@@ -0,0 +1 @@
+__version__ = "2.8.0"

models/activations.py

@@ -0,0 +1,59 @@
+import logging
+import math
+
+import torch
+import torch.nn.functional as F
+
+
+logger = logging.getLogger(__name__)
+
+
+def swish(x):
+	return x * torch.sigmoid(x)
+
+
+def _gelu_python(x):
+	""" Original Implementation of the gelu activation function in Google Bert repo when initially created.
+		For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
+		0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
+		This is now written in C in torch.nn.functional
+		Also see https://arxiv.org/abs/1606.08415
+	"""
+	return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))
+
+
+def gelu_new(x):
+	""" Implementation of the gelu activation function currently in Google Bert repo (identical to OpenAI GPT).
+		Also see https://arxiv.org/abs/1606.08415
+	"""
+	return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3.0))))
+
+
+if torch.__version__ < "1.4.0":
+	gelu = _gelu_python
+else:
+	gelu = F.gelu
+	try:
+		import torch_xla  # noqa F401
+
+		logger.warning(
+			"The torch_xla package was detected in the python environment. PyTorch/XLA and JIT is untested,"
+			" no activation function will be traced with JIT."
+		)
+	except ImportError:
+		gelu_new = torch.jit.script(gelu_new)
+
+ACT2FN = {
+	"relu": F.relu,
+	"swish": swish,
+	"gelu": gelu,
+	"tanh": torch.tanh,
+	"gelu_new": gelu_new,
+}
+
+
+def get_activation(activation_string):
+	if activation_string in ACT2FN:
+		return ACT2FN[activation_string]
+	else:
+		raise KeyError("function {} not found in ACT2FN mapping {}".format(activation_string, list(ACT2FN.keys())))

models/configuration_bert.py

@@ -0,0 +1,143 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" BERT model configuration """
+
+
+import logging
+
+from .configuration_utils import PretrainedConfig
+
+
+logger = logging.getLogger(__name__)
+
+BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+	"bert-base-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json",
+	"bert-large-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-config.json",
+	"bert-base-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-config.json",
+	"bert-large-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-config.json",
+	"bert-base-multilingual-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-config.json",
+	"bert-base-multilingual-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-config.json",
+	"bert-base-chinese": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-config.json",
+	"bert-base-german-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-config.json",
+	"bert-large-uncased-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-config.json",
+	"bert-large-cased-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-config.json",
+	"bert-large-uncased-whole-word-masking-finetuned-squad": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-config.json",
+	"bert-large-cased-whole-word-masking-finetuned-squad": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-config.json",
+	"bert-base-cased-finetuned-mrpc": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-config.json",
+	"bert-base-german-dbmdz-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-dbmdz-cased-config.json",
+	"bert-base-german-dbmdz-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-dbmdz-uncased-config.json",
+	"bert-base-japanese": "https://s3.amazonaws.com/models.huggingface.co/bert/cl-tohoku/bert-base-japanese-config.json",
+	"bert-base-japanese-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/cl-tohoku/bert-base-japanese-whole-word-masking-config.json",
+	"bert-base-japanese-char": "https://s3.amazonaws.com/models.huggingface.co/bert/cl-tohoku/bert-base-japanese-char-config.json",
+	"bert-base-japanese-char-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/cl-tohoku/bert-base-japanese-char-whole-word-masking-config.json",
+	"bert-base-finnish-cased-v1": "https://s3.amazonaws.com/models.huggingface.co/bert/TurkuNLP/bert-base-finnish-cased-v1/config.json",
+	"bert-base-finnish-uncased-v1": "https://s3.amazonaws.com/models.huggingface.co/bert/TurkuNLP/bert-base-finnish-uncased-v1/config.json",
+	"bert-base-dutch-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/wietsedv/bert-base-dutch-cased/config.json",
+}
+
+
+class BertConfig(PretrainedConfig):
+	r"""
+		This is the configuration class to store the configuration of a :class:`~transformers.BertModel`.
+		It is used to instantiate an BERT 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 BERT `bert-base-uncased <https://huggingface.co/bert-base-uncased>`__ architecture.
+
+		Configuration objects inherit from  :class:`~transformers.PretrainedConfig` and can be used
+		to control the model outputs. Read the documentation from  :class:`~transformers.PretrainedConfig`
+		for more information.
+
+
+		Args:
+			vocab_size (:obj:`int`, optional, defaults to 30522):
+				Vocabulary size of the BERT 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 "gelu"):
+				The non-linear activation function (function or string) in the encoder and pooler.
+				If string, "gelu", "relu", "swish" and "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.
+
+		Example::
+
+			from transformers import BertModel, BertConfig
+
+			# Initializing a BERT bert-base-uncased style configuration
+			configuration = BertConfig()
+
+			# Initializing a model from the bert-base-uncased style configuration
+			model = BertModel(configuration)
+
+			# Accessing the model configuration
+			configuration = model.config
+
+		Attributes:
+			pretrained_config_archive_map (Dict[str, str]):
+				A dictionary containing all the available pre-trained checkpoints.
+	"""
+	pretrained_config_archive_map = BERT_PRETRAINED_CONFIG_ARCHIVE_MAP
+	model_type = "bert"
+
+	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,
+		**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

models/configuration_roberta.py

@@ -0,0 +1,73 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" RoBERTa configuration """
+
+
+import logging
+
+from .configuration_bert import BertConfig
+
+
+logger = logging.getLogger(__name__)
+
+ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP = {
+	"roberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-config.json",
+	"roberta-large": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-config.json",
+	"roberta-large-mnli": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-mnli-config.json",
+	"distilroberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/distilroberta-base-config.json",
+	"roberta-base-openai-detector": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-openai-detector-config.json",
+	"roberta-large-openai-detector": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-openai-detector-config.json",
+}
+
+
+class RobertaConfig(BertConfig):
+	r"""
+		This is the configuration class to store the configuration of an :class:`~transformers.RobertaModel`.
+		It is used to instantiate an RoBERTa 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 BERT `bert-base-uncased <https://huggingface.co/bert-base-uncased>`__ architecture.
+
+		Configuration objects inherit from  :class:`~transformers.PretrainedConfig` and can be used
+		to control the model outputs. Read the documentation from  :class:`~transformers.PretrainedConfig`
+		for more information.
+
+		The :class:`~transformers.RobertaConfig` class directly inherits :class:`~transformers.BertConfig`.
+		It reuses the same defaults. Please check the parent class for more information.
+
+		Example::
+
+			from transformers import RobertaConfig, RobertaModel
+
+			# Initializing a RoBERTa configuration
+			configuration = RobertaConfig()
+
+			# Initializing a model from the configuration
+			model = RobertaModel(configuration)
+
+			# Accessing the model configuration
+			configuration = model.config
+
+		Attributes:
+			pretrained_config_archive_map (Dict[str, str]):
+				A dictionary containing all the available pre-trained checkpoints.
+	"""
+	pretrained_config_archive_map = ROBERTA_PRETRAINED_CONFIG_ARCHIVE_MAP
+	model_type = "roberta"
+
+	def __init__(self, pad_token_id=1, bos_token_id=0, eos_token_id=2, **kwargs):
+		"""Constructs FlaubertConfig.
+		"""
+		super().__init__(pad_token_id=pad_token_id, bos_token_id=bos_token_id, eos_token_id=eos_token_id, **kwargs)

models/configuration_utils.py

@@ -0,0 +1,432 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+""" Configuration base class and utilities."""
+
+
+import copy
+import json
+import logging
+import os
+from typing import Dict, Optional, Tuple
+
+from .file_utils import CONFIG_NAME, cached_path, hf_bucket_url, is_remote_url
+
+
+logger = logging.getLogger(__name__)
+
+
+class PretrainedConfig(object):
+	r""" Base class for all configuration classes.
+		Handles a few parameters common to all models' configurations as well as methods for loading/downloading/saving configurations.
+
+		Note:
+			A configuration file can be loaded and saved to disk. Loading the configuration file and using this file to initialize a model does **not** load the model weights.
+			It only affects the model's configuration.
+
+		Class attributes (overridden by derived classes):
+			- ``pretrained_config_archive_map``: a python ``dict`` with `shortcut names` (string) as keys and `url` (string) of associated pretrained model configurations as values.
+			- ``model_type``: a string that identifies the model type, that we serialize into the JSON file, and that we use to recreate the correct object in :class:`~transformers.AutoConfig`.
+
+		Args:
+			finetuning_task (:obj:`string` or :obj:`None`, `optional`, defaults to :obj:`None`):
+				Name of the task used to fine-tune the model. This can be used when converting from an original (TensorFlow or PyTorch) checkpoint.
+			num_labels (:obj:`int`, `optional`, defaults to `2`):
+				Number of classes to use when the model is a classification model (sequences/tokens)
+			output_attentions (:obj:`bool`, `optional`, defaults to :obj:`False`):
+				Should the model returns attentions weights.
+			output_hidden_states (:obj:`string`, `optional`, defaults to :obj:`False`):
+				Should the model returns all hidden-states.
+			torchscript (:obj:`bool`, `optional`, defaults to :obj:`False`):
+				Is the model used with Torchscript (for PyTorch models).
+	"""
+	pretrained_config_archive_map: Dict[str, str] = {}
+	model_type: str = ""
+
+	def __init__(self, **kwargs):
+		# Attributes with defaults
+		self.output_attentions = kwargs.pop("output_attentions", False)
+		self.output_hidden_states = kwargs.pop("output_hidden_states", False)
+		self.use_cache = kwargs.pop("use_cache", True)  # Not used by all models
+		self.torchscript = kwargs.pop("torchscript", False)  # Only used by PyTorch models
+		self.use_bfloat16 = kwargs.pop("use_bfloat16", False)
+		self.pruned_heads = kwargs.pop("pruned_heads", {})
+
+		# Is decoder is used in encoder-decoder models to differentiate encoder from decoder
+		self.is_encoder_decoder = kwargs.pop("is_encoder_decoder", False)
+		self.is_decoder = kwargs.pop("is_decoder", False)
+
+		# Parameters for sequence generation
+		self.max_length = kwargs.pop("max_length", 20)
+		self.min_length = kwargs.pop("min_length", 0)
+		self.do_sample = kwargs.pop("do_sample", False)
+		self.early_stopping = kwargs.pop("early_stopping", False)
+		self.num_beams = kwargs.pop("num_beams", 1)
+		self.temperature = kwargs.pop("temperature", 1.0)
+		self.top_k = kwargs.pop("top_k", 50)
+		self.top_p = kwargs.pop("top_p", 1.0)
+		self.repetition_penalty = kwargs.pop("repetition_penalty", 1.0)
+		self.length_penalty = kwargs.pop("length_penalty", 1.0)
+		self.no_repeat_ngram_size = kwargs.pop("no_repeat_ngram_size", 0)
+		self.bad_words_ids = kwargs.pop("bad_words_ids", None)
+		self.num_return_sequences = kwargs.pop("num_return_sequences", 1)
+
+		# Fine-tuning task arguments
+		self.architectures = kwargs.pop("architectures", None)
+		self.finetuning_task = kwargs.pop("finetuning_task", None)
+		self.num_labels = kwargs.pop("num_labels", 2)
+		self.empathy_num_labels = kwargs.pop("empathy_num_labels", 2)
+		self.rationale_num_labels = kwargs.pop("rationale_num_labels", 2)
+		self.id2label = kwargs.pop("id2label", {i: f"LABEL_{i}" for i in range(self.num_labels)})
+		self.id2label = dict((int(key), value) for key, value in self.id2label.items())
+		self.label2id = kwargs.pop("label2id", dict(zip(self.id2label.values(), self.id2label.keys())))
+		self.label2id = dict((key, int(value)) for key, value in self.label2id.items())
+
+		# Tokenizer arguments TODO: eventually tokenizer and models should share the same config
+		self.prefix = kwargs.pop("prefix", None)
+		self.bos_token_id = kwargs.pop("bos_token_id", None)
+		self.pad_token_id = kwargs.pop("pad_token_id", None)
+		self.eos_token_id = kwargs.pop("eos_token_id", None)
+		self.decoder_start_token_id = kwargs.pop("decoder_start_token_id", None)
+
+		# task specific arguments
+		self.task_specific_params = kwargs.pop("task_specific_params", None)
+
+		# TPU arguments
+		self.xla_device = kwargs.pop("xla_device", None)
+
+		# Additional attributes without default values
+		for key, value in kwargs.items():
+			try:
+				setattr(self, key, value)
+			except AttributeError as err:
+				logger.error("Can't set {} with value {} for {}".format(key, value, self))
+				raise err
+
+	@property
+	def num_labels(self):
+		return self._num_labels
+
+	@num_labels.setter
+	def num_labels(self, num_labels):
+		self._num_labels = num_labels
+		self.id2label = {i: "LABEL_{}".format(i) for i in range(self.num_labels)}
+		self.id2label = dict((int(key), value) for key, value in self.id2label.items())
+		self.label2id = dict(zip(self.id2label.values(), self.id2label.keys()))
+		self.label2id = dict((key, int(value)) for key, value in self.label2id.items())
+
+	def save_pretrained(self, save_directory):
+		"""
+		Save a configuration object to the directory `save_directory`, so that it
+		can be re-loaded using the :func:`~transformers.PretrainedConfig.from_pretrained` class method.
+
+		Args:
+			save_directory (:obj:`string`):
+				Directory where the configuration JSON file will be saved.
+		"""
+		assert os.path.isdir(
+			save_directory
+		), "Saving path should be a directory where the model and configuration can be saved"
+
+		# If we save using the predefined names, we can load using `from_pretrained`
+		output_config_file = os.path.join(save_directory, CONFIG_NAME)
+
+		self.to_json_file(output_config_file, use_diff=True)
+		logger.info("Configuration saved in {}".format(output_config_file))
+
+	@classmethod
+	def from_pretrained(cls, pretrained_model_name_or_path, **kwargs) -> "PretrainedConfig":
+		r"""
+
+		Instantiate a :class:`~transformers.PretrainedConfig` (or a derived class) from a pre-trained model configuration.
+
+		Args:
+			pretrained_model_name_or_path (:obj:`string`):
+				either:
+				  - a string with the `shortcut name` of a pre-trained model configuration to load from cache or
+					download, e.g.: ``bert-base-uncased``.
+				  - a string with the `identifier name` of a pre-trained model configuration that was user-uploaded to
+					our S3, e.g.: ``dbmdz/bert-base-german-cased``.
+				  - a path to a `directory` containing a configuration file saved using the
+					:func:`~transformers.PretrainedConfig.save_pretrained` method, e.g.: ``./my_model_directory/``.
+				  - a path or url to a saved configuration JSON `file`, e.g.:
+					``./my_model_directory/configuration.json``.
+			cache_dir (:obj:`string`, `optional`):
+				Path to a directory in which a downloaded pre-trained model
+				configuration should be cached if the standard cache should not be used.
+			kwargs (:obj:`Dict[str, any]`, `optional`):
+				The values in kwargs of any keys which are configuration attributes will be used to override the loaded
+				values. Behavior concerning key/value pairs whose keys are *not* configuration attributes is
+				controlled by the `return_unused_kwargs` keyword parameter.
+			force_download (:obj:`bool`, `optional`, defaults to :obj:`False`):
+				Force to (re-)download the model weights and configuration files and override the cached versions if they exist.
+			resume_download (:obj:`bool`, `optional`, defaults to :obj:`False`):
+				Do not delete incompletely recieved file. Attempt to resume the download if such a file exists.
+			proxies (:obj:`Dict`, `optional`):
+				A dictionary of proxy servers to use by protocol or endpoint, e.g.:
+				:obj:`{'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}.`
+				The proxies are used on each request.
+			return_unused_kwargs: (`optional`) bool:
+				If False, then this function returns just the final configuration object.
+				If True, then this functions returns a :obj:`Tuple(config, unused_kwargs)` where `unused_kwargs` is a
+				dictionary consisting of the key/value pairs whose keys are not configuration attributes: ie the part
+				of kwargs which has not been used to update `config` and is otherwise ignored.
+
+		Returns:
+			:class:`PretrainedConfig`: An instance of a configuration object
+
+		Examples::
+
+			# We can't instantiate directly the base class `PretrainedConfig` so let's show the examples on a
+			# derived class: BertConfig
+			config = BertConfig.from_pretrained('bert-base-uncased')    # Download configuration from S3 and cache.
+			config = BertConfig.from_pretrained('./test/saved_model/')  # E.g. config (or model) was saved using `save_pretrained('./test/saved_model/')`
+			config = BertConfig.from_pretrained('./test/saved_model/my_configuration.json')
+			config = BertConfig.from_pretrained('bert-base-uncased', output_attention=True, foo=False)
+			assert config.output_attention == True
+			config, unused_kwargs = BertConfig.from_pretrained('bert-base-uncased', output_attention=True,
+															   foo=False, return_unused_kwargs=True)
+			assert config.output_attention == True
+			assert unused_kwargs == {'foo': False}
+
+		"""
+		config_dict, kwargs = cls.get_config_dict(pretrained_model_name_or_path, **kwargs)
+		return cls.from_dict(config_dict, **kwargs)
+
+	@classmethod
+	def get_config_dict(
+		cls, pretrained_model_name_or_path: str, pretrained_config_archive_map: Optional[Dict] = None, **kwargs
+	) -> Tuple[Dict, Dict]:
+		"""
+		From a `pretrained_model_name_or_path`, resolve to a dictionary of parameters, to be used
+		for instantiating a Config using `from_dict`.
+
+		Parameters:
+			pretrained_model_name_or_path (:obj:`string`):
+				The identifier of the pre-trained checkpoint from which we want the dictionary of parameters.
+			pretrained_config_archive_map: (:obj:`Dict[str, str]`, `optional`) Dict:
+				A map of `shortcut names` to `url`. By default, will use the current class attribute.
+
+		Returns:
+			:obj:`Tuple[Dict, Dict]`: The dictionary that will be used to instantiate the configuration object.
+
+		"""
+		cache_dir = kwargs.pop("cache_dir", None)
+		force_download = kwargs.pop("force_download", False)
+		resume_download = kwargs.pop("resume_download", False)
+		proxies = kwargs.pop("proxies", None)
+		local_files_only = kwargs.pop("local_files_only", False)
+
+		if pretrained_config_archive_map is None:
+			pretrained_config_archive_map = cls.pretrained_config_archive_map
+
+		if pretrained_model_name_or_path in pretrained_config_archive_map:
+			config_file = pretrained_config_archive_map[pretrained_model_name_or_path]
+		elif os.path.isdir(pretrained_model_name_or_path):
+			config_file = os.path.join(pretrained_model_name_or_path, CONFIG_NAME)
+		elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
+			config_file = pretrained_model_name_or_path
+		else:
+			config_file = hf_bucket_url(pretrained_model_name_or_path, postfix=CONFIG_NAME)
+
+		try:
+			# Load from URL or cache if already cached
+			resolved_config_file = cached_path(
+				config_file,
+				cache_dir=cache_dir,
+				force_download=force_download,
+				proxies=proxies,
+				resume_download=resume_download,
+				local_files_only=local_files_only,
+			)
+			# Load config dict
+			if resolved_config_file is None:
+				raise EnvironmentError
+			config_dict = cls._dict_from_json_file(resolved_config_file)
+
+		except EnvironmentError:
+			if pretrained_model_name_or_path in pretrained_config_archive_map:
+				msg = "Couldn't reach server at '{}' to download pretrained model configuration file.".format(
+					config_file
+				)
+			else:
+				msg = (
+					"Can't load '{}'. Make sure that:\n\n"
+					"- '{}' is a correct model identifier listed on 'https://huggingface.co/models'\n\n"
+					"- or '{}' is the correct path to a directory containing a '{}' file\n\n".format(
+						pretrained_model_name_or_path,
+						pretrained_model_name_or_path,
+						pretrained_model_name_or_path,
+						CONFIG_NAME,
+					)
+				)
+			raise EnvironmentError(msg)
+
+		except json.JSONDecodeError:
+			msg = (
+				"Couldn't reach server at '{}' to download configuration file or "
+				"configuration file is not a valid JSON file. "
+				"Please check network or file content here: {}.".format(config_file, resolved_config_file)
+			)
+			raise EnvironmentError(msg)
+
+		if resolved_config_file == config_file:
+			logger.info("loading configuration file {}".format(config_file))
+		else:
+			logger.info("loading configuration file {} from cache at {}".format(config_file, resolved_config_file))
+
+		return config_dict, kwargs
+
+	@classmethod
+	def from_dict(cls, config_dict: Dict, **kwargs) -> "PretrainedConfig":
+		"""
+		Constructs a `Config` from a Python dictionary of parameters.
+
+		Args:
+			config_dict (:obj:`Dict[str, any]`):
+				Dictionary that will be used to instantiate the configuration object. Such a dictionary can be retrieved
+				from a pre-trained checkpoint by leveraging the :func:`~transformers.PretrainedConfig.get_config_dict`
+				method.
+			kwargs (:obj:`Dict[str, any]`):
+				Additional parameters from which to initialize the configuration object.
+
+		Returns:
+			:class:`PretrainedConfig`: An instance of a configuration object
+		"""
+		return_unused_kwargs = kwargs.pop("return_unused_kwargs", False)
+
+		config = cls(**config_dict)
+
+		if hasattr(config, "pruned_heads"):
+			config.pruned_heads = dict((int(key), value) for key, value in config.pruned_heads.items())
+
+		# Update config with kwargs if needed
+		to_remove = []
+		for key, value in kwargs.items():
+			if hasattr(config, key):
+				setattr(config, key, value)
+				to_remove.append(key)
+		for key in to_remove:
+			kwargs.pop(key, None)
+
+		logger.info("Model config %s", str(config))
+		if return_unused_kwargs:
+			return config, kwargs
+		else:
+			return config
+
+	@classmethod
+	def from_json_file(cls, json_file: str) -> "PretrainedConfig":
+		"""
+		Constructs a `Config` from the path to a json file of parameters.
+
+		Args:
+			json_file (:obj:`string`):
+				Path to the JSON file containing the parameters.
+
+		Returns:
+			:class:`PretrainedConfig`: An instance of a configuration object
+
+		"""
+		config_dict = cls._dict_from_json_file(json_file)
+		return cls(**config_dict)
+
+	@classmethod
+	def _dict_from_json_file(cls, json_file: str):
+		with open(json_file, "r", encoding="utf-8") as reader:
+			text = reader.read()
+		return json.loads(text)
+
+	def __eq__(self, other):
+		return self.__dict__ == other.__dict__
+
+	def __repr__(self):
+		return "{} {}".format(self.__class__.__name__, self.to_json_string())
+
+	def to_diff_dict(self):
+		"""
+		Removes all attributes from config which correspond to the default
+		config attributes for better readability and serializes to a Python
+		dictionary.
+
+		Returns:
+			:obj:`Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+		"""
+		config_dict = self.to_dict()
+
+		# get the default config dict
+		default_config_dict = PretrainedConfig().to_dict()
+
+		serializable_config_dict = {}
+
+		# only serialize values that differ from the default config
+		for key, value in config_dict.items():
+			if key not in default_config_dict or value != default_config_dict[key]:
+				serializable_config_dict[key] = value
+
+		return serializable_config_dict
+
+	def to_dict(self):
+		"""
+		Serializes this instance to a Python dictionary.
+
+		Returns:
+			:obj:`Dict[str, any]`: Dictionary of all the attributes that make up this configuration instance,
+		"""
+		output = copy.deepcopy(self.__dict__)
+		if hasattr(self.__class__, "model_type"):
+			output["model_type"] = self.__class__.model_type
+		return output
+
+	def to_json_string(self, use_diff=True):
+		"""
+		Serializes this instance to a JSON string.
+
+		Args:
+			use_diff (:obj:`bool`):
+				If set to True, only the difference between the config instance and the default PretrainedConfig() is serialized to JSON string.
+
+		Returns:
+			:obj:`string`: String containing all the attributes that make up this configuration instance in JSON format.
+		"""
+		if use_diff is True:
+			config_dict = self.to_diff_dict()
+		else:
+			config_dict = self.to_dict()
+		return json.dumps(config_dict, indent=2, sort_keys=True) + "\n"
+
+	def to_json_file(self, json_file_path, use_diff=True):
+		"""
+		Save this instance to a json file.
+
+		Args:
+			json_file_path (:obj:`string`):
+				Path to the JSON file in which this configuration instance's parameters will be saved.
+			use_diff (:obj:`bool`):
+				If set to True, only the difference between the config instance and the default PretrainedConfig() is serialized to JSON file.
+		"""
+		with open(json_file_path, "w", encoding="utf-8") as writer:
+			writer.write(self.to_json_string(use_diff=use_diff))
+
+	def update(self, config_dict: Dict):
+		"""
+		Updates attributes of this class
+		with attributes from `config_dict`.
+
+		Args:
+			:obj:`Dict[str, any]`: Dictionary of attributes that shall be updated for this class.
+		"""
+		for key, value in config_dict.items():
+			setattr(self, key, value)

models/file_utils.py

@@ -0,0 +1,488 @@
+"""
+Utilities for working with the local dataset cache.
+This file is adapted from the AllenNLP library at https://github.com/allenai/allennlp
+Copyright by the AllenNLP authors.
+"""
+
+import fnmatch
+import json
+import logging
+import os
+import shutil
+import sys
+import tarfile
+import tempfile
+from contextlib import contextmanager
+from functools import partial, wraps
+from hashlib import sha256
+from typing import Optional
+from urllib.parse import urlparse
+from zipfile import ZipFile, is_zipfile
+
+import requests
+from filelock import FileLock
+from tqdm.auto import tqdm
+
+# from . import __version__
+__version__ = "2.8.0"
+
+logger = logging.getLogger(__name__)  # pylint: disable=invalid-name
+
+try:
+	USE_TF = os.environ.get("USE_TF", "AUTO").upper()
+	USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()
+	if USE_TORCH in ("1", "ON", "YES", "AUTO") and USE_TF not in ("1", "ON", "YES"):
+		import torch
+
+		_torch_available = True  # pylint: disable=invalid-name
+		logger.info("PyTorch version {} available.".format(torch.__version__))
+	else:
+		logger.info("Disabling PyTorch because USE_TF is set")
+		_torch_available = False
+except ImportError:
+	_torch_available = False  # pylint: disable=invalid-name
+
+try:
+	USE_TF = os.environ.get("USE_TF", "AUTO").upper()
+	USE_TORCH = os.environ.get("USE_TORCH", "AUTO").upper()
+
+	if USE_TF in ("1", "ON", "YES", "AUTO") and USE_TORCH not in ("1", "ON", "YES"):
+		import tensorflow as tf
+
+		assert hasattr(tf, "__version__") and int(tf.__version__[0]) >= 2
+		_tf_available = True  # pylint: disable=invalid-name
+		logger.info("TensorFlow version {} available.".format(tf.__version__))
+	else:
+		logger.info("Disabling Tensorflow because USE_TORCH is set")
+		_tf_available = False
+except (ImportError, AssertionError):
+	_tf_available = False  # pylint: disable=invalid-name
+
+try:
+	from torch.hub import _get_torch_home
+
+	torch_cache_home = _get_torch_home()
+except ImportError:
+	torch_cache_home = os.path.expanduser(
+		os.getenv("TORCH_HOME", os.path.join(os.getenv("XDG_CACHE_HOME", "~/.cache"), "torch"))
+	)
+default_cache_path = os.path.join(torch_cache_home, "transformers")
+
+try:
+	from pathlib import Path
+
+	PYTORCH_PRETRAINED_BERT_CACHE = Path(
+		os.getenv("PYTORCH_TRANSFORMERS_CACHE", os.getenv("PYTORCH_PRETRAINED_BERT_CACHE", default_cache_path))
+	)
+except (AttributeError, ImportError):
+	PYTORCH_PRETRAINED_BERT_CACHE = os.getenv(
+		"PYTORCH_TRANSFORMERS_CACHE", os.getenv("PYTORCH_PRETRAINED_BERT_CACHE", default_cache_path)
+	)
+
+PYTORCH_TRANSFORMERS_CACHE = PYTORCH_PRETRAINED_BERT_CACHE  # Kept for backward compatibility
+TRANSFORMERS_CACHE = PYTORCH_PRETRAINED_BERT_CACHE  # Kept for backward compatibility
+
+WEIGHTS_NAME = "pytorch_model.bin"
+TF2_WEIGHTS_NAME = "tf_model.h5"
+TF_WEIGHTS_NAME = "model.ckpt"
+CONFIG_NAME = "config.json"
+MODEL_CARD_NAME = "modelcard.json"
+
+
+MULTIPLE_CHOICE_DUMMY_INPUTS = [[[0], [1]], [[0], [1]]]
+DUMMY_INPUTS = [[7, 6, 0, 0, 1], [1, 2, 3, 0, 0], [0, 0, 0, 4, 5]]
+DUMMY_MASK = [[1, 1, 1, 1, 1], [1, 1, 1, 0, 0], [0, 0, 0, 1, 1]]
+
+S3_BUCKET_PREFIX = "https://s3.amazonaws.com/models.huggingface.co/bert"
+CLOUDFRONT_DISTRIB_PREFIX = "https://d2ws9o8vfrpkyk.cloudfront.net"
+
+
+def is_torch_available():
+	return _torch_available
+
+
+def is_tf_available():
+	return _tf_available
+
+
+def add_start_docstrings(*docstr):
+	def docstring_decorator(fn):
+		fn.__doc__ = "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")
+		return fn
+
+	return docstring_decorator
+
+
+def add_start_docstrings_to_callable(*docstr):
+	def docstring_decorator(fn):
+		class_name = ":class:`~transformers.{}`".format(fn.__qualname__.split(".")[0])
+		intro = "   The {} forward method, overrides the :func:`__call__` special method.".format(class_name)
+		note = r"""
+
+	.. note::
+		Although the recipe for forward pass needs to be defined within
+		this function, one should call the :class:`Module` instance afterwards
+		instead of this since the former takes care of running the
+		pre and post processing steps while the latter silently ignores them.
+		"""
+		fn.__doc__ = intro + note + "".join(docstr) + (fn.__doc__ if fn.__doc__ is not None else "")
+		return fn
+
+	return docstring_decorator
+
+
+def add_end_docstrings(*docstr):
+	def docstring_decorator(fn):
+		fn.__doc__ = fn.__doc__ + "".join(docstr)
+		return fn
+
+	return docstring_decorator
+
+
+def is_remote_url(url_or_filename):
+	parsed = urlparse(url_or_filename)
+	return parsed.scheme in ("http", "https")
+
+
+def hf_bucket_url(identifier, postfix=None, cdn=False) -> str:
+	endpoint = CLOUDFRONT_DISTRIB_PREFIX if cdn else S3_BUCKET_PREFIX
+	if postfix is None:
+		return "/".join((endpoint, identifier))
+	else:
+		return "/".join((endpoint, identifier, postfix))
+
+
+def url_to_filename(url, etag=None):
+	"""
+	Convert `url` into a hashed filename in a repeatable way.
+	If `etag` is specified, append its hash to the url's, delimited
+	by a period.
+	If the url ends with .h5 (Keras HDF5 weights) adds '.h5' to the name
+	so that TF 2.0 can identify it as a HDF5 file
+	(see https://github.com/tensorflow/tensorflow/blob/00fad90125b18b80fe054de1055770cfb8fe4ba3/tensorflow/python/keras/engine/network.py#L1380)
+	"""
+	url_bytes = url.encode("utf-8")
+	url_hash = sha256(url_bytes)
+	filename = url_hash.hexdigest()
+
+	if etag:
+		etag_bytes = etag.encode("utf-8")
+		etag_hash = sha256(etag_bytes)
+		filename += "." + etag_hash.hexdigest()
+
+	if url.endswith(".h5"):
+		filename += ".h5"
+
+	return filename
+
+
+def filename_to_url(filename, cache_dir=None):
+	"""
+	Return the url and etag (which may be ``None``) stored for `filename`.
+	Raise ``EnvironmentError`` if `filename` or its stored metadata do not exist.
+	"""
+	if cache_dir is None:
+		cache_dir = TRANSFORMERS_CACHE
+	if isinstance(cache_dir, Path):
+		cache_dir = str(cache_dir)
+
+	cache_path = os.path.join(cache_dir, filename)
+	if not os.path.exists(cache_path):
+		raise EnvironmentError("file {} not found".format(cache_path))
+
+	meta_path = cache_path + ".json"
+	if not os.path.exists(meta_path):
+		raise EnvironmentError("file {} not found".format(meta_path))
+
+	with open(meta_path, encoding="utf-8") as meta_file:
+		metadata = json.load(meta_file)
+	url = metadata["url"]
+	etag = metadata["etag"]
+
+	return url, etag
+
+
+def cached_path(
+	url_or_filename,
+	cache_dir=None,
+	force_download=False,
+	proxies=None,
+	resume_download=False,
+	user_agent=None,
+	extract_compressed_file=False,
+	force_extract=False,
+	local_files_only=False,
+) -> Optional[str]:
+	"""
+	Given something that might be a URL (or might be a local path),
+	determine which. If it's a URL, download the file and cache it, and
+	return the path to the cached file. If it's already a local path,
+	make sure the file exists and then return the path.
+	Args:
+		cache_dir: specify a cache directory to save the file to (overwrite the default cache dir).
+		force_download: if True, re-dowload the file even if it's already cached in the cache dir.
+		resume_download: if True, resume the download if incompletly recieved file is found.
+		user_agent: Optional string or dict that will be appended to the user-agent on remote requests.
+		extract_compressed_file: if True and the path point to a zip or tar file, extract the compressed
+			file in a folder along the archive.
+		force_extract: if True when extract_compressed_file is True and the archive was already extracted,
+			re-extract the archive and overide the folder where it was extracted.
+
+	Return:
+		None in case of non-recoverable file (non-existent or inaccessible url + no cache on disk).
+		Local path (string) otherwise
+	"""
+	if cache_dir is None:
+		cache_dir = TRANSFORMERS_CACHE
+	if isinstance(url_or_filename, Path):
+		url_or_filename = str(url_or_filename)
+	if isinstance(cache_dir, Path):
+		cache_dir = str(cache_dir)
+
+	if is_remote_url(url_or_filename):
+		# URL, so get it from the cache (downloading if necessary)
+		output_path = get_from_cache(
+			url_or_filename,
+			cache_dir=cache_dir,
+			force_download=force_download,
+			proxies=proxies,
+			resume_download=resume_download,
+			user_agent=user_agent,
+			local_files_only=local_files_only,
+		)
+	elif os.path.exists(url_or_filename):
+		# File, and it exists.
+		output_path = url_or_filename
+	elif urlparse(url_or_filename).scheme == "":
+		# File, but it doesn't exist.
+		raise EnvironmentError("file {} not found".format(url_or_filename))
+	else:
+		# Something unknown
+		raise ValueError("unable to parse {} as a URL or as a local path".format(url_or_filename))
+
+	if extract_compressed_file:
+		if not is_zipfile(output_path) and not tarfile.is_tarfile(output_path):
+			return output_path
+
+		# Path where we extract compressed archives
+		# We avoid '.' in dir name and add "-extracted" at the end: "./model.zip" => "./model-zip-extracted/"
+		output_dir, output_file = os.path.split(output_path)
+		output_extract_dir_name = output_file.replace(".", "-") + "-extracted"
+		output_path_extracted = os.path.join(output_dir, output_extract_dir_name)
+
+		if os.path.isdir(output_path_extracted) and os.listdir(output_path_extracted) and not force_extract:
+			return output_path_extracted
+
+		# Prevent parallel extractions
+		lock_path = output_path + ".lock"
+		with FileLock(lock_path):
+			shutil.rmtree(output_path_extracted, ignore_errors=True)
+			os.makedirs(output_path_extracted)
+			if is_zipfile(output_path):
+				with ZipFile(output_path, "r") as zip_file:
+					zip_file.extractall(output_path_extracted)
+					zip_file.close()
+			elif tarfile.is_tarfile(output_path):
+				tar_file = tarfile.open(output_path)
+				tar_file.extractall(output_path_extracted)
+				tar_file.close()
+			else:
+				raise EnvironmentError("Archive format of {} could not be identified".format(output_path))
+
+		return output_path_extracted
+
+	return output_path
+
+
+def http_get(url, temp_file, proxies=None, resume_size=0, user_agent=None):
+	ua = "transformers/{}; python/{}".format(__version__, sys.version.split()[0])
+	if is_torch_available():
+		ua += "; torch/{}".format(torch.__version__)
+	if is_tf_available():
+		ua += "; tensorflow/{}".format(tf.__version__)
+	if isinstance(user_agent, dict):
+		ua += "; " + "; ".join("{}/{}".format(k, v) for k, v in user_agent.items())
+	elif isinstance(user_agent, str):
+		ua += "; " + user_agent
+	headers = {"user-agent": ua}
+	if resume_size > 0:
+		headers["Range"] = "bytes=%d-" % (resume_size,)
+	response = requests.get(url, stream=True, proxies=proxies, headers=headers)
+	if response.status_code == 416:  # Range not satisfiable
+		return
+	content_length = response.headers.get("Content-Length")
+	total = resume_size + int(content_length) if content_length is not None else None
+	progress = tqdm(
+		unit="B",
+		unit_scale=True,
+		total=total,
+		initial=resume_size,
+		desc="Downloading",
+		disable=bool(logger.getEffectiveLevel() == logging.NOTSET),
+	)
+	for chunk in response.iter_content(chunk_size=1024):
+		if chunk:  # filter out keep-alive new chunks
+			progress.update(len(chunk))
+			temp_file.write(chunk)
+	progress.close()
+
+
+def get_from_cache(
+	url,
+	cache_dir=None,
+	force_download=False,
+	proxies=None,
+	etag_timeout=10,
+	resume_download=False,
+	user_agent=None,
+	local_files_only=False,
+) -> Optional[str]:
+	"""
+	Given a URL, look for the corresponding file in the local cache.
+	If it's not there, download it. Then return the path to the cached file.
+
+	Return:
+		None in case of non-recoverable file (non-existent or inaccessible url + no cache on disk).
+		Local path (string) otherwise
+	"""
+	if cache_dir is None:
+		cache_dir = TRANSFORMERS_CACHE
+	if isinstance(cache_dir, Path):
+		cache_dir = str(cache_dir)
+
+	os.makedirs(cache_dir, exist_ok=True)
+
+	etag = None
+	if not local_files_only:
+		try:
+			response = requests.head(url, allow_redirects=True, proxies=proxies, timeout=etag_timeout)
+			if response.status_code == 200:
+				etag = response.headers.get("ETag")
+		except (EnvironmentError, requests.exceptions.Timeout):
+			# etag is already None
+			pass
+
+	filename = url_to_filename(url, etag)
+
+	# get cache path to put the file
+	cache_path = os.path.join(cache_dir, filename)
+
+	# etag is None = we don't have a connection, or url doesn't exist, or is otherwise inaccessible.
+	# try to get the last downloaded one
+	if etag is None:
+		if os.path.exists(cache_path):
+			return cache_path
+		else:
+			matching_files = [
+				file
+				for file in fnmatch.filter(os.listdir(cache_dir), filename + ".*")
+				if not file.endswith(".json") and not file.endswith(".lock")
+			]
+			if len(matching_files) > 0:
+				return os.path.join(cache_dir, matching_files[-1])
+			else:
+				# If files cannot be found and local_files_only=True,
+				# the models might've been found if local_files_only=False
+				# Notify the user about that
+				if local_files_only:
+					raise ValueError(
+						"Cannot find the requested files in the cached path and outgoing traffic has been"
+						" disabled. To enable model look-ups and downloads online, set 'local_files_only'"
+						" to False."
+					)
+				return None
+
+	# From now on, etag is not None.
+	if os.path.exists(cache_path) and not force_download:
+		return cache_path
+
+	# Prevent parallel downloads of the same file with a lock.
+	lock_path = cache_path + ".lock"
+	with FileLock(lock_path):
+
+		# If the download just completed while the lock was activated.
+		if os.path.exists(cache_path) and not force_download:
+			# Even if returning early like here, the lock will be released.
+			return cache_path
+
+		if resume_download:
+			incomplete_path = cache_path + ".incomplete"
+
+			@contextmanager
+			def _resumable_file_manager():
+				with open(incomplete_path, "a+b") as f:
+					yield f
+
+			temp_file_manager = _resumable_file_manager
+			if os.path.exists(incomplete_path):
+				resume_size = os.stat(incomplete_path).st_size
+			else:
+				resume_size = 0
+		else:
+			temp_file_manager = partial(tempfile.NamedTemporaryFile, dir=cache_dir, delete=False)
+			resume_size = 0
+
+		# Download to temporary file, then copy to cache dir once finished.
+		# Otherwise you get corrupt cache entries if the download gets interrupted.
+		with temp_file_manager() as temp_file:
+			logger.info("%s not found in cache or force_download set to True, downloading to %s", url, temp_file.name)
+
+			http_get(url, temp_file, proxies=proxies, resume_size=resume_size, user_agent=user_agent)
+
+		logger.info("storing %s in cache at %s", url, cache_path)
+		os.replace(temp_file.name, cache_path)
+
+		logger.info("creating metadata file for %s", cache_path)
+		meta = {"url": url, "etag": etag}
+		meta_path = cache_path + ".json"
+		with open(meta_path, "w") as meta_file:
+			json.dump(meta, meta_file)
+
+	return cache_path
+
+
+class cached_property(property):
+	"""
+	Descriptor that mimics @property but caches output in member variable.
+
+	From tensorflow_datasets
+
+	Built-in in functools from Python 3.8.
+	"""
+
+	def __get__(self, obj, objtype=None):
+		# See docs.python.org/3/howto/descriptor.html#properties
+		if obj is None:
+			return self
+		if self.fget is None:
+			raise AttributeError("unreadable attribute")
+		attr = "__cached_" + self.fget.__name__
+		cached = getattr(obj, attr, None)
+		if cached is None:
+			cached = self.fget(obj)
+			setattr(obj, attr, cached)
+		return cached
+
+
+def torch_required(func):
+	# Chose a different decorator name than in tests so it's clear they are not the same.
+	@wraps(func)
+	def wrapper(*args, **kwargs):
+		if is_torch_available():
+			return func(*args, **kwargs)
+		else:
+			raise ImportError(f"Method `{func.__name__}` requires PyTorch.")
+
+	return wrapper
+
+
+def tf_required(func):
+	# Chose a different decorator name than in tests so it's clear they are not the same.
+	@wraps(func)
+	def wrapper(*args, **kwargs):
+		if is_tf_available():
+			return func(*args, **kwargs)
+		else:
+			raise ImportError(f"Method `{func.__name__}` requires TF.")
+
+	return wrapper

models/modeling_bert.py

@@ -0,0 +1,1476 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch BERT model. """
+
+
+import logging
+import math
+import os
+
+import torch
+from torch import nn
+from torch.nn import CrossEntropyLoss, MSELoss
+
+from .activations import gelu, gelu_new, swish
+from .configuration_bert import BertConfig
+from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
+from .modeling_utils import PreTrainedModel, prune_linear_layer
+
+
+logger = logging.getLogger(__name__)
+
+BERT_PRETRAINED_MODEL_ARCHIVE_MAP = {
+	"bert-base-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-pytorch_model.bin",
+	"bert-large-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-pytorch_model.bin",
+	"bert-base-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-pytorch_model.bin",
+	"bert-large-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-pytorch_model.bin",
+	"bert-base-multilingual-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-pytorch_model.bin",
+	"bert-base-multilingual-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-pytorch_model.bin",
+	"bert-base-chinese": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-pytorch_model.bin",
+	"bert-base-german-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-pytorch_model.bin",
+	"bert-large-uncased-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-pytorch_model.bin",
+	"bert-large-cased-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-pytorch_model.bin",
+	"bert-large-uncased-whole-word-masking-finetuned-squad": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-pytorch_model.bin",
+	"bert-large-cased-whole-word-masking-finetuned-squad": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-pytorch_model.bin",
+	"bert-base-cased-finetuned-mrpc": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-pytorch_model.bin",
+	"bert-base-german-dbmdz-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-dbmdz-cased-pytorch_model.bin",
+	"bert-base-german-dbmdz-uncased": "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-dbmdz-uncased-pytorch_model.bin",
+	"bert-base-japanese": "https://s3.amazonaws.com/models.huggingface.co/bert/cl-tohoku/bert-base-japanese-pytorch_model.bin",
+	"bert-base-japanese-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/cl-tohoku/bert-base-japanese-whole-word-masking-pytorch_model.bin",
+	"bert-base-japanese-char": "https://s3.amazonaws.com/models.huggingface.co/bert/cl-tohoku/bert-base-japanese-char-pytorch_model.bin",
+	"bert-base-japanese-char-whole-word-masking": "https://s3.amazonaws.com/models.huggingface.co/bert/cl-tohoku/bert-base-japanese-char-whole-word-masking-pytorch_model.bin",
+	"bert-base-finnish-cased-v1": "https://s3.amazonaws.com/models.huggingface.co/bert/TurkuNLP/bert-base-finnish-cased-v1/pytorch_model.bin",
+	"bert-base-finnish-uncased-v1": "https://s3.amazonaws.com/models.huggingface.co/bert/TurkuNLP/bert-base-finnish-uncased-v1/pytorch_model.bin",
+	"bert-base-dutch-cased": "https://s3.amazonaws.com/models.huggingface.co/bert/wietsedv/bert-base-dutch-cased/pytorch_model.bin",
+}
+
+
+def load_tf_weights_in_bert(model, config, tf_checkpoint_path):
+	""" Load tf checkpoints in a pytorch model.
+	"""
+	try:
+		import re
+		import numpy as np
+		import tensorflow as tf
+	except ImportError:
+		logger.error(
+			"Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
+			"https://www.tensorflow.org/install/ for installation instructions."
+		)
+		raise
+	tf_path = os.path.abspath(tf_checkpoint_path)
+	logger.info("Converting TensorFlow checkpoint from {}".format(tf_path))
+	# Load weights from TF model
+	init_vars = tf.train.list_variables(tf_path)
+	names = []
+	arrays = []
+	for name, shape in init_vars:
+		logger.info("Loading TF weight {} with shape {}".format(name, shape))
+		array = tf.train.load_variable(tf_path, name)
+		names.append(name)
+		arrays.append(array)
+
+	for name, array in zip(names, arrays):
+		name = name.split("/")
+		# adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
+		# which are not required for using pretrained model
+		if any(
+			n in ["adam_v", "adam_m", "AdamWeightDecayOptimizer", "AdamWeightDecayOptimizer_1", "global_step"]
+			for n in name
+		):
+			logger.info("Skipping {}".format("/".join(name)))
+			continue
+		pointer = model
+		for m_name in name:
+			if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
+				scope_names = re.split(r"_(\d+)", m_name)
+			else:
+				scope_names = [m_name]
+			if scope_names[0] == "kernel" or scope_names[0] == "gamma":
+				pointer = getattr(pointer, "weight")
+			elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
+				pointer = getattr(pointer, "bias")
+			elif scope_names[0] == "output_weights":
+				pointer = getattr(pointer, "weight")
+			elif scope_names[0] == "squad":
+				pointer = getattr(pointer, "classifier")
+			else:
+				try:
+					pointer = getattr(pointer, scope_names[0])
+				except AttributeError:
+					logger.info("Skipping {}".format("/".join(name)))
+					continue
+			if len(scope_names) >= 2:
+				num = int(scope_names[1])
+				pointer = pointer[num]
+		if m_name[-11:] == "_embeddings":
+			pointer = getattr(pointer, "weight")
+		elif m_name == "kernel":
+			array = np.transpose(array)
+		try:
+			assert pointer.shape == array.shape
+		except AssertionError as e:
+			e.args += (pointer.shape, array.shape)
+			raise
+		logger.info("Initialize PyTorch weight {}".format(name))
+		pointer.data = torch.from_numpy(array)
+	return model
+
+
+def mish(x):
+	return x * torch.tanh(nn.functional.softplus(x))
+
+
+ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish, "gelu_new": gelu_new, "mish": mish}
+
+
+BertLayerNorm = torch.nn.LayerNorm
+
+
+class BertEmbeddings(nn.Module):
+	"""Construct the embeddings from word, position and token_type embeddings.
+	"""
+
+	def __init__(self, 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 = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+		self.dropout = nn.Dropout(config.hidden_dropout_prob)
+
+	def forward(self, input_ids=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 = torch.arange(seq_length, dtype=torch.long, device=device)
+			position_ids = position_ids.unsqueeze(0).expand(input_shape)
+		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)
+		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
+
+
+class BertSelfAttention(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.output_attentions = config.output_attentions
+
+		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,
+	):
+		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 BertModel 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 self.output_attentions else (context_layer,)
+		return outputs
+
+
+class BertSelfOutput(nn.Module):
+	def __init__(self, config):
+		super().__init__()
+		self.dense = nn.Linear(config.hidden_size, config.hidden_size)
+		self.LayerNorm = BertLayerNorm(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
+
+
+class BertAttention(nn.Module):
+	def __init__(self, config):
+		super().__init__()
+		self.self = BertSelfAttention(config)
+		self.output = BertSelfOutput(config)
+		self.pruned_heads = set()
+
+	def prune_heads(self, heads):
+		if len(heads) == 0:
+			return
+		mask = torch.ones(self.self.num_attention_heads, self.self.attention_head_size)
+		heads = set(heads) - self.pruned_heads  # Convert to set and remove already pruned heads
+		for head in heads:
+			# Compute how many pruned heads are before the head and move the index accordingly
+			head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
+			mask[head] = 0
+		mask = mask.view(-1).contiguous().eq(1)
+		index = torch.arange(len(mask))[mask].long()
+
+		# 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,
+	):
+		self_outputs = self.self(
+			hidden_states, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask
+		)
+		attention_output = self.output(self_outputs[0], hidden_states)
+		outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
+		return outputs
+
+
+class BertIntermediate(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
+
+
+class BertOutput(nn.Module):
+	def __init__(self, config):
+		super().__init__()
+		self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
+		self.LayerNorm = BertLayerNorm(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
+
+
+class BertLayer(nn.Module):
+	def __init__(self, config):
+		super().__init__()
+		self.attention = BertAttention(config)
+		self.is_decoder = config.is_decoder
+		if self.is_decoder:
+			self.crossattention = BertAttention(config)
+		self.intermediate = BertIntermediate(config)
+		self.output = BertOutput(config)
+
+	def forward(
+		self,
+		hidden_states,
+		attention_mask=None,
+		head_mask=None,
+		encoder_hidden_states=None,
+		encoder_attention_mask=None,
+	):
+		self_attention_outputs = self.attention(hidden_states, attention_mask, head_mask)
+		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:
+			cross_attention_outputs = self.crossattention(
+				attention_output, attention_mask, head_mask, encoder_hidden_states, encoder_attention_mask
+			)
+			attention_output = cross_attention_outputs[0]
+			outputs = outputs + cross_attention_outputs[1:]  # add cross attentions if we output attention weights
+
+		intermediate_output = self.intermediate(attention_output)
+		layer_output = self.output(intermediate_output, attention_output)
+		outputs = (layer_output,) + outputs
+		return outputs
+
+
+class BertEncoder(nn.Module):
+	def __init__(self, config):
+		super().__init__()
+		self.output_attentions = config.output_attentions
+		self.output_hidden_states = config.output_hidden_states
+		self.layer = nn.ModuleList([BertLayer(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,
+	):
+		all_hidden_states = ()
+		all_attentions = ()
+		for i, layer_module in enumerate(self.layer):
+			if self.output_hidden_states:
+				all_hidden_states = all_hidden_states + (hidden_states,)
+
+			layer_outputs = layer_module(
+				hidden_states, attention_mask, head_mask[i], encoder_hidden_states, encoder_attention_mask
+			)
+			hidden_states = layer_outputs[0]
+
+			if self.output_attentions:
+				all_attentions = all_attentions + (layer_outputs[1],)
+
+		# Add last layer
+		if self.output_hidden_states:
+			all_hidden_states = all_hidden_states + (hidden_states,)
+
+		outputs = (hidden_states,)
+		if self.output_hidden_states:
+			outputs = outputs + (all_hidden_states,)
+		if self.output_attentions:
+			outputs = outputs + (all_attentions,)
+		return outputs  # last-layer hidden state, (all hidden states), (all attentions)
+
+
+class BertPooler(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 BertPredictionHeadTransform(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 = BertLayerNorm(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
+
+
+class BertLMPredictionHead(nn.Module):
+	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 BertOnlyMLMHead(nn.Module):
+	def __init__(self, config):
+		super().__init__()
+		self.predictions = BertLMPredictionHead(config)
+
+	def forward(self, sequence_output):
+		prediction_scores = self.predictions(sequence_output)
+		return prediction_scores
+
+
+class BertOnlyNSPHead(nn.Module):
+	def __init__(self, config):
+		super().__init__()
+		self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+	def forward(self, pooled_output):
+		seq_relationship_score = self.seq_relationship(pooled_output)
+		return seq_relationship_score
+
+
+class BertPreTrainingHeads(nn.Module):
+	def __init__(self, config):
+		super().__init__()
+		self.predictions = BertLMPredictionHead(config)
+		self.seq_relationship = nn.Linear(config.hidden_size, 2)
+
+	def forward(self, sequence_output, pooled_output):
+		prediction_scores = self.predictions(sequence_output)
+		seq_relationship_score = self.seq_relationship(pooled_output)
+		return prediction_scores, seq_relationship_score
+
+
+class BertPreTrainedModel(PreTrainedModel):
+	""" An abstract class to handle weights initialization and
+		a simple interface for downloading and loading pretrained models.
+	"""
+
+	config_class = BertConfig
+	pretrained_model_archive_map = BERT_PRETRAINED_MODEL_ARCHIVE_MAP
+	load_tf_weights = load_tf_weights_in_bert
+	base_model_prefix = "bert"
+
+	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, BertLayerNorm):
+			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_()
+
+
+BERT_START_DOCSTRING = r"""
+	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.BertConfig`): 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.
+"""
+
+BERT_INPUTS_DOCSTRING = r"""
+	Args:
+		input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`):
+			Indices of input sequence tokens in the vocabulary.
+
+			Indices can be obtained using :class:`transformers.BertTokenizer`.
+			See :func:`transformers.PreTrainedTokenizer.encode` and
+			:func:`transformers.PreTrainedTokenizer.encode_plus` for details.
+
+			`What are input IDs? <../glossary.html#input-ids>`__
+		attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Mask to avoid performing attention on padding token indices.
+			Mask values selected in ``[0, 1]``:
+			``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
+
+			`What are attention masks? <../glossary.html#attention-mask>`__
+		token_type_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Segment token indices to indicate first and second portions of the inputs.
+			Indices are selected in ``[0, 1]``: ``0`` corresponds to a `sentence A` token, ``1``
+			corresponds to a `sentence B` token
+
+			`What are token type IDs? <../glossary.html#token-type-ids>`_
+		position_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Indices of positions of each input sequence tokens in the position embeddings.
+			Selected in the range ``[0, config.max_position_embeddings - 1]``.
+
+			`What are position IDs? <../glossary.html#position-ids>`_
+		head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`):
+			Mask to nullify selected heads of the self-attention modules.
+			Mask values selected in ``[0, 1]``:
+			:obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
+		inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
+			Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
+			This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+			than the model's internal embedding lookup matrix.
+		encoder_hidden_states  (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
+			Sequence of hidden-states at the output of the last layer of the encoder. Used in the cross-attention
+			if the model is configured as a decoder.
+		encoder_attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Mask to avoid performing attention on the padding token indices of the encoder input. This mask
+			is used in the cross-attention if the model is configured as a decoder.
+			Mask values selected in ``[0, 1]``:
+			``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
+"""
+
+
+@add_start_docstrings(
+	"The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
+	BERT_START_DOCSTRING,
+)
+class BertModel(BertPreTrainedModel):
+	"""
+
+	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`; an
+	:obj:`encoder_hidden_states` is expected as an input to the forward pass.
+
+	.. _`Attention is all you need`:
+		https://arxiv.org/abs/1706.03762
+
+	"""
+
+	def __init__(self, config):
+		super().__init__(config)
+		self.config = config
+
+		self.embeddings = BertEmbeddings(config)
+		self.encoder = BertEncoder(config)
+		self.pooler = BertPooler(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(BERT_INPUTS_DOCSTRING)
+	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,
+	):
+		r"""
+	Return:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+		last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
+			Sequence of hidden-states at the output of the last layer of the model.
+		pooler_output (:obj:`torch.FloatTensor`: of shape :obj:`(batch_size, hidden_size)`):
+			Last layer hidden-state of the first token of the sequence (classification token)
+			further processed by a Linear layer and a Tanh activation function. The Linear
+			layer weights are trained from the next sentence prediction (classification)
+			objective during pre-training.
+
+			This output is usually *not* a good summary
+			of the semantic content of the input, you're often better with averaging or pooling
+			the sequence of hidden-states for the whole input sequence.
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+
+	Examples::
+
+		from transformers import BertModel, BertTokenizer
+		import torch
+
+		tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+		model = BertModel.from_pretrained('bert-base-uncased')
+
+		input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids)
+
+		last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple
+
+		"""
+
+		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, self.device
+		)
+
+		# If a 2D ou 3D attention mask is provided for the cross-attention
+		# we need to make broadcastabe 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,
+		)
+		sequence_output = encoder_outputs[0]
+		pooled_output = self.pooler(sequence_output)
+
+		outputs = (sequence_output, pooled_output,) + encoder_outputs[
+			1:
+		]  # add hidden_states and attentions if they are here
+		return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+	"""Bert Model with two heads on top as done during the pre-training: a `masked language modeling` head and
+	a `next sentence prediction (classification)` head. """,
+	BERT_START_DOCSTRING,
+)
+class BertForPreTraining(BertPreTrainedModel):
+	def __init__(self, config):
+		super().__init__(config)
+
+		self.bert = BertModel(config)
+		self.cls = BertPreTrainingHeads(config)
+
+		self.init_weights()
+
+	def get_output_embeddings(self):
+		return self.cls.predictions.decoder
+
+	@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		masked_lm_labels=None,
+		next_sentence_label=None,
+	):
+		r"""
+		masked_lm_labels (``torch.LongTensor`` of shape ``(batch_size, sequence_length)``, `optional`, defaults to :obj:`None`):
+			Labels for computing the masked language modeling loss.
+			Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
+			Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
+			in ``[0, ..., config.vocab_size]``
+		next_sentence_label (``torch.LongTensor`` of shape ``(batch_size,)``, `optional`, defaults to :obj:`None`):
+			Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see :obj:`input_ids` docstring)
+			Indices should be in ``[0, 1]``.
+			``0`` indicates sequence B is a continuation of sequence A,
+			``1`` indicates sequence B is a random sequence.
+
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+		loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+			Total loss as the sum of the masked language modeling loss and the next sequence prediction (classification) loss.
+		prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
+			Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+		seq_relationship_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, 2)`):
+			Prediction scores of the next sequence prediction (classification) head (scores of True/False
+			continuation before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when :obj:`config.output_hidden_states=True`):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+
+
+	Examples::
+
+		from transformers import BertTokenizer, BertForPreTraining
+		import torch
+
+		tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+		model = BertForPreTraining.from_pretrained('bert-base-uncased')
+
+		input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids)
+
+		prediction_scores, seq_relationship_scores = outputs[:2]
+
+		"""
+
+		outputs = self.bert(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		sequence_output, pooled_output = outputs[:2]
+		prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)
+
+		outputs = (prediction_scores, seq_relationship_score,) + outputs[
+			2:
+		]  # add hidden states and attention if they are here
+
+		if masked_lm_labels is not None and next_sentence_label is not None:
+			loss_fct = CrossEntropyLoss()
+			masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+			next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
+			total_loss = masked_lm_loss + next_sentence_loss
+			outputs = (total_loss,) + outputs
+
+		return outputs  # (loss), prediction_scores, seq_relationship_score, (hidden_states), (attentions)
+
+
+@add_start_docstrings("""Bert Model with a `language modeling` head on top. """, BERT_START_DOCSTRING)
+class BertForMaskedLM(BertPreTrainedModel):
+	def __init__(self, config):
+		super().__init__(config)
+
+		self.bert = BertModel(config)
+		self.cls = BertOnlyMLMHead(config)
+
+		self.init_weights()
+
+	def get_output_embeddings(self):
+		return self.cls.predictions.decoder
+
+	@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		masked_lm_labels=None,
+		encoder_hidden_states=None,
+		encoder_attention_mask=None,
+		lm_labels=None,
+	):
+		r"""
+		masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the masked language modeling loss.
+			Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
+			Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
+			in ``[0, ..., config.vocab_size]``
+		lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the left-to-right language modeling loss (next word prediction).
+			Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
+			Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
+			in ``[0, ..., config.vocab_size]``
+
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+		masked_lm_loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+			Masked language modeling loss.
+		ltr_lm_loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`lm_labels` is provided):
+				Next token prediction loss.
+		prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
+			Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+
+		Examples::
+
+			from transformers import BertTokenizer, BertForMaskedLM
+			import torch
+
+			tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+			model = BertForMaskedLM.from_pretrained('bert-base-uncased')
+
+			input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+			outputs = model(input_ids, masked_lm_labels=input_ids)
+
+			loss, prediction_scores = outputs[:2]
+
+		"""
+
+		outputs = self.bert(
+			input_ids,
+			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,
+		)
+
+		sequence_output = outputs[0]
+		prediction_scores = self.cls(sequence_output)
+
+		outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention if they are here
+
+		# Although this may seem awkward, BertForMaskedLM supports two scenarios:
+		# 1. If a tensor that contains the indices of masked labels is provided,
+		#    the cross-entropy is the MLM cross-entropy that measures the likelihood
+		#    of predictions for masked words.
+		# 2. If `lm_labels` is provided we are in a causal scenario where we
+		#    try to predict the next token for each input in the decoder.
+		if masked_lm_labels is not None:
+			loss_fct = CrossEntropyLoss()  # -100 index = padding token
+			masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+			outputs = (masked_lm_loss,) + outputs
+
+		if lm_labels is not None:
+			# we are doing next-token prediction; shift prediction scores and input ids by one
+			prediction_scores = prediction_scores[:, :-1, :].contiguous()
+			lm_labels = lm_labels[:, 1:].contiguous()
+			loss_fct = CrossEntropyLoss()
+			ltr_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), lm_labels.view(-1))
+			outputs = (ltr_lm_loss,) + outputs
+
+		return outputs  # (ltr_lm_loss), (masked_lm_loss), prediction_scores, (hidden_states), (attentions)
+
+	def prepare_inputs_for_generation(self, input_ids, attention_mask=None, **model_kwargs):
+		input_shape = input_ids.shape
+		effective_batch_size = input_shape[0]
+
+		# if model is used as a decoder in encoder-decoder model, the decoder attention mask is created on the fly
+		if attention_mask is None:
+			attention_mask = input_ids.new_ones(input_shape)
+
+		# if model is does not use a causal mask then add a dummy token
+		if self.config.is_decoder is False:
+			assert self.config.pad_token_id is not None, "The PAD token should be defined for generation"
+			attention_mask = torch.cat(
+				[attention_mask, attention_mask.new_zeros((attention_mask.shape[0], 1))], dim=-1
+			)
+
+			dummy_token = torch.full(
+				(effective_batch_size, 1), self.config.pad_token_id, dtype=torch.long, device=input_ids.device
+			)
+			input_ids = torch.cat([input_ids, dummy_token], dim=1)
+
+		return {"input_ids": input_ids, "attention_mask": attention_mask}
+
+
+@add_start_docstrings(
+	"""Bert Model with a `next sentence prediction (classification)` head on top. """, BERT_START_DOCSTRING,
+)
+class BertForNextSentencePrediction(BertPreTrainedModel):
+	def __init__(self, config):
+		super().__init__(config)
+
+		self.bert = BertModel(config)
+		self.cls = BertOnlyNSPHead(config)
+
+		self.init_weights()
+
+	@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		next_sentence_label=None,
+	):
+		r"""
+		next_sentence_label (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the next sequence prediction (classification) loss. Input should be a sequence pair (see ``input_ids`` docstring)
+			Indices should be in ``[0, 1]``.
+			``0`` indicates sequence B is a continuation of sequence A,
+			``1`` indicates sequence B is a random sequence.
+
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`next_sentence_label` is provided):
+			Next sequence prediction (classification) loss.
+		seq_relationship_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, 2)`):
+			Prediction scores of the next sequence prediction (classification) head (scores of True/False continuation before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+
+	Examples::
+
+		from transformers import BertTokenizer, BertForNextSentencePrediction
+		import torch
+
+		tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+		model = BertForNextSentencePrediction.from_pretrained('bert-base-uncased')
+
+		input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids)
+
+		seq_relationship_scores = outputs[0]
+
+		"""
+
+		outputs = self.bert(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		pooled_output = outputs[1]
+
+		seq_relationship_score = self.cls(pooled_output)
+
+		outputs = (seq_relationship_score,) + outputs[2:]  # add hidden states and attention if they are here
+		if next_sentence_label is not None:
+			loss_fct = CrossEntropyLoss()
+			next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
+			outputs = (next_sentence_loss,) + outputs
+
+		return outputs  # (next_sentence_loss), seq_relationship_score, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+	"""Bert Model transformer with a sequence classification/regression head on top (a linear layer on top of
+	the pooled output) e.g. for GLUE tasks. """,
+	BERT_START_DOCSTRING,
+)
+class BertForSequenceClassification(BertPreTrainedModel):
+	def __init__(self, config):
+		super().__init__(config)
+		self.num_labels = config.num_labels
+
+		self.bert = BertModel(config)
+		self.dropout = nn.Dropout(config.hidden_dropout_prob)
+		self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
+
+		self.init_weights()
+
+	@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		labels=None,
+	):
+		r"""
+		labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the sequence classification/regression loss.
+			Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
+			If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+			If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
+			Classification (or regression if config.num_labels==1) loss.
+		logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
+			Classification (or regression if config.num_labels==1) scores (before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+
+	Examples::
+
+		from transformers import BertTokenizer, BertForSequenceClassification
+		import torch
+
+		tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+		model = BertForSequenceClassification.from_pretrained('bert-base-uncased')
+
+		input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+		labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids, labels=labels)
+
+		loss, logits = outputs[:2]
+
+		"""
+
+		outputs = self.bert(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		pooled_output = outputs[1]
+
+		pooled_output = self.dropout(pooled_output)
+		logits = self.classifier(pooled_output)
+
+		outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
+
+		if labels is not None:
+			if self.num_labels == 1:
+				#  We are doing regression
+				loss_fct = MSELoss()
+				loss = loss_fct(logits.view(-1), labels.view(-1))
+			else:
+				loss_fct = CrossEntropyLoss()
+				loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+			outputs = (loss,) + outputs
+
+		return outputs  # (loss), logits, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+	"""Bert Model with a multiple choice classification head on top (a linear layer on top of
+	the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
+	BERT_START_DOCSTRING,
+)
+class BertForMultipleChoice(BertPreTrainedModel):
+	def __init__(self, config):
+		super().__init__(config)
+
+		self.bert = BertModel(config)
+		self.dropout = nn.Dropout(config.hidden_dropout_prob)
+		self.classifier = nn.Linear(config.hidden_size, 1)
+
+		self.init_weights()
+
+	@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		labels=None,
+	):
+		r"""
+		labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the multiple choice classification loss.
+			Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
+			of the input tensors. (see `input_ids` above)
+
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor` of shape `(1,)`, `optional`, returned when :obj:`labels` is provided):
+			Classification loss.
+		classification_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices)`):
+			`num_choices` is the second dimension of the input tensors. (see `input_ids` above).
+
+			Classification scores (before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+
+	Examples::
+
+		from transformers import BertTokenizer, BertForMultipleChoice
+		import torch
+
+		tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+		model = BertForMultipleChoice.from_pretrained('bert-base-uncased')
+		choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]
+
+		input_ids = torch.tensor([tokenizer.encode(s, add_special_tokens=True) for s in choices]).unsqueeze(0)  # Batch size 1, 2 choices
+		labels = torch.tensor(1).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids, labels=labels)
+
+		loss, classification_scores = outputs[:2]
+
+		"""
+		num_choices = input_ids.shape[1]
+
+		input_ids = input_ids.view(-1, input_ids.size(-1))
+		attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+		token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+		position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+
+		outputs = self.bert(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		pooled_output = outputs[1]
+
+		pooled_output = self.dropout(pooled_output)
+		logits = self.classifier(pooled_output)
+		reshaped_logits = logits.view(-1, num_choices)
+
+		outputs = (reshaped_logits,) + outputs[2:]  # add hidden states and attention if they are here
+
+		if labels is not None:
+			loss_fct = CrossEntropyLoss()
+			loss = loss_fct(reshaped_logits, labels)
+			outputs = (loss,) + outputs
+
+		return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+	"""Bert 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. """,
+	BERT_START_DOCSTRING,
+)
+class BertForTokenClassification(BertPreTrainedModel):
+	def __init__(self, config):
+		super().__init__(config)
+		self.num_labels = config.num_labels
+
+		self.bert = BertModel(config)
+		self.dropout = nn.Dropout(config.hidden_dropout_prob)
+		self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+		self.init_weights()
+
+	@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		labels=None,
+	):
+		r"""
+		labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the token classification loss.
+			Indices should be in ``[0, ..., config.num_labels - 1]``.
+
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided) :
+			Classification loss.
+		scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`)
+			Classification scores (before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+
+	Examples::
+
+		from transformers import BertTokenizer, BertForTokenClassification
+		import torch
+
+		tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+		model = BertForTokenClassification.from_pretrained('bert-base-uncased')
+
+		input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+		labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids, labels=labels)
+
+		loss, scores = outputs[:2]
+
+		"""
+
+		outputs = self.bert(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		sequence_output = outputs[0]
+
+		sequence_output = self.dropout(sequence_output)
+		logits = self.classifier(sequence_output)
+
+		outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
+		if labels is not None:
+			loss_fct = CrossEntropyLoss()
+			# Only keep active parts of the loss
+			if attention_mask is not None:
+				active_loss = attention_mask.view(-1) == 1
+				active_logits = logits.view(-1, self.num_labels)
+				active_labels = torch.where(
+					active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+				)
+				loss = loss_fct(active_logits, active_labels)
+			else:
+				loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+			outputs = (loss,) + outputs
+
+		return outputs  # (loss), scores, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+	"""Bert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear
+	layers on top of the hidden-states output to compute `span start logits` and `span end logits`). """,
+	BERT_START_DOCSTRING,
+)
+class BertForQuestionAnswering(BertPreTrainedModel):
+	def __init__(self, config):
+		super().__init__(config)
+		self.num_labels = config.num_labels
+
+		self.bert = BertModel(config)
+		self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+		self.init_weights()
+
+	@add_start_docstrings_to_callable(BERT_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		start_positions=None,
+		end_positions=None,
+	):
+		r"""
+		start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for position (index) of the start of the labelled span for computing the token classification loss.
+			Positions are clamped to the length of the sequence (`sequence_length`).
+			Position outside of the sequence are not taken into account for computing the loss.
+		end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for position (index) of the end of the labelled span for computing the token classification loss.
+			Positions are clamped to the length of the sequence (`sequence_length`).
+			Position outside of the sequence are not taken into account for computing the loss.
+
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.BertConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
+			Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
+		start_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
+			Span-start scores (before SoftMax).
+		end_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
+			Span-end scores (before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+
+	Examples::
+
+		from transformers import BertTokenizer, BertForQuestionAnswering
+		import torch
+
+		tokenizer = BertTokenizer.from_pretrained('bert-base-uncased')
+		model = BertForQuestionAnswering.from_pretrained('bert-large-uncased-whole-word-masking-finetuned-squad')
+
+		question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+		encoding = tokenizer.encode_plus(question, text)
+		input_ids, token_type_ids = encoding["input_ids"], encoding["token_type_ids"]
+		start_scores, end_scores = model(torch.tensor([input_ids]), token_type_ids=torch.tensor([token_type_ids]))
+
+		all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
+		answer = ' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1])
+
+		assert answer == "a nice puppet"
+
+		"""
+
+		outputs = self.bert(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		sequence_output = outputs[0]
+
+		logits = self.qa_outputs(sequence_output)
+		start_logits, end_logits = logits.split(1, dim=-1)
+		start_logits = start_logits.squeeze(-1)
+		end_logits = end_logits.squeeze(-1)
+
+		outputs = (start_logits, end_logits,) + outputs[2:]
+		if start_positions is not None and end_positions is not None:
+			# If we are on multi-GPU, split add a dimension
+			if len(start_positions.size()) > 1:
+				start_positions = start_positions.squeeze(-1)
+			if len(end_positions.size()) > 1:
+				end_positions = end_positions.squeeze(-1)
+			# sometimes the start/end positions are outside our model inputs, we ignore these terms
+			ignored_index = start_logits.size(1)
+			start_positions.clamp_(0, ignored_index)
+			end_positions.clamp_(0, ignored_index)
+
+			loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+			start_loss = loss_fct(start_logits, start_positions)
+			end_loss = loss_fct(end_logits, end_positions)
+			total_loss = (start_loss + end_loss) / 2
+			outputs = (total_loss,) + outputs
+
+		return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)

models/modeling_utils.py

@@ -0,0 +1,2098 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors, Facebook AI Research authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch BERT model."""
+
+import logging
+import os
+from typing import Callable, Tuple
+
+import torch
+from torch import Tensor, device, dtype, nn
+from torch.nn import CrossEntropyLoss
+from torch.nn import functional as F
+
+from .activations import get_activation
+from .configuration_utils import PretrainedConfig
+from .file_utils import (
+	DUMMY_INPUTS,
+	TF2_WEIGHTS_NAME,
+	TF_WEIGHTS_NAME,
+	WEIGHTS_NAME,
+	cached_path,
+	hf_bucket_url,
+	is_remote_url,
+)
+
+
+logger = logging.getLogger(__name__)
+
+
+try:
+	from torch.nn import Identity
+except ImportError:
+	# Older PyTorch compatibility
+	class Identity(nn.Module):
+		r"""A placeholder identity operator that is argument-insensitive.
+		"""
+
+		def __init__(self, *args, **kwargs):
+			super().__init__()
+
+		def forward(self, input):
+			return input
+
+
+class ModuleUtilsMixin:
+	"""
+	A few utilities for torch.nn.Modules, to be used as a mixin.
+	"""
+
+	def num_parameters(self, only_trainable: bool = False) -> int:
+		"""
+		Get number of (optionally, trainable) parameters in the module.
+		"""
+		params = filter(lambda x: x.requires_grad, self.parameters()) if only_trainable else self.parameters()
+		return sum(p.numel() for p in params)
+
+	@staticmethod
+	def _hook_rss_memory_pre_forward(module, *args, **kwargs):
+		try:
+			import psutil
+		except (ImportError):
+			raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.")
+
+		process = psutil.Process(os.getpid())
+		mem = process.memory_info()
+		module.mem_rss_pre_forward = mem.rss
+		return None
+
+	@staticmethod
+	def _hook_rss_memory_post_forward(module, *args, **kwargs):
+		try:
+			import psutil
+		except (ImportError):
+			raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.")
+
+		process = psutil.Process(os.getpid())
+		mem = process.memory_info()
+		module.mem_rss_post_forward = mem.rss
+		mem_rss_diff = module.mem_rss_post_forward - module.mem_rss_pre_forward
+		module.mem_rss_diff = mem_rss_diff + (module.mem_rss_diff if hasattr(module, "mem_rss_diff") else 0)
+		return None
+
+	def add_memory_hooks(self):
+		""" Add a memory hook before and after each sub-module forward pass to record increase in memory consumption.
+			Increase in memory consumption is stored in a `mem_rss_diff` attribute for each module and can be reset to zero with `model.reset_memory_hooks_state()`
+		"""
+		for module in self.modules():
+			module.register_forward_pre_hook(self._hook_rss_memory_pre_forward)
+			module.register_forward_hook(self._hook_rss_memory_post_forward)
+		self.reset_memory_hooks_state()
+
+	def reset_memory_hooks_state(self):
+		for module in self.modules():
+			module.mem_rss_diff = 0
+			module.mem_rss_post_forward = 0
+			module.mem_rss_pre_forward = 0
+
+	@property
+	def device(self) -> device:
+		return next(self.parameters()).device
+
+	@property
+	def dtype(self) -> dtype:
+		return next(self.parameters()).dtype
+
+	def invert_attention_mask(self, encoder_attention_mask: Tensor) -> Tensor:
+		"""type: torch.Tensor -> torch.Tensor"""
+		if encoder_attention_mask.dim() == 3:
+			encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
+		if encoder_attention_mask.dim() == 2:
+			encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
+		# T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
+		# Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
+		# /transformer/transformer_layers.py#L270
+		# encoder_extended_attention_mask = (encoder_extended_attention_mask ==
+		# encoder_extended_attention_mask.transpose(-1, -2))
+		encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+		encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -1e9
+		return encoder_extended_attention_mask
+
+	def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: tuple, device: device):
+		"""Makes broadcastable attention mask and causal mask so that future and maked tokens are ignored.
+
+		Arguments:
+			attention_mask: torch.Tensor with 1 indicating tokens to ATTEND to
+			input_shape: tuple, shape of input_ids
+			device: torch.Device, usually self.device
+
+		Returns:
+			torch.Tensor with dtype of attention_mask.dtype
+		"""
+		# 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.
+		if attention_mask.dim() == 3:
+			extended_attention_mask = attention_mask[:, None, :, :]
+		elif attention_mask.dim() == 2:
+			# Provided a padding mask of dimensions [batch_size, seq_length]
+			# - if the model is a decoder, apply a causal mask in addition to the padding mask
+			# - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
+			if self.config.is_decoder:
+				batch_size, seq_length = input_shape
+				seq_ids = torch.arange(seq_length, device=device)
+				causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
+				# causal and attention masks must have same type with pytorch version < 1.3
+				causal_mask = causal_mask.to(attention_mask.dtype)
+				extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
+			else:
+				extended_attention_mask = attention_mask[:, None, None, :]
+		else:
+			raise ValueError(
+				"Wrong shape for input_ids (shape {}) or attention_mask (shape {})".format(
+					input_shape, attention_mask.shape
+				)
+			)
+
+		# Since attention_mask is 1.0 for positions we want to attend and 0.0 for
+		# masked positions, this operation will create a tensor which is 0.0 for
+		# positions we want to attend and -10000.0 for masked positions.
+		# Since we are adding it to the raw scores before the softmax, this is
+		# effectively the same as removing these entirely.
+		extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
+		extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
+		return extended_attention_mask
+
+	def get_head_mask(self, head_mask, num_hidden_layers):
+		"""
+		# 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
+		Arguments:
+			head_mask: torch.Tensor or None: has shape [num_heads] or [num_hidden_layers x num_heads]
+			num_hidden_layers: int
+		Returns:
+			 Tensor of shape shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
+			 or list with [None] for each layer
+		"""
+		if head_mask is not None:
+			head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
+		else:
+			head_mask = [None] * num_hidden_layers
+
+		return head_mask
+
+	def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers):
+		"""-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]"""
+		if head_mask.dim() == 1:
+			head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
+			head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1)
+		elif head_mask.dim() == 2:
+			head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
+		assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}"
+		head_mask = head_mask.to(dtype=self.dtype)  # switch to fload if need + fp16 compatibility
+		return head_mask
+
+
+class PreTrainedModel(nn.Module, ModuleUtilsMixin):
+	r""" Base class for all models.
+
+		:class:`~transformers.PreTrainedModel` takes care of storing the configuration of the models and handles methods for loading/downloading/saving models
+		as well as a few methods common to all models to (i) resize the input embeddings and (ii) prune heads in the self-attention heads.
+
+		Class attributes (overridden by derived classes):
+			- ``config_class``: a class derived from :class:`~transformers.PretrainedConfig` to use as configuration class for this model architecture.
+			- ``pretrained_model_archive_map``: a python ``dict`` of with `short-cut-names` (string) as keys and `url` (string) of associated pretrained weights as values.
+			- ``load_tf_weights``: a python ``method`` for loading a TensorFlow checkpoint in a PyTorch model, taking as arguments:
+
+				- ``model``: an instance of the relevant subclass of :class:`~transformers.PreTrainedModel`,
+				- ``config``: an instance of the relevant subclass of :class:`~transformers.PretrainedConfig`,
+				- ``path``: a path (string) to the TensorFlow checkpoint.
+
+			- ``base_model_prefix``: a string indicating the attribute associated to the base model in derived classes of the same architecture adding modules on top of the base model.
+	"""
+	config_class = None
+	pretrained_model_archive_map = {}
+	base_model_prefix = ""
+
+	@property
+	def dummy_inputs(self):
+		""" Dummy inputs to do a forward pass in the network.
+
+		Returns:
+			torch.Tensor with dummy inputs
+		"""
+		return {"input_ids": torch.tensor(DUMMY_INPUTS)}
+
+	def __init__(self, config, *inputs, **kwargs):
+		super().__init__()
+		if not isinstance(config, PretrainedConfig):
+			raise ValueError(
+				"Parameter config in `{}(config)` should be an instance of class `PretrainedConfig`. "
+				"To create a model from a pretrained model use "
+				"`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
+					self.__class__.__name__, self.__class__.__name__
+				)
+			)
+		# Save config in model
+		self.config = config
+
+	@property
+	def base_model(self):
+		return getattr(self, self.base_model_prefix, self)
+
+	def get_input_embeddings(self):
+		"""
+		Returns the model's input embeddings.
+
+		Returns:
+			:obj:`nn.Module`:
+				A torch module mapping vocabulary to hidden states.
+		"""
+		base_model = getattr(self, self.base_model_prefix, self)
+		if base_model is not self:
+			return base_model.get_input_embeddings()
+		else:
+			raise NotImplementedError
+
+	def set_input_embeddings(self, value):
+		"""
+		Set model's input embeddings
+
+		Args:
+			value (:obj:`nn.Module`):
+				A module mapping vocabulary to hidden states.
+		"""
+		base_model = getattr(self, self.base_model_prefix, self)
+		if base_model is not self:
+			base_model.set_input_embeddings(value)
+		else:
+			raise NotImplementedError
+
+	def get_output_embeddings(self):
+		"""
+		Returns the model's output embeddings.
+
+		Returns:
+			:obj:`nn.Module`:
+				A torch module mapping hidden states to vocabulary.
+		"""
+		return None  # Overwrite for models with output embeddings
+
+	def tie_weights(self):
+		"""
+		Tie the weights between the input embeddings and the output embeddings.
+		If the `torchscript` flag is set in the configuration, can't handle parameter sharing so we are cloning
+		the weights instead.
+		"""
+		output_embeddings = self.get_output_embeddings()
+		if output_embeddings is not None:
+			self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings())
+
+	def _tie_or_clone_weights(self, output_embeddings, input_embeddings):
+		""" Tie or clone module weights depending of whether we are using TorchScript or not
+		"""
+		if self.config.torchscript:
+			output_embeddings.weight = nn.Parameter(input_embeddings.weight.clone())
+		else:
+			output_embeddings.weight = input_embeddings.weight
+
+		if getattr(output_embeddings, "bias", None) is not None:
+			output_embeddings.bias.data = torch.nn.functional.pad(
+				output_embeddings.bias.data,
+				(0, output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0],),
+				"constant",
+				0,
+			)
+		if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"):
+			output_embeddings.out_features = input_embeddings.num_embeddings
+
+	def resize_token_embeddings(self, new_num_tokens=None):
+		""" Resize input token embeddings matrix of the model if new_num_tokens != config.vocab_size.
+		Take care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.
+
+		Arguments:
+
+			new_num_tokens: (`optional`) int:
+				New number of tokens in the embedding matrix. Increasing the size will add newly initialized vectors at the end. Reducing the size will remove vectors from the end.
+				If not provided or None: does nothing and just returns a pointer to the input tokens ``torch.nn.Embeddings`` Module of the model.
+
+		Return: ``torch.nn.Embeddings``
+			Pointer to the input tokens Embeddings Module of the model
+		"""
+		base_model = getattr(self, self.base_model_prefix, self)  # get the base model if needed
+		model_embeds = base_model._resize_token_embeddings(new_num_tokens)
+		if new_num_tokens is None:
+			return model_embeds
+
+		# Update base model and current model config
+		self.config.vocab_size = new_num_tokens
+		base_model.vocab_size = new_num_tokens
+
+		# Tie weights again if needed
+		self.tie_weights()
+
+		return model_embeds
+
+	def _resize_token_embeddings(self, new_num_tokens):
+		old_embeddings = self.get_input_embeddings()
+		new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens)
+		self.set_input_embeddings(new_embeddings)
+		return self.get_input_embeddings()
+
+	def _get_resized_embeddings(self, old_embeddings, new_num_tokens=None):
+		""" Build a resized Embedding Module from a provided token Embedding Module.
+			Increasing the size will add newly initialized vectors at the end
+			Reducing the size will remove vectors from the end
+
+		Args:
+			new_num_tokens: (`optional`) int
+				New number of tokens in the embedding matrix.
+				Increasing the size will add newly initialized vectors at the end
+				Reducing the size will remove vectors from the end
+				If not provided or None: return the provided token Embedding Module.
+		Return: ``torch.nn.Embeddings``
+			Pointer to the resized Embedding Module or the old Embedding Module if new_num_tokens is None
+		"""
+		if new_num_tokens is None:
+			return old_embeddings
+
+		old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
+		if old_num_tokens == new_num_tokens:
+			return old_embeddings
+
+		# Build new embeddings
+		new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim)
+		new_embeddings.to(old_embeddings.weight.device)
+
+		# initialize all new embeddings (in particular added tokens)
+		self._init_weights(new_embeddings)
+
+		# Copy token embeddings from the previous weights
+		num_tokens_to_copy = min(old_num_tokens, new_num_tokens)
+		new_embeddings.weight.data[:num_tokens_to_copy, :] = old_embeddings.weight.data[:num_tokens_to_copy, :]
+
+		return new_embeddings
+
+	def init_weights(self):
+		""" Initialize and prunes weights if needed. """
+		# Initialize weights
+		self.apply(self._init_weights)
+
+		# Prune heads if needed
+		if self.config.pruned_heads:
+			self.prune_heads(self.config.pruned_heads)
+
+		# Tie weights if needed
+		self.tie_weights()
+
+	def prune_heads(self, heads_to_prune):
+		""" Prunes heads of the base model.
+
+			Arguments:
+
+				heads_to_prune: dict with keys being selected layer indices (`int`) and associated values being the list of heads to prune in said layer (list of `int`).
+				E.g. {1: [0, 2], 2: [2, 3]} will prune heads 0 and 2 on layer 1 and heads 2 and 3 on layer 2.
+		"""
+		# save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads
+		for layer, heads in heads_to_prune.items():
+			union_heads = set(self.config.pruned_heads.get(layer, [])) | set(heads)
+			self.config.pruned_heads[layer] = list(union_heads)  # Unfortunately we have to store it as list for JSON
+
+		self.base_model._prune_heads(heads_to_prune)
+
+	def save_pretrained(self, save_directory):
+		""" Save a model and its configuration file to a directory, so that it
+			can be re-loaded using the `:func:`~transformers.PreTrainedModel.from_pretrained`` class method.
+
+			Arguments:
+				save_directory: directory to which to save.
+		"""
+		assert os.path.isdir(
+			save_directory
+		), "Saving path should be a directory where the model and configuration can be saved"
+
+		# Only save the model itself if we are using distributed training
+		model_to_save = self.module if hasattr(self, "module") else self
+
+		# Attach architecture to the config
+		model_to_save.config.architectures = [model_to_save.__class__.__name__]
+
+		# If we save using the predefined names, we can load using `from_pretrained`
+		output_model_file = os.path.join(save_directory, WEIGHTS_NAME)
+
+		if getattr(self.config, "xla_device", False):
+			import torch_xla.core.xla_model as xm
+
+			if xm.is_master_ordinal():
+				# Save configuration file
+				model_to_save.config.save_pretrained(save_directory)
+			# xm.save takes care of saving only from master
+			xm.save(model_to_save.state_dict(), output_model_file)
+		else:
+			model_to_save.config.save_pretrained(save_directory)
+			torch.save(model_to_save.state_dict(), output_model_file)
+
+		logger.info("Model weights saved in {}".format(output_model_file))
+
+	@classmethod
+	def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
+		r"""Instantiate a pretrained pytorch model from a pre-trained model configuration.
+
+		The model is set in evaluation mode by default using ``model.eval()`` (Dropout modules are deactivated)
+		To train the model, you should first set it back in training mode with ``model.train()``
+
+		The warning ``Weights from XXX not initialized from pretrained model`` means that the weights of XXX do not come pre-trained with the rest of the model.
+		It is up to you to train those weights with a downstream fine-tuning task.
+
+		The warning ``Weights from XXX not used in YYY`` means that the layer XXX is not used by YYY, therefore those weights are discarded.
+
+		Parameters:
+			pretrained_model_name_or_path: either:
+			  - a string with the `shortcut name` of a pre-trained model to load from cache or download, e.g.: ``bert-base-uncased``.
+			  - a string with the `identifier name` of a pre-trained model that was user-uploaded to our S3, e.g.: ``dbmdz/bert-base-german-cased``.
+			  - a path to a `directory` containing model weights saved using :func:`~transformers.PreTrainedModel.save_pretrained`, e.g.: ``./my_model_directory/``.
+			  - a path or url to a `tensorflow index checkpoint file` (e.g. `./tf_model/model.ckpt.index`). In this case, ``from_tf`` should be set to True and a configuration object should be provided as ``config`` argument. This loading path is slower than converting the TensorFlow checkpoint in a PyTorch model using the provided conversion scripts and loading the PyTorch model afterwards.
+			  - None if you are both providing the configuration and state dictionary (resp. with keyword arguments ``config`` and ``state_dict``)
+
+			model_args: (`optional`) Sequence of positional arguments:
+				All remaning positional arguments will be passed to the underlying model's ``__init__`` method
+
+			config: (`optional`) one of:
+				- an instance of a class derived from :class:`~transformers.PretrainedConfig`, or
+				- a string valid as input to :func:`~transformers.PretrainedConfig.from_pretrained()`
+				Configuration for the model to use instead of an automatically loaded configuation. Configuration can be automatically loaded when:
+					- the model is a model provided by the library (loaded with the ``shortcut-name`` string of a pretrained model), or
+					- the model was saved using :func:`~transformers.PreTrainedModel.save_pretrained` and is reloaded by suppling the save directory.
+					- the model is loaded by suppling a local directory as ``pretrained_model_name_or_path`` and a configuration JSON file named `config.json` is found in the directory.
+
+			state_dict: (`optional`) dict:
+				an optional state dictionnary for the model to use instead of a state dictionary loaded from saved weights file.
+				This option can be used if you want to create a model from a pretrained configuration but load your own weights.
+				In this case though, you should check if using :func:`~transformers.PreTrainedModel.save_pretrained` and :func:`~transformers.PreTrainedModel.from_pretrained` is not a simpler option.
+
+			cache_dir: (`optional`) string:
+				Path to a directory in which a downloaded pre-trained model
+				configuration should be cached if the standard cache should not be used.
+
+			force_download: (`optional`) boolean, default False:
+				Force to (re-)download the model weights and configuration files and override the cached versions if they exists.
+
+			resume_download: (`optional`) boolean, default False:
+				Do not delete incompletely recieved file. Attempt to resume the download if such a file exists.
+
+			proxies: (`optional`) dict, default None:
+				A dictionary of proxy servers to use by protocol or endpoint, e.g.: {'http': 'foo.bar:3128', 'http://hostname': 'foo.bar:4012'}.
+				The proxies are used on each request.
+
+			output_loading_info: (`optional`) boolean:
+				Set to ``True`` to also return a dictionnary containing missing keys, unexpected keys and error messages.
+
+			kwargs: (`optional`) Remaining dictionary of keyword arguments:
+				Can be used to update the configuration object (after it being loaded) and initiate the model. (e.g. ``output_attention=True``). Behave differently depending on whether a `config` is provided or automatically loaded:
+
+				- If a configuration is provided with ``config``, ``**kwargs`` will be directly passed to the underlying model's ``__init__`` method (we assume all relevant updates to the configuration have already been done)
+				- If a configuration is not provided, ``kwargs`` will be first passed to the configuration class initialization function (:func:`~transformers.PretrainedConfig.from_pretrained`). Each key of ``kwargs`` that corresponds to a configuration attribute will be used to override said attribute with the supplied ``kwargs`` value. Remaining keys that do not correspond to any configuration attribute will be passed to the underlying model's ``__init__`` function.
+
+		Examples::
+
+			# For example purposes. Not runnable.
+			model = BertModel.from_pretrained('bert-base-uncased')    # Download model and configuration from S3 and cache.
+			model = BertModel.from_pretrained('./test/saved_model/')  # E.g. model was saved using `save_pretrained('./test/saved_model/')`
+			model = BertModel.from_pretrained('bert-base-uncased', output_attention=True)  # Update configuration during loading
+			assert model.config.output_attention == True
+			# Loading from a TF checkpoint file instead of a PyTorch model (slower)
+			config = BertConfig.from_json_file('./tf_model/my_tf_model_config.json')
+			model = BertModel.from_pretrained('./tf_model/my_tf_checkpoint.ckpt.index', from_tf=True, config=config)
+
+		"""
+		config = kwargs.pop("config", None)
+		state_dict = kwargs.pop("state_dict", None)
+		cache_dir = kwargs.pop("cache_dir", None)
+		from_tf = kwargs.pop("from_tf", False)
+		force_download = kwargs.pop("force_download", False)
+		resume_download = kwargs.pop("resume_download", False)
+		proxies = kwargs.pop("proxies", None)
+		output_loading_info = kwargs.pop("output_loading_info", False)
+		local_files_only = kwargs.pop("local_files_only", False)
+
+		# Load config if we don't provide a configuration
+		if not isinstance(config, PretrainedConfig):
+			config_path = config if config is not None else pretrained_model_name_or_path
+			config, model_kwargs = cls.config_class.from_pretrained(
+				config_path,
+				*model_args,
+				cache_dir=cache_dir,
+				return_unused_kwargs=True,
+				force_download=force_download,
+				resume_download=resume_download,
+				proxies=proxies,
+				local_files_only=local_files_only,
+				**kwargs,
+			)
+		else:
+			model_kwargs = kwargs
+
+		# Load model
+		if pretrained_model_name_or_path is not None:
+			if pretrained_model_name_or_path in cls.pretrained_model_archive_map:
+				archive_file = cls.pretrained_model_archive_map[pretrained_model_name_or_path]
+			elif os.path.isdir(pretrained_model_name_or_path):
+				if from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")):
+					# Load from a TF 1.0 checkpoint
+					archive_file = os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
+				elif from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)):
+					# Load from a TF 2.0 checkpoint
+					archive_file = os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)
+				elif os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
+					# Load from a PyTorch checkpoint
+					archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
+				else:
+					raise EnvironmentError(
+						"Error no file named {} found in directory {} or `from_tf` set to False".format(
+							[WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME + ".index"],
+							pretrained_model_name_or_path,
+						)
+					)
+			elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
+				archive_file = pretrained_model_name_or_path
+			elif os.path.isfile(pretrained_model_name_or_path + ".index"):
+				assert (
+					from_tf
+				), "We found a TensorFlow checkpoint at {}, please set from_tf to True to load from this checkpoint".format(
+					pretrained_model_name_or_path + ".index"
+				)
+				archive_file = pretrained_model_name_or_path + ".index"
+			else:
+				archive_file = hf_bucket_url(
+					pretrained_model_name_or_path, postfix=(TF2_WEIGHTS_NAME if from_tf else WEIGHTS_NAME),
+				)
+
+			# redirect to the cache, if necessary
+			try:
+				resolved_archive_file = cached_path(
+					archive_file,
+					cache_dir=cache_dir,
+					force_download=force_download,
+					proxies=proxies,
+					resume_download=resume_download,
+					local_files_only=local_files_only,
+				)
+			except EnvironmentError:
+				if pretrained_model_name_or_path in cls.pretrained_model_archive_map:
+					msg = "Couldn't reach server at '{}' to download pretrained weights.".format(archive_file)
+				else:
+					msg = (
+						"Model name '{}' was not found in model name list ({}). "
+						"We assumed '{}' was a path or url to model weight files named one of {} but "
+						"couldn't find any such file at this path or url.".format(
+							pretrained_model_name_or_path,
+							", ".join(cls.pretrained_model_archive_map.keys()),
+							archive_file,
+							[WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME],
+						)
+					)
+				raise EnvironmentError(msg)
+
+			if resolved_archive_file == archive_file:
+				logger.info("loading weights file {}".format(archive_file))
+			else:
+				logger.info("loading weights file {} from cache at {}".format(archive_file, resolved_archive_file))
+		else:
+			resolved_archive_file = None
+
+		# Instantiate model.
+		model = cls(config, *model_args, **model_kwargs)
+
+		if state_dict is None and not from_tf:
+			try:
+				state_dict = torch.load(resolved_archive_file, map_location="cpu")
+			except Exception:
+				raise OSError(
+					"Unable to load weights from pytorch checkpoint file. "
+					"If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True. "
+				)
+
+		missing_keys = []
+		unexpected_keys = []
+		error_msgs = []
+
+		if from_tf:
+			if resolved_archive_file.endswith(".index"):
+				# Load from a TensorFlow 1.X checkpoint - provided by original authors
+				model = cls.load_tf_weights(model, config, resolved_archive_file[:-6])  # Remove the '.index'
+			else:
+				# Load from our TensorFlow 2.0 checkpoints
+				try:
+					from transformers import load_tf2_checkpoint_in_pytorch_model
+
+					model = load_tf2_checkpoint_in_pytorch_model(model, resolved_archive_file, allow_missing_keys=True)
+				except ImportError:
+					logger.error(
+						"Loading a TensorFlow model in PyTorch, requires both PyTorch and TensorFlow to be installed. Please see "
+						"https://pytorch.org/ and https://www.tensorflow.org/install/ for installation instructions."
+					)
+					raise
+		else:
+			# Convert old format to new format if needed from a PyTorch state_dict
+			old_keys = []
+			new_keys = []
+			for key in state_dict.keys():
+				new_key = None
+				if "gamma" in key:
+					new_key = key.replace("gamma", "weight")
+				if "beta" in key:
+					new_key = key.replace("beta", "bias")
+				if new_key:
+					old_keys.append(key)
+					new_keys.append(new_key)
+			for old_key, new_key in zip(old_keys, new_keys):
+				state_dict[new_key] = state_dict.pop(old_key)
+
+			# copy state_dict so _load_from_state_dict can modify it
+			metadata = getattr(state_dict, "_metadata", None)
+			state_dict = state_dict.copy()
+			if metadata is not None:
+				state_dict._metadata = metadata
+
+			# PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
+			# so we need to apply the function recursively.
+			def load(module: nn.Module, prefix=""):
+				local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
+				module._load_from_state_dict(
+					state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs,
+				)
+				for name, child in module._modules.items():
+					if child is not None:
+						load(child, prefix + name + ".")
+
+			# Make sure we are able to load base models as well as derived models (with heads)
+			start_prefix = ""
+			model_to_load = model
+			has_prefix_module = any(s.startswith(cls.base_model_prefix) for s in state_dict.keys())
+			if not hasattr(model, cls.base_model_prefix) and has_prefix_module:
+				start_prefix = cls.base_model_prefix + "."
+			if hasattr(model, cls.base_model_prefix) and not has_prefix_module:
+				model_to_load = getattr(model, cls.base_model_prefix)
+
+			load(model_to_load, prefix=start_prefix)
+
+			if model.__class__.__name__ != model_to_load.__class__.__name__:
+				base_model_state_dict = model_to_load.state_dict().keys()
+				head_model_state_dict_without_base_prefix = [
+					key.split(cls.base_model_prefix + ".")[-1] for key in model.state_dict().keys()
+				]
+
+				missing_keys.extend(head_model_state_dict_without_base_prefix - base_model_state_dict)
+
+			if len(missing_keys) > 0:
+				logger.info(
+					"Weights of {} not initialized from pretrained model: {}".format(
+						model.__class__.__name__, missing_keys
+					)
+				)
+			if len(unexpected_keys) > 0:
+				logger.info(
+					"Weights from pretrained model not used in {}: {}".format(
+						model.__class__.__name__, unexpected_keys
+					)
+				)
+			if len(error_msgs) > 0:
+				raise RuntimeError(
+					"Error(s) in loading state_dict for {}:\n\t{}".format(
+						model.__class__.__name__, "\n\t".join(error_msgs)
+					)
+				)
+		model.tie_weights()  # make sure token embedding weights are still tied if needed
+
+		# Set model in evaluation mode to deactivate DropOut modules by default
+		model.eval()
+
+		if output_loading_info:
+			loading_info = {
+				"missing_keys": missing_keys,
+				"unexpected_keys": unexpected_keys,
+				"error_msgs": error_msgs,
+			}
+			return model, loading_info
+
+		if hasattr(config, "xla_device") and config.xla_device:
+			import torch_xla.core.xla_model as xm
+
+			model = xm.send_cpu_data_to_device(model, xm.xla_device())
+			model = model.to(xm.xla_device())
+
+		return model
+
+	def prepare_inputs_for_generation(self, input_ids, **kwargs):
+		return {"input_ids": input_ids}
+
+	def prepare_scores_for_generation(self, scores, **kwargs):
+		return scores
+
+	def _use_cache(self, outputs, use_cache):
+		"""During generation, decide whether to pass the `past` variable to the next forward pass."""
+		if len(outputs) <= 1 or use_cache is False:
+			return False
+		if hasattr(self.config, "mem_len") and self.config.mem_len == 0:
+			return False
+		return True
+
+	def enforce_repetition_penalty_(self, lprobs, batch_size, num_beams, prev_output_tokens, repetition_penalty):
+		"""repetition penalty (from CTRL paper https://arxiv.org/abs/1909.05858). """
+		for i in range(batch_size * num_beams):
+			for previous_token in set(prev_output_tokens[i].tolist()):
+				# if score < 0 then repetition penalty has to multiplied to reduce the previous token probability
+				if lprobs[i, previous_token] < 0:
+					lprobs[i, previous_token] *= repetition_penalty
+				else:
+					lprobs[i, previous_token] /= repetition_penalty
+
+	@torch.no_grad()
+	def generate(
+		self,
+		input_ids=None,
+		max_length=None,
+		min_length=None,
+		do_sample=None,
+		early_stopping=None,
+		num_beams=None,
+		temperature=None,
+		top_k=None,
+		top_p=None,
+		repetition_penalty=None,
+		bad_words_ids=None,
+		bos_token_id=None,
+		pad_token_id=None,
+		eos_token_id=None,
+		length_penalty=None,
+		no_repeat_ngram_size=None,
+		num_return_sequences=None,
+		attention_mask=None,
+		decoder_start_token_id=None,
+		use_cache=None,
+	):
+		r""" Generates sequences for models with a LM head. The method currently supports greedy decoding, beam-search decoding, sampling with temperature, sampling with top-k or nucleus sampling.
+
+		Adapted in part from `Facebook's XLM beam search code`_.
+
+		.. _`Facebook's XLM beam search code`:
+		   https://github.com/facebookresearch/XLM/blob/9e6f6814d17be4fe5b15f2e6c43eb2b2d76daeb4/src/model/transformer.py#L529
+
+
+		Parameters:
+
+			input_ids: (`optional`) `torch.LongTensor` of shape `(batch_size, sequence_length)`
+				The sequence used as a prompt for the generation. If `None` the method initializes
+				it as an empty `torch.LongTensor` of shape `(1,)`.
+
+			max_length: (`optional`) int
+				The max length of the sequence to be generated.  Between `min_length` and infinity. Default to 20.
+
+			min_length: (`optional`) int
+				The min length of the sequence to be generated.  Between 0 and infinity. Default to 0.
+
+			do_sample: (`optional`) bool
+				If set to `False` greedy decoding is used. Otherwise sampling is used. Defaults to `False` as defined in `configuration_utils.PretrainedConfig`.
+
+			early_stopping: (`optional`) bool
+				if set to `True` beam search is stopped when at least `num_beams` sentences finished per batch. Defaults to `False` as defined in `configuration_utils.PretrainedConfig`.
+
+			num_beams: (`optional`) int
+				Number of beams for beam search. Must be between 1 and infinity. 1 means no beam search. Default to 1.
+
+			temperature: (`optional`) float
+				The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
+
+			top_k: (`optional`) int
+				The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
+
+			top_p: (`optional`) float
+				The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
+
+			repetition_penalty: (`optional`) float
+				The parameter for repetition penalty. Between 1.0 and infinity. 1.0 means no penalty. Default to 1.0.
+
+			pad_token_id: (`optional`) int
+				Padding token. Default to specicic model pad_token_id or None if it does not exist.
+
+			bos_token_id: (`optional`) int
+				BOS token. Defaults to `bos_token_id` as defined in the models config.
+
+			eos_token_id: (`optional`) int
+				EOS token. Defaults to `eos_token_id` as defined in the models config.
+
+			length_penalty: (`optional`) float
+				Exponential penalty to the length. Default to 1.
+
+			no_repeat_ngram_size: (`optional`) int
+				If set to int > 0, all ngrams of size `no_repeat_ngram_size` can only occur once.
+			bad_words_ids: (`optional`) list of lists of int
+				`bad_words_ids` contains tokens that are not allowed to be generated. In order to get the tokens of the words that should not appear in the generated text, use `tokenizer.encode(bad_word, add_prefix_space=True)`.
+
+			num_return_sequences: (`optional`) int
+				The number of independently computed returned sequences for each element in the batch. Default to 1.
+
+			attention_mask (`optional`) obj: `torch.LongTensor` of same shape as `input_ids`
+				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.
+				Defaults to `None`.
+
+			`What are attention masks? <../glossary.html#attention-mask>`__
+
+			decoder_start_token_id=None: (`optional`) int
+				If an encoder-decoder model starts decoding with a different token than BOS.
+				Defaults to `None` and is changed to `BOS` later.
+
+			use_cache: (`optional`) bool
+				If `use_cache` is True, past key values are used to speed up decoding if applicable to model. Defaults to `True`.
+
+		Return:
+
+			output: `torch.LongTensor` of shape `(batch_size * num_return_sequences, sequence_length)`
+				sequence_length is either equal to max_length or shorter if all batches finished early due to the `eos_token_id`
+
+		Examples::
+
+			tokenizer = AutoTokenizer.from_pretrained('distilgpt2')   # Initialize tokenizer
+			model = AutoModelWithLMHead.from_pretrained('distilgpt2')    # Download model and configuration from S3 and cache.
+			outputs = model.generate(max_length=40)  # do greedy decoding
+			print('Generated: {}'.format(tokenizer.decode(outputs[0], skip_special_tokens=True)))
+
+			tokenizer = AutoTokenizer.from_pretrained('openai-gpt')   # Initialize tokenizer
+			model = AutoModelWithLMHead.from_pretrained('openai-gpt')    # Download model and configuration from S3 and cache.
+			input_context = 'The dog'
+			input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
+			outputs = model.generate(input_ids=input_ids, num_beams=5, num_return_sequences=3, temperature=1.5)  # generate 3 independent sequences using beam search decoding (5 beams) with sampling from initial context 'The dog'
+			for i in range(3): #  3 output sequences were generated
+				print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))
+
+			tokenizer = AutoTokenizer.from_pretrained('distilgpt2')   # Initialize tokenizer
+			model = AutoModelWithLMHead.from_pretrained('distilgpt2')    # Download model and configuration from S3 and cache.
+			input_context = 'The dog'
+			input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
+			outputs = model.generate(input_ids=input_ids, max_length=40, temperature=0.7, num_return_sequences=3)  # 3 generate sequences using by sampling
+			for i in range(3): #  3 output sequences were generated
+				print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))
+
+			tokenizer = AutoTokenizer.from_pretrained('ctrl')   # Initialize tokenizer
+			model = AutoModelWithLMHead.from_pretrained('ctrl')    # Download model and configuration from S3 and cache.
+			input_context = 'Legal My neighbor is'  # "Legal" is one of the control codes for ctrl
+			input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
+			outputs = model.generate(input_ids=input_ids, max_length=50, temperature=0.7, repetition_penalty=1.2)  # generate sequences
+			print('Generated: {}'.format(tokenizer.decode(outputs[0], skip_special_tokens=True)))
+
+			tokenizer = AutoTokenizer.from_pretrained('gpt2')   # Initialize tokenizer
+			model = AutoModelWithLMHead.from_pretrained('gpt2')    # Download model and configuration from S3 and cache.
+			input_context = 'My cute dog'  # "Legal" is one of the control codes for ctrl
+			bad_words_ids = [tokenizer.encode(bad_word, add_prefix_space=True) for bad_word in ['idiot', 'stupid', 'shut up']]
+			input_ids = tokenizer.encode(input_context, return_tensors='pt')  # encode input context
+			outputs = model.generate(input_ids=input_ids, max_length=100, do_sample=True, bad_words_ids=bad_words_ids)  # generate sequences without allowing bad_words to be generated
+		"""
+
+		# We cannot generate if the model does not have a LM head
+		if self.get_output_embeddings() is None:
+			raise AttributeError(
+				"You tried to generate sequences with a model that does not have a LM Head."
+				"Please use another model class (e.g. `OpenAIGPTLMHeadModel`, `XLNetLMHeadModel`, `GPT2LMHeadModel`, `CTRLLMHeadModel`, `T5WithLMHeadModel`, `TransfoXLLMHeadModel`, `XLMWithLMHeadModel`, `BartForConditionalGeneration` )"
+			)
+
+		max_length = max_length if max_length is not None else self.config.max_length
+		min_length = min_length if min_length is not None else self.config.min_length
+		do_sample = do_sample if do_sample is not None else self.config.do_sample
+		early_stopping = early_stopping if early_stopping is not None else self.config.early_stopping
+		use_cache = use_cache if use_cache is not None else self.config.use_cache
+		num_beams = num_beams if num_beams is not None else self.config.num_beams
+		temperature = temperature if temperature is not None else self.config.temperature
+		top_k = top_k if top_k is not None else self.config.top_k
+		top_p = top_p if top_p is not None else self.config.top_p
+		repetition_penalty = repetition_penalty if repetition_penalty is not None else self.config.repetition_penalty
+		bos_token_id = bos_token_id if bos_token_id is not None else self.config.bos_token_id
+		pad_token_id = pad_token_id if pad_token_id is not None else self.config.pad_token_id
+		eos_token_id = eos_token_id if eos_token_id is not None else self.config.eos_token_id
+		length_penalty = length_penalty if length_penalty is not None else self.config.length_penalty
+		no_repeat_ngram_size = (
+			no_repeat_ngram_size if no_repeat_ngram_size is not None else self.config.no_repeat_ngram_size
+		)
+		bad_words_ids = bad_words_ids if bad_words_ids is not None else self.config.bad_words_ids
+		num_return_sequences = (
+			num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
+		)
+		decoder_start_token_id = (
+			decoder_start_token_id if decoder_start_token_id is not None else self.config.decoder_start_token_id
+		)
+
+		if input_ids is not None:
+			batch_size = input_ids.shape[0]  # overriden by the input batch_size
+		else:
+			batch_size = 1
+
+		assert isinstance(max_length, int) and max_length > 0, "`max_length` should be a strictly positive integer."
+		assert isinstance(min_length, int) and min_length >= 0, "`min_length` should be a positive integer."
+		assert isinstance(do_sample, bool), "`do_sample` should be a boolean."
+		assert isinstance(early_stopping, bool), "`early_stopping` should be a boolean."
+		assert isinstance(use_cache, bool), "`use_cache` should be a boolean."
+		assert isinstance(num_beams, int) and num_beams > 0, "`num_beams` should be a strictly positive integer."
+		assert temperature > 0, "`temperature` should be strictly positive."
+		assert isinstance(top_k, int) and top_k >= 0, "`top_k` should be a positive integer."
+		assert 0 <= top_p <= 1, "`top_p` should be between 0 and 1."
+		assert repetition_penalty >= 1.0, "`repetition_penalty` should be >= 1."
+		assert input_ids is not None or (
+			isinstance(bos_token_id, int) and bos_token_id >= 0
+		), "If input_ids is not defined, `bos_token_id` should be a positive integer."
+		assert pad_token_id is None or (
+			isinstance(pad_token_id, int) and (pad_token_id >= 0)
+		), "`pad_token_id` should be a positive integer."
+		assert (eos_token_id is None) or (
+			isinstance(eos_token_id, int) and (eos_token_id >= 0)
+		), "`eos_token_id` should be a positive integer."
+		assert length_penalty > 0, "`length_penalty` should be strictly positive."
+		assert (
+			isinstance(no_repeat_ngram_size, int) and no_repeat_ngram_size >= 0
+		), "`no_repeat_ngram_size` should be a positive integer."
+		assert (
+			isinstance(num_return_sequences, int) and num_return_sequences > 0
+		), "`num_return_sequences` should be a strictly positive integer."
+		assert (
+			bad_words_ids is None or isinstance(bad_words_ids, list) and isinstance(bad_words_ids[0], list)
+		), "`bad_words_ids` is either `None` or a list of lists of tokens that should not be generated"
+
+		if input_ids is None:
+			assert isinstance(bos_token_id, int) and bos_token_id >= 0, (
+				"you should either supply a context to complete as `input_ids` input "
+				"or a `bos_token_id` (integer >= 0) as a first token to start the generation."
+			)
+			input_ids = torch.full(
+				(batch_size, 1), bos_token_id, dtype=torch.long, device=next(self.parameters()).device,
+			)
+		else:
+			assert input_ids.dim() == 2, "Input prompt should be of shape (batch_size, sequence length)."
+
+		# not allow to duplicate outputs when greedy decoding
+		if do_sample is False:
+			if num_beams == 1:
+				# no_beam_search greedy generation conditions
+				assert (
+					num_return_sequences == 1
+				), "Greedy decoding will always produce the same output for num_beams == 1 and num_return_sequences > 1. Please set num_return_sequences = 1"
+
+			else:
+				# beam_search greedy generation conditions
+				assert (
+					num_beams >= num_return_sequences
+				), "Greedy beam search decoding cannot return more sequences than it has beams. Please set num_beams >= num_return_sequences"
+
+		# create attention mask if necessary
+		# TODO (PVP): this should later be handled by the forward fn() in each model in the future see PR 3140
+		if (attention_mask is None) and (pad_token_id is not None) and (pad_token_id in input_ids):
+			attention_mask = input_ids.ne(pad_token_id).long()
+		elif attention_mask is None:
+			attention_mask = input_ids.new_ones(input_ids.shape)
+
+		# set pad_token_id to eos_token_id if not set. Important that this is done after
+		# attention_mask is created
+		if pad_token_id is None and eos_token_id is not None:
+			logger.warning(
+				"Setting `pad_token_id` to {} (first `eos_token_id`) to generate sequence".format(eos_token_id)
+			)
+			pad_token_id = eos_token_id
+
+		# current position and vocab size
+		if hasattr(self.config, "vocab_size"):
+			vocab_size = self.config.vocab_size
+		elif (
+			self.config.is_encoder_decoder
+			and hasattr(self.config, "decoder")
+			and hasattr(self.config.decoder, "vocab_size")
+		):
+			vocab_size = self.config.decoder.vocab_size
+
+		# set effective batch size and effective batch multiplier according to do_sample
+		if do_sample:
+			effective_batch_size = batch_size * num_return_sequences
+			effective_batch_mult = num_return_sequences
+		else:
+			effective_batch_size = batch_size
+			effective_batch_mult = 1
+
+		if self.config.is_encoder_decoder:
+			if decoder_start_token_id is None:
+				decoder_start_token_id = bos_token_id
+
+			assert (
+				decoder_start_token_id is not None
+			), "decoder_start_token_id or bos_token_id has to be defined for encoder-decoder generation"
+			assert hasattr(self, "get_encoder"), "{} should have a 'get_encoder' function defined".format(self)
+			assert callable(self.get_encoder), "{} should be a method".format(self.get_encoder)
+
+			# get encoder and store encoder outputs
+			encoder = self.get_encoder()
+
+			encoder_outputs: tuple = encoder(input_ids, attention_mask=attention_mask)
+
+		# Expand input ids if num_beams > 1 or num_return_sequences > 1
+		if num_return_sequences > 1 or num_beams > 1:
+			input_ids_len = input_ids.shape[-1]
+			input_ids = input_ids.unsqueeze(1).expand(batch_size, effective_batch_mult * num_beams, input_ids_len)
+			attention_mask = attention_mask.unsqueeze(1).expand(
+				batch_size, effective_batch_mult * num_beams, input_ids_len
+			)
+
+			input_ids = input_ids.contiguous().view(
+				effective_batch_size * num_beams, input_ids_len
+			)  # shape: (batch_size * num_return_sequences * num_beams, cur_len)
+			attention_mask = attention_mask.contiguous().view(
+				effective_batch_size * num_beams, input_ids_len
+			)  # shape: (batch_size * num_return_sequences * num_beams, cur_len)
+
+		if self.config.is_encoder_decoder:
+			# create empty decoder_input_ids
+			input_ids = torch.full(
+				(effective_batch_size * num_beams, 1),
+				decoder_start_token_id,
+				dtype=torch.long,
+				device=next(self.parameters()).device,
+			)
+			cur_len = 1
+
+			assert (
+				batch_size == encoder_outputs[0].shape[0]
+			), f"expected encoder_outputs[0] to have 1st dimension bs={batch_size}, got {encoder_outputs[0].shape[0]} "
+
+			# expand batch_idx to assign correct encoder output for expanded input_ids (due to num_beams > 1 and num_return_sequences > 1)
+			expanded_batch_idxs = (
+				torch.arange(batch_size)
+				.view(-1, 1)
+				.repeat(1, num_beams * effective_batch_mult)
+				.view(-1)
+				.to(input_ids.device)
+			)
+			# expand encoder_outputs
+			encoder_outputs = (encoder_outputs[0].index_select(0, expanded_batch_idxs), *encoder_outputs[1:])
+
+		else:
+			encoder_outputs = None
+			cur_len = input_ids.shape[-1]
+
+		if num_beams > 1:
+			output = self._generate_beam_search(
+				input_ids,
+				cur_len=cur_len,
+				max_length=max_length,
+				min_length=min_length,
+				do_sample=do_sample,
+				early_stopping=early_stopping,
+				temperature=temperature,
+				top_k=top_k,
+				top_p=top_p,
+				repetition_penalty=repetition_penalty,
+				no_repeat_ngram_size=no_repeat_ngram_size,
+				bad_words_ids=bad_words_ids,
+				bos_token_id=bos_token_id,
+				pad_token_id=pad_token_id,
+				decoder_start_token_id=decoder_start_token_id,
+				eos_token_id=eos_token_id,
+				batch_size=effective_batch_size,
+				num_return_sequences=num_return_sequences,
+				length_penalty=length_penalty,
+				num_beams=num_beams,
+				vocab_size=vocab_size,
+				encoder_outputs=encoder_outputs,
+				attention_mask=attention_mask,
+				use_cache=use_cache,
+			)
+		else:
+			output = self._generate_no_beam_search(
+				input_ids,
+				cur_len=cur_len,
+				max_length=max_length,
+				min_length=min_length,
+				do_sample=do_sample,
+				temperature=temperature,
+				top_k=top_k,
+				top_p=top_p,
+				repetition_penalty=repetition_penalty,
+				no_repeat_ngram_size=no_repeat_ngram_size,
+				bad_words_ids=bad_words_ids,
+				bos_token_id=bos_token_id,
+				pad_token_id=pad_token_id,
+				decoder_start_token_id=decoder_start_token_id,
+				eos_token_id=eos_token_id,
+				batch_size=effective_batch_size,
+				encoder_outputs=encoder_outputs,
+				attention_mask=attention_mask,
+				use_cache=use_cache,
+			)
+
+		return output
+
+	def _generate_no_beam_search(
+		self,
+		input_ids,
+		cur_len,
+		max_length,
+		min_length,
+		do_sample,
+		temperature,
+		top_k,
+		top_p,
+		repetition_penalty,
+		no_repeat_ngram_size,
+		bad_words_ids,
+		bos_token_id,
+		pad_token_id,
+		eos_token_id,
+		decoder_start_token_id,
+		batch_size,
+		encoder_outputs,
+		attention_mask,
+		use_cache,
+	):
+		""" Generate sequences for each example without beam search (num_beams == 1).
+			All returned sequence are generated independantly.
+		"""
+		# length of generated sentences / unfinished sentences
+		unfinished_sents = input_ids.new(batch_size).fill_(1)
+		sent_lengths = input_ids.new(batch_size).fill_(max_length)
+
+		past = encoder_outputs  # defined for encoder-decoder models, None for decoder-only models
+
+		while cur_len < max_length:
+			model_inputs = self.prepare_inputs_for_generation(
+				input_ids, past=past, attention_mask=attention_mask, use_cache=use_cache
+			)
+
+			outputs = self(**model_inputs)
+			next_token_logits = outputs[0][:, -1, :]
+
+			# if model has past, then set the past variable to speed up decoding
+			if self._use_cache(outputs, use_cache):
+				past = outputs[1]
+
+			# repetition penalty from CTRL paper (https://arxiv.org/abs/1909.05858)
+			if repetition_penalty != 1.0:
+				self.enforce_repetition_penalty_(next_token_logits, batch_size, 1, input_ids, repetition_penalty)
+
+			if no_repeat_ngram_size > 0:
+				# calculate a list of banned tokens to prevent repetitively generating the same ngrams
+				# from fairseq: https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345
+				banned_tokens = calc_banned_ngram_tokens(input_ids, batch_size, no_repeat_ngram_size, cur_len)
+				for batch_idx in range(batch_size):
+					next_token_logits[batch_idx, banned_tokens[batch_idx]] = -float("inf")
+
+			if bad_words_ids is not None:
+				# calculate a list of banned tokens according to bad words
+				banned_tokens = calc_banned_bad_words_ids(input_ids, bad_words_ids)
+
+				for batch_idx in range(batch_size):
+					next_token_logits[batch_idx, banned_tokens[batch_idx]] = -float("inf")
+
+			# set eos token prob to zero if min_length is not reached
+			if eos_token_id is not None and cur_len < min_length:
+				next_token_logits[:, eos_token_id] = -float("inf")
+
+			if do_sample:
+				# Temperature (higher temperature => more likely to sample low probability tokens)
+				if temperature != 1.0:
+					next_token_logits = next_token_logits / temperature
+				# Top-p/top-k filtering
+				next_token_logits = top_k_top_p_filtering(next_token_logits, top_k=top_k, top_p=top_p)
+				# Sample
+				probs = F.softmax(next_token_logits, dim=-1)
+				next_token = torch.multinomial(probs, num_samples=1).squeeze(1)
+			else:
+				# Greedy decoding
+				next_token = torch.argmax(next_token_logits, dim=-1)
+
+			# update generations and finished sentences
+			if eos_token_id is not None:
+				# pad finished sentences if eos_token_id exist
+				tokens_to_add = next_token * unfinished_sents + (pad_token_id) * (1 - unfinished_sents)
+			else:
+				tokens_to_add = next_token
+
+			input_ids = torch.cat([input_ids, tokens_to_add.unsqueeze(-1)], dim=-1)
+
+			if eos_token_id is not None:
+				eos_in_sents = tokens_to_add == eos_token_id
+				# if sentence is unfinished and the token to add is eos, sent_lengths is filled with current length
+				is_sents_unfinished_and_token_to_add_is_eos = unfinished_sents.mul(eos_in_sents.long()).bool()
+				sent_lengths.masked_fill_(is_sents_unfinished_and_token_to_add_is_eos, cur_len + 1)
+				# unfinished_sents is set to zero if eos in sentence
+				unfinished_sents.mul_((~eos_in_sents).long())
+
+			# stop when there is a </s> in each sentence, or if we exceed the maximul length
+			if unfinished_sents.max() == 0:
+				break
+
+			# extend attention_mask for new generated input if only decoder
+			if self.config.is_encoder_decoder is False:
+				attention_mask = torch.cat(
+					[attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
+				)
+
+			cur_len = cur_len + 1
+
+		# if there are different sentences lengths in the batch, some batches have to be padded
+		if sent_lengths.min().item() != sent_lengths.max().item():
+			assert pad_token_id is not None, "`Pad_token_id` has to be defined if batches have different lengths"
+			# finished sents are filled with pad_token
+			decoded = input_ids.new(batch_size, sent_lengths.max().item()).fill_(pad_token_id)
+		else:
+			decoded = input_ids
+
+		for hypo_idx, hypo in enumerate(input_ids):
+			decoded[hypo_idx, : sent_lengths[hypo_idx]] = hypo[: sent_lengths[hypo_idx]]
+
+		return decoded
+
+	def _generate_beam_search(
+		self,
+		input_ids,
+		cur_len,
+		max_length,
+		min_length,
+		do_sample,
+		early_stopping,
+		temperature,
+		top_k,
+		top_p,
+		repetition_penalty,
+		no_repeat_ngram_size,
+		bad_words_ids,
+		bos_token_id,
+		pad_token_id,
+		eos_token_id,
+		decoder_start_token_id,
+		batch_size,
+		num_return_sequences,
+		length_penalty,
+		num_beams,
+		vocab_size,
+		encoder_outputs,
+		attention_mask,
+		use_cache,
+	):
+		""" Generate sequences for each example with beam search.
+		"""
+
+		# generated hypotheses
+		generated_hyps = [
+			BeamHypotheses(num_beams, max_length, length_penalty, early_stopping=early_stopping)
+			for _ in range(batch_size)
+		]
+
+		# scores for each sentence in the beam
+		beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
+
+		# for greedy decoding it is made sure that only tokens of the first beam are considered to avoid sampling the exact same tokens three times
+		if do_sample is False:
+			beam_scores[:, 1:] = -1e9
+		beam_scores = beam_scores.view(-1)  # shape (batch_size * num_beams,)
+
+		# cache compute states
+		past = encoder_outputs  # defined for encoder-decoder models, None for decoder-only models
+
+		# done sentences
+		done = [False for _ in range(batch_size)]
+
+		while cur_len < max_length:
+			model_inputs = self.prepare_inputs_for_generation(
+				input_ids, past=past, attention_mask=attention_mask, use_cache=use_cache
+			)
+			outputs = self(**model_inputs)  # (batch_size * num_beams, cur_len, vocab_size)
+			next_token_logits = outputs[0][:, -1, :]  # (batch_size * num_beams, vocab_size)
+
+			# if model has past, then set the past variable to speed up decoding
+			if self._use_cache(outputs, use_cache):
+				past = outputs[1]
+
+			# repetition penalty (from CTRL paper https://arxiv.org/abs/1909.05858)
+			if repetition_penalty != 1.0:
+				self.enforce_repetition_penalty_(
+					next_token_logits, batch_size, num_beams, input_ids, repetition_penalty,
+				)
+
+			if temperature != 1.0:
+				next_token_logits = next_token_logits / temperature
+
+			scores = F.log_softmax(next_token_logits, dim=-1)  # (batch_size * num_beams, vocab_size)
+			if self.config.is_encoder_decoder and do_sample is False:
+				# TODO (PVP) still a bit hacky here - there might be a better solutino
+				scores = self.prepare_scores_for_generation(scores, cur_len=cur_len, max_length=max_length)
+
+			# set eos token prob to zero if min_length is not reached
+			if eos_token_id is not None and cur_len < min_length:
+				scores[:, eos_token_id] = -float("inf")
+
+			if no_repeat_ngram_size > 0:
+				# calculate a list of banned tokens to prevent repetitively generating the same ngrams
+				num_batch_hypotheses = batch_size * num_beams
+				# from fairseq: https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345
+				banned_batch_tokens = calc_banned_ngram_tokens(
+					input_ids, num_batch_hypotheses, no_repeat_ngram_size, cur_len
+				)
+				for i, banned_tokens in enumerate(banned_batch_tokens):
+					scores[i, banned_tokens] = -float("inf")
+
+			if bad_words_ids is not None:
+				# calculate a list of banned tokens according to bad words
+				banned_tokens = calc_banned_bad_words_ids(input_ids, bad_words_ids)
+
+				for i, banned_tokens in enumerate(banned_tokens):
+					scores[i, banned_tokens] = -float("inf")
+
+			assert scores.shape == (batch_size * num_beams, vocab_size), "Shapes of scores: {} != {}".format(
+				scores.shape, (batch_size * num_beams, vocab_size)
+			)
+
+			if do_sample:
+				_scores = scores + beam_scores[:, None].expand_as(scores)  # (batch_size * num_beams, vocab_size)
+				# Top-p/top-k filtering
+				_scores = top_k_top_p_filtering(
+					_scores, top_k=top_k, top_p=top_p, min_tokens_to_keep=2
+				)  # (batch_size * num_beams, vocab_size)
+				# re-organize to group the beam together to sample from all beam_idxs
+				_scores = _scores.contiguous().view(
+					batch_size, num_beams * vocab_size
+				)  # (batch_size, num_beams * vocab_size)
+
+				# Sample 2 next tokens for each beam (so we have some spare tokens and match output of greedy beam search)
+				probs = F.softmax(_scores, dim=-1)
+				next_tokens = torch.multinomial(probs, num_samples=2 * num_beams)  # (batch_size, num_beams * 2)
+				# Compute next scores
+				next_scores = torch.gather(_scores, -1, next_tokens)  # (batch_size, num_beams * 2)
+				# sort the sampled vector to make sure that the first num_beams samples are the best
+				next_scores, next_scores_indices = torch.sort(next_scores, descending=True, dim=1)
+				next_tokens = torch.gather(next_tokens, -1, next_scores_indices)  # (batch_size, num_beams * 2)
+
+			else:
+				next_scores = scores + beam_scores[:, None].expand_as(scores)  # (batch_size * num_beams, vocab_size)
+
+				# re-organize to group the beam together (we are keeping top hypothesis accross beams)
+				next_scores = next_scores.view(
+					batch_size, num_beams * vocab_size
+				)  # (batch_size, num_beams * vocab_size)
+
+				next_scores, next_tokens = torch.topk(next_scores, 2 * num_beams, dim=1, largest=True, sorted=True)
+
+			assert next_scores.size() == next_tokens.size() == (batch_size, 2 * num_beams)
+
+			# next batch beam content
+			next_batch_beam = []
+
+			# for each sentence
+			for batch_idx in range(batch_size):
+
+				# if we are done with this sentence
+				if done[batch_idx]:
+					assert (
+						len(generated_hyps[batch_idx]) >= num_beams
+					), "Batch can only be done if at least {} beams have been generated".format(num_beams)
+					assert (
+						eos_token_id is not None and pad_token_id is not None
+					), "generated beams >= num_beams -> eos_token_id and pad_token have to be defined"
+					next_batch_beam.extend([(0, pad_token_id, 0)] * num_beams)  # pad the batch
+					continue
+
+				# next sentence beam content
+				next_sent_beam = []
+
+				# next tokens for this sentence
+				for beam_token_rank, (beam_token_id, beam_token_score) in enumerate(
+					zip(next_tokens[batch_idx], next_scores[batch_idx])
+				):
+					# get beam and token IDs
+					beam_id = beam_token_id // vocab_size
+					token_id = beam_token_id % vocab_size
+
+					effective_beam_id = batch_idx * num_beams + beam_id
+					# add to generated hypotheses if end of sentence or last iteration
+					if (eos_token_id is not None) and (token_id.item() == eos_token_id):
+						# if beam_token does not belong to top num_beams tokens, it should not be added
+						is_beam_token_worse_than_top_num_beams = beam_token_rank >= num_beams
+						if is_beam_token_worse_than_top_num_beams:
+							continue
+						generated_hyps[batch_idx].add(
+							input_ids[effective_beam_id].clone(), beam_token_score.item(),
+						)
+					else:
+						# add next predicted token if it is not eos_token
+						next_sent_beam.append((beam_token_score, token_id, effective_beam_id))
+
+					# the beam for next step is full
+					if len(next_sent_beam) == num_beams:
+						break
+
+				# Check if were done so that we can save a pad step if all(done)
+				done[batch_idx] = done[batch_idx] or generated_hyps[batch_idx].is_done(
+					next_scores[batch_idx].max().item(), cur_len=cur_len
+				)
+
+				# update next beam content
+				assert len(next_sent_beam) == num_beams, "Beam should always be full"
+				next_batch_beam.extend(next_sent_beam)
+				assert len(next_batch_beam) == num_beams * (batch_idx + 1)
+
+			# stop when we are done with each sentence
+			if all(done):
+				break
+
+			# sanity check / prepare next batch
+			assert len(next_batch_beam) == batch_size * num_beams
+			beam_scores = beam_scores.new([x[0] for x in next_batch_beam])
+			beam_tokens = input_ids.new([x[1] for x in next_batch_beam])
+			beam_idx = input_ids.new([x[2] for x in next_batch_beam])
+
+			# re-order batch
+			input_ids = input_ids[beam_idx, :]
+			input_ids = torch.cat([input_ids, beam_tokens.unsqueeze(1)], dim=-1)
+			# re-order internal states
+			if past is not None:
+				past = self._reorder_cache(past, beam_idx)
+
+			# extend attention_mask for new generated input if only decoder
+			if self.config.is_encoder_decoder is False:
+				attention_mask = torch.cat(
+					[attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
+				)
+
+			# update current length
+			cur_len = cur_len + 1
+
+		# finalize all open beam hypotheses and end to generated hypotheses
+		for batch_idx in range(batch_size):
+			if done[batch_idx]:
+				continue
+
+			# test that beam scores match previously calculated scores if not eos and batch_idx not done
+			if eos_token_id is not None and all(
+				(token_id % vocab_size).item() is not eos_token_id for token_id in next_tokens[batch_idx]
+			):
+				assert torch.all(
+					next_scores[batch_idx, :num_beams] == beam_scores.view(batch_size, num_beams)[batch_idx]
+				), "If batch_idx is not done, final next scores: {} have to equal to accumulated beam_scores: {}".format(
+					next_scores[:, :num_beams][batch_idx], beam_scores.view(batch_size, num_beams)[batch_idx],
+				)
+
+			# need to add best num_beams hypotheses to generated hyps
+			for beam_id in range(num_beams):
+				effective_beam_id = batch_idx * num_beams + beam_id
+				final_score = beam_scores[effective_beam_id].item()
+				final_tokens = input_ids[effective_beam_id]
+				generated_hyps[batch_idx].add(final_tokens, final_score)
+
+		# depending on whether greedy generation is wanted or not define different output_batch_size and output_num_return_sequences_per_batch
+		output_batch_size = batch_size if do_sample else batch_size * num_return_sequences
+		output_num_return_sequences_per_batch = 1 if do_sample else num_return_sequences
+
+		# select the best hypotheses
+		sent_lengths = input_ids.new(output_batch_size)
+		best = []
+
+		# retrieve best hypotheses
+		for i, hypotheses in enumerate(generated_hyps):
+			sorted_hyps = sorted(hypotheses.beams, key=lambda x: x[0])
+			for j in range(output_num_return_sequences_per_batch):
+				effective_batch_idx = output_num_return_sequences_per_batch * i + j
+				best_hyp = sorted_hyps.pop()[1]
+				sent_lengths[effective_batch_idx] = len(best_hyp)
+				best.append(best_hyp)
+
+		# shorter batches are filled with pad_token
+		if sent_lengths.min().item() != sent_lengths.max().item():
+			assert pad_token_id is not None, "`Pad_token_id` has to be defined"
+			sent_max_len = min(sent_lengths.max().item() + 1, max_length)
+			decoded = input_ids.new(output_batch_size, sent_max_len).fill_(pad_token_id)
+
+			# fill with hypothesis and eos_token_id if necessary
+			for i, hypo in enumerate(best):
+				decoded[i, : sent_lengths[i]] = hypo
+				if sent_lengths[i] < max_length:
+					decoded[i, sent_lengths[i]] = eos_token_id
+		else:
+			# none of the hypotheses have an eos_token
+			assert (len(hypo) == max_length for hypo in best)
+			decoded = torch.stack(best).type(torch.long).to(next(self.parameters()).device)
+
+		return decoded
+
+	# force one of token_ids to be generated by setting prob of all other tokens to 0.
+	def _force_token_ids_generation(self, scores, token_ids) -> None:
+		if isinstance(token_ids, int):
+			token_ids = [token_ids]
+		all_but_token_ids_mask = torch.tensor(
+			[x for x in range(self.config.vocab_size) if x not in token_ids],
+			dtype=torch.long,
+			device=next(self.parameters()).device,
+		)
+		assert len(scores.shape) == 2, "scores should be of rank 2 with shape: [batch_size, vocab_size]"
+		scores[:, all_but_token_ids_mask] = -float("inf")
+
+	@staticmethod
+	def _reorder_cache(past: Tuple, beam_idx: Tensor) -> Tuple[Tensor]:
+		return tuple(layer_past.index_select(1, beam_idx) for layer_past in past)
+
+
+def calc_banned_ngram_tokens(prev_input_ids: Tensor, num_hypos: int, no_repeat_ngram_size: int, cur_len: int) -> None:
+	"""Copied from fairseq for no_repeat_ngram in beam_search"""
+	if cur_len + 1 < no_repeat_ngram_size:
+		# return no banned tokens if we haven't generated no_repeat_ngram_size tokens yet
+		return [[] for _ in range(num_hypos)]
+	generated_ngrams = [{} for _ in range(num_hypos)]
+	for idx in range(num_hypos):
+		gen_tokens = prev_input_ids[idx].tolist()
+		generated_ngram = generated_ngrams[idx]
+		for ngram in zip(*[gen_tokens[i:] for i in range(no_repeat_ngram_size)]):
+			prev_ngram_tuple = tuple(ngram[:-1])
+			generated_ngram[prev_ngram_tuple] = generated_ngram.get(prev_ngram_tuple, []) + [ngram[-1]]
+
+	def _get_generated_ngrams(hypo_idx):
+		# Before decoding the next token, prevent decoding of ngrams that have already appeared
+		start_idx = cur_len + 1 - no_repeat_ngram_size
+		ngram_idx = tuple(prev_input_ids[hypo_idx, start_idx:cur_len].tolist())
+		return generated_ngrams[hypo_idx].get(ngram_idx, [])
+
+	banned_tokens = [_get_generated_ngrams(hypo_idx) for hypo_idx in range(num_hypos)]
+	return banned_tokens
+
+
+def calc_banned_bad_words_ids(prev_input_ids, bad_words_ids):
+	banned_tokens = []
+
+	def _tokens_match(prev_tokens, tokens):
+		if len(tokens) == 0:
+			# if bad word tokens is just one token always ban it
+			return True
+		if len(tokens) > len(prev_input_ids):
+			# if bad word tokens are longer then prev input_ids they can't be equal
+			return False
+
+		if prev_tokens[-len(tokens) :] == tokens:
+			# if tokens match
+			return True
+		else:
+			return False
+
+	for prev_input_ids_slice in prev_input_ids:
+		banned_tokens_slice = []
+
+		for banned_token_seq in bad_words_ids:
+			assert len(banned_token_seq) > 0, "Banned words token sequences {} cannot have an empty list".format(
+				bad_words_ids
+			)
+
+			if _tokens_match(prev_input_ids_slice.tolist(), banned_token_seq[:-1]) is False:
+				# if tokens do not match continue
+				continue
+
+			banned_tokens_slice.append(banned_token_seq[-1])
+
+		banned_tokens.append(banned_tokens_slice)
+
+	return banned_tokens
+
+
+def top_k_top_p_filtering(logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1):
+	""" Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
+		Args:
+			logits: logits distribution shape (batch size, vocabulary size)
+			if top_k > 0: keep only top k tokens with highest probability (top-k filtering).
+			if top_p < 1.0: keep the top tokens with cumulative probability >= top_p (nucleus filtering).
+				Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
+			Make sure we keep at least min_tokens_to_keep per batch example in the output
+		From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
+	"""
+	if top_k > 0:
+		top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
+		# Remove all tokens with a probability less than the last token of the top-k
+		indices_to_remove = logits < torch.topk(logits, top_k)[0][..., -1, None]
+		logits[indices_to_remove] = filter_value
+
+	if top_p < 1.0:
+		sorted_logits, sorted_indices = torch.sort(logits, descending=True)
+		cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)
+
+		# Remove tokens with cumulative probability above the threshold (token with 0 are kept)
+		sorted_indices_to_remove = cumulative_probs > top_p
+		if min_tokens_to_keep > 1:
+			# Keep at least min_tokens_to_keep (set to min_tokens_to_keep-1 because we add the first one below)
+			sorted_indices_to_remove[..., :min_tokens_to_keep] = 0
+		# Shift the indices to the right to keep also the first token above the threshold
+		sorted_indices_to_remove[..., 1:] = sorted_indices_to_remove[..., :-1].clone()
+		sorted_indices_to_remove[..., 0] = 0
+
+		# scatter sorted tensors to original indexing
+		indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
+		logits[indices_to_remove] = filter_value
+	return logits
+
+
+class BeamHypotheses(object):
+	def __init__(self, num_beams, max_length, length_penalty, early_stopping):
+		"""
+		Initialize n-best list of hypotheses.
+		"""
+		self.max_length = max_length - 1  # ignoring bos_token
+		self.length_penalty = length_penalty
+		self.early_stopping = early_stopping
+		self.num_beams = num_beams
+		self.beams = []
+		self.worst_score = 1e9
+
+	def __len__(self):
+		"""
+		Number of hypotheses in the list.
+		"""
+		return len(self.beams)
+
+	def add(self, hyp, sum_logprobs):
+		"""
+		Add a new hypothesis to the list.
+		"""
+		score = sum_logprobs / len(hyp) ** self.length_penalty
+		if len(self) < self.num_beams or score > self.worst_score:
+			self.beams.append((score, hyp))
+			if len(self) > self.num_beams:
+				sorted_scores = sorted([(s, idx) for idx, (s, _) in enumerate(self.beams)])
+				del self.beams[sorted_scores[0][1]]
+				self.worst_score = sorted_scores[1][0]
+			else:
+				self.worst_score = min(score, self.worst_score)
+
+	def is_done(self, best_sum_logprobs, cur_len=None):
+		"""
+		If there are enough hypotheses and that none of the hypotheses being generated
+		can become better than the worst one in the heap, then we are done with this sentence.
+		"""
+
+		if len(self) < self.num_beams:
+			return False
+		elif self.early_stopping:
+			return True
+		else:
+			if cur_len is None:
+				cur_len = self.max_length
+			cur_score = best_sum_logprobs / cur_len ** self.length_penalty
+			ret = self.worst_score >= cur_score
+			return ret
+
+
+class Conv1D(nn.Module):
+	def __init__(self, nf, nx):
+		""" Conv1D layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2)
+			Basically works like a Linear layer but the weights are transposed
+		"""
+		super().__init__()
+		self.nf = nf
+		w = torch.empty(nx, nf)
+		nn.init.normal_(w, std=0.02)
+		self.weight = nn.Parameter(w)
+		self.bias = nn.Parameter(torch.zeros(nf))
+
+	def forward(self, x):
+		size_out = x.size()[:-1] + (self.nf,)
+		x = torch.addmm(self.bias, x.view(-1, x.size(-1)), self.weight)
+		x = x.view(*size_out)
+		return x
+
+
+class PoolerStartLogits(nn.Module):
+	""" Compute SQuAD start_logits from sequence hidden states. """
+
+	def __init__(self, config):
+		super().__init__()
+		self.dense = nn.Linear(config.hidden_size, 1)
+
+	def forward(self, hidden_states, p_mask=None):
+		""" Args:
+			**p_mask**: (`optional`) ``torch.FloatTensor`` of shape `(batch_size, seq_len)`
+				invalid position mask such as query and special symbols (PAD, SEP, CLS)
+				1.0 means token should be masked.
+		"""
+		x = self.dense(hidden_states).squeeze(-1)
+
+		if p_mask is not None:
+			if next(self.parameters()).dtype == torch.float16:
+				x = x * (1 - p_mask) - 65500 * p_mask
+			else:
+				x = x * (1 - p_mask) - 1e30 * p_mask
+
+		return x
+
+
+class PoolerEndLogits(nn.Module):
+	""" Compute SQuAD end_logits from sequence hidden states and start token hidden state.
+	"""
+
+	def __init__(self, config):
+		super().__init__()
+		self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size)
+		self.activation = nn.Tanh()
+		self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
+		self.dense_1 = nn.Linear(config.hidden_size, 1)
+
+	def forward(self, hidden_states, start_states=None, start_positions=None, p_mask=None):
+		""" Args:
+			One of ``start_states``, ``start_positions`` should be not None.
+			If both are set, ``start_positions`` overrides ``start_states``.
+
+			**start_states**: ``torch.LongTensor`` of shape identical to hidden_states
+				hidden states of the first tokens for the labeled span.
+			**start_positions**: ``torch.LongTensor`` of shape ``(batch_size,)``
+				position of the first token for the labeled span:
+			**p_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, seq_len)``
+				Mask of invalid position such as query and special symbols (PAD, SEP, CLS)
+				1.0 means token should be masked.
+		"""
+		assert (
+			start_states is not None or start_positions is not None
+		), "One of start_states, start_positions should be not None"
+		if start_positions is not None:
+			slen, hsz = hidden_states.shape[-2:]
+			start_positions = start_positions[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
+			start_states = hidden_states.gather(-2, start_positions)  # shape (bsz, 1, hsz)
+			start_states = start_states.expand(-1, slen, -1)  # shape (bsz, slen, hsz)
+
+		x = self.dense_0(torch.cat([hidden_states, start_states], dim=-1))
+		x = self.activation(x)
+		x = self.LayerNorm(x)
+		x = self.dense_1(x).squeeze(-1)
+
+		if p_mask is not None:
+			if next(self.parameters()).dtype == torch.float16:
+				x = x * (1 - p_mask) - 65500 * p_mask
+			else:
+				x = x * (1 - p_mask) - 1e30 * p_mask
+
+		return x
+
+
+class PoolerAnswerClass(nn.Module):
+	""" Compute SQuAD 2.0 answer class from classification and start tokens hidden states. """
+
+	def __init__(self, config):
+		super().__init__()
+		self.dense_0 = nn.Linear(config.hidden_size * 2, config.hidden_size)
+		self.activation = nn.Tanh()
+		self.dense_1 = nn.Linear(config.hidden_size, 1, bias=False)
+
+	def forward(self, hidden_states, start_states=None, start_positions=None, cls_index=None):
+		"""
+		Args:
+			One of ``start_states``, ``start_positions`` should be not None.
+			If both are set, ``start_positions`` overrides ``start_states``.
+
+			**start_states**: ``torch.LongTensor`` of shape identical to ``hidden_states``.
+				hidden states of the first tokens for the labeled span.
+			**start_positions**: ``torch.LongTensor`` of shape ``(batch_size,)``
+				position of the first token for the labeled span.
+			**cls_index**: torch.LongTensor of shape ``(batch_size,)``
+				position of the CLS token. If None, take the last token.
+
+			note(Original repo):
+				no dependency on end_feature so that we can obtain one single `cls_logits`
+				for each sample
+		"""
+		hsz = hidden_states.shape[-1]
+		assert (
+			start_states is not None or start_positions is not None
+		), "One of start_states, start_positions should be not None"
+		if start_positions is not None:
+			start_positions = start_positions[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
+			start_states = hidden_states.gather(-2, start_positions).squeeze(-2)  # shape (bsz, hsz)
+
+		if cls_index is not None:
+			cls_index = cls_index[:, None, None].expand(-1, -1, hsz)  # shape (bsz, 1, hsz)
+			cls_token_state = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, hsz)
+		else:
+			cls_token_state = hidden_states[:, -1, :]  # shape (bsz, hsz)
+
+		x = self.dense_0(torch.cat([start_states, cls_token_state], dim=-1))
+		x = self.activation(x)
+		x = self.dense_1(x).squeeze(-1)
+
+		return x
+
+
+class SQuADHead(nn.Module):
+	r""" A SQuAD head inspired by XLNet.
+
+	Parameters:
+		config (:class:`~transformers.XLNetConfig`): Model configuration class with all the parameters of the model.
+
+	Inputs:
+		**hidden_states**: ``torch.FloatTensor`` of shape ``(batch_size, seq_len, hidden_size)``
+			hidden states of sequence tokens
+		**start_positions**: ``torch.LongTensor`` of shape ``(batch_size,)``
+			position of the first token for the labeled span.
+		**end_positions**: ``torch.LongTensor`` of shape ``(batch_size,)``
+			position of the last token for the labeled span.
+		**cls_index**: torch.LongTensor of shape ``(batch_size,)``
+			position of the CLS token. If None, take the last token.
+		**is_impossible**: ``torch.LongTensor`` of shape ``(batch_size,)``
+			Whether the question has a possible answer in the paragraph or not.
+		**p_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, seq_len)``
+			Mask of invalid position such as query and special symbols (PAD, SEP, CLS)
+			1.0 means token should be masked.
+
+	Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
+		**loss**: (`optional`, returned if both ``start_positions`` and ``end_positions`` are provided) ``torch.FloatTensor`` of shape ``(1,)``:
+			Classification loss as the sum of start token, end token (and is_impossible if provided) classification losses.
+		**start_top_log_probs**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
+			``torch.FloatTensor`` of shape ``(batch_size, config.start_n_top)``
+			Log probabilities for the top config.start_n_top start token possibilities (beam-search).
+		**start_top_index**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
+			``torch.LongTensor`` of shape ``(batch_size, config.start_n_top)``
+			Indices for the top config.start_n_top start token possibilities (beam-search).
+		**end_top_log_probs**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
+			``torch.FloatTensor`` of shape ``(batch_size, config.start_n_top * config.end_n_top)``
+			Log probabilities for the top ``config.start_n_top * config.end_n_top`` end token possibilities (beam-search).
+		**end_top_index**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
+			``torch.LongTensor`` of shape ``(batch_size, config.start_n_top * config.end_n_top)``
+			Indices for the top ``config.start_n_top * config.end_n_top`` end token possibilities (beam-search).
+		**cls_logits**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
+			``torch.FloatTensor`` of shape ``(batch_size,)``
+			Log probabilities for the ``is_impossible`` label of the answers.
+	"""
+
+	def __init__(self, config):
+		super().__init__()
+		self.start_n_top = config.start_n_top
+		self.end_n_top = config.end_n_top
+
+		self.start_logits = PoolerStartLogits(config)
+		self.end_logits = PoolerEndLogits(config)
+		self.answer_class = PoolerAnswerClass(config)
+
+	def forward(
+		self, hidden_states, start_positions=None, end_positions=None, cls_index=None, is_impossible=None, p_mask=None,
+	):
+		outputs = ()
+
+		start_logits = self.start_logits(hidden_states, p_mask=p_mask)
+
+		if start_positions is not None and end_positions is not None:
+			# If we are on multi-GPU, let's remove the dimension added by batch splitting
+			for x in (start_positions, end_positions, cls_index, is_impossible):
+				if x is not None and x.dim() > 1:
+					x.squeeze_(-1)
+
+			# during training, compute the end logits based on the ground truth of the start position
+			end_logits = self.end_logits(hidden_states, start_positions=start_positions, p_mask=p_mask)
+
+			loss_fct = CrossEntropyLoss()
+			start_loss = loss_fct(start_logits, start_positions)
+			end_loss = loss_fct(end_logits, end_positions)
+			total_loss = (start_loss + end_loss) / 2
+
+			if cls_index is not None and is_impossible is not None:
+				# Predict answerability from the representation of CLS and START
+				cls_logits = self.answer_class(hidden_states, start_positions=start_positions, cls_index=cls_index)
+				loss_fct_cls = nn.BCEWithLogitsLoss()
+				cls_loss = loss_fct_cls(cls_logits, is_impossible)
+
+				# note(zhiliny): by default multiply the loss by 0.5 so that the scale is comparable to start_loss and end_loss
+				total_loss += cls_loss * 0.5
+
+			outputs = (total_loss,) + outputs
+
+		else:
+			# during inference, compute the end logits based on beam search
+			bsz, slen, hsz = hidden_states.size()
+			start_log_probs = F.softmax(start_logits, dim=-1)  # shape (bsz, slen)
+
+			start_top_log_probs, start_top_index = torch.topk(
+				start_log_probs, self.start_n_top, dim=-1
+			)  # shape (bsz, start_n_top)
+			start_top_index_exp = start_top_index.unsqueeze(-1).expand(-1, -1, hsz)  # shape (bsz, start_n_top, hsz)
+			start_states = torch.gather(hidden_states, -2, start_top_index_exp)  # shape (bsz, start_n_top, hsz)
+			start_states = start_states.unsqueeze(1).expand(-1, slen, -1, -1)  # shape (bsz, slen, start_n_top, hsz)
+
+			hidden_states_expanded = hidden_states.unsqueeze(2).expand_as(
+				start_states
+			)  # shape (bsz, slen, start_n_top, hsz)
+			p_mask = p_mask.unsqueeze(-1) if p_mask is not None else None
+			end_logits = self.end_logits(hidden_states_expanded, start_states=start_states, p_mask=p_mask)
+			end_log_probs = F.softmax(end_logits, dim=1)  # shape (bsz, slen, start_n_top)
+
+			end_top_log_probs, end_top_index = torch.topk(
+				end_log_probs, self.end_n_top, dim=1
+			)  # shape (bsz, end_n_top, start_n_top)
+			end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top)
+			end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top)
+
+			start_states = torch.einsum("blh,bl->bh", hidden_states, start_log_probs)
+			cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index)
+
+			outputs = (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits,) + outputs
+
+		# return start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits
+		# or (if labels are provided) (total_loss,)
+		return outputs
+
+
+class SequenceSummary(nn.Module):
+	r""" Compute a single vector summary of a sequence hidden states according to various possibilities:
+		Args of the config class:
+			summary_type:
+				- 'last' => [default] take the last token hidden state (like XLNet)
+				- 'first' => take the first token hidden state (like Bert)
+				- 'mean' => take the mean of all tokens hidden states
+				- 'cls_index' => supply a Tensor of classification token position (GPT/GPT-2)
+				- 'attn' => Not implemented now, use multi-head attention
+			summary_use_proj: Add a projection after the vector extraction
+			summary_proj_to_labels: If True, the projection outputs to config.num_labels classes (otherwise to hidden_size). Default: False.
+			summary_activation: 'tanh' or another string => add an activation to the output, Other => no activation. Default
+			summary_first_dropout: Add a dropout before the projection and activation
+			summary_last_dropout: Add a dropout after the projection and activation
+	"""
+
+	def __init__(self, config: PretrainedConfig):
+		super().__init__()
+
+		self.summary_type = getattr(config, "summary_type", "last")
+		if self.summary_type == "attn":
+			# We should use a standard multi-head attention module with absolute positional embedding for that.
+			# Cf. https://github.com/zihangdai/xlnet/blob/master/modeling.py#L253-L276
+			# We can probably just use the multi-head attention module of PyTorch >=1.1.0
+			raise NotImplementedError
+
+		self.summary = Identity()
+		if hasattr(config, "summary_use_proj") and config.summary_use_proj:
+			if hasattr(config, "summary_proj_to_labels") and config.summary_proj_to_labels and config.num_labels > 0:
+				num_classes = config.num_labels
+			else:
+				num_classes = config.hidden_size
+			self.summary = nn.Linear(config.hidden_size, num_classes)
+
+		activation_string = getattr(config, "summary_activation", None)
+		self.activation: Callable = (get_activation(activation_string) if activation_string else Identity())
+
+		self.first_dropout = Identity()
+		if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0:
+			self.first_dropout = nn.Dropout(config.summary_first_dropout)
+
+		self.last_dropout = Identity()
+		if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0:
+			self.last_dropout = nn.Dropout(config.summary_last_dropout)
+
+	def forward(self, hidden_states, cls_index=None):
+		""" hidden_states: float Tensor in shape [bsz, ..., seq_len, hidden_size], the hidden-states of the last layer.
+			cls_index: [optional] position of the classification token if summary_type == 'cls_index',
+				shape (bsz,) or more generally (bsz, ...) where ... are optional leading dimensions of hidden_states.
+				if summary_type == 'cls_index' and cls_index is None:
+					we take the last token of the sequence as classification token
+		"""
+		if self.summary_type == "last":
+			output = hidden_states[:, -1]
+		elif self.summary_type == "first":
+			output = hidden_states[:, 0]
+		elif self.summary_type == "mean":
+			output = hidden_states.mean(dim=1)
+		elif self.summary_type == "cls_index":
+			if cls_index is None:
+				cls_index = torch.full_like(hidden_states[..., :1, :], hidden_states.shape[-2] - 1, dtype=torch.long,)
+			else:
+				cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)
+				cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))
+			# shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states
+			output = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, XX, hidden_size)
+		elif self.summary_type == "attn":
+			raise NotImplementedError
+
+		output = self.first_dropout(output)
+		output = self.summary(output)
+		output = self.activation(output)
+		output = self.last_dropout(output)
+
+		return output
+
+
+def create_position_ids_from_input_ids(input_ids, padding_idx):
+	""" Replace non-padding symbols with their position numbers. Position numbers begin at
+	padding_idx+1. Padding symbols are ignored. This is modified from fairseq's
+	`utils.make_positions`.
+
+	:param torch.Tensor x:
+	:return torch.Tensor:
+	"""
+	# The series of casts and type-conversions here are carefully balanced to both work with ONNX export and XLA.
+	mask = input_ids.ne(padding_idx).int()
+	incremental_indices = torch.cumsum(mask, dim=1).type_as(mask) * mask
+	return incremental_indices.long() + padding_idx
+
+
+def prune_linear_layer(layer, index, dim=0):
+	""" Prune a linear layer (a model parameters) to keep only entries in index.
+		Return the pruned layer as a new layer with requires_grad=True.
+		Used to remove heads.
+	"""
+	index = index.to(layer.weight.device)
+	W = layer.weight.index_select(dim, index).clone().detach()
+	if layer.bias is not None:
+		if dim == 1:
+			b = layer.bias.clone().detach()
+		else:
+			b = layer.bias[index].clone().detach()
+	new_size = list(layer.weight.size())
+	new_size[dim] = len(index)
+	new_layer = nn.Linear(new_size[1], new_size[0], bias=layer.bias is not None).to(layer.weight.device)
+	new_layer.weight.requires_grad = False
+	new_layer.weight.copy_(W.contiguous())
+	new_layer.weight.requires_grad = True
+	if layer.bias is not None:
+		new_layer.bias.requires_grad = False
+		new_layer.bias.copy_(b.contiguous())
+		new_layer.bias.requires_grad = True
+	return new_layer
+
+
+def prune_conv1d_layer(layer, index, dim=1):
+	""" Prune a Conv1D layer (a model parameters) to keep only entries in index.
+		A Conv1D work as a Linear layer (see e.g. BERT) but the weights are transposed.
+		Return the pruned layer as a new layer with requires_grad=True.
+		Used to remove heads.
+	"""
+	index = index.to(layer.weight.device)
+	W = layer.weight.index_select(dim, index).clone().detach()
+	if dim == 0:
+		b = layer.bias.clone().detach()
+	else:
+		b = layer.bias[index].clone().detach()
+	new_size = list(layer.weight.size())
+	new_size[dim] = len(index)
+	new_layer = Conv1D(new_size[1], new_size[0]).to(layer.weight.device)
+	new_layer.weight.requires_grad = False
+	new_layer.weight.copy_(W.contiguous())
+	new_layer.weight.requires_grad = True
+	new_layer.bias.requires_grad = False
+	new_layer.bias.copy_(b.contiguous())
+	new_layer.bias.requires_grad = True
+	return new_layer
+
+
+def prune_layer(layer, index, dim=None):
+	""" Prune a Conv1D or nn.Linear layer (a model parameters) to keep only entries in index.
+		Return the pruned layer as a new layer with requires_grad=True.
+		Used to remove heads.
+	"""
+	if isinstance(layer, nn.Linear):
+		return prune_linear_layer(layer, index, dim=0 if dim is None else dim)
+	elif isinstance(layer, Conv1D):
+		return prune_conv1d_layer(layer, index, dim=1 if dim is None else dim)
+	else:
+		raise ValueError("Can't prune layer of class {}".format(layer.__class__))

models/models.py

@@ -0,0 +1,263 @@
+from .roberta import RobertaForTokenClassification, RobertaModel
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss, MSELoss
+import math
+import torch.nn.functional as F
+
+from .modeling_bert import BertEmbeddings, BertLayerNorm, BertModel, BertPreTrainedModel, gelu
+from .configuration_roberta import RobertaConfig
+from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
+
+from transformers import GPT2Model
+from transformers import AutoModelWithLMHead, AutoTokenizer
+
+
+ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP = {
+	"roberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-pytorch_model.bin",
+	"roberta-large": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-pytorch_model.bin",
+	"roberta-large-mnli": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-mnli-pytorch_model.bin",
+	"distilroberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/distilroberta-base-pytorch_model.bin",
+	"roberta-base-openai-detector": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-openai-detector-pytorch_model.bin",
+	"roberta-large-openai-detector": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-openai-detector-pytorch_model.bin",
+}
+
+
+class Norm(nn.Module):
+	def __init__(self, d_model, eps = 1e-6):
+		super().__init__()
+	
+		self.size = d_model
+		# create two learnable parameters to calibrate normalisation
+		self.alpha = nn.Parameter(torch.ones(self.size))
+		self.bias = nn.Parameter(torch.zeros(self.size))
+		self.eps = eps
+	def forward(self, x):
+		norm = self.alpha * (x - x.mean(dim=-1, keepdim=True)) \
+		/ (x.std(dim=-1, keepdim=True) + self.eps) + self.bias
+		return norm
+
+
+class MultiHeadAttention(nn.Module):
+	def __init__(self, heads, d_model, dropout = 0.1):
+		super().__init__()
+		
+		self.d_model = d_model
+		self.d_k = d_model // heads
+		self.h = heads
+		
+		self.q_linear = nn.Linear(d_model, d_model)
+		self.v_linear = nn.Linear(d_model, d_model)
+		self.k_linear = nn.Linear(d_model, d_model)
+		self.dropout = nn.Dropout(dropout)
+		self.out = nn.Linear(d_model, d_model)
+	
+	def forward(self, q, k, v, mask=None):
+		
+		bs = q.size(0)
+				
+		k = self.k_linear(k).view(bs, -1, self.h, self.d_k)
+		q = self.q_linear(q).view(bs, -1, self.h, self.d_k)
+		v = self.v_linear(v).view(bs, -1, self.h, self.d_k)
+			   
+		k = k.transpose(1,2)
+		q = q.transpose(1,2)
+		v = v.transpose(1,2)
+
+		scores = self.attention(q, k, v, self.d_k, mask, self.dropout)
+		
+		concat = scores.transpose(1,2).contiguous()\
+		.view(bs, -1, self.d_model)
+		
+		output = self.out(concat)
+	
+		return output
+	
+	def attention(self, q, k, v, d_k, mask=None, dropout=None):
+	
+		scores = torch.matmul(q, k.transpose(-2, -1)) /  math.sqrt(d_k)
+		
+		if mask is not None:
+			mask = mask.unsqueeze(1)
+			scores = scores.masked_fill(mask == 0, -1e9)
+		
+		scores = F.softmax(scores, dim=-1)
+			
+		if dropout is not None:
+			scores = dropout(scores)
+			
+		output = torch.matmul(scores, v)
+		
+		return output
+
+
+class SeekerEncoder(BertPreTrainedModel):
+	config_class = RobertaConfig
+	pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
+	base_model_prefix = "roberta"
+
+	def __init__(self, config):
+		super().__init__(config)
+		self.num_labels = config.num_labels
+		self.roberta = RobertaModel(config)
+		self.init_weights()
+	
+	def get_input_embeddings(self):
+		return self.roberta.embeddings.word_embeddings
+
+	def set_input_embeddings(self, value):
+		self.roberta.embeddings.word_embeddings = value
+	
+
+class ResponderEncoder(BertPreTrainedModel):
+	config_class = RobertaConfig
+	pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
+	base_model_prefix = "roberta"
+
+	def __init__(self, config):
+		super().__init__(config)
+		self.roberta = RobertaModel(config)
+		self.init_weights()
+	
+	def get_input_embeddings(self):
+		return self.roberta.embeddings.word_embeddings
+
+	def set_input_embeddings(self, value):
+		self.roberta.embeddings.word_embeddings = value
+
+class BiEncoderAttentionWithRationaleClassification(nn.Module):
+
+	def __init__(self, hidden_dropout_prob=0.2, rationale_num_labels=2, empathy_num_labels=3, hidden_size=768, attn_heads = 1):
+		super().__init__()
+
+		self.dropout = nn.Dropout(hidden_dropout_prob)
+		self.rationale_classifier = nn.Linear(hidden_size, rationale_num_labels)
+		self.attn = MultiHeadAttention(attn_heads, hidden_size)
+		self.norm = Norm(hidden_size)
+		self.rationale_num_labels = rationale_num_labels
+		self.empathy_num_labels = empathy_num_labels
+		self.empathy_classifier = RobertaClassificationHead(hidden_size = 768)
+
+		self.apply(self._init_weights)
+
+		self.seeker_encoder = SeekerEncoder.from_pretrained(
+								"roberta-base", # Use the 12-layer BERT model, with an uncased vocab.
+								output_attentions = False, # Whether the model returns attentions weights.
+								output_hidden_states = False)
+
+		self.responder_encoder = ResponderEncoder.from_pretrained(
+								"roberta-base", # Use the 12-layer BERT model, with an uncased vocab.
+								output_attentions = False, # Whether the model returns attentions weights.
+								output_hidden_states = False)
+
+	
+	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
+			initializer_range=0.02
+			module.weight.data.normal_(mean=0.0, std=initializer_range)
+		elif isinstance(module, BertLayerNorm):
+			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_()
+
+
+	# @add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids_SP=None,
+		input_ids_RP=None,
+		attention_mask_SP=None,
+		attention_mask_RP=None,
+		token_type_ids_SP=None,
+		token_type_ids_RP=None,
+		position_ids_SP=None,
+		position_ids_RP=None,
+		head_mask_SP=None,
+		head_mask_RP=None,
+		inputs_embeds_SP=None,
+		inputs_embeds_RP=None,
+		empathy_labels=None,
+		rationale_labels=None,
+		lambda_EI=1,
+		lambda_RE=0.1
+	):
+		outputs_SP = self.seeker_encoder.roberta(
+			input_ids_SP,
+			attention_mask=attention_mask_SP,
+			token_type_ids=token_type_ids_SP,
+			position_ids=position_ids_SP,
+			head_mask=head_mask_SP,
+			inputs_embeds=inputs_embeds_SP,
+		)
+
+
+		outputs_RP = self.responder_encoder.roberta(
+			input_ids_RP,
+			attention_mask=attention_mask_RP,
+			token_type_ids=token_type_ids_RP,
+			position_ids=position_ids_RP,
+			head_mask=head_mask_RP,
+			inputs_embeds=inputs_embeds_RP,
+		)
+
+		sequence_output_SP = outputs_SP[0]
+		sequence_output_RP = outputs_RP[0]
+
+		sequence_output_RP = sequence_output_RP + self.dropout(self.attn(sequence_output_RP, sequence_output_SP, sequence_output_SP))
+
+		logits_empathy = self.empathy_classifier(sequence_output_RP[:, 0, :]) # (sequence_output_RP[:, 0, :]) #(torch.tanh(concat_tensor))
+		
+		sequence_output = self.dropout(sequence_output_RP)
+		logits_rationales = self.rationale_classifier(sequence_output)
+		outputs = (logits_empathy,logits_rationales) + outputs_RP[2:]
+
+		loss_rationales = 0.0
+		loss_empathy = 0.0
+
+		if rationale_labels is not None:
+			loss_fct = CrossEntropyLoss()
+			# Only keep active parts of the loss
+			if attention_mask_RP is not None:
+				active_loss = attention_mask_RP.view(-1) == 1
+				active_logits = logits_rationales.view(-1, self.rationale_num_labels)
+				active_labels = torch.where(
+					active_loss, rationale_labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(rationale_labels)
+				)
+				loss_rationales = loss_fct(active_logits, active_labels)
+			else:
+				loss_rationales = loss_fct(logits_rationales.view(-1, self.rationale_num_labels), rationale_labels.view(-1))
+
+		if empathy_labels is not None:
+			loss_fct = CrossEntropyLoss()
+			loss_empathy = loss_fct(logits_empathy.view(-1, self.empathy_num_labels), empathy_labels.view(-1))
+
+			loss = lambda_EI * loss_empathy + lambda_RE * loss_rationales
+
+			outputs = (loss, loss_empathy, loss_rationales) + outputs
+
+		return outputs  # (loss), (scores_empathy, scores_rationales), (hidden_states), (attentions)
+
+
+
+class RobertaClassificationHead(nn.Module):
+	"""Head for sentence-level classification tasks."""
+
+	def __init__(self, hidden_dropout_prob=0.1, hidden_size=768, empathy_num_labels=3):
+		super().__init__()
+
+		self.dense = nn.Linear(hidden_size, hidden_size)
+		self.dropout = nn.Dropout(hidden_dropout_prob)
+		self.out_proj = nn.Linear(hidden_size, empathy_num_labels)
+
+	def forward(self, features, **kwargs):
+		x = features[:, :]  # take <s> token (equiv. to [CLS])
+		x = self.dropout(x)
+		x = self.dense(x)
+		x = torch.relu(x)
+		x = self.dropout(x)
+		x = self.out_proj(x)
+		return x

models/roberta.py

@@ -0,0 +1,685 @@
+# coding=utf-8
+# Copyright 2018 The Google AI Language Team Authors and The HuggingFace Inc. team.
+# Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+# Licensed under the Apache License, Version 2.0 (the "License");
+# you may not use this file except in compliance with the License.
+# You may obtain a copy of the License at
+#
+#     http://www.apache.org/licenses/LICENSE-2.0
+#
+# Unless required by applicable law or agreed to in writing, software
+# distributed under the License is distributed on an "AS IS" BASIS,
+# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+# See the License for the specific language governing permissions and
+# limitations under the License.
+"""PyTorch RoBERTa model. """
+
+
+import logging
+
+import torch
+import torch.nn as nn
+from torch.nn import CrossEntropyLoss, MSELoss
+
+from .configuration_roberta import RobertaConfig
+from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
+from .modeling_bert import BertEmbeddings, BertLayerNorm, BertModel, BertPreTrainedModel, gelu
+from .modeling_utils import create_position_ids_from_input_ids
+
+
+logger = logging.getLogger(__name__)
+
+ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP = {
+	"roberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-pytorch_model.bin",
+	"roberta-large": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-pytorch_model.bin",
+	"roberta-large-mnli": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-mnli-pytorch_model.bin",
+	"distilroberta-base": "https://s3.amazonaws.com/models.huggingface.co/bert/distilroberta-base-pytorch_model.bin",
+	"roberta-base-openai-detector": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-base-openai-detector-pytorch_model.bin",
+	"roberta-large-openai-detector": "https://s3.amazonaws.com/models.huggingface.co/bert/roberta-large-openai-detector-pytorch_model.bin",
+}
+
+
+class RobertaEmbeddings(BertEmbeddings):
+	"""
+	Same as BertEmbeddings with a tiny tweak for positional embeddings indexing.
+	"""
+
+	def __init__(self, config):
+		super().__init__(config)
+		self.padding_idx = 1
+		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
+		)
+
+	def forward(self, input_ids=None, token_type_ids=None, position_ids=None, inputs_embeds=None):
+		if position_ids is None:
+			if input_ids is not None:
+				# Create the position ids from the input token ids. Any padded tokens remain padded.
+				position_ids = create_position_ids_from_input_ids(input_ids, self.padding_idx).to(input_ids.device)
+			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
+		)
+
+	def create_position_ids_from_inputs_embeds(self, inputs_embeds):
+		""" We are provided embeddings directly. We cannot infer which are padded so just generate
+		sequential position ids.
+		:param torch.Tensor inputs_embeds:
+		:return torch.Tensor:
+		"""
+		input_shape = inputs_embeds.size()[:-1]
+		sequence_length = input_shape[1]
+
+		position_ids = torch.arange(
+			self.padding_idx + 1, sequence_length + self.padding_idx + 1, dtype=torch.long, device=inputs_embeds.device
+		)
+		return position_ids.unsqueeze(0).expand(input_shape)
+
+
+ROBERTA_START_DOCSTRING = r"""
+	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.RobertaConfig`): 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.
+"""
+
+ROBERTA_INPUTS_DOCSTRING = r"""
+	Args:
+		input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`):
+			Indices of input sequence tokens in the vocabulary.
+			Indices can be obtained using :class:`transformers.RobertaTokenizer`.
+			See :func:`transformers.PreTrainedTokenizer.encode` and
+			:func:`transformers.PreTrainedTokenizer.encode_plus` for details.
+			`What are input IDs? <../glossary.html#input-ids>`__
+		attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Mask to avoid performing attention on padding token indices.
+			Mask values selected in ``[0, 1]``:
+			``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
+			`What are attention masks? <../glossary.html#attention-mask>`__
+		token_type_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Segment token indices to indicate first and second portions of the inputs.
+			Indices are selected in ``[0, 1]``: ``0`` corresponds to a `sentence A` token, ``1``
+			corresponds to a `sentence B` token
+			`What are token type IDs? <../glossary.html#token-type-ids>`_
+		position_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Indices of positions of each input sequence tokens in the position embeddings.
+			Selected in the range ``[0, config.max_position_embeddings - 1]``.
+			`What are position IDs? <../glossary.html#position-ids>`_
+		head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`):
+			Mask to nullify selected heads of the self-attention modules.
+			Mask values selected in ``[0, 1]``:
+			:obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
+		inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
+			Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
+			This is useful if you want more control over how to convert `input_ids` indices into associated vectors
+			than the model's internal embedding lookup matrix.
+"""
+
+
+@add_start_docstrings(
+	"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.
+	"""
+
+	config_class = RobertaConfig
+	pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
+	base_model_prefix = "roberta"
+
+	def __init__(self, config):
+		super().__init__(config)
+
+		self.embeddings = RobertaEmbeddings(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
+
+
+@add_start_docstrings("""RoBERTa Model with a `language modeling` head on top. """, ROBERTA_START_DOCSTRING)
+class RobertaForMaskedLM(BertPreTrainedModel):
+	config_class = RobertaConfig
+	pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
+	base_model_prefix = "roberta"
+
+	def __init__(self, config):
+		super().__init__(config)
+
+		self.roberta = RobertaModel(config)
+		self.lm_head = RobertaLMHead(config)
+
+		self.init_weights()
+
+	def get_output_embeddings(self):
+		return self.lm_head.decoder
+
+	@add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		masked_lm_labels=None,
+	):
+		r"""
+		masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the masked language modeling loss.
+			Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
+			Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with labels
+			in ``[0, ..., config.vocab_size]``
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
+		masked_lm_loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
+			Masked language modeling loss.
+		prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
+			Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+	Examples::
+		from transformers import RobertaTokenizer, RobertaForMaskedLM
+		import torch
+		tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+		model = RobertaForMaskedLM.from_pretrained('roberta-base')
+		input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids, masked_lm_labels=input_ids)
+		loss, prediction_scores = outputs[:2]
+		"""
+		outputs = self.roberta(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+		sequence_output = outputs[0]
+		prediction_scores = self.lm_head(sequence_output)
+
+		outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention if they are here
+
+		if masked_lm_labels is not None:
+			loss_fct = CrossEntropyLoss()
+			masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
+			outputs = (masked_lm_loss,) + outputs
+
+		return outputs  # (masked_lm_loss), prediction_scores, (hidden_states), (attentions)
+
+
+class RobertaLMHead(nn.Module):
+	"""Roberta Head for masked language modeling."""
+
+	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.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, features, **kwargs):
+		x = self.dense(features)
+		x = gelu(x)
+		x = self.layer_norm(x)
+
+		# project back to size of vocabulary with bias
+		x = self.decoder(x)
+
+		return x
+
+
+@add_start_docstrings(
+	"""RoBERTa Model transformer with a sequence classification/regression head on top (a linear layer
+	on top of the pooled output) e.g. for GLUE tasks. """,
+	ROBERTA_START_DOCSTRING,
+)
+class RobertaForSequenceClassification(BertPreTrainedModel):
+	config_class = RobertaConfig
+	pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
+	base_model_prefix = "roberta"
+
+	def __init__(self, config):
+
+		print(config)
+
+
+		super().__init__(config)
+		self.num_labels = config.num_labels
+
+		self.roberta = RobertaModel(config)
+		self.classifier_ER = RobertaClassificationHead(config)
+		self.classifier_IP = RobertaClassificationHead(config)
+		self.classifier_EX = RobertaClassificationHead(config)
+
+	@add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		labels_ER=None,
+		labels_IP=None,
+		labels_EX=None,
+		lambda_ER=0.1,
+		lambda_IP=0.4,
+		lambda_EX=0.5
+	):
+		r"""
+		labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the sequence classification/regression loss.
+			Indices should be in :obj:`[0, ..., config.num_labels - 1]`.
+			If :obj:`config.num_labels == 1` a regression loss is computed (Mean-Square loss),
+			If :obj:`config.num_labels > 1` a classification loss is computed (Cross-Entropy).
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`label` is provided):
+			Classification (or regression if config.num_labels==1) loss.
+		logits (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, config.num_labels)`):
+			Classification (or regression if config.num_labels==1) scores (before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+	Examples::
+		from transformers import RobertaTokenizer, RobertaForSequenceClassification
+		import torch
+		tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+		model = RobertaForSequenceClassification.from_pretrained('roberta-base')
+		input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+		labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids, labels=labels)
+		loss, logits = outputs[:2]
+		"""
+		
+		outputs = self.roberta(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		sequence_output = outputs[0]
+
+
+		logits_ER = self.classifier_ER(sequence_output)
+		logits_IP = self.classifier_IP(sequence_output)
+		logits_EX = self.classifier_EX(sequence_output)
+
+
+		outputs = (logits_ER, logits_IP, logits_EX) + outputs[2:]
+
+
+		if labels_ER is not None:
+			loss_fct = CrossEntropyLoss()
+			# loss_fct_IP = CrossEntropyLoss()
+			# loss_fct_EX = CrossEntropyLoss()
+			
+			
+			loss_ER = loss_fct(logits_ER.view(-1, self.num_labels), labels_ER.view(-1))
+			loss_IP = loss_fct(logits_IP.view(-1, self.num_labels), labels_IP.view(-1))
+			loss_EX = loss_fct(logits_EX.view(-1, self.num_labels), labels_EX.view(-1))
+			
+
+			loss = (lambda_ER * loss_ER) + (lambda_IP * loss_IP) + (lambda_EX * loss_EX)
+			
+			outputs = (loss,) + outputs
+			
+			
+			# if self.num_labels == 1:
+			# 	#  We are doing regression
+			# 	loss_fct = MSELoss()
+			# 	loss = loss_fct(logits.view(-1), labels.view(-1))
+			# else:
+			# 	loss_fct = CrossEntropyLoss()
+			# 	loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+			# outputs = (loss,) + outputs
+
+		return outputs  # (loss), logits, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+	"""Roberta Model with a multiple choice classification head on top (a linear layer on top of
+	the pooled output and a softmax) e.g. for RocStories/SWAG tasks. """,
+	ROBERTA_START_DOCSTRING,
+)
+class RobertaForMultipleChoice(BertPreTrainedModel):
+	config_class = RobertaConfig
+	pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
+	base_model_prefix = "roberta"
+
+	def __init__(self, config):
+		super().__init__(config)
+
+		self.roberta = RobertaModel(config)
+		self.dropout = nn.Dropout(config.hidden_dropout_prob)
+		self.classifier = nn.Linear(config.hidden_size, 1)
+
+		self.init_weights()
+
+	@add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		token_type_ids=None,
+		attention_mask=None,
+		labels=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+	):
+		r"""
+		labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the multiple choice classification loss.
+			Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
+			of the input tensors. (see `input_ids` above)
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor`` of shape ``(1,)`, `optional`, returned when :obj:`labels` is provided):
+			Classification loss.
+		classification_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, num_choices)`):
+			`num_choices` is the second dimension of the input tensors. (see `input_ids` above).
+			Classification scores (before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+	Examples::
+		from transformers import RobertaTokenizer, RobertaForMultipleChoice
+		import torch
+		tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+		model = RobertaForMultipleChoice.from_pretrained('roberta-base')
+		choices = ["Hello, my dog is cute", "Hello, my cat is amazing"]
+		input_ids = torch.tensor([tokenizer.encode(s, add_special_tokens=True) for s in choices]).unsqueeze(0)  # Batch size 1, 2 choices
+		labels = torch.tensor(1).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids, labels=labels)
+		loss, classification_scores = outputs[:2]
+		"""
+		num_choices = input_ids.shape[1]
+
+		flat_input_ids = input_ids.view(-1, input_ids.size(-1))
+		flat_position_ids = position_ids.view(-1, position_ids.size(-1)) if position_ids is not None else None
+		flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
+		flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
+		outputs = self.roberta(
+			flat_input_ids,
+			position_ids=flat_position_ids,
+			token_type_ids=flat_token_type_ids,
+			attention_mask=flat_attention_mask,
+			head_mask=head_mask,
+		)
+		pooled_output = outputs[1]
+
+		pooled_output = self.dropout(pooled_output)
+		logits = self.classifier(pooled_output)
+		reshaped_logits = logits.view(-1, num_choices)
+
+		outputs = (reshaped_logits,) + outputs[2:]  # add hidden states and attention if they are here
+
+		if labels is not None:
+			loss_fct = CrossEntropyLoss()
+			loss = loss_fct(reshaped_logits, labels)
+			outputs = (loss,) + outputs
+
+		return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)
+
+
+@add_start_docstrings(
+	"""Roberta 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. """,
+	ROBERTA_START_DOCSTRING,
+)
+class RobertaForTokenClassification(BertPreTrainedModel):
+	config_class = RobertaConfig
+	pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
+	base_model_prefix = "roberta"
+
+	def __init__(self, config):
+		super().__init__(config)
+		self.num_labels = config.num_labels
+
+		self.roberta = RobertaModel(config)
+		self.dropout = nn.Dropout(config.hidden_dropout_prob)
+		self.classifier = nn.Linear(config.hidden_size, config.num_labels)
+
+		self.init_weights()
+
+	@add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids=None,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		labels=None,
+	):
+		r"""
+		labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
+			Labels for computing the token classification loss.
+			Indices should be in ``[0, ..., config.num_labels - 1]``.
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided) :
+			Classification loss.
+		scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`)
+			Classification scores (before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+	Examples::
+		from transformers import RobertaTokenizer, RobertaForTokenClassification
+		import torch
+		tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+		model = RobertaForTokenClassification.from_pretrained('roberta-base')
+		input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
+		labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0)  # Batch size 1
+		outputs = model(input_ids, labels=labels)
+		loss, scores = outputs[:2]
+		"""
+
+		outputs = self.roberta(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		sequence_output = outputs[0]
+
+		sequence_output = self.dropout(sequence_output)
+		logits = self.classifier(sequence_output)
+
+		outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
+
+		if labels is not None:
+			loss_fct = CrossEntropyLoss()
+			# Only keep active parts of the loss
+			if attention_mask is not None:
+				active_loss = attention_mask.view(-1) == 1
+				active_logits = logits.view(-1, self.num_labels)
+				active_labels = torch.where(
+					active_loss, labels.view(-1), torch.tensor(loss_fct.ignore_index).type_as(labels)
+				)
+				loss = loss_fct(active_logits, active_labels)
+			else:
+				loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
+			outputs = (loss,) + outputs
+
+		return outputs  # (loss), scores, (hidden_states), (attentions)
+
+
+class RobertaClassificationHead(nn.Module):
+	"""Head for sentence-level classification tasks."""
+
+	def __init__(self, config):
+		super().__init__()
+		self.dense = nn.Linear(config.hidden_size, 128)
+		self.dropout = nn.Dropout(config.hidden_dropout_prob)
+		self.out_proj = nn.Linear(128, config.num_labels)
+
+	def forward(self, features, **kwargs):
+		x = features[:, 0, :]  # take <s> token (equiv. to [CLS])
+		x = self.dropout(x)
+		x = self.dense(x)
+		x = torch.tanh(x)
+		x = self.dropout(x)
+		x = self.out_proj(x)
+		return x
+
+
+@add_start_docstrings(
+	"""Roberta Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
+	the hidden-states output to compute `span start logits` and `span end logits`). """,
+	ROBERTA_START_DOCSTRING,
+)
+class RobertaForQuestionAnswering(BertPreTrainedModel):
+	config_class = RobertaConfig
+	pretrained_model_archive_map = ROBERTA_PRETRAINED_MODEL_ARCHIVE_MAP
+	base_model_prefix = "roberta"
+
+	def __init__(self, config):
+		super().__init__(config)
+		self.num_labels = config.num_labels
+
+		self.roberta = RobertaModel(config)
+		self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
+
+		self.init_weights()
+
+	@add_start_docstrings_to_callable(ROBERTA_INPUTS_DOCSTRING)
+	def forward(
+		self,
+		input_ids,
+		attention_mask=None,
+		token_type_ids=None,
+		position_ids=None,
+		head_mask=None,
+		inputs_embeds=None,
+		start_positions=None,
+		end_positions=None,
+	):
+		r"""
+		start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for position (index) of the start of the labelled span for computing the token classification loss.
+			Positions are clamped to the length of the sequence (`sequence_length`).
+			Position outside of the sequence are not taken into account for computing the loss.
+		end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
+			Labels for position (index) of the end of the labelled span for computing the token classification loss.
+			Positions are clamped to the length of the sequence (`sequence_length`).
+			Position outside of the sequence are not taken into account for computing the loss.
+	Returns:
+		:obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.RobertaConfig`) and inputs:
+		loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when :obj:`labels` is provided):
+			Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
+		start_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
+			Span-start scores (before SoftMax).
+		end_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length,)`):
+			Span-end scores (before SoftMax).
+		hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
+			of shape :obj:`(batch_size, sequence_length, hidden_size)`.
+			Hidden-states of the model at the output of each layer plus the initial embedding outputs.
+		attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
+			Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
+			:obj:`(batch_size, num_heads, sequence_length, sequence_length)`.
+			Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
+			heads.
+	Examples::
+		# The checkpoint roberta-large is not fine-tuned for question answering. Please see the
+		# examples/run_squad.py example to see how to fine-tune a model to a question answering task.
+		from transformers import RobertaTokenizer, RobertaForQuestionAnswering
+		import torch
+		tokenizer = RobertaTokenizer.from_pretrained('roberta-base')
+		model = RobertaForQuestionAnswering.from_pretrained('roberta-base')
+		question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
+		input_ids = tokenizer.encode(question, text)
+		start_scores, end_scores = model(torch.tensor([input_ids]))
+		all_tokens = tokenizer.convert_ids_to_tokens(input_ids)
+		answer = ' '.join(all_tokens[torch.argmax(start_scores) : torch.argmax(end_scores)+1])
+		"""
+
+		outputs = self.roberta(
+			input_ids,
+			attention_mask=attention_mask,
+			token_type_ids=token_type_ids,
+			position_ids=position_ids,
+			head_mask=head_mask,
+			inputs_embeds=inputs_embeds,
+		)
+
+		sequence_output = outputs[0]
+
+		logits = self.qa_outputs(sequence_output)
+		start_logits, end_logits = logits.split(1, dim=-1)
+		start_logits = start_logits.squeeze(-1)
+		end_logits = end_logits.squeeze(-1)
+
+		outputs = (start_logits, end_logits,) + outputs[2:]
+		if start_positions is not None and end_positions is not None:
+			# If we are on multi-GPU, split add a dimension
+			if len(start_positions.size()) > 1:
+				start_positions = start_positions.squeeze(-1)
+			if len(end_positions.size()) > 1:
+				end_positions = end_positions.squeeze(-1)
+			# sometimes the start/end positions are outside our model inputs, we ignore these terms
+			ignored_index = start_logits.size(1)
+			start_positions.clamp_(0, ignored_index)
+			end_positions.clamp_(0, ignored_index)
+
+			loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
+			start_loss = loss_fct(start_logits, start_positions)
+			end_loss = loss_fct(end_logits, end_positions)
+			total_loss = (start_loss + end_loss) / 2
+			outputs = (total_loss,) + outputs
+
+		return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)

requirements.txt

@@ -1,3 +1,4 @@
 transformers>=4.30.0
 torch>=1.9.0
 numpy>=1.19.0
+sentencepiece