modeling_utils.py 78.2 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
20
# 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
21
import typing
22
23
24

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

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

Aymeric Augustin's avatar
Aymeric Augustin committed
40

41
42
logger = logging.getLogger(__name__)

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

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

        def forward(self, input):
            return input

57

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


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

74
        :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
75
        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.
76
77

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

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

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

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

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

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

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

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

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

        Args:
            value (:obj:`nn.Module`):
                A module mapping vocabulary to hidden states.
thomwolf's avatar
thomwolf committed
139
140
141
142
143
144
        """
        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
145

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

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

156
    def tie_weights(self):
157
158
159
160
        """
        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
161
        """
thomwolf's avatar
thomwolf committed
162
163
164
        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
165

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

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

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

188
189
190
        Arguments:

            new_num_tokens: (`optional`) int:
191
                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.
192
                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
193

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

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

thomwolf's avatar
thomwolf committed
209
210
        return model_embeds

211
    def _resize_token_embeddings(self, new_num_tokens):
thomwolf's avatar
thomwolf committed
212
213
214
215
        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()
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
244
245
246
247
248
249
250

    def _get_resized_embeddings(self, old_embeddings, new_num_tokens=None):
        """ Build a resized Embedding Module from a provided token Embedding Module.
            Increasing the size will add newly initialized vectors at the end
            Reducing the size will remove vectors from the end

        Args:
            new_num_tokens: (`optional`) int
                New number of tokens in the embedding matrix.
                Increasing the size will add newly initialized vectors at the end
                Reducing the size will remove vectors from the end
                If not provided or None: return the provided token Embedding Module.
        Return: ``torch.nn.Embeddings``
            Pointer to the resized Embedding Module or the old Embedding Module if new_num_tokens is None
        """
        if new_num_tokens is None:
            return old_embeddings

        old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
        if old_num_tokens == new_num_tokens:
            return old_embeddings

        # Build new embeddings
        new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim)
        new_embeddings.to(old_embeddings.weight.device)

        # initialize all new embeddings (in particular added tokens)
        self._init_weights(new_embeddings)

        # Copy word embeddings from the previous weights
        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)
        new_embeddings.weight.data[:num_tokens_to_copy, :] = old_embeddings.weight.data[:num_tokens_to_copy, :]

        return new_embeddings

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

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

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

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

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

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

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

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

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

295
296
297
        # 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
298
        logger.info("Model weights saved in {}".format(output_model_file))
299

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

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

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

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

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

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

            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
333
                This option can be used if you want to create a model from a pretrained configuration but load your own weights.
334
                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.
335
336

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

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

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

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

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

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

        Examples::
thomwolf's avatar
thomwolf committed
360

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

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

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

thomwolf's avatar
thomwolf committed
398
        # Load model
thomwolf's avatar
thomwolf committed
399
        if pretrained_model_name_or_path is not None:
400
            if pretrained_model_name_or_path in cls.pretrained_model_archive_map:
thomwolf's avatar
thomwolf committed
401
402
                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
403
404
                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
405
                    archive_file = os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
thomwolf's avatar
thomwolf committed
406
407
408
409
410
                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
411
                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
thomwolf's avatar
thomwolf committed
412
                else:
413
414
415
416
417
                    raise EnvironmentError(
                        "Error no file named {} found in directory {} or `from_tf` set to False".format(
                            [WEIGHTS_NAME, TF2_WEIGHTS_NAME, TF_WEIGHTS_NAME + ".index"], pretrained_model_name_or_path
                        )
                    )
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
429
430
                archive_file = hf_bucket_url(
                    pretrained_model_name_or_path, postfix=(TF2_WEIGHTS_NAME if from_tf else WEIGHTS_NAME)
                )
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(
524
525
                    state_dict, prefix, local_metadata, True, missing_keys, unexpected_keys, error_msgs
                )
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
543
                model_to_load = getattr(model, cls.base_model_prefix)

            load(model_to_load, prefix=start_prefix)
            if len(missing_keys) > 0:
544
545
546
547
548
                logger.info(
                    "Weights of {} not initialized from pretrained model: {}".format(
                        model.__class__.__name__, missing_keys
                    )
                )
549
            if len(unexpected_keys) > 0:
550
551
552
553
554
                logger.info(
                    "Weights from pretrained model not used in {}: {}".format(
                        model.__class__.__name__, unexpected_keys
                    )
                )
555
            if len(error_msgs) > 0:
556
557
558
559
560
                raise RuntimeError(
                    "Error(s) in loading state_dict for {}:\n\t{}".format(
                        model.__class__.__name__, "\n\t".join(error_msgs)
                    )
                )
561
        model.tie_weights()  # make sure word embedding weights are still tied if needed
562

563
564
565
        # Set model in evaluation mode to desactivate DropOut modules by default
        model.eval()

thomwolf's avatar
thomwolf committed
566
567
568
569
        if output_loading_info:
            loading_info = {"missing_keys": missing_keys, "unexpected_keys": unexpected_keys, "error_msgs": error_msgs}
            return model, loading_info

570
571
        return model

thomwolf's avatar
thomwolf committed
572
573
574
    def prepare_inputs_for_generation(self, input_ids, **kwargs):
        return {"input_ids": input_ids}

575
    def _do_output_past(self, outputs):
Sam Shleifer's avatar
Sam Shleifer committed
576
577
578
579
580
581
        """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:
