modeling_utils.py 88.6 KB
Newer Older
1
# coding=utf-8
thomwolf's avatar
thomwolf committed
2
# Copyright 2018 The Google AI Language Team Authors, Facebook AI Research authors and The HuggingFace Inc. team.
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
# 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
20
import typing
21
22
23

import torch
from torch import nn
24
25
from torch.nn import CrossEntropyLoss
from torch.nn import functional as F
26

27
from .activations import get_activation
28
from .configuration_utils import PretrainedConfig
29
from .file_utils import (
Aymeric Augustin's avatar
Aymeric Augustin committed
30
    DUMMY_INPUTS,
31
32
33
34
35
36
37
    TF2_WEIGHTS_NAME,
    TF_WEIGHTS_NAME,
    WEIGHTS_NAME,
    cached_path,
    hf_bucket_url,
    is_remote_url,
)
38

Aymeric Augustin's avatar
Aymeric Augustin committed
39

40
41
logger = logging.getLogger(__name__)

thomwolf's avatar
thomwolf committed
42
43
44
45
46
47
48
try:
    from torch.nn import Identity
except ImportError:
    # Older PyTorch compatibility
    class Identity(nn.Module):
        r"""A placeholder identity operator that is argument-insensitive.
        """
49

thomwolf's avatar
thomwolf committed
50
        def __init__(self, *args, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
51
            super().__init__()
thomwolf's avatar
thomwolf committed
52
53
54
55

        def forward(self, input):
            return input

56

57
class ModuleUtilsMixin:
Julien Chaumond's avatar
Julien Chaumond committed
58
59
60
61
62
63
64
65
66
67
68
69
    """
    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)


70
class PreTrainedModel(nn.Module, ModuleUtilsMixin):
71
72
    r""" Base class for all models.

73
        :class:`~transformers.PreTrainedModel` takes care of storing the configuration of the models and handles methods for loading/downloading/saving models
Julien Chaumond's avatar
Julien Chaumond committed
74
        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.
75
76

        Class attributes (overridden by derived classes):
77
            - ``config_class``: a class derived from :class:`~transformers.PretrainedConfig` to use as configuration class for this model architecture.
78
79
80
            - ``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:

81
82
                - ``model``: an instance of the relevant subclass of :class:`~transformers.PreTrainedModel`,
                - ``config``: an instance of the relevant subclass of :class:`~transformers.PretrainedConfig`,
83
84
85
                - ``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.
86
    """
87
    config_class = None
88
89
90
    pretrained_model_archive_map = {}
    base_model_prefix = ""

91
92
93
94
95
96
97
    @property
    def dummy_inputs(self):
        """ Dummy inputs to do a forward pass in the network.

        Returns:
            torch.Tensor with dummy inputs
        """
98
        return {"input_ids": torch.tensor(DUMMY_INPUTS)}
99

100
    def __init__(self, config, *inputs, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
101
        super().__init__()
102
103
104
105
106
107
        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__
108
109
                )
            )
thomwolf's avatar
thomwolf committed
110
        # Save config in model
111
112
        self.config = config

113
114
115
    @property
    def base_model(self):
        return getattr(self, self.base_model_prefix, self)
thomwolf's avatar
thomwolf committed
116

thomwolf's avatar
thomwolf committed
117
    def get_input_embeddings(self):
118
119
120
121
122
123
        """
        Returns the model's input embeddings.

        Returns:
            :obj:`nn.Module`:
                A torch module mapping vocabulary to hidden states.
thomwolf's avatar
thomwolf committed
124
        """
125
        base_model = getattr(self, self.base_model_prefix, self)
thomwolf's avatar
thomwolf committed
126
127
128
129
        if base_model is not self:
            return base_model.get_input_embeddings()
        else:
            raise NotImplementedError
thomwolf's avatar
thomwolf committed
130

thomwolf's avatar
thomwolf committed
131
    def set_input_embeddings(self, value):
132
133
134
135
136
137
        """
        Set model's input embeddings

        Args:
            value (:obj:`nn.Module`):
                A module mapping vocabulary to hidden states.
thomwolf's avatar
thomwolf committed
138
139
140
141
142
143
        """
        base_model = getattr(self, self.base_model_prefix, self)
        if base_model is not self:
            base_model.set_input_embeddings(value)
        else:
            raise NotImplementedError
thomwolf's avatar
thomwolf committed
144

thomwolf's avatar
thomwolf committed
145
    def get_output_embeddings(self):
146
147
148
149
150
151
        """
        Returns the model's output embeddings.

        Returns:
            :obj:`nn.Module`:
                A torch module mapping hidden states to vocabulary.
thomwolf's avatar
thomwolf committed
152
        """
153
        return None  # Overwrite for models with output embeddings
thomwolf's avatar
thomwolf committed
154

155
    def tie_weights(self):
156
157
158
159
        """
        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.
thomwolf's avatar
thomwolf committed
160
        """
thomwolf's avatar
thomwolf committed
161
162
163
        output_embeddings = self.get_output_embeddings()
        if output_embeddings is not None:
            self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings())
thomwolf's avatar
thomwolf committed
164

165
    def _tie_or_clone_weights(self, output_embeddings, input_embeddings):
thomwolf's avatar
thomwolf committed
166
167
168
        """ Tie or clone module weights depending of weither we are using TorchScript or not
        """
        if self.config.torchscript:
169
            output_embeddings.weight = nn.Parameter(input_embeddings.weight.clone())
thomwolf's avatar
thomwolf committed
170
        else:
171
            output_embeddings.weight = input_embeddings.weight
thomwolf's avatar
thomwolf committed
172

Sam Shleifer's avatar
Sam Shleifer committed
173
        if getattr(output_embeddings, "bias", None) is not None:
174
175
            output_embeddings.bias.data = torch.nn.functional.pad(
                output_embeddings.bias.data,
Patrick von Platen's avatar
Patrick von Platen committed
176
                (0, output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0],),
177
178
                "constant",
                0,
179
            )
180
        if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"):
181
            output_embeddings.out_features = input_embeddings.num_embeddings
182

thomwolf's avatar
thomwolf committed
183
184
    def resize_token_embeddings(self, new_num_tokens=None):
        """ Resize input token embeddings matrix of the model if new_num_tokens != config.vocab_size.
185
        Take care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.
thomwolf's avatar
thomwolf committed
186

187
188
189
        Arguments:

            new_num_tokens: (`optional`) int:
190
                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.
191
                If not provided or None: does nothing and just returns a pointer to the input tokens ``torch.nn.Embeddings`` Module of the model.
thomwolf's avatar
thomwolf committed
192

thomwolf's avatar
thomwolf committed
193
        Return: ``torch.nn.Embeddings``
194
            Pointer to the input tokens Embeddings Module of the model
thomwolf's avatar
thomwolf committed
195
196
        """
        base_model = getattr(self, self.base_model_prefix, self)  # get the base model if needed
thomwolf's avatar
thomwolf committed
197
198
199
        model_embeds = base_model._resize_token_embeddings(new_num_tokens)
        if new_num_tokens is None:
            return model_embeds
thomwolf's avatar
thomwolf committed
200
201
202
203
204
205

        # 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
206
        self.tie_weights()
thomwolf's avatar
thomwolf committed
207

thomwolf's avatar
thomwolf committed
208
209
        return model_embeds

210
    def _resize_token_embeddings(self, new_num_tokens):
thomwolf's avatar
thomwolf committed
211
212
213
214
        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()
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243

    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)

244
        # Copy token embeddings from the previous weights
245
246
247
248
249
        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

250
251
252
253
254
255
256
257
258
    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)

259
260
261
        # Tie weights if needed
        self.tie_weights()

thomwolf's avatar
thomwolf committed
262
263
    def prune_heads(self, heads_to_prune):
        """ Prunes heads of the base model.
264
265
266
267

            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`).
268
                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.
thomwolf's avatar
thomwolf committed
269
        """
270
        # save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads
271
        for layer, heads in heads_to_prune.items():
272
273
274
            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

275
        self.base_model._prune_heads(heads_to_prune)
thomwolf's avatar
thomwolf committed
276

277
    def save_pretrained(self, save_directory):
278
        """ Save a model and its configuration file to a directory, so that it
279
            can be re-loaded using the `:func:`~transformers.PreTrainedModel.from_pretrained`` class method.
280
        """
281
282
283
        assert os.path.isdir(
            save_directory
        ), "Saving path should be a directory where the model and configuration can be saved"
284

Julien Chaumond's avatar
Julien Chaumond committed
285
        # Only save the model itself if we are using distributed training
286
        model_to_save = self.module if hasattr(self, "module") else self
287

Julien Chaumond's avatar
Julien Chaumond committed
288
289
290
        # Attach architecture to the config
        model_to_save.config.architectures = [model_to_save.__class__.__name__]

thomwolf's avatar
thomwolf committed
291
292
293
        # Save configuration file
        model_to_save.config.save_pretrained(save_directory)

294
295
296
        # If we save using the predefined names, we can load using `from_pretrained`
        output_model_file = os.path.join(save_directory, WEIGHTS_NAME)
        torch.save(model_to_save.state_dict(), output_model_file)
thomwolf's avatar
thomwolf committed
297
        logger.info("Model weights saved in {}".format(output_model_file))
298

299
    @classmethod
300
    def from_pretrained(cls, pretrained_model_name_or_path, *model_args, **kwargs):
301
302
        r"""Instantiate a pretrained pytorch model from a pre-trained model configuration.

303
304
305
        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()``

306
307
308
309
310
        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.

311
312
        Parameters:
            pretrained_model_name_or_path: either:
Lysandre's avatar
Fixes  
Lysandre committed
313
314
315
316
317
              - 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``)
318
319
320
321

            model_args: (`optional`) Sequence of positional arguments:
                All remaning positional arguments will be passed to the underlying model's ``__init__`` method

322
            config: (`optional`) one of:
Lysandre's avatar
Fixes  
Lysandre committed
323
324
                - an instance of a class derived from :class:`~transformers.PretrainedConfig`, or
                - a string valid as input to :func:`~transformers.PretrainedConfig.from_pretrained()`
325
                Configuration for the model to use instead of an automatically loaded configuation. Configuration can be automatically loaded when:
Lysandre's avatar
Fixes  
Lysandre committed
326
327
328
                    - 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.
329
330
331

            state_dict: (`optional`) dict:
                an optional state dictionnary for the model to use instead of a state dictionary loaded from saved weights file.
thomwolf's avatar
typos  
thomwolf committed
332
                This option can be used if you want to create a model from a pretrained configuration but load your own weights.
333
                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.
334
335

            cache_dir: (`optional`) string:
thomwolf's avatar
thomwolf committed
336
337
                Path to a directory in which a downloaded pre-trained model
                configuration should be cached if the standard cache should not be used.
338

339
340
341
            force_download: (`optional`) boolean, default False:
                Force to (re-)download the model weights and configuration files and override the cached versions if they exists.

342
343
344
            resume_download: (`optional`) boolean, default False:
                Do not delete incompletely recieved file. Attempt to resume the download if such a file exists.

345
346
347
348
            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.

349
            output_loading_info: (`optional`) boolean:
thomwolf's avatar
thomwolf committed
350
                Set to ``True`` to also return a dictionnary containing missing keys, unexpected keys and error messages.
351
352
353
354
355

            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)
356
                - 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.
357
358

        Examples::
thomwolf's avatar
thomwolf committed
359

Lysandre's avatar
Lysandre committed
360
            # For example purposes. Not runnable.
thomwolf's avatar
thomwolf committed
361
362
363
364
365
366
367
            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)
thomwolf's avatar
thomwolf committed
368

369
        """
370
371
372
373
374
375
376
377
        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)
378
        local_files_only = kwargs.pop("local_files_only", False)
thomwolf's avatar
thomwolf committed
379

380
381
382
        # 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
383
            config, model_kwargs = cls.config_class.from_pretrained(
384
385
386
387
                config_path,
                *model_args,
                cache_dir=cache_dir,
                return_unused_kwargs=True,
388
                force_download=force_download,
389
                resume_download=resume_download,
390
                proxies=proxies,
391
                local_files_only=local_files_only,
392
                **kwargs,
393
394
395
            )
        else:
            model_kwargs = kwargs
396

thomwolf's avatar
thomwolf committed
397
        # Load model
thomwolf's avatar
thomwolf committed
398
        if pretrained_model_name_or_path is not None:
399
            if pretrained_model_name_or_path in cls.pretrained_model_archive_map:
thomwolf's avatar
thomwolf committed
400
401
                archive_file = cls.pretrained_model_archive_map[pretrained_model_name_or_path]
            elif os.path.isdir(pretrained_model_name_or_path):
thomwolf's avatar
thomwolf committed
402
403
                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
thomwolf's avatar
thomwolf committed
404
                    archive_file = os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
thomwolf's avatar
thomwolf committed
405
406
407
408
409
                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
thomwolf's avatar
thomwolf committed
410
                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
thomwolf's avatar
thomwolf committed
411
                else:
412
413
                    raise EnvironmentError(
                        "Error no file named {} found in directory {} or `from_tf` set to False".format(
Patrick von Platen's avatar
Patrick von Platen committed
414
                            [WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME + ".index"],
Patrick von Platen's avatar
Patrick von Platen committed
415
                            pretrained_model_name_or_path,
416
417
                        )
                    )
418
            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
419
                archive_file = pretrained_model_name_or_path
420
            elif os.path.isfile(pretrained_model_name_or_path + ".index"):
421
422
423
424
425
                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"
                )
426
                archive_file = pretrained_model_name_or_path + ".index"
427
            else:
thomwolf's avatar
thomwolf committed
428
                archive_file = hf_bucket_url(
Patrick von Platen's avatar
Patrick von Platen committed
429
                    pretrained_model_name_or_path, postfix=(TF2_WEIGHTS_NAME if from_tf else WEIGHTS_NAME),
thomwolf's avatar
thomwolf committed
430
                )
431

thomwolf's avatar
thomwolf committed
432
433
            # redirect to the cache, if necessary
            try:
434
435
436
437
438
439
                resolved_archive_file = cached_path(
                    archive_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
440
                    local_files_only=local_files_only,
441
                )
thomwolf's avatar
thomwolf committed
442
            except EnvironmentError:
thomwolf's avatar
thomwolf committed
443
                if pretrained_model_name_or_path in cls.pretrained_model_archive_map:
444
                    msg = "Couldn't reach server at '{}' to download pretrained weights.".format(archive_file)
thomwolf's avatar
thomwolf committed
445
                else:
446
447
448
                    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 "
thomwolf's avatar
thomwolf committed
449
                        "couldn't find any such file at this path or url.".format(
thomwolf's avatar
thomwolf committed
450
                            pretrained_model_name_or_path,
451
                            ", ".join(cls.pretrained_model_archive_map.keys()),
thomwolf's avatar
thomwolf committed
452
                            archive_file,
453
454
455
                            [WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME],
                        )
                    )
thomwolf's avatar
thomwolf committed
456
457
                raise EnvironmentError(msg)

thomwolf's avatar
thomwolf committed
458
459
            if resolved_archive_file == archive_file:
                logger.info("loading weights file {}".format(archive_file))
460
            else:
461
                logger.info("loading weights file {} from cache at {}".format(archive_file, resolved_archive_file))
462
        else:
thomwolf's avatar
thomwolf committed
463
            resolved_archive_file = None
464
465

        # Instantiate model.
466
        model = cls(config, *model_args, **model_kwargs)
thomwolf's avatar
thomwolf committed
467

468
        if state_dict is None and not from_tf:
469
            try:
470
                state_dict = torch.load(resolved_archive_file, map_location="cpu")
471
            except Exception:
472
473
474
475
                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. "
                )
476

477
478
479
        missing_keys = []
        unexpected_keys = []
        error_msgs = []
480
481

        if from_tf:
482
            if resolved_archive_file.endswith(".index"):
483
484
485
486
487
                # 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:
488
                    from transformers import load_tf2_checkpoint_in_pytorch_model
489

490
                    model = load_tf2_checkpoint_in_pytorch_model(model, resolved_archive_file, allow_missing_keys=True)
491
                except ImportError:
492
493
494
495
                    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."
                    )
496
                    raise
497
498
499
500
501
502
        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
503
504
505
506
                if "gamma" in key:
                    new_key = key.replace("gamma", "weight")
                if "beta" in key:
                    new_key = key.replace("beta", "bias")
507
508
509
510
511
512
513
                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
514
            metadata = getattr(state_dict, "_metadata", None)
515
516
517
518
            state_dict = state_dict.copy()
            if metadata is not None:
                state_dict._metadata = metadata

519
520
            # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
            # so we need to apply the function recursively.
Julien Chaumond's avatar
Julien Chaumond committed
521
            def load(module: nn.Module, prefix=""):
522
523
                local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
                module._load_from_state_dict(
Patrick von Platen's avatar
Patrick von Platen committed
524
                    state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs,
525
                )
526
527
                for name, child in module._modules.items():
                    if child is not None:
528
                        load(child, prefix + name + ".")
529
530

            # Make sure we are able to load base models as well as derived models (with heads)
531
            start_prefix = ""
532
            model_to_load = model
533
534
535
536
537
538
539
            if not hasattr(model, cls.base_model_prefix) and any(
                s.startswith(cls.base_model_prefix) for s in state_dict.keys()
            ):
                start_prefix = cls.base_model_prefix + "."
            if hasattr(model, cls.base_model_prefix) and not any(
                s.startswith(cls.base_model_prefix) for s in state_dict.keys()
            ):
540
541
542
                model_to_load = getattr(model, cls.base_model_prefix)

            load(model_to_load, prefix=start_prefix)
543
544
545
546
547
548
549
550
551

            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)

552
            if len(missing_keys) > 0:
553
554
555
556
557
                logger.info(
                    "Weights of {} not initialized from pretrained model: {}".format(
                        model.__class__.__name__, missing_keys
                    )
                )
558
            if len(unexpected_keys) > 0:
559
560
561
562
563
                logger.info(
                    "Weights from pretrained model not used in {}: {}".format(
                        model.__class__.__name__, unexpected_keys
                    )
                )
564
            if len(error_msgs) > 0:
565
566
567
568
569
                raise RuntimeError(
                    "Error(s) in loading state_dict for {}:\n\t{}".format(
                        model.__class__.__name__, "\n\t".join(error_msgs)
                    )
                )
570
        model.tie_weights()  # make sure token embedding weights are still tied if needed
571

572
573
574
        # Set model in evaluation mode to desactivate DropOut modules by default
        model.eval()

thomwolf's avatar
thomwolf committed
575
        if output_loading_info:
576
577
578
579
580
            loading_info = {
                "missing_keys": missing_keys,
                "unexpected_keys": unexpected_keys,
                "error_msgs": error_msgs,
            }
thomwolf's avatar
thomwolf committed
581
582
            return model, loading_info

583
584
        return model

thomwolf's avatar
thomwolf committed
585
586
587
    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}

patrickvonplaten's avatar
patrickvonplaten committed
588
589
590
    def prepare_scores_for_generation(self, scores, **kwargs):
        return scores

591
    def _do_output_past(self, outputs):
Sam Shleifer's avatar
Sam Shleifer committed
592
593
594
595
596
597
        """During generation, decide whether to pass the `past` variable to the next forward pass."""
        has_output_past = getattr(self.config, "output_past", False)
        mem_len = getattr(self.config, "mem_len", 0)
        if len(outputs) <= 1:
            return False
        if mem_len > 0 or has_output_past:
598
            return True
599
600
        return False

Sam Shleifer's avatar
Sam Shleifer committed
601
602
603
604
605
606
607
608
609
610
    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

thomwolf's avatar
thomwolf committed
611
    @torch.no_grad()
612
613
614
615
    def generate(
        self,
        input_ids=None,
        max_length=None,
616
        min_length=None,
617
        do_sample=True,
Patrick von Platen's avatar
Patrick von Platen committed
618
        early_stopping=False,
619
620
621
622
623
624
625
626
627
        num_beams=None,
        temperature=None,
        top_k=None,
        top_p=None,
        repetition_penalty=None,
        bos_token_id=None,
        pad_token_id=None,
        eos_token_ids=None,
        length_penalty=None,
Patrick von Platen's avatar
Patrick von Platen committed
628
        no_repeat_ngram_size=None,
629
        num_return_sequences=None,
Patrick von Platen's avatar
Patrick von Platen committed
630
        attention_mask=None,
631
    ):
632
        r""" Generates sequences for models with a LM head. The method currently supports greedy or penalized greedy decoding, sampling with top-k or nucleus sampling
thomwolf's avatar
thomwolf committed
633
        and beam-search.
thomwolf's avatar
thomwolf committed
634

635
636
637
638
639
640
641
        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:
thomwolf's avatar
thomwolf committed
642

643
            input_ids: (`optional`) `torch.LongTensor` of shape `(batch_size, sequence_length)`
thomwolf's avatar
thomwolf committed
644
                The sequence used as a prompt for the generation. If `None` the method initializes
645
646
647
                it as an empty `torch.LongTensor` of shape `(1,)`.

            max_length: (`optional`) int
thomwolf's avatar
thomwolf committed
648
                The max length of the sequence to be generated.  Between 1 and infinity. Default to 20.
649
650

            do_sample: (`optional`) bool
651
                If set to `False` greedy decoding is used. Otherwise sampling is used. Defaults to `True`.
652
653
654
655
656

            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
Sam Shleifer's avatar
Sam Shleifer committed
657
                The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
658
659

            top_k: (`optional`) int
thomwolf's avatar
thomwolf committed
660
                The number of highest probability vocabulary tokens to keep for top-k-filtering. Between 1 and infinity. Default to 50.
661
662

            top_p: (`optional`) float
thomwolf's avatar
thomwolf committed
663
                The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
664
665
666
667
668

            repetition_penalty: (`optional`) float
                The parameter for repetition penalty. Between 1.0 and infinity. 1.0 means no penalty. Default to 1.0.

            bos_token_id: (`optional`) int
thomwolf's avatar
thomwolf committed
669
                Beginning of sentence token if no prompt is provided. Default to 0.
670
671

            eos_token_ids: (`optional`) int or list of int
thomwolf's avatar
thomwolf committed
672
                End of sequence token or list of tokens to stop the generation. Default to 0.
673
            length_penalty: (`optional`) float
thomwolf's avatar
thomwolf committed
674
                Exponential penalty to the length. Default to 1.
675
676
677
678

            num_return_sequences: (`optional`) int
                The number of independently computed returned sequences for each element in the batch. Default to 1.

679
680
681
682
683
        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`

684
685
686
687
        Examples::

            tokenizer = AutoTokenizer.from_pretrained('distilgpt2')   # Initialize tokenizer
            model = AutoModelWithLMHead.from_pretrained('distilgpt2')    # Download model and configuration from S3 and cache.
688
            outputs = model.generate(max_length=40, do_sample=False)  # do greedy decoding
689
690
691
692
693
694
            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 = torch.tensor(tokenizer.encode(input_context)).unsqueeze(0)  # encode input context
695
            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'
696
            for i in range(3): #  3 output sequences were generated
697
                print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))
698
699
700
701
702

            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 = torch.tensor(tokenizer.encode(input_context)).unsqueeze(0)  # encode input context
703
            outputs = model.generate(input_ids=input_ids, max_length=40, temperature=0.7, num_return_sequences=3)  # 3 generate sequences using by sampling
704
705
            for i in range(3): #  3 output sequences were generated
                print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))
706
707
708
709
710

            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 = torch.tensor(tokenizer.encode(input_context)).unsqueeze(0)  # encode input context
711
            outputs = model.generate(input_ids=input_ids, max_length=50, temperature=0.7, repetition_penalty=1.2)  # generate sequences
712
713
            print('Generated: {}'.format(tokenizer.decode(outputs[0], skip_special_tokens=True)))

thomwolf's avatar
thomwolf committed
714
715
716
717
        """

        # We cannot generate if the model does not have a LM head
        if self.get_output_embeddings() is None:
718
719
            raise AttributeError(
                "You tried to generate sequences with a model that does not have a LM Head."
720
                "Please use another model class (e.g. `OpenAIGPTLMHeadModel`, `XLNetLMHeadModel`, `GPT2LMHeadModel`, `CTRLLMHeadModel`, `T5WithLMHeadModel`, `TransfoXLLMHeadModel`)"
721
            )
thomwolf's avatar
thomwolf committed
722

723
        max_length = max_length if max_length is not None else self.config.max_length
724
        min_length = min_length if min_length is not None else self.config.min_length
725
        do_sample = do_sample if do_sample is not None else self.config.do_sample
Patrick von Platen's avatar
Patrick von Platen committed
726
        early_stopping = early_stopping if early_stopping is not None else self.config.early_stopping
727
728
729
730
731
732
733
734
735
        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_ids = eos_token_ids if eos_token_ids is not None else self.config.eos_token_ids
        length_penalty = length_penalty if length_penalty is not None else self.config.length_penalty
Patrick von Platen's avatar
Patrick von Platen committed
736
737
738
        no_repeat_ngram_size = (
            no_repeat_ngram_size if no_repeat_ngram_size is not None else self.config.no_repeat_ngram_size
        )
739
740
741
        num_return_sequences = (
            num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
        )
thomwolf's avatar
thomwolf committed
742
743
744

        if input_ids is not None:
            batch_size = input_ids.shape[0]  # overriden by the input batch_size
thomwolf's avatar
thomwolf committed
745
746
        else:
            batch_size = 1
thomwolf's avatar
thomwolf committed
747
748
749
        if isinstance(eos_token_ids, int):
            eos_token_ids = [eos_token_ids]

Sam Shleifer's avatar
Sam Shleifer committed
750
        assert isinstance(max_length, int) and max_length > 0, "`max_length` should be a strictly positive integer."
751
        assert isinstance(min_length, int) and min_length >= 0, "`min_length` should be a positive integer."
thomwolf's avatar
thomwolf committed
752
        assert isinstance(do_sample, bool), "`do_sample` should be a boolean."
Patrick von Platen's avatar
Patrick von Platen committed
753
        assert isinstance(early_stopping, bool), "`early_stopping` should be a boolean."
Sam Shleifer's avatar
Sam Shleifer committed
754
755
        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."
756
        assert isinstance(top_k, int) and top_k >= 0, "`top_k` should be a positive integer."
thomwolf's avatar
thomwolf committed
757
        assert 0 <= top_p <= 1, "`top_p` should be between 0 and 1."
thomwolf's avatar
thomwolf committed
758
        assert repetition_penalty >= 1.0, "`repetition_penalty` should be >= 1."
759
760
761
762
763
764
765
766
        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_ids is None) or (
            isinstance(eos_token_ids, (list, tuple)) and ((isinstance(e, int) and e >= 0) for e in eos_token_ids)
767
        ), "`eos_token_ids` should be a positive integer or a list/tuple of positive integers."
Sam Shleifer's avatar
Sam Shleifer committed
768
        assert length_penalty > 0, "`length_penalty` should be strictly positive."
Patrick von Platen's avatar
Patrick von Platen committed
769
770
771
        assert (
            isinstance(no_repeat_ngram_size, int) and no_repeat_ngram_size >= 0
        ), "`no_repeat_ngram_size` should be a positive integer."
772
773
        assert (
            isinstance(num_return_sequences, int) and num_return_sequences > 0
Sam Shleifer's avatar
Sam Shleifer committed
774
        ), "`num_return_sequences` should be a strictly positive integer."
thomwolf's avatar
thomwolf committed
775
776

        if input_ids is None:
777
778
779
780
            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."
            )
781
            input_ids = torch.full(
Patrick von Platen's avatar
Patrick von Platen committed
782
                (batch_size, 1), bos_token_id, dtype=torch.long, device=next(self.parameters()).device,
783
            )
thomwolf's avatar
thomwolf committed
784
        else:
785
            assert input_ids.dim() == 2, "Input prompt should be of shape (batch_size, sequence length)."
thomwolf's avatar
thomwolf committed
786

787
        # not allow to duplicate outputs when greedy decoding
788
789
790
791
792
793
794
795
796
797
798
799
800
        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"

Patrick von Platen's avatar
Patrick von Platen committed
801
        # create attention mask if necessary
patrickvonplaten's avatar
patrickvonplaten committed
802
        # TODO (PVP): this should later be handled by the forward fn() in each model in the future see PR 3140
Patrick von Platen's avatar
Patrick von Platen committed
803
804
805
806
807
808
809
        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_ids if not set. Important that this is done after
        # attention_mask is created
810
811
812
813
814
815
        if pad_token_id is None and eos_token_ids is not None:
            logger.warning(
                "Setting `pad_token_id` to {} (first `eos_token_id`) to generate sequence".format(eos_token_ids[0])
            )
            pad_token_id = eos_token_ids[0]

thomwolf's avatar
thomwolf committed
816
817
818
        # current position and vocab size
        vocab_size = self.config.vocab_size

819
820
        # set effective batch size and effective batch multiplier according to do_sample
        if do_sample:
thomwolf's avatar
thomwolf committed
821
            effective_batch_size = batch_size * num_return_sequences
822
            effective_batch_mult = num_return_sequences
thomwolf's avatar
thomwolf committed
823
824
        else:
            effective_batch_size = batch_size
825
826
827
828
829
830
            effective_batch_mult = 1

        # 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)
831
832
833
            attention_mask = attention_mask.unsqueeze(1).expand(
                batch_size, effective_batch_mult * num_beams, input_ids_len
            )
Patrick von Platen's avatar
Patrick von Platen committed
834

835
836
837
            input_ids = input_ids.contiguous().view(
                effective_batch_size * num_beams, input_ids_len
            )  # shape: (batch_size * num_return_sequences * num_beams, cur_len)
Patrick von Platen's avatar
Patrick von Platen committed
838
839
840
            attention_mask = attention_mask.contiguous().view(
                effective_batch_size * num_beams, input_ids_len
            )  # shape: (batch_size * num_return_sequences * num_beams, cur_len)
thomwolf's avatar
thomwolf committed
841

Patrick von Platen's avatar
Patrick von Platen committed
842
        if self.config.is_encoder_decoder:
Patrick von Platen's avatar
Patrick von Platen committed
843
            eos_token_id = eos_token_ids[0]
Patrick von Platen's avatar
Patrick von Platen committed
844
845
            assert bos_token_id is not None, "Encoder Decoder Models need to have a bos_token_id"
            assert eos_token_id is not None, "Encoder Decoder Models need to have a eos_token_id"
Patrick von Platen's avatar
Patrick von Platen committed
846
847
848
849
            # encoder decoder need to start with empty input_ids and copy the input_ids to encoder_inputs
            encoder_inputs = input_ids
            input_ids = torch.full(
                (effective_batch_size * num_beams, 1),
850
851
                #                eos_token_id,  # TODO (PVP): to check if this is the only solution -> quite hacky to do this
                bos_token_id,
Patrick von Platen's avatar
Patrick von Platen committed
852
853
854
                dtype=torch.long,
                device=next(self.parameters()).device,
            )
855
            cur_len = 1
856
857
858
859

            # put model in generation mode if it has one
            if hasattr(self.model, "generation_mode"):
                self.model.decoder.generation_mode = True
Patrick von Platen's avatar
Patrick von Platen committed
860
861
862
863
        else:
            encoder_inputs = None
            cur_len = input_ids.shape[-1]

thomwolf's avatar
thomwolf committed
864
        if num_beams > 1:
865
866
            output = self._generate_beam_search(
                input_ids,
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
                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,
                bos_token_id=bos_token_id,
                pad_token_id=pad_token_id,
                eos_token_ids=eos_token_ids,
                batch_size=effective_batch_size,
                num_return_sequences=num_return_sequences,
                length_penalty=length_penalty,
                num_beams=num_beams,
                vocab_size=vocab_size,
                encoder_inputs=encoder_inputs,
                attention_mask=attention_mask,
887
            )
thomwolf's avatar
thomwolf committed
888
        else:
889
890
            output = self._generate_no_beam_search(
                input_ids,
891
892
893
894
895
896
897
898
899
900
901
902
903
904
                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,
                pad_token_id=pad_token_id,
                eos_token_ids=eos_token_ids,
                batch_size=effective_batch_size,
                encoder_inputs=encoder_inputs,
                attention_mask=attention_mask,
905
            )
thomwolf's avatar
thomwolf committed
906
907

        return output
thomwolf's avatar
thomwolf committed
908

909
910
911
912
913
    def _generate_no_beam_search(
        self,
        input_ids,
        cur_len,
        max_length,
914
        min_length,
915
916
917
918
919
        do_sample,
        temperature,
        top_k,
        top_p,
        repetition_penalty,
Patrick von Platen's avatar
Patrick von Platen committed
920
        no_repeat_ngram_size,
921
922
923
        pad_token_id,
        eos_token_ids,
        batch_size,
Patrick von Platen's avatar
Patrick von Platen committed
924
        encoder_inputs,
925
        attention_mask,
926
    ):
thomwolf's avatar
thomwolf committed
927
        """ Generate sequences for each example without beam search (num_beams == 1).