582
            return True
583
584
        return False

Sam Shleifer's avatar
Sam Shleifer committed
585
586
587
588
589
590
591
592
593
594
    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
595
    @torch.no_grad()
596
597
598
599
    def generate(
        self,
        input_ids=None,
        max_length=None,
600
        do_sample=True,
601
602
603
604
605
606
607
608
609
610
611
        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,
        num_return_sequences=None,
    ):
612
        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
613
        and beam-search.
thomwolf's avatar
thomwolf committed
614

615
616
617
618
619
620
621
        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
622

623
            input_ids: (`optional`) `torch.LongTensor` of shape `(batch_size, sequence_length)`
thomwolf's avatar
thomwolf committed
624
                The sequence used as a prompt for the generation. If `None` the method initializes
625
626
627
                it as an empty `torch.LongTensor` of shape `(1,)`.

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

            do_sample: (`optional`) bool
631
                If set to `False` greedy decoding is used. Otherwise sampling is used. Defaults to `True`.
632
633
634
635
636

            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
637
                The value used to module the next token probabilities. Must be strictly positive. Default to 1.0.
638
639

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

            top_p: (`optional`) float
thomwolf's avatar
thomwolf committed
643
                The cumulative probability of parameter highest probability vocabulary tokens to keep for nucleus sampling. Must be between 0 and 1. Default to 1.
644
645
646
647
648

            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
649
                Beginning of sentence token if no prompt is provided. Default to 0.
650
651

            eos_token_ids: (`optional`) int or list of int
thomwolf's avatar
thomwolf committed
652
                End of sequence token or list of tokens to stop the generation. Default to 0.
653
            length_penalty: (`optional`) float
thomwolf's avatar
thomwolf committed
654
                Exponential penalty to the length. Default to 1.
655
656
657
658

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