928
929
            All returned sequence are generated independantly.
        """
930
        # length of generated sentences / unfinished sentences
thomwolf's avatar
thomwolf committed
931
        unfinished_sents = input_ids.new(batch_size).fill_(1)
932
        sent_lengths = input_ids.new(batch_size).fill_(max_length)
thomwolf's avatar
thomwolf committed
933

934
        past = None
thomwolf's avatar
thomwolf committed
935
        while cur_len < max_length:
936
937
938
            model_inputs = self.prepare_inputs_for_generation(
                input_ids, past=past, encoder_inputs=encoder_inputs, attention_mask=attention_mask
            )
Sam Shleifer's avatar
Sam Shleifer committed
939

thomwolf's avatar
thomwolf committed
940
941
942
            outputs = self(**model_inputs)
            next_token_logits = outputs[0][:, -1, :]

patrickvonplaten's avatar
patrickvonplaten committed
943
            # if model has past, then set the past variable to speed up decoding
944
            if self._do_output_past(outputs):
945
946
                past = outputs[1]

thomwolf's avatar
thomwolf committed
947
948
            # repetition penalty from CTRL paper (https://arxiv.org/abs/1909.05858)
            if repetition_penalty != 1.0:
Sam Shleifer's avatar
Sam Shleifer committed
949
                self.enforce_repetition_penalty_(next_token_logits, batch_size, 1, input_ids, repetition_penalty)
thomwolf's avatar
thomwolf committed
950

Patrick von Platen's avatar
Patrick von Platen committed
951
            if no_repeat_ngram_size > 0:
patrickvonplaten's avatar
patrickvonplaten committed
952
                # calculate a list of banned tokens to prevent repetitively generating the same ngrams
Patrick von Platen's avatar
Patrick von Platen committed
953
                # from fairseq: https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345
patrickvonplaten's avatar
patrickvonplaten committed
954
                banned_tokens = calc_banned_tokens(input_ids, batch_size, no_repeat_ngram_size, cur_len)
Patrick von Platen's avatar
Patrick von Platen committed
955
                for batch_idx in range(batch_size):
956
                    next_token_logits[batch_idx, banned_tokens[batch_idx]] = -float("inf")
Patrick von Platen's avatar
Patrick von Platen committed
957

patrickvonplaten's avatar
patrickvonplaten committed
958
            # set eos token prob to zero if min_length is not reached
959
960
            if eos_token_ids is not None and cur_len < min_length:
                for eos_token_id in eos_token_ids:
961
                    next_token_logits[:, eos_token_id] = -float("inf")
962

thomwolf's avatar
thomwolf committed
963
964
            if do_sample:
                # Temperature (higher temperature => more likely to sample low probability tokens)
965
                if temperature != 1.0:
thomwolf's avatar
thomwolf committed
966
967
968
969
                    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
970
971
                probs = F.softmax(next_token_logits, dim=-1)
                next_token = torch.multinomial(probs, num_samples=1).squeeze(1)
thomwolf's avatar
thomwolf committed
972
973
            else:
                # Greedy decoding
974
                next_token = torch.argmax(next_token_logits, dim=-1)
thomwolf's avatar
thomwolf committed
975
976

            # update generations and finished sentences
977
978
979
980
981
982
            if eos_token_ids is not None:
                # pad finished sentences if eos_token_ids exist
                tokens_to_add = next_token * unfinished_sents + (pad_token_id) * (1 - unfinished_sents)
            else:
                tokens_to_add = next_token

983
            input_ids = torch.cat([input_ids, tokens_to_add.unsqueeze(-1)], dim=-1)
984
985
986
987
988
989
990
991
992
993

            if eos_token_ids is not None:
                for eos_token_id in eos_token_ids:
                    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())

thomwolf's avatar
thomwolf committed
994
995
996
997
            # stop when there is a </s> in each sentence, or if we exceed the maximul length
            if unfinished_sents.max() == 0:
                break

patrickvonplaten's avatar
patrickvonplaten committed
998
            # extend attention_mask for new generated input if only decoder
Patrick von Platen's avatar
Patrick von Platen committed
999
            if self.config.is_encoder_decoder is False:
1000
                attention_mask = torch.cat(
patrickvonplaten's avatar
patrickvonplaten committed
1001
                    [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
1002
                )
Patrick von Platen's avatar
Patrick von Platen committed
1003

1004
1005
            cur_len = cur_len + 1

1006
1007
1008
1009
1010
1011
1012
1013
1014
1015
        # 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]]
1016

1017
        return decoded
thomwolf's avatar
thomwolf committed
1018

1019
1020
1021
1022
1023
    def _generate_beam_search(
        self,
        input_ids,
        cur_len,
        max_length,
1024
        min_length,
1025
        do_sample,
Patrick von Platen's avatar
Patrick von Platen committed
1026
        early_stopping,
1027
1028
1029
1030
        temperature,
        top_k,
        top_p,
        repetition_penalty,
Patrick von Platen's avatar
Patrick von Platen committed
1031
        no_repeat_ngram_size,
Patrick von Platen's avatar
Patrick von Platen committed
1032
        bos_token_id,
1033
1034
1035
        pad_token_id,
        eos_token_ids,
        batch_size,
1036
        num_return_sequences,
1037
1038
1039
        length_penalty,
        num_beams,
        vocab_size,
Patrick von Platen's avatar
Patrick von Platen committed
1040
        encoder_inputs,
Patrick von Platen's avatar
Patrick von Platen committed
1041
        attention_mask,
1042
    ):
thomwolf's avatar
thomwolf committed
1043
        """ Generate sequences for each example with beam search.
1044
        """
thomwolf's avatar
thomwolf committed
1045
1046

        # generated hypotheses
1047
        generated_hyps = [
1048
1049
            BeamHypotheses(num_beams, max_length, length_penalty, early_stopping=early_stopping)
            for _ in range(batch_size)
1050
        ]
thomwolf's avatar
thomwolf committed
1051
1052
1053

        # scores for each sentence in the beam
        beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
patrickvonplaten's avatar
patrickvonplaten committed
1054
1055

        # 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
Patrick von Platen's avatar
Patrick von Platen committed
1056
1057
        if do_sample is False:
            beam_scores[:, 1:] = -1e9
1058
        beam_scores = beam_scores.view(-1)  # shape (batch_size * num_beams,)
thomwolf's avatar
thomwolf committed
1059
1060

        # cache compute states
1061
        past = None
thomwolf's avatar
thomwolf committed
1062
1063
1064
1065
1066

        # done sentences
        done = [False for _ in range(batch_size)]

        while cur_len < max_length:
1067
1068
1069
            model_inputs = self.prepare_inputs_for_generation(
                input_ids, past=past, encoder_inputs=encoder_inputs, attention_mask=attention_mask
            )
1070
            outputs = self(**model_inputs)  # (batch_size * num_beams, cur_len, vocab_size)
1071
            next_token_logits = outputs[0][:, -1, :]  # (batch_size * num_beams, vocab_size)
1072

patrickvonplaten's avatar
patrickvonplaten committed
1073
            # if model has past, then set the past variable to speed up decoding
1074
            if self._do_output_past(outputs):
1075
                past = outputs[1]
thomwolf's avatar
thomwolf committed
1076

1077
1078
            # repetition penalty (from CTRL paper https://arxiv.org/abs/1909.05858)
            if repetition_penalty != 1.0:
1079
                self.enforce_repetition_penalty_(
Patrick von Platen's avatar
Patrick von Platen committed
1080
                    next_token_logits, batch_size, num_beams, input_ids, repetition_penalty,
1081
                )
thomwolf's avatar
thomwolf committed
1082

patrickvonplaten's avatar
patrickvonplaten committed
1083
1084
1085
1086
            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)
1087
1088
1089
1090
            #            if (
            #                self.config.is_encoder_decoder and do_sample is False
            #            ):  # TODO(PVP) to be refactored later - do we need this boolean flag here? Also Only add for beam_search or also for no_beam_search? The prepare scores fn is ugly here
            #                scores = self.prepare_scores_for_generation(scores, cur_len, max_length)
patrickvonplaten's avatar
patrickvonplaten committed
1091
1092
1093

            # set eos token prob to zero if min_length is not reached
            if eos_token_ids is not None and cur_len < min_length:
Patrick von Platen's avatar
Patrick von Platen committed
1094
                for eos_token_id in eos_token_ids:
1095
                    scores[:, eos_token_id] = -float("inf")
Patrick von Platen's avatar
Patrick von Platen committed
1096

Patrick von Platen's avatar
Patrick von Platen committed
1097
            if no_repeat_ngram_size > 0:
patrickvonplaten's avatar
patrickvonplaten committed
1098
1099
                # calculate a list of banned tokens to prevent repetitively generating the same ngrams
                num_batch_hypotheses = batch_size * num_beams
Patrick von Platen's avatar
Patrick von Platen committed
1100
                # from fairseq: https://github.com/pytorch/fairseq/blob/a07cb6f40480928c9e0548b737aadd36ee66ac76/fairseq/sequence_generator.py#L345
1101
1102
1103
                banned_batch_tokens = calc_banned_tokens(
                    input_ids, num_batch_hypotheses, no_repeat_ngram_size, cur_len
                )
patrickvonplaten's avatar
patrickvonplaten committed
1104
                for i, banned_tokens in enumerate(banned_batch_tokens):
1105
                    scores[i, banned_tokens] = -float("inf")
Patrick von Platen's avatar
Patrick von Platen committed
1106

1107
1108
1109
            assert scores.shape == (batch_size * num_beams, vocab_size), "Shapes of scores: {} != {}".format(
                scores.shape, (batch_size * num_beams, vocab_size)
            )
1110

1111
            if do_sample:
1112
                _scores = scores + beam_scores[:, None].expand_as(scores)  # (batch_size * num_beams, vocab_size)
1113
                # Top-p/top-k filtering
1114
1115
                _scores = top_k_top_p_filtering(
                    _scores, top_k=top_k, top_p=top_p, min_tokens_to_keep=2
1116
                )  # (batch_size * num_beams, vocab_size)
1117
1118
1119
1120
1121
                # 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)

1122
                # Sample 2 next tokens for each beam (so we have some spare tokens and match output of greedy beam search)
1123
1124
                probs = F.softmax(_scores, dim=-1)
                next_tokens = torch.multinomial(probs, num_samples=2 * num_beams)  # (batch_size, num_beams * 2)
1125
                # Compute next scores
1126
                next_scores = torch.gather(_scores, -1, next_tokens)  # (batch_size, num_beams * 2)
1127
1128
1129
                # 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)
patrickvonplaten's avatar
patrickvonplaten committed
1130

1131
            else:
1132
                next_scores = scores + beam_scores[:, None].expand_as(scores)  # (batch_size * num_beams, vocab_size)
patrickvonplaten's avatar
patrickvonplaten committed
1133

1134
                # re-organize to group the beam together (we are keeping top hypothesis accross beams)
1135
1136
1137
                next_scores = next_scores.view(
                    batch_size, num_beams * vocab_size
                )  # (batch_size, num_beams * vocab_size)
1138

1139
                next_scores, next_tokens = torch.topk(next_scores, 2 * num_beams, dim=1, largest=True, sorted=True)
thomwolf's avatar
thomwolf committed
1140

1141
            assert next_scores.size() == next_tokens.size() == (batch_size, 2 * num_beams)
thomwolf's avatar
thomwolf committed
1142
1143
1144
1145
1146

            # next batch beam content
            next_batch_beam = []

            # for each sentence
1147
            for batch_idx in range(batch_size):
thomwolf's avatar
thomwolf committed
1148
1149

                # if we are done with this sentence
1150
1151
1152
1153
1154
1155
1156
                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_ids is not None and pad_token_id is not None
                    ), "generated beams >= num_beams -> eos_token_id and pad_token have to be defined"
thomwolf's avatar
thomwolf committed
1157
1158
1159
1160
1161
1162
                    next_batch_beam.extend([(0, pad_token_id, 0)] * num_beams)  # pad the batch
                    continue

                # next sentence beam content
                next_sent_beam = []

1163
                # next tokens for this sentence
1164
1165
1166
                for beam_token_rank, (beam_token_id, beam_token_score) in enumerate(
                    zip(next_tokens[batch_idx], next_scores[batch_idx])
                ):
thomwolf's avatar
thomwolf committed
1167
                    # get beam and word IDs
1168
1169
                    beam_id = beam_token_id // vocab_size
                    token_id = beam_token_id % vocab_size
thomwolf's avatar
thomwolf committed
1170

1171
                    effective_beam_id = batch_idx * num_beams + beam_id
patrickvonplaten's avatar
patrickvonplaten committed
1172

1173
                    # add to generated hypotheses if end of sentence
patrickvonplaten's avatar
patrickvonplaten committed
1174
                    if (eos_token_ids is not None) and (token_id.item() in eos_token_ids):
1175
1176
1177
                        # 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:
patrickvonplaten's avatar
patrickvonplaten committed
1178
                            continue
1179
                        generated_hyps[batch_idx].add(
1180
                            input_ids[effective_beam_id].clone(), beam_token_score.item(),
1181
                        )
thomwolf's avatar
thomwolf committed
1182
                    else:
1183
                        # add next predicted word if it is not eos_token
1184
                        next_sent_beam.append((beam_token_score, token_id, effective_beam_id))
thomwolf's avatar
thomwolf committed
1185
1186
1187
1188
1189

                    # the beam for next step is full
                    if len(next_sent_beam) == num_beams:
                        break

patrickvonplaten's avatar
patrickvonplaten committed
1190
1191
1192
1193
1194
                # 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
                )

thomwolf's avatar
thomwolf committed
1195
                # update next beam content
1196
                assert len(next_sent_beam) == num_beams, "Beam should always be full"
thomwolf's avatar
thomwolf committed
1197
                next_batch_beam.extend(next_sent_beam)
1198
                assert len(next_batch_beam) == num_beams * (batch_idx + 1)
thomwolf's avatar
thomwolf committed
1199

patrickvonplaten's avatar
patrickvonplaten committed
1200
1201
1202
1203
            # stop when we are done with each sentence
            if all(done):
                break

thomwolf's avatar
thomwolf committed
1204
1205
1206
            # 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])
1207
            beam_tokens = input_ids.new([x[1] for x in next_batch_beam])
thomwolf's avatar
thomwolf committed
1208
1209
            beam_idx = input_ids.new([x[2] for x in next_batch_beam])

1210
            # re-order batch
thomwolf's avatar
thomwolf committed
1211
            input_ids = input_ids[beam_idx, :]
1212
            input_ids = torch.cat([input_ids, beam_tokens.unsqueeze(1)], dim=-1)
1213
1214
1215

            # re-order internal states
            if past:
Sam Shleifer's avatar
Sam Shleifer committed
1216
                past = self._reorder_cache(past, beam_idx)
thomwolf's avatar
thomwolf committed
1217

patrickvonplaten's avatar
patrickvonplaten committed
1218
            # extend attention_mask for new generated input if only decoder
Patrick von Platen's avatar
Patrick von Platen committed
1219
            if self.config.is_encoder_decoder is False:
1220
                attention_mask = torch.cat(
patrickvonplaten's avatar
patrickvonplaten committed
1221
                    [attention_mask, attention_mask.new_ones((attention_mask.shape[0], 1))], dim=-1
1222
                )
Patrick von Platen's avatar
Patrick von Platen committed
1223

1224
1225
1226
1227
            # update current length
            cur_len = cur_len + 1

        # finalize all open beam hypotheses and end to generated hypotheses
1228
        for batch_idx in range(batch_size):
1229
1230
            if done[batch_idx]:
                continue
1231

1232
1233
1234
1235
1236
1237
1238
            # test that beam scores match previously calculated scores if not eos and batch_idx not done
            if eos_token_ids is not None and all(
                (token_id % vocab_size).item() not in eos_token_ids 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(
Patrick von Platen's avatar
Patrick von Platen committed
1239
                    next_scores[:, :num_beams][batch_idx], beam_scores.view(batch_size, num_beams)[batch_idx],
1240
1241
1242
1243
1244
1245
1246
1247
                )

            # 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)
thomwolf's avatar
thomwolf committed
1248

1249
1250
1251
        # 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
thomwolf's avatar
thomwolf committed
1252
1253

        # select the best hypotheses
1254
        sent_lengths = input_ids.new(output_batch_size)
thomwolf's avatar
thomwolf committed
1255
        best = []
thomwolf's avatar
thomwolf committed
1256

1257
        # retrieve best hypotheses
thomwolf's avatar
thomwolf committed
1258
        for i, hypotheses in enumerate(generated_hyps):
1259
1260
1261
1262
1263
1264
            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)
thomwolf's avatar
thomwolf committed
1265

1266
1267
1268
        # 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"
1269
            sent_max_len = min(sent_lengths.max().item() + 1, max_length)
1270
            decoded = input_ids.new(output_batch_size, sent_max_len).fill_(pad_token_id)
1271
1272
1273
1274
1275
1276
1277
1278
1279
1280

            # 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_ids[0]
        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)
thomwolf's avatar
thomwolf committed
1281

1282
1283
1284
        #        if self.config.is_encoder_decoder:
        # do not return first <EOS> token
        #            return decoded[:, 1:]
Patrick von Platen's avatar
Patrick von Platen committed
1285
        return decoded
thomwolf's avatar
thomwolf committed
1286

Patrick von Platen's avatar
Patrick von Platen committed
1287
    # force one of token_ids to be generated by setting prob of all other tokens to 0.
patrickvonplaten's avatar
patrickvonplaten committed
1288
    def _force_token_ids_generation(self, scores, token_ids):
Patrick von Platen's avatar
Patrick von Platen committed
1289
1290
1291
1292
1293
1294
1295
        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,
        )
patrickvonplaten's avatar
patrickvonplaten committed
1296
        assert len(scores.shape) == 2, "scores should be of rank 2 with shape: [batch_size, vocab_size]"
1297
        scores[:, all_but_token_ids_mask] = -float("inf")
Patrick von Platen's avatar
Patrick von Platen committed
1298

Sam Shleifer's avatar
Sam Shleifer committed
1299
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
1311
1312
    @staticmethod
    def _reorder_cache(past, beam_idx):
        reordered_past = []
        for layer_past in past:
            # get the correct batch idx from layer past batch dim
            # batch dim of `past` and `mems` is at 2nd position
            reordered_layer_past = [layer_past[:, i].unsqueeze(1).clone().detach() for i in beam_idx]
            reordered_layer_past = torch.cat(reordered_layer_past, dim=1)
            # check that shape matches
            assert reordered_layer_past.shape == layer_past.shape
            reordered_past.append(reordered_layer_past)
        past = tuple(reordered_past)
        return past

thomwolf's avatar
thomwolf committed
1313

patrickvonplaten's avatar
patrickvonplaten committed
1314
def calc_banned_tokens(prev_input_ids, num_hypos, no_repeat_ngram_size, cur_len):
Patrick von Platen's avatar
Patrick von Platen committed
1315
    # Copied from fairseq for no_repeat_ngram in beam_search"""