659
660
661
662
663
        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`

664
665
666
667
        Examples::

            tokenizer = AutoTokenizer.from_pretrained('distilgpt2')   # Initialize tokenizer
            model = AutoModelWithLMHead.from_pretrained('distilgpt2')    # Download model and configuration from S3 and cache.
668
            outputs = model.generate(max_length=40, bos_token_id=tokenizer.bos_token_id, eos_token_ids=tokenizer.eos_token_id, do_sample=False)  # do greedy decoding
669
670
671
672
673
674
            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
675
            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'
676
            for i in range(3): #  3 output sequences were generated
677
                print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))
678
679
680
681
682

            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
683
684
685
            outputs = model.generate(input_ids=input_ids, max_length=40, temperature=0.7, bos_token_id=tokenizer.bos_token_id, pad_token_id=tokenizer.pad_token_id, eos_token_ids=tokenizer.eos_token_id, num_return_sequences=3)  # 3 generate sequences using by sampling
            for i in range(3): #  3 output sequences were generated
                print('Generated {}: {}'.format(i, tokenizer.decode(outputs[i], skip_special_tokens=True)))
686
687
688
689
690

            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
691
            outputs = model.generate(input_ids=input_ids, max_length=50, temperature=0.7, repetition_penalty=1.2)  # generate sequences
692
693
            print('Generated: {}'.format(tokenizer.decode(outputs[0], skip_special_tokens=True)))

thomwolf's avatar
thomwolf committed
694
695
696
697
        """

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

703
704
705
706
707
708
709
710
711
712
713
        max_length = max_length if max_length is not None else self.config.max_length
        do_sample = do_sample if do_sample is not None else self.config.do_sample
        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
714
715
716
        num_return_sequences = (
            num_return_sequences if num_return_sequences is not None else self.config.num_return_sequences
        )
thomwolf's avatar
thomwolf committed
717
718
719

        if input_ids is not None:
            batch_size = input_ids.shape[0]  # overriden by the input batch_size
thomwolf's avatar
thomwolf committed
720
721
        else:
            batch_size = 1
thomwolf's avatar
thomwolf committed
722
723
724
        if isinstance(eos_token_ids, int):
            eos_token_ids = [eos_token_ids]

Sam Shleifer's avatar
Sam Shleifer committed
725
        assert isinstance(max_length, int) and max_length > 0, "`max_length` should be a strictly positive integer."
thomwolf's avatar
thomwolf committed
726
        assert isinstance(do_sample, bool), "`do_sample` should be a boolean."
Sam Shleifer's avatar
Sam Shleifer committed
727
728
        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."
729
        assert isinstance(top_k, int) and top_k >= 0, "`top_k` should be a positive integer."
thomwolf's avatar
thomwolf committed
730
        assert 0 <= top_p <= 1, "`top_p` should be between 0 and 1."
thomwolf's avatar
thomwolf committed
731
        assert repetition_penalty >= 1.0, "`repetition_penalty` should be >= 1."
732
733
734
735
736
737
738
739
        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)
740
        ), "`eos_token_ids` should be a positive integer or a list/tuple of positive integers."
Sam Shleifer's avatar
Sam Shleifer committed
741
        assert length_penalty > 0, "`length_penalty` should be strictly positive."
742
743
        assert (
            isinstance(num_return_sequences, int) and num_return_sequences > 0
Sam Shleifer's avatar
Sam Shleifer committed
744
        ), "`num_return_sequences` should be a strictly positive integer."
thomwolf's avatar
thomwolf committed
745
746

        if input_ids is None:
747
748
749
750
            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."
            )
751
752
753
            input_ids = torch.full(
                (batch_size, 1), bos_token_id, dtype=torch.long, device=next(self.parameters()).device
            )
thomwolf's avatar
thomwolf committed
754
        else:
755
            assert input_ids.dim() == 2, "Input prompt should be of shape (batch_size, sequence length)."
thomwolf's avatar
thomwolf committed
756

757
758
759
760
761
762
763
764
765
766
767
768
769
        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"

770
771
772
773
774
775
        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
776
        # current position and vocab size
thomwolf's avatar
thomwolf committed
777
        cur_len = input_ids.shape[1]
thomwolf's avatar
thomwolf committed
778
779
        vocab_size = self.config.vocab_size

780
        if num_return_sequences != 1 and do_sample:
thomwolf's avatar
thomwolf committed
781
782
            # Expand input to num return sequences
            input_ids = input_ids.unsqueeze(1).expand(batch_size, num_return_sequences, cur_len)
783
784
            input_ids = input_ids.contiguous().view(
                batch_size * num_return_sequences, cur_len
Sam Shleifer's avatar
Sam Shleifer committed
785
            )  # shape: (batch_size * num_return_sequences, cur_len)
thomwolf's avatar
thomwolf committed
786
787
788
789
            effective_batch_size = batch_size * num_return_sequences
        else:
            effective_batch_size = batch_size

thomwolf's avatar
thomwolf committed
790
        if num_beams > 1:
791
792
793
794
795
796
797
798
799
800
801
802
            output = self._generate_beam_search(
                input_ids,
                cur_len,
                max_length,
                do_sample,
                temperature,
                top_k,
                top_p,
                repetition_penalty,
                pad_token_id,
                eos_token_ids,
                effective_batch_size,
803
                num_return_sequences,
804
805
806
807
                length_penalty,
                num_beams,
                vocab_size,
            )
thomwolf's avatar
thomwolf committed
808
        else:
809
810
811
812
813
814
815
816
817
818
819
820
821
            output = self._generate_no_beam_search(
                input_ids,
                cur_len,
                max_length,
                do_sample,
                temperature,
                top_k,
                top_p,
                repetition_penalty,
                pad_token_id,
                eos_token_ids,
                effective_batch_size,
            )
thomwolf's avatar
thomwolf committed
822
823

        return output
thomwolf's avatar
thomwolf committed
824

825
826
827
828
829
830
831
832
833
834
835
836
837
838
    def _generate_no_beam_search(
        self,
        input_ids,
        cur_len,
        max_length,
        do_sample,
        temperature,
        top_k,
        top_p,
        repetition_penalty,
        pad_token_id,
        eos_token_ids,
        batch_size,
    ):
thomwolf's avatar
thomwolf committed
839
        """ Generate sequences for each example without beam search (num_beams == 1).
840
841
            All returned sequence are generated independantly.
        """
thomwolf's avatar
thomwolf committed
842
        # current position / max lengths / length of generated sentences / unfinished sentences
843

thomwolf's avatar
thomwolf committed
844
        unfinished_sents = input_ids.new(batch_size).fill_(1)
845
        sent_lengths = input_ids.new(batch_size).fill_(max_length)
thomwolf's avatar
thomwolf committed
846

847
        past = None
thomwolf's avatar
thomwolf committed
848
        while cur_len < max_length:
849
            model_inputs = self.prepare_inputs_for_generation(input_ids, past=past)
Sam Shleifer's avatar
Sam Shleifer committed
850

thomwolf's avatar
thomwolf committed
851
852
853
            outputs = self(**model_inputs)
            next_token_logits = outputs[0][:, -1, :]

patrickvonplaten's avatar
patrickvonplaten committed
854
            # if model has past, then set the past variable to speed up decoding
855
            if self._do_output_past(outputs):
856
857
                past = outputs[1]

thomwolf's avatar
thomwolf committed
858
859
            # repetition penalty from CTRL paper (https://arxiv.org/abs/1909.05858)
            if repetition_penalty != 1.0:
Sam Shleifer's avatar
Sam Shleifer committed
860
                self.enforce_repetition_penalty_(next_token_logits, batch_size, 1, input_ids, repetition_penalty)
thomwolf's avatar
thomwolf committed
861
862
863

            if do_sample:
                # Temperature (higher temperature => more likely to sample low probability tokens)
864
                if temperature != 1.0:
thomwolf's avatar
thomwolf committed
865
866
867
868
                    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
869
                next_token = torch.multinomial(F.softmax(next_token_logits, dim=-1), num_samples=1).squeeze(1)
thomwolf's avatar
thomwolf committed
870
871
            else:
                # Greedy decoding
872
                next_token = torch.argmax(next_token_logits, dim=-1)
thomwolf's avatar
thomwolf committed
873
874

            # update generations and finished sentences
875
876
877
878
879
880
            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

881
            input_ids = torch.cat([input_ids, tokens_to_add.unsqueeze(-1)], dim=-1)
882
883
884
885
886
887
888
889
890
891

            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
892
893
894
895
896
897
            cur_len = cur_len + 1

            # stop when there is a </s> in each sentence, or if we exceed the maximul length
            if unfinished_sents.max() == 0:
                break

898
899
900
901
902
903
904
905
906
907
        # 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]]
908

909
        return decoded
thomwolf's avatar
thomwolf committed
910

911
912
913
914
915
916
917
918
919
920
921
922
923
    def _generate_beam_search(
        self,
        input_ids,
        cur_len,
        max_length,
        do_sample,
        temperature,
        top_k,
        top_p,
        repetition_penalty,
        pad_token_id,
        eos_token_ids,
        batch_size,
924
        num_return_sequences,
925
926
927
928
        length_penalty,
        num_beams,
        vocab_size,
    ):
thomwolf's avatar
thomwolf committed
929
        """ Generate sequences for each example with beam search.
930
        """
931

thomwolf's avatar
thomwolf committed
932
        # Expand input to num beams
Sam Shleifer's avatar
Sam Shleifer committed
933
        # assert input_ids.shape == (batch_size * num_beams, cur_len)
thomwolf's avatar
thomwolf committed
934
        input_ids = input_ids.unsqueeze(1).expand(batch_size, num_beams, cur_len)
935
        input_ids = input_ids.contiguous().view(batch_size * num_beams, cur_len)  # (batch_size * num_beams, cur_len)
thomwolf's avatar
thomwolf committed
936
937

        # generated hypotheses
938
939
940
        generated_hyps = [
            BeamHypotheses(num_beams, max_length, length_penalty, early_stopping=False) for _ in range(batch_size)
        ]
thomwolf's avatar
thomwolf committed
941
942
943
944

        # scores for each sentence in the beam
        beam_scores = torch.zeros((batch_size, num_beams), dtype=torch.float, device=input_ids.device)
        beam_scores[:, 1:] = -1e9
945
        beam_scores = beam_scores.view(-1)  # shape (batch_size * num_beams,)
thomwolf's avatar
thomwolf committed
946
947

        # cache compute states
948
        past = None
thomwolf's avatar
thomwolf committed
949
950
951
952
953

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

        while cur_len < max_length:
954
955
956
957
            model_inputs = self.prepare_inputs_for_generation(input_ids, past=past)
            outputs = self(**model_inputs)  # (batch_size * num_beams, cur_len, vocab_size)
            scores = outputs[0][:, -1, :]  # (batch_size * num_beams, vocab_size)

patrickvonplaten's avatar
patrickvonplaten committed
958
            # if model has past, then set the past variable to speed up decoding
959
            if self._do_output_past(outputs):
960
                past = outputs[1]
thomwolf's avatar
thomwolf committed
961

962
963
            # repetition penalty (from CTRL paper https://arxiv.org/abs/1909.05858)
            if repetition_penalty != 1.0:
Sam Shleifer's avatar
Sam Shleifer committed
964
                self.enforce_repetition_penalty_(scores, batch_size, num_beams, input_ids, repetition_penalty)
thomwolf's avatar
thomwolf committed
965

966
967
            if do_sample:
                # Temperature (higher temperature => more likely to sample low probability tokens)
968
                if temperature != 1.0:
969
970
                    scores = scores / temperature
                # Top-p/top-k filtering
971
972
973
                scores = top_k_top_p_filtering(
                    scores, top_k=top_k, top_p=top_p, min_tokens_to_keep=2
                )  # (batch_size * num_beams, vocab_size)
974
                # Sample 2 next words for each beam (so we have some spare tokens and match output of greedy beam search)
975
                next_words = torch.multinomial(F.softmax(scores, dim=-1), num_samples=2)  # (batch_size * num_beams, 2)
976
                # Compute next scores
977
978
979
                _scores = F.log_softmax(scores, dim=-1)  # (batch_size * num_beams, vocab_size)
                _scores = torch.gather(_scores, -1, next_words)  # (batch_size * num_beams, 2)
                next_scores = _scores + beam_scores[:, None].expand_as(_scores)  # (batch_size * num_beams, 2)
980
                # Match shape of greedy beam search
981
982
                next_words = next_words.view(batch_size, 2 * num_beams)  # (batch_size, 2 * num_beams)
                next_scores = next_scores.view(batch_size, 2 * num_beams)  # (batch_size, 2 * num_beams)
983
984
            else:
                # do greedy beam search
985
                scores = F.log_softmax(scores, dim=-1)  # (batch_size * num_beams, vocab_size)
986
987
                assert scores.size() == (batch_size * num_beams, vocab_size)
                # Add the log prob of the new beams to the log prob of the beginning of the sequence (sum of logs == log of the product)
988
                _scores = scores + beam_scores[:, None].expand_as(scores)  # (batch_size * num_beams, vocab_size)
989
                # re-organize to group the beam together (we are keeping top hypothesis accross beams)
990
991
                _scores = _scores.view(batch_size, num_beams * vocab_size)  # (batch_size, num_beams * vocab_size)
                next_scores, next_words = torch.topk(_scores, 2 * num_beams, dim=1, largest=True, sorted=True)
thomwolf's avatar
thomwolf committed
992
993
994
995
996
997
998
999

            assert next_scores.size() == next_words.size() == (batch_size, 2 * num_beams)

            # next batch beam content
            # list of (batch_size * num_beams) tuple(next hypothesis score, next word, current position in the batch)
            next_batch_beam = []

            # for each sentence
1000
            for batch_idx in range(batch_size):
thomwolf's avatar
thomwolf committed
1001
1002

                # if we are done with this sentence
1003
1004
1005
1006
1007
1008
1009
1010
1011
1012
                done[batch_idx] = done[batch_idx] or generated_hyps[batch_idx].is_done(
                    next_scores[batch_idx].max().item()
                )
                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
1013
1014
1015
1016
1017
1018
1019
                    next_batch_beam.extend([(0, pad_token_id, 0)] * num_beams)  # pad the batch
                    continue

                # next sentence beam content
                next_sent_beam = []

                # next words for this sentence
1020
                for idx, score in zip(next_words[batch_idx], next_scores[batch_idx]):
thomwolf's avatar
thomwolf committed
1021
1022
1023
1024
1025

                    # get beam and word IDs
                    beam_id = idx // vocab_size
                    word_id = idx % vocab_size

1026
1027
1028
1029
                    # add to generated hypotheses if end of sentence or last iteration
                    if eos_token_ids is not None and word_id.item() in eos_token_ids:
                        generated_hyps[batch_idx].add(
                            input_ids[batch_idx * num_beams + beam_id, :cur_len].clone(), score.item()
1030
                        )
thomwolf's avatar
thomwolf committed
1031
                    else:
1032
1033
                        # add next predicted word if it is not eos_token
                        next_sent_beam.append((score, word_id, batch_idx * num_beams + beam_id))
thomwolf's avatar
thomwolf committed
1034
1035
1036
1037
1038
1039

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

                # update next beam content
1040
                assert len(next_sent_beam) == num_beams, "Beam should always be full"
thomwolf's avatar
thomwolf committed
1041
                next_batch_beam.extend(next_sent_beam)
1042
                assert len(next_batch_beam) == num_beams * (batch_idx + 1)
thomwolf's avatar
thomwolf committed
1043
1044
1045
1046
1047
1048
1049

            # sanity check / prepare next batch
            assert len(next_batch_beam) == batch_size * num_beams
            beam_scores = beam_scores.new([x[0] for x in next_batch_beam])
            beam_words = input_ids.new([x[1] for x in next_batch_beam])
            beam_idx = input_ids.new([x[2] for x in next_batch_beam])

1050
            # re-order batch
thomwolf's avatar
thomwolf committed
1051
1052
            input_ids = input_ids[beam_idx, :]
            input_ids = torch.cat([input_ids, beam_words.unsqueeze(1)], dim=-1)
1053
1054
1055

            # re-order internal states
            if past:
Sam Shleifer's avatar
Sam Shleifer committed
1056
                past = self._reorder_cache(past, beam_idx)
thomwolf's avatar
thomwolf committed
1057
1058
1059
1060
1061
1062
1063
1064

            # update current length
            cur_len = cur_len + 1

            # stop when we are done with each sentence
            if all(done):
                break

1065
1066
1067
1068
1069
1070
1071
1072
1073
1074
1075
        for batch_idx in range(batch_size):
            # Add all open beam hypothesis to generated_hyps
            if not done[batch_idx]:
                for idx, score in zip(next_words[batch_idx], next_scores[batch_idx]):

                    # get beam and word IDs
                    beam_id = idx // vocab_size
                    word_id = idx % vocab_size
                    generated_hyps[batch_idx].add(
                        input_ids[batch_idx * num_beams + beam_id, :cur_len].clone(), score.item()
                    )
thomwolf's avatar
thomwolf committed
1076

1077
1078
1079
1080
        # 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
1081
        # select the best hypotheses
1082
        sent_lengths = input_ids.new(output_batch_size)
thomwolf's avatar
thomwolf committed
1083
        best = []
thomwolf's avatar
thomwolf committed
1084

1085
        # retrieve best hypotheses
thomwolf's avatar
thomwolf committed
1086
        for i, hypotheses in enumerate(generated_hyps):
1087
1088
1089
1090
1091
1092
            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
1093

1094
1095
1096
1097
        # shorter batches are filled with pad_token
        if sent_lengths.min().item() != sent_lengths.max().item():
            assert pad_token_id is not None, "`Pad_token_id` has to be defined"
            sent_max_len = min(sent_lengths.max().item() + 1, max_length)
1098
            decoded = input_ids.new(output_batch_size, sent_max_len).fill_(pad_token_id)
1099
1100
1101
1102
1103
1104
1105
1106
1107
1108

            # 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
1109

thomwolf's avatar
thomwolf committed
1110
1111
        return decoded

Sam Shleifer's avatar
Sam Shleifer committed
1112
1113
1114
1115
1116
1117
1118
1119
1120
1121
1122
1123
1124
1125
    @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
1126

1127
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
1128
1129
    """ Filter a distribution of logits using top-k and/or nucleus (top-p) filtering
        Args:
thomwolf's avatar
thomwolf committed
1130
            logits: logits distribution shape (batch size, vocabulary size)
1131
1132
            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
1133
                Nucleus filtering is described in Holtzman et al. (http://arxiv.org/abs/1904.09751)
thomwolf's avatar
thomwolf committed
1134
            Make sure we keep at least min_tokens_to_keep per batch example in the output
thomwolf's avatar
thomwolf committed
1135
1136
1137
        From: https://gist.github.com/thomwolf/1a5a29f6962089e871b94cbd09daf317
    """
    if top_k > 0:
thomwolf's avatar
thomwolf committed
1138
        top_k = min(max(top_k, min_tokens_to_keep), logits.size(-1))  # Safety check
thomwolf's avatar
thomwolf committed
1139
1140
1141
1142
        # 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

1143
    if top_p < 1.0:
thomwolf's avatar
thomwolf committed
1144
1145
1146
        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
1147
        # Remove tokens with cumulative probability above the threshold (token with 0 are kept)
thomwolf's avatar
thomwolf committed
1148
        sorted_indices_to_remove = cumulative_probs > top_p
thomwolf's avatar
thomwolf committed
1149
1150
1151
        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
1152
1153
1154
1155
1156
        # 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
1157
        indices_to_remove = sorted_indices_to_remove.scatter(1, sorted_indices, sorted_indices_to_remove)
thomwolf's avatar
thomwolf committed
1158
1159
        logits[indices_to_remove] = filter_value
    return logits
thomwolf's avatar
thomwolf committed
1160
1161
1162


class BeamHypotheses(object):
1163
    def __init__(self, num_beams, max_length, length_penalty, early_stopping):
thomwolf's avatar
thomwolf committed
1164
1165
1166
1167
1168
1169
        """
        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
1170
1171
        self.num_beams = num_beams
        self.beams = []
thomwolf's avatar
thomwolf committed
1172
1173
1174
1175
1176
1177
        self.worst_score = 1e9

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

thomwolf's avatar
thomwolf committed
1180
1181
1182
1183
1184
    def add(self, hyp, sum_logprobs):
        """
        Add a new hypothesis to the list.
        """
        score = sum_logprobs / len(hyp) ** self.length_penalty
1185
1186
1187
1188
1189
        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
1190
1191
1192
                self.worst_score = sorted_scores[1][0]
            else:
                self.worst_score = min(score, self.worst_score)
thomwolf's avatar
thomwolf committed
1193

Sam Shleifer's avatar
Sam Shleifer committed
1194
    def is_done(self, best_sum_logprobs, cur_len=None):
thomwolf's avatar
thomwolf committed
1195
1196
1197
1198
        """
        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
1199

1200
        if len(self) < self.num_beams:
thomwolf's avatar
thomwolf committed
1201
1202
1203
1204
            return False
        elif self.early_stopping:
            return True
        else:
Sam Shleifer's avatar
Sam Shleifer committed
1205
1206
1207
1208
1209
            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
1210
1211


thomwolf's avatar
thomwolf committed
1212
1213
class Conv1D(nn.Module):
    def __init__(self, nf, nx):
thomwolf's avatar
thomwolf committed
1214
        """ Conv1D layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2)
thomwolf's avatar
thomwolf committed
1215
1216
            Basically works like a Linear layer but the weights are transposed
        """
Julien Chaumond's avatar
Julien Chaumond committed
1217
        super().__init__()
thomwolf's avatar
thomwolf committed
1218
1219
1220
1221
1222
1223
1224
1225
1226
1227
1228
1229
1230
        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
1231
1232
class PoolerStartLogits(nn.Module):
    """ Compute SQuAD start_logits from sequence hidden states. """
1233

thomwolf's avatar
thomwolf committed
1234
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1235
        super().__init__()
thomwolf's avatar
thomwolf committed
1236
1237
1238
1239
        self.dense = nn.Linear(config.hidden_size, 1)

    def forward(self, hidden_states, p_mask=None):
        """ Args:
1240
1241
1242
            **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
1243
        """
thomwolf's avatar
thomwolf committed
1244
1245
1246
        x = self.dense(hidden_states).squeeze(-1)

        if p_mask is not None:
1247
1248
1249
1250
            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
1251
1252
1253
1254
1255
1256
1257

        return x


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

thomwolf's avatar
thomwolf committed
1259
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1260
        super().__init__()
thomwolf's avatar
thomwolf committed
1261
1262
1263
1264
1265
1266
1267
        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:
1268
1269
1270
1271
1272
1273
            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,)``
1274
                position of the first token for the labeled span:
1275
1276
1277
            **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
1278
        """
1279
1280
1281
        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
1282
        if start_positions is not None:
1283
            slen, hsz = hidden_states.shape[-2:]
1284
1285
1286
            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
1287
1288
1289
1290
1291
1292
1293

        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:
1294
1295
1296
1297
            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
1298
1299
1300
1301
1302
1303

        return x


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

thomwolf's avatar
thomwolf committed
1305
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1306
        super().__init__()
thomwolf's avatar
thomwolf committed
1307
1308
1309
1310
1311
        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):
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
1322
1323
1324
1325
1326
        """
        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
1327
        """
1328
        hsz = hidden_states.shape[-1]
1329
1330
1331
        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
1332
        if start_positions is not None:
1333
1334
            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
1335
1336

        if cls_index is not None:
1337
1338
            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
1339
        else:
1340
            cls_token_state = hidden_states[:, -1, :]  # shape (bsz, hsz)
thomwolf's avatar
thomwolf committed
1341
1342
1343
1344
1345
1346
1347
1348
1349

        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):
1350
1351
1352
    r""" A SQuAD head inspired by XLNet.

    Parameters:
1353
        config (:class:`~transformers.XLNetConfig`): Model configuration class with all the parameters of the model.
1354
1355
1356
1357
1358
1359
1360
1361
1362
1363
1364
1365
1366
1367
1368
1369
1370
1371
1372

    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
1373
        **start_top_log_probs**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1374
1375
            ``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
1376
        **start_top_index**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1377
1378
            ``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
1379
        **end_top_log_probs**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1380
1381
            ``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
1382
        **end_top_index**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1383
1384
            ``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
1385
        **cls_logits**: (`optional`, returned if ``start_positions`` or ``end_positions`` is not provided)
1386
1387
            ``torch.FloatTensor`` of shape ``(batch_size,)``
            Log probabilities for the ``is_impossible`` label of the answers.
thomwolf's avatar
thomwolf committed
1388
    """
1389

thomwolf's avatar
thomwolf committed
1390
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1391
        super().__init__()
thomwolf's avatar
thomwolf committed
1392
1393
1394
1395
1396
1397
1398
        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)

1399
1400
1401
    def forward(
        self, hidden_states, start_positions=None, end_positions=None, cls_index=None, is_impossible=None, p_mask=None
    ):
thomwolf's avatar
thomwolf committed
1402
1403
        outputs = ()

thomwolf's avatar
thomwolf committed
1404
        start_logits = self.start_logits(hidden_states, p_mask=p_mask)
thomwolf's avatar
thomwolf committed
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419
1420
1421
1422
1423
1424
1425
1426
1427

        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
1428
1429

            outputs = (total_loss,) + outputs
thomwolf's avatar
thomwolf committed
1430
1431
1432
1433

        else:
            # during inference, compute the end logits based on beam search
            bsz, slen, hsz = hidden_states.size()
1434
1435
1436
1437
1438
1439
1440
1441
1442
1443
1444
1445
            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
1446
1447
            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)
1448
            end_log_probs = F.softmax(end_logits, dim=1)  # shape (bsz, slen, start_n_top)
thomwolf's avatar
thomwolf committed
1449

1450
1451
1452
            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
1453
1454
1455
1456
1457
1458
1459
1460
1461
            end_top_log_probs = end_top_log_probs.view(-1, self.start_n_top * self.end_n_top)
            end_top_index = end_top_index.view(-1, self.start_n_top * self.end_n_top)

            start_states = torch.einsum("blh,bl->bh", hidden_states, start_log_probs)
            cls_logits = self.answer_class(hidden_states, start_states=start_states, cls_index=cls_index)

            outputs = (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits) + outputs

        # return start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits
1462
        # or (if labels are provided) (total_loss,)
thomwolf's avatar
thomwolf committed
1463
1464
1465
1466
        return outputs


class SequenceSummary(nn.Module):
thomwolf's avatar
thomwolf committed
1467
    r""" Compute a single vector summary of a sequence hidden states according to various possibilities:
thomwolf's avatar
thomwolf committed
1468
1469
1470
1471
1472
        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
1473
                - 'cls_index' => supply a Tensor of classification token position (GPT/GPT-2)
thomwolf's avatar
thomwolf committed
1474
1475
                - 'attn' => Not implemented now, use multi-head attention
            summary_use_proj: Add a projection after the vector extraction
1476
            summary_proj_to_labels: If True, the projection outputs to config.num_labels classes (otherwise to hidden_size). Default: False.
1477
            summary_activation: 'tanh' or another string => add an activation to the output, Other => no activation. Default
1478
1479
            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
1480
    """
1481

1482
    def __init__(self, config: PretrainedConfig):
Julien Chaumond's avatar
Julien Chaumond committed
1483
        super().__init__()
thomwolf's avatar
thomwolf committed
1484

1485
        self.summary_type = getattr(config, "summary_type", "last")
1486
        if self.summary_type == "attn":
thomwolf's avatar
thomwolf committed
1487
1488
1489
1490
1491
            # 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
1492
        self.summary = Identity()
1493
1494
        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:
1495
                num_classes = config.num_labels
thomwolf's avatar
thomwolf committed
1496
1497
1498
1499
            else:
                num_classes = config.hidden_size
            self.summary = nn.Linear(config.hidden_size, num_classes)

1500
1501
1502
1503
        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
1504

thomwolf's avatar
thomwolf committed
1505
        self.first_dropout = Identity()
1506
        if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0:
1507
1508
            self.first_dropout = nn.Dropout(config.summary_first_dropout)

thomwolf's avatar
thomwolf committed
1509
        self.last_dropout = Identity()
1510
        if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0:
1511
            self.last_dropout = nn.Dropout(config.summary_last_dropout)
thomwolf's avatar
thomwolf committed
1512

thomwolf's avatar
thomwolf committed
1513
    def forward(self, hidden_states, cls_index=None):
1514
        """ hidden_states: float Tensor in shape [bsz, ..., seq_len, hidden_size], the hidden-states of the last layer.
thomwolf's avatar
thomwolf committed
1515
            cls_index: [optional] position of the classification token if summary_type == 'cls_index',
thomwolf's avatar
thomwolf committed
1516
                shape (bsz,) or more generally (bsz, ...) where ... are optional leading dimensions of hidden_states.
thomwolf's avatar
thomwolf committed
1517
                if summary_type == 'cls_index' and cls_index is None:
thomwolf's avatar
thomwolf committed
1518
1519
                    we take the last token of the sequence as classification token
        """
1520
        if self.summary_type == "last":
thomwolf's avatar
thomwolf committed
1521
            output = hidden_states[:, -1]
1522
        elif self.summary_type == "first":
thomwolf's avatar
thomwolf committed
1523
            output = hidden_states[:, 0]
1524
        elif self.summary_type == "mean":
thomwolf's avatar
thomwolf committed
1525
            output = hidden_states.mean(dim=1)
1526
        elif self.summary_type == "cls_index":
thomwolf's avatar
thomwolf committed
1527
            if cls_index is None:
1528
                cls_index = torch.full_like(hidden_states[..., :1, :], hidden_states.shape[-2] - 1, dtype=torch.long)
thomwolf's avatar
thomwolf committed
1529
            else:
thomwolf's avatar
thomwolf committed
1530
                cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)
1531
                cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))
thomwolf's avatar
thomwolf committed
1532
            # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states
1533
1534
            output = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, XX, hidden_size)
        elif self.summary_type == "attn":
thomwolf's avatar
thomwolf committed
1535
1536
            raise NotImplementedError

1537
        output = self.first_dropout(output)
thomwolf's avatar
thomwolf committed
1538
1539
        output = self.summary(output)
        output = self.activation(output)
1540
        output = self.last_dropout(output)
thomwolf's avatar
thomwolf committed
1541
1542
1543
1544

        return output


Sam Shleifer's avatar
Sam Shleifer committed
1545
1546
1547
1548
1549
1550
1551
1552
1553
1554
1555
1556
1557
1558
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


1559
1560
1561
1562
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579
1580
1581
1582
1583
1584
1585
1586
1587
1588
1589
1590
1591
1592
1593
1594
1595
1596
1597
1598
1599
1600
1601
1602
1603
1604
1605
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
1606
1607
1608
1609
1610
1611
1612
1613
1614
1615
1616
1617
1618


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