patrickvonplaten's avatar
patrickvonplaten committed
1316
    if cur_len + 1 < no_repeat_ngram_size:
Patrick von Platen's avatar
Patrick von Platen committed
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
1327
1328
        # 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
patrickvonplaten's avatar
patrickvonplaten committed
1329
        start_idx = cur_len + 1 - no_repeat_ngram_size
1330
        ngram_idx = tuple(prev_input_ids[hypo_idx, start_idx:cur_len].tolist())
Patrick von Platen's avatar
Patrick von Platen committed
1331
1332
1333
1334
1335
1336
        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


1337
def top_k_top_p_filtering(logits, top_k=0, top_p=1.0, filter_value=-float("Inf"), min_tokens_to_keep=1):
thomwolf's avatar
thomwolf committed
1338
1339
    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
        Args:
thomwolf's avatar
thomwolf committed
1340
            logits: logits distribution shape (batch size, vocabulary size)
1341
1342
            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).
thomwolf's avatar
thomwolf committed
1343
                Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
thomwolf's avatar
thomwolf committed
1344
            Make sure we keep at least min_tokens_to_keep per batch example in the output
thomwolf's avatar
thomwolf committed
1345
1346
1347
        From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
    """
    if top_k > 0:
thomwolf's avatar
thomwolf committed
1348
        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
thomwolf's avatar
thomwolf committed
1349
1350
1351
1352
        # 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

1353
    if top_p < 1.0:
thomwolf's avatar
thomwolf committed
1354
1355
1356
        sorted_logits, sorted_indices = torch.sort(logits, descending=True)
        cumulative_probs = torch.cumsum(F.softmax(sorted_logits, dim=-1), dim=-1)

thomwolf's avatar
thomwolf committed
1357
        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
thomwolf's avatar
thomwolf committed
1358
        sorted_indices_to_remove = cumulative_probs > top_p
thomwolf's avatar
thomwolf committed
1359
1360
1361
        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
thomwolf's avatar
thomwolf committed
1362
1363
1364
1365
1366
        # 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
1367
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
thomwolf's avatar
thomwolf committed
1368
1369
        logits[indices_to_remove] = filter_value
    return logits
thomwolf's avatar
thomwolf committed
1370
1371
1372


class BeamHypotheses(object):
1373
    def __init__(self, num_beams, max_length, length_penalty, early_stopping):
thomwolf's avatar
thomwolf committed
1374
1375
1376
1377
1378
1379
        """
        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
1380
1381
        self.num_beams = num_beams
        self.beams = []
thomwolf's avatar
thomwolf committed
1382
1383
1384
1385
1386
1387
        self.worst_score = 1e9

    def __len__(self):
        """
        Number of hypotheses in the list.
        """
1388
        return len(self.beams)
thomwolf's avatar
thomwolf committed
1389

thomwolf's avatar
thomwolf committed
1390
1391
1392
1393
1394
    def add(self, hyp, sum_logprobs):
        """
        Add a new hypothesis to the list.
        """
        score = sum_logprobs / len(hyp) ** self.length_penalty
1395
1396
1397
1398
1399
        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]]
thomwolf's avatar
thomwolf committed
1400
1401
1402
                self.worst_score = sorted_scores[1][0]
            else:
                self.worst_score = min(score, self.worst_score)
thomwolf's avatar
thomwolf committed
1403

Sam Shleifer's avatar
Sam Shleifer committed
1404
    def is_done(self, best_sum_logprobs, cur_len=None):
thomwolf's avatar
thomwolf committed
1405
1406
1407
1408
        """
        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.
        """
Sam Shleifer's avatar
Sam Shleifer committed
1409

1410
        if len(self) < self.num_beams:
thomwolf's avatar
thomwolf committed
1411
1412
1413
1414
            return False
        elif self.early_stopping:
            return True
        else:
Sam Shleifer's avatar
Sam Shleifer committed
1415
1416
1417
1418
1419
            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
thomwolf's avatar
thomwolf committed
1420
1421


thomwolf's avatar
thomwolf committed
1422
1423
class Conv1D(nn.Module):
    def __init__(self, nf, nx):
thomwolf's avatar
thomwolf committed
1424
        """ Conv1D layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2)
thomwolf's avatar
thomwolf committed
1425
1426
            Basically works like a Linear layer but the weights are transposed
        """
Julien Chaumond's avatar
Julien Chaumond committed
1427
        super().__init__()
thomwolf's avatar
thomwolf committed
1428
1429
1430
1431
1432
1433
1434
1435
1436
1437
1438
1439
1440
        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


thomwolf's avatar
thomwolf committed
1441
1442
class PoolerStartLogits(nn.Module):
    """ Compute SQuAD start_logits from sequence hidden states. """
1443

thomwolf's avatar
thomwolf committed
1444
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1445
        super().__init__()
thomwolf's avatar
thomwolf committed
1446
1447
1448
1449
        self.dense = nn.Linear(config.hidden_size, 1)

    def forward(self, hidden_states, p_mask=None):
        """ Args:
1450
1451
1452
            **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.
thomwolf's avatar
thomwolf committed
1453
        """
thomwolf's avatar
thomwolf committed
1454
1455
1456
        x = self.dense(hidden_states).squeeze(-1)

        if p_mask is not None:
1457
1458
1459
1460
            if next(self.parameters()).dtype == torch.float16:
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e30 * p_mask
thomwolf's avatar
thomwolf committed
1461
1462
1463
1464
1465
1466
1467

        return x


class PoolerEndLogits(nn.Module):
    """ Compute SQuAD end_logits from sequence hidden states and start token hidden state.
    """
1468

thomwolf's avatar
thomwolf committed
1469
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1470
        super().__init__()
thomwolf's avatar
thomwolf committed
1471
1472
1473
1474
1475
1476
1477
        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:
1478
1479
1480
1481
1482
1483
            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,)``
1484
                position of the first token for the labeled span:
1485
1486
1487
            **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.
thomwolf's avatar
thomwolf committed
1488
        """
1489
1490
1491
        assert (
            start_states is not None or start_positions is not None
        ), "One of start_states, start_positions should be not None"
thomwolf's avatar
thomwolf committed
1492
        if start_positions is not None:
1493
            slen, hsz = hidden_states.shape[-2:]
1494
1495
1496
            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)
thomwolf's avatar
thomwolf committed
1497
1498
1499
1500
1501
1502
1503

        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:
1504
1505
1506
1507
            if next(self.parameters()).dtype == torch.float16:
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e30 * p_mask
thomwolf's avatar
thomwolf committed
1508
1509
1510
1511
1512
1513

        return x


class PoolerAnswerClass(nn.Module):
    """ Compute SQuAD 2.0 answer class from classification and start tokens hidden states. """
1514

thomwolf's avatar
thomwolf committed
1515
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1516
        super().__init__()
thomwolf's avatar
thomwolf committed
1517
1518
1519
1520
1521
        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):
1522
1523
1524
1525
1526
1527
1528
1529
1530
1531
1532
1533
1534
1535
1536
        """
        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
thomwolf's avatar
thomwolf committed
1537
        """
1538
        hsz = hidden_states.shape[-1]
1539
1540
1541
        assert (
            start_states is not None or start_positions is not None
        ), "One of start_states, start_positions should be not None"
thomwolf's avatar
thomwolf committed
1542
        if start_positions is not None:
1543
1544
            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)
thomwolf's avatar
thomwolf committed
1545
1546

        if cls_index is not None:
1547
1548
            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)
thomwolf's avatar
thomwolf committed
1549
        else:
1550
            cls_token_state = hidden_states[:, -1, :]  # shape (bsz, hsz)
thomwolf's avatar
thomwolf committed
1551
1552
1553
1554
1555
1556
1557
1558
1559

        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):
1560
1561
1562
    r""" A SQuAD head inspired by XLNet.

    Parameters:
1563
        config (:class:`~transformers.XLNetConfig`): Model configuration class with all the parameters of the model.
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582

    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.
thomwolf's avatar
thomwolf committed
1583
        **start_top_log_probs**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1584
1585
            ``torch.FloatTensor`` of shape ``(batch_size, config.start_n_top)``
            Log probabilities for the top config.start_n_top start token possibilities (beam-search).
thomwolf's avatar
thomwolf committed
1586
        **start_top_index**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1587
1588
            ``torch.LongTensor`` of shape ``(batch_size, config.start_n_top)``
            Indices for the top config.start_n_top start token possibilities (beam-search).
thomwolf's avatar
thomwolf committed
1589
        **end_top_log_probs**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1590
1591
            ``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).
thomwolf's avatar
thomwolf committed
1592
        **end_top_index**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1593
1594
            ``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).
thomwolf's avatar
thomwolf committed
1595
        **cls_logits**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1596
1597
            ``torch.FloatTensor`` of shape ``(batch_size,)``
            Log probabilities for the ``is_impossible`` label of the answers.
thomwolf's avatar
thomwolf committed
1598
    """
1599

thomwolf's avatar
thomwolf committed
1600
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1601
        super().__init__()
thomwolf's avatar
thomwolf committed
1602
1603
1604
1605
1606
1607
1608
        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)

1609
    def forward(
Patrick von Platen's avatar
Patrick von Platen committed
1610
        self, hidden_states, start_positions=None, end_positions=None, cls_index=None, is_impossible=None, p_mask=None,
1611
    ):
thomwolf's avatar
thomwolf committed
1612
1613
        outputs = ()

thomwolf's avatar
thomwolf committed
1614
        start_logits = self.start_logits(hidden_states, p_mask=p_mask)
thomwolf's avatar
thomwolf committed
1615
1616
1617
1618
1619
1620
1621
1622
1623
1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
1634
1635
1636
1637

        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
1638
1639

            outputs = (total_loss,) + outputs
thomwolf's avatar
thomwolf committed
1640
1641
1642
1643

        else:
            # during inference, compute the end logits based on beam search
            bsz, slen, hsz = hidden_states.size()
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
            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)
thomwolf's avatar
thomwolf committed
1656
1657
            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)
1658
            end_log_probs = F.softmax(end_logits, dim=1)  # shape (bsz, slen, start_n_top)
thomwolf's avatar
thomwolf committed
1659

1660
1661
1662
            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)
thomwolf's avatar
thomwolf committed
1663
1664
1665
1666
1667
1668
            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)

Patrick von Platen's avatar
Patrick von Platen committed
1669
            outputs = (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits,) + outputs
thomwolf's avatar
thomwolf committed
1670
1671

        # return start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits
1672
        # or (if labels are provided) (total_loss,)
thomwolf's avatar
thomwolf committed
1673
1674
1675
1676
        return outputs


class SequenceSummary(nn.Module):
thomwolf's avatar
thomwolf committed
1677
    r""" Compute a single vector summary of a sequence hidden states according to various possibilities:
thomwolf's avatar
thomwolf committed
1678
1679
1680
1681
1682
        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
thomwolf's avatar
thomwolf committed
1683
                - 'cls_index' => supply a Tensor of classification token position (GPT/GPT-2)
thomwolf's avatar
thomwolf committed
1684
1685
                - 'attn' => Not implemented now, use multi-head attention
            summary_use_proj: Add a projection after the vector extraction
1686
            summary_proj_to_labels: If True, the projection outputs to config.num_labels classes (otherwise to hidden_size). Default: False.
1687
            summary_activation: 'tanh' or another string => add an activation to the output, Other => no activation. Default
1688
1689
            summary_first_dropout: Add a dropout before the projection and activation
            summary_last_dropout: Add a dropout after the projection and activation
thomwolf's avatar
thomwolf committed
1690
    """
1691

1692
    def __init__(self, config: PretrainedConfig):
Julien Chaumond's avatar
Julien Chaumond committed
1693
        super().__init__()
thomwolf's avatar
thomwolf committed
1694

1695
        self.summary_type = getattr(config, "summary_type", "last")
1696
        if self.summary_type == "attn":
thomwolf's avatar
thomwolf committed
1697
1698
1699
1700
1701
            # 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

thomwolf's avatar
thomwolf committed
1702
        self.summary = Identity()
1703
1704
        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:
1705
                num_classes = config.num_labels
thomwolf's avatar
thomwolf committed
1706
1707
1708
1709
            else:
                num_classes = config.hidden_size
            self.summary = nn.Linear(config.hidden_size, num_classes)

1710
1711
1712
1713
        activation_string = getattr(config, "summary_activation", None)
        self.activation = (
            get_activation(activation_string) if activation_string else Identity()
        )  # type: typing.Callable
thomwolf's avatar
thomwolf committed
1714

thomwolf's avatar
thomwolf committed
1715
        self.first_dropout = Identity()
1716
        if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0:
1717
1718
            self.first_dropout = nn.Dropout(config.summary_first_dropout)

thomwolf's avatar
thomwolf committed
1719
        self.last_dropout = Identity()
1720
        if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0:
1721
            self.last_dropout = nn.Dropout(config.summary_last_dropout)
thomwolf's avatar
thomwolf committed
1722

thomwolf's avatar
thomwolf committed
1723
    def forward(self, hidden_states, cls_index=None):
1724
        """ hidden_states: float Tensor in shape [bsz, ..., seq_len, hidden_size], the hidden-states of the last layer.
thomwolf's avatar
thomwolf committed
1725
            cls_index: [optional] position of the classification token if summary_type == 'cls_index',
thomwolf's avatar
thomwolf committed
1726
                shape (bsz,) or more generally (bsz, ...) where ... are optional leading dimensions of hidden_states.
thomwolf's avatar
thomwolf committed
1727
                if summary_type == 'cls_index' and cls_index is None:
thomwolf's avatar
thomwolf committed
1728
1729
                    we take the last token of the sequence as classification token
        """
1730
        if self.summary_type == "last":
thomwolf's avatar
thomwolf committed
1731
            output = hidden_states[:, -1]
1732
        elif self.summary_type == "first":
thomwolf's avatar
thomwolf committed
1733
            output = hidden_states[:, 0]
1734
        elif self.summary_type == "mean":
thomwolf's avatar
thomwolf committed
1735
            output = hidden_states.mean(dim=1)
1736
        elif self.summary_type == "cls_index":
thomwolf's avatar
thomwolf committed
1737
            if cls_index is None:
Patrick von Platen's avatar
Patrick von Platen committed
1738
                cls_index = torch.full_like(hidden_states[..., :1, :], hidden_states.shape[-2] - 1, dtype=torch.long,)
thomwolf's avatar
thomwolf committed
1739
            else:
thomwolf's avatar
thomwolf committed
1740
                cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)
1741
                cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))
thomwolf's avatar
thomwolf committed
1742
            # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states
1743
1744
            output = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, XX, hidden_size)
        elif self.summary_type == "attn":
thomwolf's avatar
thomwolf committed
1745
1746
            raise NotImplementedError

1747
        output = self.first_dropout(output)
thomwolf's avatar
thomwolf committed
1748
1749
        output = self.summary(output)
        output = self.activation(output)
1750
        output = self.last_dropout(output)
thomwolf's avatar
thomwolf committed
1751
1752
1753
1754

        return output


Sam Shleifer's avatar
Sam Shleifer committed
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
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_indicies = torch.cumsum(mask, dim=1).type_as(mask) * mask
    return incremental_indicies.long() + padding_idx


1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
1814
1815
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
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828


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__))