"vscode:/vscode.git/clone" did not exist on "662b143b71eb5ef775e27a8f79798bb28b3283bd"
modeling_utils.py 115 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
# 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.

Patrick von Platen's avatar
Patrick von Platen committed
17
import inspect
18
import os
19
import re
20
from contextlib import contextmanager
21
from dataclasses import dataclass
22
from functools import partial
23
from typing import Any, Callable, Dict, List, Optional, Set, Tuple, Union
24
25

import torch
26
from torch import Tensor, device, nn
27
from torch.nn import CrossEntropyLoss
28

29
30
from requests import HTTPError

31
from .activations import get_activation
32
from .configuration_utils import PretrainedConfig
33
from .deepspeed import deepspeed_config, is_deepspeed_zero3_enabled
34
from .dynamic_module_utils import custom_object_save
35
from .file_utils import (
Aymeric Augustin's avatar
Aymeric Augustin committed
36
    DUMMY_INPUTS,
37
    FLAX_WEIGHTS_NAME,
38
39
40
    TF2_WEIGHTS_NAME,
    TF_WEIGHTS_NAME,
    WEIGHTS_NAME,
41
    EntryNotFoundError,
42
    ModelOutput,
Sylvain Gugger's avatar
Sylvain Gugger committed
43
    PushToHubMixin,
44
45
    RepositoryNotFoundError,
    RevisionNotFoundError,
46
    cached_path,
47
    copy_func,
48
    has_file,
49
    hf_bucket_url,
50
    is_offline_mode,
51
    is_remote_url,
Sylvain Gugger's avatar
Sylvain Gugger committed
52
    replace_return_docstrings,
53
)
54
from .generation_utils import GenerationMixin
Lysandre Debut's avatar
Lysandre Debut committed
55
from .utils import logging
56
from .utils.versions import require_version_core
57

Aymeric Augustin's avatar
Aymeric Augustin committed
58

Lysandre Debut's avatar
Lysandre Debut committed
59
logger = logging.get_logger(__name__)
60

61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80

_init_weights = True


@contextmanager
def no_init_weights(_enable=True):
    """
    Context manager to globally disable weight initialization to speed up loading large models.

    TODO(Patrick): Delete safety argument `_enable=True` at next major version. .
    """
    global _init_weights
    if _enable:
        _init_weights = False
    try:
        yield
    finally:
        _init_weights = True


thomwolf's avatar
thomwolf committed
81
82
83
84
85
try:
    from torch.nn import Identity
except ImportError:
    # Older PyTorch compatibility
    class Identity(nn.Module):
Lysandre's avatar
Lysandre committed
86
        r"""A placeholder identity operator that is argument-insensitive."""
87

thomwolf's avatar
thomwolf committed
88
        def __init__(self, *args, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
89
            super().__init__()
thomwolf's avatar
thomwolf committed
90
91
92
93

        def forward(self, input):
            return input

94

95
def find_pruneable_heads_and_indices(
Sylvain Gugger's avatar
Sylvain Gugger committed
96
97
98
    heads: List[int], n_heads: int, head_size: int, already_pruned_heads: Set[int]
) -> Tuple[Set[int], torch.LongTensor]:
    """
99
    Finds the heads and their indices taking `already_pruned_heads` into account.
Sylvain Gugger's avatar
Sylvain Gugger committed
100
101

    Args:
102
103
104
105
        heads (`List[int]`): List of the indices of heads to prune.
        n_heads (`int`): The number of heads in the model.
        head_size (`int`): The size of each head.
        already_pruned_heads (`Set[int]`): A set of already pruned heads.
Sylvain Gugger's avatar
Sylvain Gugger committed
106
107

    Returns:
108
        `Tuple[Set[int], torch.LongTensor]`: A tuple with the remaining heads and their corresponding indices.
Sylvain Gugger's avatar
Sylvain Gugger committed
109
    """
110
111
112
113
114
115
116
117
118
119
120
    mask = torch.ones(n_heads, head_size)
    heads = set(heads) - already_pruned_heads  # Convert to set and remove already pruned heads
    for head in heads:
        # Compute how many pruned heads are before the head and move the index accordingly
        head = head - sum(1 if h < head else 0 for h in already_pruned_heads)
        mask[head] = 0
    mask = mask.view(-1).contiguous().eq(1)
    index: torch.LongTensor = torch.arange(len(mask))[mask].long()
    return heads, index


Lysandre Debut's avatar
Lysandre Debut committed
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
def get_parameter_device(parameter: Union[nn.Module, GenerationMixin, "ModuleUtilsMixin"]):
    try:
        return next(parameter.parameters()).device
    except StopIteration:
        # For nn.DataParallel compatibility in PyTorch 1.5

        def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]:
            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
            return tuples

        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
        first_tuple = next(gen)
        return first_tuple[1].device


def get_parameter_dtype(parameter: Union[nn.Module, GenerationMixin, "ModuleUtilsMixin"]):
    try:
        return next(parameter.parameters()).dtype
    except StopIteration:
        # For nn.DataParallel compatibility in PyTorch 1.5

        def find_tensor_attributes(module: nn.Module) -> List[Tuple[str, Tensor]]:
            tuples = [(k, v) for k, v in module.__dict__.items() if torch.is_tensor(v)]
            return tuples

        gen = parameter._named_members(get_members_fn=find_tensor_attributes)
        first_tuple = next(gen)
        return first_tuple[1].dtype


151
class ModuleUtilsMixin:
Julien Chaumond's avatar
Julien Chaumond committed
152
    """
153
    A few utilities for `torch.nn.Modules`, to be used as a mixin.
Julien Chaumond's avatar
Julien Chaumond committed
154
155
    """

156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
    @staticmethod
    def _hook_rss_memory_pre_forward(module, *args, **kwargs):
        try:
            import psutil
        except (ImportError):
            raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.")

        process = psutil.Process(os.getpid())
        mem = process.memory_info()
        module.mem_rss_pre_forward = mem.rss
        return None

    @staticmethod
    def _hook_rss_memory_post_forward(module, *args, **kwargs):
        try:
            import psutil
        except (ImportError):
            raise ImportError("You need to install psutil (pip install psutil) to use memory tracing.")

        process = psutil.Process(os.getpid())
        mem = process.memory_info()
        module.mem_rss_post_forward = mem.rss
        mem_rss_diff = module.mem_rss_post_forward - module.mem_rss_pre_forward
        module.mem_rss_diff = mem_rss_diff + (module.mem_rss_diff if hasattr(module, "mem_rss_diff") else 0)
        return None

    def add_memory_hooks(self):
Sylvain Gugger's avatar
Sylvain Gugger committed
183
184
185
        """
        Add a memory hook before and after each sub-module forward pass to record increase in memory consumption.

Sylvain Gugger's avatar
Sylvain Gugger committed
186
187
        Increase in memory consumption is stored in a `mem_rss_diff` attribute for each module and can be reset to zero
        with `model.reset_memory_hooks_state()`.
188
189
190
191
192
193
194
        """
        for module in self.modules():
            module.register_forward_pre_hook(self._hook_rss_memory_pre_forward)
            module.register_forward_hook(self._hook_rss_memory_post_forward)
        self.reset_memory_hooks_state()

    def reset_memory_hooks_state(self):
Sylvain Gugger's avatar
Sylvain Gugger committed
195
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
196
        Reset the `mem_rss_diff` attribute of each module (see [`~modeling_utils.ModuleUtilsMixin.add_memory_hooks`]).
Sylvain Gugger's avatar
Sylvain Gugger committed
197
        """
198
199
200
201
202
        for module in self.modules():
            module.mem_rss_diff = 0
            module.mem_rss_post_forward = 0
            module.mem_rss_pre_forward = 0

203
    @property
204
    def device(self) -> device:
205
        """
206
        `torch.device`: The device on which the module is (assuming that all the module parameters are on the same
207
        device).
208
        """
Lysandre Debut's avatar
Lysandre Debut committed
209
        return get_parameter_device(self)
210

211
    @property
212
    def dtype(self) -> torch.dtype:
213
        """
214
        `torch.dtype`: The dtype of the module (assuming that all the module parameters have the same dtype).
215
        """
Lysandre Debut's avatar
Lysandre Debut committed
216
        return get_parameter_dtype(self)
217
218

    def invert_attention_mask(self, encoder_attention_mask: Tensor) -> Tensor:
Sylvain Gugger's avatar
Sylvain Gugger committed
219
220
221
222
        """
        Invert an attention mask (e.g., switches 0. and 1.).

        Args:
223
            encoder_attention_mask (`torch.Tensor`): An attention mask.
Sylvain Gugger's avatar
Sylvain Gugger committed
224
225

        Returns:
226
            `torch.Tensor`: The inverted attention mask.
Sylvain Gugger's avatar
Sylvain Gugger committed
227
        """
228
229
230
231
232
233
234
235
236
237
        if encoder_attention_mask.dim() == 3:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, :, :]
        if encoder_attention_mask.dim() == 2:
            encoder_extended_attention_mask = encoder_attention_mask[:, None, None, :]
        # T5 has a mask that can compare sequence ids, we can simulate this here with this transposition
        # Cf. https://github.com/tensorflow/mesh/blob/8d2465e9bc93129b913b5ccc6a59aa97abd96ec6/mesh_tensorflow
        # /transformer/transformer_layers.py#L270
        # encoder_extended_attention_mask = (encoder_extended_attention_mask ==
        # encoder_extended_attention_mask.transpose(-1, -2))
        encoder_extended_attention_mask = encoder_extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
238
239
240

        if self.dtype == torch.float16:
            encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -1e4
241
        elif self.dtype in [torch.bfloat16, torch.float32]:
242
243
244
            encoder_extended_attention_mask = (1.0 - encoder_extended_attention_mask) * -1e9
        else:
            raise ValueError(
245
                f"{self.dtype} not recognized. `dtype` should be set to either `torch.float32` or `torch.float16`"
246
247
            )

248
249
        return encoder_extended_attention_mask

Sylvain Gugger's avatar
Sylvain Gugger committed
250
251
252
    def get_extended_attention_mask(self, attention_mask: Tensor, input_shape: Tuple[int], device: device) -> Tensor:
        """
        Makes broadcastable attention and causal masks so that future and masked tokens are ignored.
253
254

        Arguments:
255
            attention_mask (`torch.Tensor`):
Sylvain Gugger's avatar
Sylvain Gugger committed
256
                Mask with ones indicating tokens to attend to, zeros for tokens to ignore.
257
            input_shape (`Tuple[int]`):
Sylvain Gugger's avatar
Sylvain Gugger committed
258
                The shape of the input to the model.
259
            device: (`torch.device`):
Sylvain Gugger's avatar
Sylvain Gugger committed
260
                The device of the input to the model.
261
262

        Returns:
263
            `torch.Tensor` The extended attention mask, with a the same dtype as `attention_mask.dtype`.
264
265
266
267
268
269
270
271
272
273
274
275
276
        """
        # We can provide a self-attention mask of dimensions [batch_size, from_seq_length, to_seq_length]
        # ourselves in which case we just need to make it broadcastable to all heads.
        if attention_mask.dim() == 3:
            extended_attention_mask = attention_mask[:, None, :, :]
        elif attention_mask.dim() == 2:
            # Provided a padding mask of dimensions [batch_size, seq_length]
            # - if the model is a decoder, apply a causal mask in addition to the padding mask
            # - if the model is an encoder, make the mask broadcastable to [batch_size, num_heads, seq_length, seq_length]
            if self.config.is_decoder:
                batch_size, seq_length = input_shape
                seq_ids = torch.arange(seq_length, device=device)
                causal_mask = seq_ids[None, None, :].repeat(batch_size, seq_length, 1) <= seq_ids[None, :, None]
277
                # in case past_key_values are used we need to add a prefix ones mask to the causal mask
Patrick von Platen's avatar
Patrick von Platen committed
278
279
280
                # causal and attention masks must have same type with pytorch version < 1.3
                causal_mask = causal_mask.to(attention_mask.dtype)

281
282
283
                if causal_mask.shape[1] < attention_mask.shape[1]:
                    prefix_seq_len = attention_mask.shape[1] - causal_mask.shape[1]
                    causal_mask = torch.cat(
Patrick von Platen's avatar
Patrick von Platen committed
284
285
286
287
288
289
290
                        [
                            torch.ones(
                                (batch_size, seq_length, prefix_seq_len), device=device, dtype=causal_mask.dtype
                            ),
                            causal_mask,
                        ],
                        axis=-1,
291
292
                    )

293
294
295
296
297
                extended_attention_mask = causal_mask[:, None, :, :] * attention_mask[:, None, None, :]
            else:
                extended_attention_mask = attention_mask[:, None, None, :]
        else:
            raise ValueError(
298
                f"Wrong shape for input_ids (shape {input_shape}) or attention_mask (shape {attention_mask.shape})"
299
300
301
302
303
304
305
306
307
308
309
            )

        # Since attention_mask is 1.0 for positions we want to attend and 0.0 for
        # masked positions, this operation will create a tensor which is 0.0 for
        # positions we want to attend and -10000.0 for masked positions.
        # Since we are adding it to the raw scores before the softmax, this is
        # effectively the same as removing these entirely.
        extended_attention_mask = extended_attention_mask.to(dtype=self.dtype)  # fp16 compatibility
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
        return extended_attention_mask

Sylvain Gugger's avatar
Sylvain Gugger committed
310
311
312
    def get_head_mask(
        self, head_mask: Optional[Tensor], num_hidden_layers: int, is_attention_chunked: bool = False
    ) -> Tensor:
313
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
314
315
316
        Prepare the head mask if needed.

        Args:
317
            head_mask (`torch.Tensor` with shape `[num_heads]` or `[num_hidden_layers x num_heads]`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
318
                The mask indicating if we should keep the heads or not (1.0 for keep, 0.0 for discard).
319
            num_hidden_layers (`int`):
Sylvain Gugger's avatar
Sylvain Gugger committed
320
                The number of hidden layers in the model.
321
            is_attention_chunked: (`bool`, *optional*, defaults to `False`):
Sylvain Gugger's avatar
Sylvain Gugger committed
322
323
                Whether or not the attentions scores are computed by chunks or not.

324
        Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
325
326
            `torch.Tensor` with shape `[num_hidden_layers x batch x num_heads x seq_length x seq_length]` or list with
            `[None]` for each layer.
327
328
329
        """
        if head_mask is not None:
            head_mask = self._convert_head_mask_to_5d(head_mask, num_hidden_layers)
Patrick von Platen's avatar
Patrick von Platen committed
330
331
            if is_attention_chunked is True:
                head_mask = head_mask.unsqueeze(-1)
332
333
334
335
336
337
338
339
340
341
342
343
344
        else:
            head_mask = [None] * num_hidden_layers

        return head_mask

    def _convert_head_mask_to_5d(self, head_mask, num_hidden_layers):
        """-> [num_hidden_layers x batch x num_heads x seq_length x seq_length]"""
        if head_mask.dim() == 1:
            head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
            head_mask = head_mask.expand(num_hidden_layers, -1, -1, -1, -1)
        elif head_mask.dim() == 2:
            head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
        assert head_mask.dim() == 5, f"head_mask.dim != 5, instead {head_mask.dim()}"
345
        head_mask = head_mask.to(dtype=self.dtype)  # switch to float if need + fp16 compatibility
346
347
        return head_mask

348
349
350
351
352
    def num_parameters(self, only_trainable: bool = False, exclude_embeddings: bool = False) -> int:
        """
        Get number of (optionally, trainable or non-embeddings) parameters in the module.

        Args:
353
            only_trainable (`bool`, *optional*, defaults to `False`):
354
355
                Whether or not to return only the number of trainable parameters

356
            exclude_embeddings (`bool`, *optional*, defaults to `False`):
357
358
359
                Whether or not to return only the number of non-embeddings parameters

        Returns:
360
            `int`: The number of parameters.
361
362
        """

363
364
365
366
367
368
369
370
371
372
        if exclude_embeddings:
            embedding_param_names = [
                f"{name}.weight" for name, module_type in self.named_modules() if isinstance(module_type, nn.Embedding)
            ]
            non_embedding_parameters = [
                parameter for name, parameter in self.named_parameters() if name not in embedding_param_names
            ]
            return sum(p.numel() for p in non_embedding_parameters if p.requires_grad or not only_trainable)
        else:
            return sum(p.numel() for p in self.parameters() if p.requires_grad or not only_trainable)
373
374
375
376
377
378

    def estimate_tokens(self, input_dict: Dict[str, Union[torch.Tensor, Any]]) -> int:
        """
        Helper function to estimate the total number of tokens from the model inputs.

        Args:
379
            inputs (`dict`): The model inputs.
380
381

        Returns:
382
            `int`: The total number of tokens.
383
        """
384
385
        if self.main_input_name in input_dict:
            return input_dict[self.main_input_name].numel()
386
        else:
387
            logger.warn(
388
389
390
391
392
393
394
395
396
397
                "Could not estimate the number of tokens of the input, floating-point operations will not be computed"
            )
            return 0

    def floating_point_ops(
        self, input_dict: Dict[str, Union[torch.Tensor, Any]], exclude_embeddings: bool = True
    ) -> int:
        """
        Get number of (optionally, non-embeddings) floating-point operations for the forward and backward passes of a
        batch with this transformer model. Default approximation neglects the quadratic dependency on the number of
Sylvain Gugger's avatar
Sylvain Gugger committed
398
399
        tokens (valid if `12 * d_model << sequence_length`) as laid out in [this
        paper](https://arxiv.org/pdf/2001.08361.pdf) section 2.1. Should be overridden for transformers with parameter
Sylvain Gugger's avatar
Sylvain Gugger committed
400
        re-use e.g. Albert or Universal Transformers, or if doing long-range modeling with very high sequence lengths.
401
402

        Args:
403
            batch_size (`int`):
404
405
                The batch size for the forward pass.

406
            sequence_length (`int`):
407
408
                The number of tokens in each line of the batch.

409
            exclude_embeddings (`bool`, *optional*, defaults to `True`):
410
411
412
                Whether or not to count embedding and softmax operations.

        Returns:
413
            `int`: The number of floating-point operations.
414
415
416
417
        """

        return 6 * self.estimate_tokens(input_dict) * self.num_parameters(exclude_embeddings=exclude_embeddings)

Julien Chaumond's avatar
Julien Chaumond committed
418

Sylvain Gugger's avatar
Sylvain Gugger committed
419
class PreTrainedModel(nn.Module, ModuleUtilsMixin, GenerationMixin, PushToHubMixin):
420
421
    r"""
    Base class for all models.
422

Sylvain Gugger's avatar
Sylvain Gugger committed
423
424
    [`PreTrainedModel`] takes care of storing the configuration of the models and handles methods for loading,
    downloading and saving models as well as a few methods common to all models to:
425

426
427
        - resize the input embeddings,
        - prune heads in the self-attention heads.
428

429
    Class attributes (overridden by derived classes):
Sylvain Gugger's avatar
Sylvain Gugger committed
430

Sylvain Gugger's avatar
Sylvain Gugger committed
431
432
433
434
        - **config_class** ([`PretrainedConfig`]) -- A subclass of [`PretrainedConfig`] to use as configuration class
          for this model architecture.
        - **load_tf_weights** (`Callable`) -- A python *method* for loading a TensorFlow checkpoint in a PyTorch model,
          taking as arguments:
435

Sylvain Gugger's avatar
Sylvain Gugger committed
436
437
            - **model** ([`PreTrainedModel`]) -- An instance of the model on which to load the TensorFlow checkpoint.
            - **config** ([`PreTrainedConfig`]) -- An instance of the configuration associated to the model.
438
            - **path** (`str`) -- A path to the TensorFlow checkpoint.
439

Sylvain Gugger's avatar
Sylvain Gugger committed
440
441
        - **base_model_prefix** (`str`) -- 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.
442
        - **is_parallelizable** (`bool`) -- A flag indicating whether this model supports model parallelization.
Sylvain Gugger's avatar
Sylvain Gugger committed
443
444
        - **main_input_name** (`str`) -- The name of the principal input to the model (often `input_ids` for NLP
          models, `pixel_values` for vision models and `input_values` for speech models).
445
    """
446
    config_class = None
447
    base_model_prefix = ""
448
    main_input_name = "input_ids"
449
    _auto_class = None
450

451
452
453
454
455
456
457
    # a list of re pattern of tensor names to ignore from the model when loading the model weights
    # (and avoid unnecessary warnings).
    _keys_to_ignore_on_load_missing = None
    # a list of re pattern of tensor names to ignore from the weights when loading the model weights
    # (and avoid unnecessary warnings).
    _keys_to_ignore_on_load_unexpected = None
    # a list of of tensor names to ignore when saving the model (useful for keys that aren't
458
    # trained, but which are deterministic, or tied variables)
459
    _keys_to_ignore_on_save = None
460

461
    is_parallelizable = False
462
    supports_gradient_checkpointing = False
463

464
    @property
465
    def dummy_inputs(self) -> Dict[str, torch.Tensor]:
466
        """
467
        `Dict[str, torch.Tensor]`: Dummy inputs to do a forward pass in the network.
468
        """
469
        return {"input_ids": torch.tensor(DUMMY_INPUTS)}
470

471
472
473
474
475
476
477
    @property
    def framework(self) -> str:
        """
        :str: Identifies that this is a PyTorch model.
        """
        return "pt"

478
    def __init__(self, config: PretrainedConfig, *inputs, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
479
        super().__init__()
480
481
        if not isinstance(config, PretrainedConfig):
            raise ValueError(
482
483
484
                f"Parameter config in `{self.__class__.__name__}(config)` should be an instance of class "
                "`PretrainedConfig`. To create a model from a pretrained model use "
                f"`model = {self.__class__.__name__}.from_pretrained(PRETRAINED_MODEL_NAME)`"
485
            )
486
        # Save config and origin of the pretrained weights if given in model
487
        self.config = config
488
        self.name_or_path = config.name_or_path
489
490
491
492
493
494
495
496
497
498
499
500
501
502

    def post_init(self):
        """
        A method executed at the end of each Transformer model initialization, to execute code that needs the model's
        modules properly initialized (such as weight initialization).
        """
        self.init_weights()
        self._backward_compatibility_gradient_checkpointing()

    def _backward_compatibility_gradient_checkpointing(self):
        if self.supports_gradient_checkpointing and getattr(self.config, "gradient_checkpointing", False):
            self.gradient_checkpointing_enable()
            # Remove the attribute now that is has been consumed, so it's no saved in the config.
            delattr(self.config, "gradient_checkpointing")
503

504
505
506
507
508
509
    @classmethod
    def _from_config(cls, config, **kwargs):
        """
        All context managers that the model should be initialized under go here.

        Args:
510
511
            torch_dtype (`torch.dtype`, *optional*):
                Override the default `torch.dtype` and load the model under this dtype.
512
513
514
515
516
517
518
519
520
521
522
523
524
525
        """
        torch_dtype = kwargs.pop("torch_dtype", None)

        # override default dtype if needed
        dtype_orig = None
        if torch_dtype is not None:
            dtype_orig = cls._set_default_torch_dtype(torch_dtype)

        if is_deepspeed_zero3_enabled():
            import deepspeed

            logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
            # this immediately partitions the model across all gpus, to avoid the overhead in time
            # and memory copying it on CPU or each GPU first
526
            with deepspeed.zero.Init(config_dict_or_path=deepspeed_config()):
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
                model = cls(config, **kwargs)
        else:
            model = cls(config, **kwargs)

        # restore default dtype if it was modified
        if dtype_orig is not None:
            torch.set_default_dtype(dtype_orig)

        return model

    @classmethod
    def _set_default_torch_dtype(cls, dtype: torch.dtype) -> torch.dtype:
        """
        Change the default dtype and return the previous one. This is needed when wanting to instantiate the model
        under specific dtype.

        Args:
544
            dtype (`torch.dtype`):
545
546
547
                a floating dtype to set to.

        Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
548
549
            `torch.dtype`: the original `dtype` that can be used to restore `torch.set_default_dtype(dtype)` if it was
            modified. If it wasn't, returns `None`.
550

551
552
        Note `set_default_dtype` currently only works with floating-point types and asserts if for example,
        `torch.int64` is passed. So if a non-float `dtype` is passed this functions will throw an exception.
553
554
555
556
557
558
559
560
561
562
563
        """
        if not dtype.is_floating_point:
            raise ValueError(
                f"Can't instantiate {cls.__name__} model under dtype={dtype} since it is not a floating point dtype"
            )

        logger.info(f"Instantiating {cls.__name__} model under default dtype {dtype}.")
        dtype_orig = torch.get_default_dtype()
        torch.set_default_dtype(dtype)
        return dtype_orig

564
    @property
565
566
    def base_model(self) -> nn.Module:
        """
567
        `torch.nn.Module`: The main body of the model.
568
        """
569
        return getattr(self, self.base_model_prefix, self)
thomwolf's avatar
thomwolf committed
570

571
    def get_input_embeddings(self) -> nn.Module:
572
573
574
575
        """
        Returns the model's input embeddings.

        Returns:
576
            `nn.Module`: A torch module mapping vocabulary to hidden states.
thomwolf's avatar
thomwolf committed
577
        """
578
        base_model = getattr(self, self.base_model_prefix, self)
thomwolf's avatar
thomwolf committed
579
580
581
582
        if base_model is not self:
            return base_model.get_input_embeddings()
        else:
            raise NotImplementedError
thomwolf's avatar
thomwolf committed
583

584
    def set_input_embeddings(self, value: nn.Module):
585
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
586
        Set model's input embeddings.
587
588

        Args:
589
            value (`nn.Module`): A module mapping vocabulary to hidden states.
thomwolf's avatar
thomwolf committed
590
591
592
593
594
595
        """
        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
596

597
    def get_output_embeddings(self) -> nn.Module:
598
599
600
601
        """
        Returns the model's output embeddings.

        Returns:
602
            `nn.Module`: A torch module mapping hidden states to vocabulary.
thomwolf's avatar
thomwolf committed
603
        """
604
        return None  # Overwrite for models with output embeddings
thomwolf's avatar
thomwolf committed
605

606
607
608
609
    def _init_weights(self, module):
        """
        Initialize the weights. This method should be overridden by derived class.
        """
610
        raise NotImplementedError(f"Make sure `_init_weights` is implemented for {self.__class__}")
611

612
    def tie_weights(self):
613
614
        """
        Tie the weights between the input embeddings and the output embeddings.
615

Sylvain Gugger's avatar
Sylvain Gugger committed
616
617
        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
618
        """
thomwolf's avatar
thomwolf committed
619
        output_embeddings = self.get_output_embeddings()
620
        if output_embeddings is not None and self.config.tie_word_embeddings:
thomwolf's avatar
thomwolf committed
621
            self._tie_or_clone_weights(output_embeddings, self.get_input_embeddings())
thomwolf's avatar
thomwolf committed
622

623
        if self.config.is_encoder_decoder and self.config.tie_encoder_decoder:
Weizhen's avatar
Weizhen committed
624
625
            if hasattr(self, self.base_model_prefix):
                self = getattr(self, self.base_model_prefix)
626
627
            self._tie_encoder_decoder_weights(self.encoder, self.decoder, self.base_model_prefix)

Sylvain Gugger's avatar
Sylvain Gugger committed
628
629
630
631
        for module in self.modules():
            if hasattr(module, "_tie_weights"):
                module._tie_weights()

632
633
634
    @staticmethod
    def _tie_encoder_decoder_weights(encoder: nn.Module, decoder: nn.Module, base_model_prefix: str):
        uninitialized_encoder_weights: List[str] = []
Weizhen's avatar
Weizhen committed
635
636
637
638
        if decoder.__class__ != encoder.__class__:
            logger.info(
                f"{decoder.__class__} and {encoder.__class__} are not equal. In this case make sure that all encoder weights are correctly initialized."
            )
639
640
641
642
643
644
645
646
647
648

        def tie_encoder_to_decoder_recursively(
            decoder_pointer: nn.Module,
            encoder_pointer: nn.Module,
            module_name: str,
            uninitialized_encoder_weights: List[str],
            depth=0,
        ):
            assert isinstance(decoder_pointer, nn.Module) and isinstance(
                encoder_pointer, nn.Module
649
            ), f"{decoder_pointer} and {encoder_pointer} have to be of type nn.Module"
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
            if hasattr(decoder_pointer, "weight"):
                assert hasattr(encoder_pointer, "weight")
                encoder_pointer.weight = decoder_pointer.weight
                if hasattr(decoder_pointer, "bias"):
                    assert hasattr(encoder_pointer, "bias")
                    encoder_pointer.bias = decoder_pointer.bias
                return

            encoder_modules = encoder_pointer._modules
            decoder_modules = decoder_pointer._modules
            if len(decoder_modules) > 0:
                assert (
                    len(encoder_modules) > 0
                ), f"Encoder module {encoder_pointer} does not match decoder module {decoder_pointer}"

                all_encoder_weights = set([module_name + "/" + sub_name for sub_name in encoder_modules.keys()])
                encoder_layer_pos = 0
                for name, module in decoder_modules.items():
                    if name.isdigit():
                        encoder_name = str(int(name) + encoder_layer_pos)
                        decoder_name = name
Weizhen's avatar
Weizhen committed
671
672
673
                        if not isinstance(decoder_modules[decoder_name], type(encoder_modules[encoder_name])) and len(
                            encoder_modules
                        ) != len(decoder_modules):
674
675
                            # this can happen if the name corresponds to the position in a list module list of layers
                            # in this case the decoder has added a cross-attention that the encoder does not have
676
                            # thus skip this step and subtract one layer pos from encoder
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
                            encoder_layer_pos -= 1
                            continue
                    elif name not in encoder_modules:
                        continue
                    elif depth > 500:
                        raise ValueError(
                            "Max depth of recursive function `tie_encoder_to_decoder` reached. It seems that there is a circular dependency between two or more `nn.Modules` of your model."
                        )
                    else:
                        decoder_name = encoder_name = name
                    tie_encoder_to_decoder_recursively(
                        decoder_modules[decoder_name],
                        encoder_modules[encoder_name],
                        module_name + "/" + name,
                        uninitialized_encoder_weights,
                        depth=depth + 1,
                    )
                    all_encoder_weights.remove(module_name + "/" + encoder_name)

                uninitialized_encoder_weights += list(all_encoder_weights)

        # tie weights recursively
        tie_encoder_to_decoder_recursively(decoder, encoder, base_model_prefix, uninitialized_encoder_weights)
        if len(uninitialized_encoder_weights) > 0:
            logger.warning(
                f"The following encoder weights were not tied to the decoder {uninitialized_encoder_weights}"
            )

705
    def _tie_or_clone_weights(self, output_embeddings, input_embeddings):
Lysandre's avatar
Lysandre committed
706
        """Tie or clone module weights depending of whether we are using TorchScript or not"""
thomwolf's avatar
thomwolf committed
707
        if self.config.torchscript:
708
            output_embeddings.weight = nn.Parameter(input_embeddings.weight.clone())
thomwolf's avatar
thomwolf committed
709
        else:
710
            output_embeddings.weight = input_embeddings.weight
thomwolf's avatar
thomwolf committed
711

Sam Shleifer's avatar
Sam Shleifer committed
712
        if getattr(output_embeddings, "bias", None) is not None:
713
            output_embeddings.bias.data = nn.functional.pad(
714
                output_embeddings.bias.data,
Lysandre's avatar
Lysandre committed
715
716
717
718
                (
                    0,
                    output_embeddings.weight.shape[0] - output_embeddings.bias.shape[0],
                ),
719
720
                "constant",
                0,
721
            )
722
        if hasattr(output_embeddings, "out_features") and hasattr(input_embeddings, "num_embeddings"):
723
            output_embeddings.out_features = input_embeddings.num_embeddings
724

725
    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None) -> nn.Embedding:
726
        """
727
        Resizes input token embeddings matrix of the model if `new_num_tokens != config.vocab_size`.
728

729
        Takes care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.
thomwolf's avatar
thomwolf committed
730

731
        Arguments:
732
            new_num_tokens (`int`, *optional*):
733
                The number of new tokens in the embedding matrix. Increasing the size will add newly initialized
Sylvain Gugger's avatar
Sylvain Gugger committed
734
735
                vectors at the end. Reducing the size will remove vectors from the end. If not provided or `None`, just
                returns a pointer to the input tokens `torch.nn.Embedding` module of the model without doing anything.
736
737

        Return:
738
            `torch.nn.Embedding`: Pointer to the input tokens Embeddings Module of the model.
thomwolf's avatar
thomwolf committed
739
        """
740
        model_embeds = self._resize_token_embeddings(new_num_tokens)
thomwolf's avatar
thomwolf committed
741
742
        if new_num_tokens is None:
            return model_embeds
thomwolf's avatar
thomwolf committed
743
744
745

        # Update base model and current model config
        self.config.vocab_size = new_num_tokens
746
        self.vocab_size = new_num_tokens
thomwolf's avatar
thomwolf committed
747
748

        # Tie weights again if needed
749
        self.tie_weights()
thomwolf's avatar
thomwolf committed
750

thomwolf's avatar
thomwolf committed
751
752
        return model_embeds

753
    def _resize_token_embeddings(self, new_num_tokens):
thomwolf's avatar
thomwolf committed
754
755
756
        old_embeddings = self.get_input_embeddings()
        new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens)
        self.set_input_embeddings(new_embeddings)
757
758
759
760
761
762
763

        # if word embeddings are not tied, make sure that lm head is resized as well
        if self.get_output_embeddings() is not None and not self.config.tie_word_embeddings:
            old_lm_head = self.get_output_embeddings()
            new_lm_head = self._get_resized_lm_head(old_lm_head, new_num_tokens)
            self.set_output_embeddings(new_lm_head)

thomwolf's avatar
thomwolf committed
764
        return self.get_input_embeddings()
765

766
    def _get_resized_embeddings(
767
768
        self, old_embeddings: nn.Embedding, new_num_tokens: Optional[int] = None
    ) -> nn.Embedding:
769
770
771
        """
        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
772
773

        Args:
774
            old_embeddings (`torch.nn.Embedding`):
775
                Old embeddings to be resized.
776
            new_num_tokens (`int`, *optional*):
777
                New number of tokens in the embedding matrix.
778
779

                Increasing the size will add newly initialized vectors at the end. Reducing the size will remove
780
781
                vectors from the end. If not provided or `None`, just returns a pointer to the input tokens
                ``torch.nn.Embedding``` module of the model without doing anything.
782
783

        Return:
784
785
            `torch.nn.Embedding`: Pointer to the resized Embedding Module or the old Embedding Module if
            `new_num_tokens` is `None`
786
787
788
789
        """
        if new_num_tokens is None:
            return old_embeddings

790
791
792
793
794
795
796
797
        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_embeddings.weight, modifier_rank=None):
                old_num_tokens, old_embedding_dim = old_embeddings.weight.size()
        else:
            old_num_tokens, old_embedding_dim = old_embeddings.weight.size()

798
799
800
        if old_num_tokens == new_num_tokens:
            return old_embeddings

801
802
        if not isinstance(old_embeddings, nn.Embedding):
            raise TypeError(
803
                f"Old embeddings are of type {type(old_embeddings)}, which is not an instance of {nn.Embedding}. "
804
805
806
                f"You should either use a different resize function or make sure that `old_embeddings` are an instance of {nn.Embedding}."
            )

807
        # Build new embeddings
808
809
        new_embeddings = nn.Embedding(new_num_tokens, old_embedding_dim)
        new_embeddings.to(self.device, dtype=old_embeddings.weight.dtype)
810
811
812
813

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

814
        # Copy token embeddings from the previous weights
815
816
817
818
819
820
821
822
823
824
825

        # numbers of tokens to copy
        n = min(old_num_tokens, new_num_tokens)
        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_embeddings.weight, modifier_rank=0):
                if torch.distributed.get_rank() == 0:
                    new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :]
        else:
            new_embeddings.weight.data[:n, :] = old_embeddings.weight.data[:n, :]
826
827
828

        return new_embeddings

829
    def _get_resized_lm_head(
830
831
        self, old_lm_head: nn.Linear, new_num_tokens: Optional[int] = None, transposed: Optional[bool] = False
    ) -> nn.Linear:
832
833
834
835
836
        """
        Build a resized Linear Module from a provided old Linear Module. Increasing the size will add newly initialized
        vectors at the end. Reducing the size will remove vectors from the end

        Args:
837
            old_lm_head (`torch.nn.Linear`):
838
                Old lm head liner layer to be resized.
839
            new_num_tokens (`int`, *optional*):
840
841
842
                New number of tokens in the linear matrix.

                Increasing the size will add newly initialized vectors at the end. Reducing the size will remove
843
                vectors from the end. If not provided or `None`, just returns a pointer to the input tokens
Sylvain Gugger's avatar
Sylvain Gugger committed
844
845
846
                ``torch.nn.Linear``` module of the model without doing anything. transposed (`bool`, *optional*,
                defaults to `False`): Whether `old_lm_head` is transposed or not. If True `old_lm_head.size()` is
                `lm_head_dim, vocab_size` else `vocab_size, lm_head_dim`.
847
848

        Return:
Sylvain Gugger's avatar
Sylvain Gugger committed
849
850
            `torch.nn.Linear`: Pointer to the resized Linear Module or the old Linear Module if `new_num_tokens` is
            `None`
851
852
853
854
        """
        if new_num_tokens is None:
            return old_lm_head

855
856
857
858
859
860
861
862
863
864
865
        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_lm_head.weight, modifier_rank=None):
                old_num_tokens, old_lm_head_dim = (
                    old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size()
                )
        else:
            old_num_tokens, old_lm_head_dim = (
                old_lm_head.weight.size() if not transposed else old_lm_head.weight.t().size()
            )
866
867
868
869
870
871

        if old_num_tokens == new_num_tokens:
            return old_lm_head

        if not isinstance(old_lm_head, nn.Linear):
            raise TypeError(
872
                f"Old language model head is of type {type(old_lm_head)}, which is not an instance of {nn.Linear}. "
873
                f"You should either use a different resize function or make sure that `old_lm_head` are an instance of {nn.Linear}."
874
875
876
877
878
            )

        # Build new lm head
        new_lm_head_shape = (old_lm_head_dim, new_num_tokens) if not transposed else (new_num_tokens, old_lm_head_dim)
        has_new_lm_head_bias = old_lm_head.bias is not None
879
880
        new_lm_head = nn.Linear(*new_lm_head_shape, bias=has_new_lm_head_bias)
        new_lm_head = new_lm_head.to(self.device, dtype=old_lm_head.weight.dtype)
881
882
883
884
885
886

        # initialize new lm head (in particular added tokens)
        self._init_weights(new_lm_head)

        num_tokens_to_copy = min(old_num_tokens, new_num_tokens)

887
888
889
890
891
892
893
894
895
896
897
898
899
900
901
902
903
904
905
        # XXX: put the long block of code in a wrapper
        if is_deepspeed_zero3_enabled():
            import deepspeed

            with deepspeed.zero.GatheredParameters(old_lm_head.weight, modifier_rank=0):
                if torch.distributed.get_rank() == 0:
                    # Copy old lm head weights to new lm head
                    if not transposed:
                        new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[
                            :num_tokens_to_copy, :
                        ]
                    else:
                        new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[
                            :, :num_tokens_to_copy
                        ]

                    # Copy bias weights to new lm head
                    if has_new_lm_head_bias:
                        new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]
906
        else:
907
908
909
910
911
            # Copy old lm head weights to new lm head
            if not transposed:
                new_lm_head.weight.data[:num_tokens_to_copy, :] = old_lm_head.weight.data[:num_tokens_to_copy, :]
            else:
                new_lm_head.weight.data[:, :num_tokens_to_copy] = old_lm_head.weight.data[:, :num_tokens_to_copy]
912

913
914
915
            # Copy bias weights to new lm head
            if has_new_lm_head_bias:
                new_lm_head.bias.data[:num_tokens_to_copy] = old_lm_head.bias.data[:num_tokens_to_copy]
916
917
918

        return new_lm_head

919
920
921
922
923
924
925
926
927
928
929
930
    def resize_position_embeddings(self, new_num_position_embeddings: int):
        raise NotImplementedError(
            f"`resize_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should "
            f"overwrite this method in the class {self.__class__} in `modeling_{self.__class__.__module__}.py`"
        )

    def get_position_embeddings(self) -> Union[nn.Embedding, Tuple[nn.Embedding]]:
        raise NotImplementedError(
            f"`get_position_embeddings` is not implemented for {self.__class__}`. To implement it, you should "
            f"overwrite this method in the class {self.__class__} in `modeling_{self.__class__.__module__}.py`"
        )

931
    def init_weights(self):
932
        """
933
        If needed prunes and maybe initializes weights.
934
        """
935
936
937
938
        # Prune heads if needed
        if self.config.pruned_heads:
            self.prune_heads(self.config.pruned_heads)

939
940
941
942
943
944
945
        if _init_weights:
            # Initialize weights
            self.apply(self._init_weights)

            # Tie weights should be skipped when not initializing all weights
            # since from_pretrained(...) calls tie weights anyways
            self.tie_weights()
946

947
948
949
    def prune_heads(self, heads_to_prune: Dict[int, List[int]]):
        """
        Prunes heads of the base model.
950

951
        Arguments:
952
            heads_to_prune (`Dict[int, List[int]]`):
Sylvain Gugger's avatar
Sylvain Gugger committed
953
954
955
                Dictionary with keys being selected layer indices (`int`) and associated values being the list of heads
                to prune in said layer (list of `int`). For instance {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
956
        """
957
        # save new sets of pruned heads as union of previously stored pruned heads and newly pruned heads
958
        for layer, heads in heads_to_prune.items():
959
960
961
            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

962
        self.base_model._prune_heads(heads_to_prune)
thomwolf's avatar
thomwolf committed
963

964
    def gradient_checkpointing_enable(self):
965
966
967
968
969
970
971
972
973
974
        """
        Activates gradient checkpointing for the current model.

        Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
        activations".
        """
        if not self.supports_gradient_checkpointing:
            raise ValueError(f"{self.__class__.__name__} does not support gradient checkpointing.")
        self.apply(partial(self._set_gradient_checkpointing, value=True))

975
    def gradient_checkpointing_disable(self):
976
977
978
979
980
981
982
983
984
        """
        Deactivates gradient checkpointing for the current model.

        Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
        activations".
        """
        if self.supports_gradient_checkpointing:
            self.apply(partial(self._set_gradient_checkpointing, value=False))

985
986
987
988
989
990
991
992
993
994
    @property
    def is_gradient_checkpointing(self) -> bool:
        """
        Whether gradient checkpointing is activated for this model or not.

        Note that in other frameworks this feature can be referred to as "activation checkpointing" or "checkpoint
        activations".
        """
        return any(hasattr(m, "gradient_checkpointing") and m.gradient_checkpointing for m in self.modules())

995
996
997
998
999
1000
    def save_pretrained(
        self,
        save_directory: Union[str, os.PathLike],
        save_config: bool = True,
        state_dict: Optional[dict] = None,
        save_function: Callable = torch.save,
Sylvain Gugger's avatar
Sylvain Gugger committed
1001
1002
        push_to_hub: bool = False,
        **kwargs,
1003
    ):
1004
1005
        """
        Save a model and its configuration file to a directory, so that it can be re-loaded using the
1006
        `[`~PreTrainedModel.from_pretrained`]` class method.
1007

1008
        Arguments:
1009
            save_directory (`str` or `os.PathLike`):
1010
                Directory to which to save. Will be created if it doesn't exist.
1011
            save_config (`bool`, *optional*, defaults to `True`):
1012
                Whether or not to save the config of the model. Useful when in distributed training like TPUs and need
Sylvain Gugger's avatar
Sylvain Gugger committed
1013
1014
                to call this function on all processes. In this case, set `save_config=True` only on the main process
                to avoid race conditions.
1015
            state_dict (nested dictionary of `torch.Tensor`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1016
1017
1018
                The state dictionary of the model to save. Will default to `self.state_dict()`, but can be used to only
                save parts of the model or if special precautions need to be taken when recovering the state dictionary
                of a model (like when using model parallelism).
1019
            save_function (`Callable`):
1020
                The function to use to save the state dictionary. Useful on distributed training like TPUs when one
1021
1022
                need to replace `torch.save` by another method.
            push_to_hub (`bool`, *optional*, defaults to `False`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1023
                Whether or not to push your model to the Hugging Face model hub after saving it.
1024

1025
                <Tip warning={true}>
1026

Sylvain Gugger's avatar
Sylvain Gugger committed
1027
1028
1029
                Using `push_to_hub=True` will synchronize the repository you are pushing to with `save_directory`,
                which requires `save_directory` to be a local clone of the repo you are pushing to if it's an existing
                folder. Pass along `temp_dir=True` to use a temporary directory instead.
1030
1031

                </Tip>
1032

Sylvain Gugger's avatar
Sylvain Gugger committed
1033
            kwargs:
Sylvain Gugger's avatar
Sylvain Gugger committed
1034
                Additional key word arguments passed along to the [`~file_utils.PushToHubMixin.push_to_hub`] method.
1035
        """
1036
        if os.path.isfile(save_directory):
1037
            logger.error(f"Provided path ({save_directory}) should be a directory, not a file")
1038
            return
1039
1040
1041
1042
1043

        if push_to_hub:
            commit_message = kwargs.pop("commit_message", None)
            repo = self._create_or_get_repo(save_directory, **kwargs)

1044
        os.makedirs(save_directory, exist_ok=True)
1045

Julien Chaumond's avatar
Julien Chaumond committed
1046
        # Only save the model itself if we are using distributed training
1047
        model_to_save = unwrap_model(self)
1048

1049
1050
1051
1052
1053
        # save the string version of dtype to the config, e.g. convert torch.float32 => "float32"
        # we currently don't use this setting automatically, but may start to use with v5
        dtype = get_parameter_dtype(model_to_save)
        model_to_save.config.torch_dtype = str(dtype).split(".")[1]

Julien Chaumond's avatar
Julien Chaumond committed
1054
1055
1056
        # Attach architecture to the config
        model_to_save.config.architectures = [model_to_save.__class__.__name__]

1057
1058
1059
1060
1061
        # If we have a custom model, we copy the file defining it in the folder and set the attributes so it can be
        # loaded from the Hub.
        if self._auto_class is not None:
            custom_object_save(self, save_directory, config=self.config)

1062
1063
1064
1065
1066
1067
1068
        # Save the config
        if save_config:
            model_to_save.config.save_pretrained(save_directory)

        # Save the model
        if state_dict is None:
            state_dict = model_to_save.state_dict()
1069
1070

        # Handle the case where some state_dict keys shouldn't be saved
1071
        if self._keys_to_ignore_on_save is not None:
1072
            for ignore_key in self._keys_to_ignore_on_save:
1073
1074
                if ignore_key in state_dict.keys():
                    del state_dict[ignore_key]
1075

1076
1077
        # If we save using the predefined names, we can load using `from_pretrained`
        output_model_file = os.path.join(save_directory, WEIGHTS_NAME)
1078
        save_function(state_dict, output_model_file)
1079

1080
        logger.info(f"Model weights saved in {output_model_file}")
1081

Sylvain Gugger's avatar
Sylvain Gugger committed
1082
        if push_to_hub:
1083
            url = self._push_to_hub(repo, commit_message=commit_message)
Sylvain Gugger's avatar
Sylvain Gugger committed
1084
1085
            logger.info(f"Model pushed to the hub in this commit: {url}")

1086
    @classmethod
1087
    def from_pretrained(cls, pretrained_model_name_or_path: Optional[Union[str, os.PathLike]], *model_args, **kwargs):
1088
1089
        r"""
        Instantiate a pretrained pytorch model from a pre-trained model configuration.
1090

Sylvain Gugger's avatar
Sylvain Gugger committed
1091
1092
        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()`.
1093

1094
        The warning *Weights from XXX not initialized from pretrained model* means that the weights of XXX do not come
1095
1096
        pretrained with the rest of the model. It is up to you to train those weights with a downstream fine-tuning
        task.
1097

1098
        The warning *Weights from XXX not used in YYY* means that the layer XXX is not used by YYY, therefore those
1099
        weights are discarded.
1100

1101
        Parameters:
1102
            pretrained_model_name_or_path (`str` or `os.PathLike`, *optional*):
1103
1104
                Can be either:

1105
                    - A string, the *model id* of a pretrained model hosted inside a model repo on huggingface.co.
Sylvain Gugger's avatar
Sylvain Gugger committed
1106
1107
                      Valid model ids can be located at the root-level, like `bert-base-uncased`, or namespaced under a
                      user or organization name, like `dbmdz/bert-base-german-cased`.
1108
1109
1110
                    - A path to a *directory* containing model weights saved using
                      [`~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
Sylvain Gugger's avatar
Sylvain Gugger committed
1111
1112
1113
                      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.
1114
                    - A path or url to a model folder containing a *flax checkpoint file* in *.msgpack* format (e.g,
Sylvain Gugger's avatar
Sylvain Gugger committed
1115
1116
                      `./flax_model/` containing `flax_model.msgpack`). In this case, `from_flax` should be set to
                      `True`.
1117
1118
1119
1120
1121
                    - `None` if you are both providing the configuration and state dictionary (resp. with keyword
                      arguments `config` and `state_dict`).
            model_args (sequence of positional arguments, *optional*):
                All remaining positional arguments will be passed to the underlying model's `__init__` method.
            config (`Union[PretrainedConfig, str, os.PathLike]`, *optional*):
1122
1123
                Can be either:

1124
1125
                    - an instance of a class derived from [`PretrainedConfig`],
                    - a string or path valid as input to [`~PretrainedConfig.from_pretrained`].
1126

1127
                Configuration for the model to use instead of an automatically loaded configuration. Configuration can
1128
1129
                be automatically loaded when:

1130
                    - The model is a model provided by the library (loaded with the *model id* string of a pretrained
1131
                      model).
Sylvain Gugger's avatar
Sylvain Gugger committed
1132
1133
                    - The model was saved using [`~PreTrainedModel.save_pretrained`] and is reloaded by supplying the
                      save directory.
1134
1135
1136
                    - The model is loaded by supplying a local directory as `pretrained_model_name_or_path` and a
                      configuration JSON file named *config.json* is found in the directory.
            state_dict (`Dict[str, torch.Tensor]`, *optional*):
1137
1138
1139
                A state dictionary to use instead of a state dictionary loaded from saved weights file.

                This option can be used if you want to create a model from a pretrained configuration but load your own
Sylvain Gugger's avatar
Sylvain Gugger committed
1140
                weights. In this case though, you should check if using [`~PreTrainedModel.save_pretrained`] and
1141
1142
                [`~PreTrainedModel.from_pretrained`] is not a simpler option.
            cache_dir (`Union[str, os.PathLike]`, *optional*):
1143
1144
                Path to a directory in which a downloaded pretrained model configuration should be cached if the
                standard cache should not be used.
1145
            from_tf (`bool`, *optional*, defaults to `False`):
1146
                Load the model weights from a TensorFlow checkpoint save file (see docstring of
1147
1148
                `pretrained_model_name_or_path` argument).
            from_flax (`bool`, *optional*, defaults to `False`):
1149
                Load the model weights from a Flax checkpoint save file (see docstring of
1150
1151
                `pretrained_model_name_or_path` argument).
            ignore_mismatched_sizes (`bool`, *optional*, defaults to `False`):
1152
1153
1154
                Whether or not to raise an error if some of the weights from the checkpoint do not have the same size
                as the weights of the model (if for instance, you are instantiating a model with 10 labels from a
                checkpoint with 3 labels).
1155
            force_download (`bool`, *optional*, defaults to `False`):
1156
1157
                Whether or not to force the (re-)download of the model weights and configuration files, overriding the
                cached versions if they exist.
1158
            resume_download (`bool`, *optional*, defaults to `False`):
1159
1160
                Whether or not to delete incompletely received files. Will attempt to resume the download if such a
                file exists.
1161
            proxies (`Dict[str, str]`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1162
1163
                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.
1164
            output_loading_info(`bool`, *optional*, defaults to `False`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1165
                Whether ot not to also return a dictionary containing missing keys, unexpected keys and error messages.
1166
            local_files_only(`bool`, *optional*, defaults to `False`):
Stas Bekman's avatar
Stas Bekman committed
1167
                Whether or not to only look at local files (i.e., do not try to download the model).
1168
            use_auth_token (`str` or *bool*, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1169
1170
                The token to use as HTTP bearer authorization for remote files. If `True`, will use the token generated
                when running `transformers-cli login` (stored in `~/.huggingface`).
1171
            revision(`str`, *optional*, defaults to `"main"`):
Julien Chaumond's avatar
Julien Chaumond committed
1172
                The specific model version to use. It can be a branch name, a tag name, or a commit id, since we use a
1173
                git-based system for storing models and other artifacts on huggingface.co, so `revision` can be any
Julien Chaumond's avatar
Julien Chaumond committed
1174
                identifier allowed by git.
1175
            mirror(`str`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1176
1177
1178
                Mirror source to accelerate downloads in China. If you are from China and have an accessibility
                problem, you can set this option to resolve it. Note that we do not guarantee the timeliness or safety.
                Please refer to the mirror site for more information.
1179
            _fast_init(`bool`, *optional*, defaults to ```True`):
1180
                Whether or not to disable fast initialization.
1181
            low_cpu_mem_usage(`bool``, *optional*, defaults to ```False`):
1182
1183
                Tries to not use more than 1x model size in CPU memory (including peak memory) while loading the model.
                This is an experimental feature and a subject to change at any moment.
1184
            torch_dtype (`str` or `torch.dtype`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1185
1186
                Override the default `torch.dtype` and load the model under this dtype. If `"auto"` is passed the dtype
                will be automatically derived from the model's weights.
1187

1188
1189
                <Tip warning={true}>

Sylvain Gugger's avatar
Sylvain Gugger committed
1190
1191
1192
                One should only disable *_fast_init* to ensure backwards compatibility with `transformers.__version__ <
                4.6.0` for seeded model initialization. This argument will be removed at the next major version. See
                [pull request 11471](https://github.com/huggingface/transformers/pull/11471) for more information.
1193

1194
                </Tip>
1195

1196
            kwargs (remaining dictionary of keyword arguments, *optional*):
1197
                Can be used to update the configuration object (after it being loaded) and initiate the model (e.g.,
1198
                `output_attentions=True`). Behaves differently depending on whether a `config` is provided or
1199
1200
                automatically loaded:

1201
1202
                    - 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
1203
                      already been done)
1204
                    - If a configuration is not provided, `kwargs` will be first passed to the configuration class
Sylvain Gugger's avatar
Sylvain Gugger committed
1205
1206
1207
1208
                      initialization function ([`~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.
1209
1210
1211
1212
1213
1214
1215
1216
1217

        <Tip>

        Passing `use_auth_token=True`` is required when you want to use a private model.

        </Tip>

        <Tip>

Sylvain Gugger's avatar
Sylvain Gugger committed
1218
1219
        Activate the special ["offline-mode"](https://huggingface.co/transformers/installation.html#offline-mode) to
        use this method in a firewalled environment.
1220
1221
1222
1223
1224
1225
1226

        </Tip>

        Examples:

        ```python
        >>> from transformers import BertConfig, BertModel
Sylvain Gugger's avatar
Sylvain Gugger committed
1227

1228
        >>> # Download model and configuration from huggingface.co and cache.
Sylvain Gugger's avatar
Sylvain Gugger committed
1229
        >>> model = BertModel.from_pretrained("bert-base-uncased")
1230
        >>> # Model was saved using *save_pretrained('./test/saved_model/')* (for example purposes, not runnable).
Sylvain Gugger's avatar
Sylvain Gugger committed
1231
        >>> model = BertModel.from_pretrained("./test/saved_model/")
1232
        >>> # Update configuration during loading.
Sylvain Gugger's avatar
Sylvain Gugger committed
1233
        >>> model = BertModel.from_pretrained("bert-base-uncased", output_attentions=True)
1234
1235
        >>> assert model.config.output_attentions == True
        >>> # Loading from a TF checkpoint file instead of a PyTorch model (slower, for example purposes, not runnable).
Sylvain Gugger's avatar
Sylvain Gugger committed
1236
1237
        >>> 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)
1238
        >>> # Loading from a Flax checkpoint file instead of a PyTorch model (slower)
Sylvain Gugger's avatar
Sylvain Gugger committed
1239
        >>> model = BertModel.from_pretrained("bert-base-uncased", from_flax=True)
1240
        ```"""
1241
1242
1243
1244
        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)
1245
        from_flax = kwargs.pop("from_flax", False)
1246
        ignore_mismatched_sizes = kwargs.pop("ignore_mismatched_sizes", False)
1247
1248
1249
1250
        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)
1251
        local_files_only = kwargs.pop("local_files_only", False)
1252
        use_auth_token = kwargs.pop("use_auth_token", None)
Julien Chaumond's avatar
Julien Chaumond committed
1253
        revision = kwargs.pop("revision", None)
1254
        mirror = kwargs.pop("mirror", None)
1255
1256
        from_pipeline = kwargs.pop("_from_pipeline", None)
        from_auto_class = kwargs.pop("_from_auto", False)
1257
        _fast_init = kwargs.pop("_fast_init", True)
1258
        torch_dtype = kwargs.pop("torch_dtype", None)
1259
        low_cpu_mem_usage = kwargs.pop("low_cpu_mem_usage", False)
1260
1261

        from_pt = not (from_tf | from_flax)
1262
1263
1264
1265

        user_agent = {"file_type": "model", "framework": "pytorch", "from_auto_class": from_auto_class}
        if from_pipeline is not None:
            user_agent["using_pipeline"] = from_pipeline
thomwolf's avatar
thomwolf committed
1266

1267
1268
1269
1270
        if is_offline_mode() and not local_files_only:
            logger.info("Offline mode: forcing local_files_only=True")
            local_files_only = True

1271
1272
1273
        # 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
1274
            config, model_kwargs = cls.config_class.from_pretrained(
1275
1276
1277
                config_path,
                cache_dir=cache_dir,
                return_unused_kwargs=True,
1278
                force_download=force_download,
1279
                resume_download=resume_download,
1280
                proxies=proxies,
1281
                local_files_only=local_files_only,
1282
                use_auth_token=use_auth_token,
Julien Chaumond's avatar
Julien Chaumond committed
1283
                revision=revision,
1284
1285
                _from_auto=from_auto_class,
                _from_pipeline=from_pipeline,
1286
                **kwargs,
1287
1288
1289
            )
        else:
            model_kwargs = kwargs
1290

thomwolf's avatar
thomwolf committed
1291
        # Load model
thomwolf's avatar
thomwolf committed
1292
        if pretrained_model_name_or_path is not None:
1293
            pretrained_model_name_or_path = str(pretrained_model_name_or_path)
1294
            if os.path.isdir(pretrained_model_name_or_path):
thomwolf's avatar
thomwolf committed
1295
                if from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")):
1296
                    # Load from a TF 1.0 checkpoint in priority if from_tf
thomwolf's avatar
thomwolf committed
1297
                    archive_file = os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
thomwolf's avatar
thomwolf committed
1298
                elif from_tf and os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)):
1299
                    # Load from a TF 2.0 checkpoint in priority if from_tf
thomwolf's avatar
thomwolf committed
1300
                    archive_file = os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)
1301
1302
1303
                elif from_flax and os.path.isfile(os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)):
                    # Load from a Flax checkpoint in priority if from_flax
                    archive_file = os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME)
thomwolf's avatar
thomwolf committed
1304
1305
                elif os.path.isfile(os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)):
                    # Load from a PyTorch checkpoint
thomwolf's avatar
thomwolf committed
1306
                    archive_file = os.path.join(pretrained_model_name_or_path, WEIGHTS_NAME)
1307
1308
1309
1310
1311
1312
1313
1314
1315
1316
1317
1318
1319
1320
1321
                # At this stage we don't have a weight file so we will raise an error.
                elif os.path.isfile(
                    os.path.join(pretrained_model_name_or_path, TF_WEIGHTS_NAME + ".index")
                ) or os.path.isfile(os.path.join(pretrained_model_name_or_path, TF2_WEIGHTS_NAME)):
                    raise EnvironmentError(
                        f"Error no file named {WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} but "
                        "there is a file for TensorFlow weights. Use `from_tf=True` to load this model from those "
                        "weights."
                    )
                elif os.path.join(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME):
                    raise EnvironmentError(
                        f"Error no file named {WEIGHTS_NAME} found in directory {pretrained_model_name_or_path} but "
                        "there is a file for Flax weights. Use `from_flax=True` to load this model from those "
                        "weights."
                    )
thomwolf's avatar
thomwolf committed
1322
                else:
1323
                    raise EnvironmentError(
1324
1325
                        f"Error no file named {WEIGHTS_NAME}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME + '.index'} or "
                        f"{FLAX_WEIGHTS_NAME} found in directory {pretrained_model_name_or_path}."
1326
                    )
1327
            elif os.path.isfile(pretrained_model_name_or_path) or is_remote_url(pretrained_model_name_or_path):
1328
                archive_file = pretrained_model_name_or_path
1329
            elif os.path.isfile(pretrained_model_name_or_path + ".index"):
1330
1331
1332
1333
1334
                if not from_tf:
                    raise ValueError(
                        f"We found a TensorFlow checkpoint at {pretrained_model_name_or_path + '.index'}, please set "
                        "from_tf to True to load from this checkpoint."
                    )
1335
                archive_file = pretrained_model_name_or_path + ".index"
1336
            else:
1337
1338
1339
1340
1341
1342
1343
1344
                # set correct filename
                if from_tf:
                    filename = TF2_WEIGHTS_NAME
                elif from_flax:
                    filename = FLAX_WEIGHTS_NAME
                else:
                    filename = WEIGHTS_NAME

thomwolf's avatar
thomwolf committed
1345
                archive_file = hf_bucket_url(
Julien Chaumond's avatar
Julien Chaumond committed
1346
                    pretrained_model_name_or_path,
1347
                    filename=filename,
Julien Chaumond's avatar
Julien Chaumond committed
1348
                    revision=revision,
1349
                    mirror=mirror,
thomwolf's avatar
thomwolf committed
1350
                )
1351

thomwolf's avatar
thomwolf committed
1352
            try:
1353
                # Load from URL or cache if already cached
1354
1355
1356
1357
1358
1359
                resolved_archive_file = cached_path(
                    archive_file,
                    cache_dir=cache_dir,
                    force_download=force_download,
                    proxies=proxies,
                    resume_download=resume_download,
1360
                    local_files_only=local_files_only,
1361
                    use_auth_token=use_auth_token,
1362
                    user_agent=user_agent,
1363
                )
1364
1365
1366
1367
1368
1369
1370
1371
1372
1373
1374
1375
1376
1377
1378
1379
1380
1381
1382
1383
1384
1385
1386
1387
1388
1389
1390
1391
1392
1393
1394
1395
1396
1397
1398
1399
1400
1401
1402
1403
1404
1405
1406
1407
1408
1409
1410
1411
1412
1413
1414
1415
1416
1417
1418
1419

            except RepositoryNotFoundError as err:
                logger.error(err)
                raise EnvironmentError(
                    f"{pretrained_model_name_or_path} is not a local folder and is not a valid model identifier "
                    "listed on 'https://huggingface.co/models'\nIf this is a private repository, make sure to pass a "
                    "token having permission to this repo with `use_auth_token` or log in with `huggingface-cli "
                    "login` and pass `use_auth_token=True`."
                )
            except RevisionNotFoundError as err:
                logger.error(err)
                raise EnvironmentError(
                    f"{revision} is not a valid git identifier (branch name, tag name or commit id) that exists for "
                    "this model name. Check the model page at "
                    f"'https://huggingface.co/{pretrained_model_name_or_path}' for available revisions."
                )
            except EntryNotFoundError as err:
                logger.error(err)
                if filename == WEIGHTS_NAME:
                    has_file_kwargs = {
                        "revision": revision,
                        "mirror": mirror,
                        "proxies": proxies,
                        "use_auth_token": use_auth_token,
                    }
                    if has_file(pretrained_model_name_or_path, TF2_WEIGHTS_NAME, **has_file_kwargs):
                        raise EnvironmentError(
                            f"{pretrained_model_name_or_path} does not appear to have a file named {WEIGHTS_NAME} but "
                            "there is a file for TensorFlow weights. Use `from_tf=True` to load this model from those "
                            "weights."
                        )
                    elif has_file(pretrained_model_name_or_path, FLAX_WEIGHTS_NAME, **has_file_kwargs):
                        raise EnvironmentError(
                            f"{pretrained_model_name_or_path} does not appear to have a file named {WEIGHTS_NAME} but "
                            "there is a file for Flax weights. Use `from_flax=True` to load this model from those "
                            "weights."
                        )
                    else:
                        logger.error(err)
                        raise EnvironmentError(
                            f"{pretrained_model_name_or_path} does not appear to have a file named {WEIGHTS_NAME}, "
                            f"{TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or {FLAX_WEIGHTS_NAME}."
                        )
                else:
                    raise EnvironmentError(
                        f"{pretrained_model_name_or_path} does not appear to have a file named {filename}."
                    )
            except HTTPError as err:
                logger.error(err)
                raise EnvironmentError(
                    "We couldn't connect to 'https://huggingface.co/' to load this model and it looks like "
                    f"{pretrained_model_name_or_path} is not the path to a directory conaining a a file named "
                    f"{WEIGHTS_NAME}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or {FLAX_WEIGHTS_NAME}.\n"
                    "Checkout your internet connection or see how to run the library in offline mode at "
                    "'https://huggingface.co/docs/transformers/installation#offline-mode'."
                )
Julien Chaumond's avatar
Julien Chaumond committed
1420
1421
            except EnvironmentError as err:
                logger.error(err)
1422
1423
1424
1425
1426
1427
                raise EnvironmentError(
                    f"Can't load the model for '{pretrained_model_name_or_path}'. If you were trying to load it from "
                    "'https://huggingface.co/models', make sure you don't have a local directory with the same name. "
                    f"Otherwise, make sure '{pretrained_model_name_or_path}' is the correct path to a directory "
                    f"containing a file named {WEIGHTS_NAME}, {TF2_WEIGHTS_NAME}, {TF_WEIGHTS_NAME} or "
                    f"{FLAX_WEIGHTS_NAME}."
1428
                )
1429

thomwolf's avatar
thomwolf committed
1430
            if resolved_archive_file == archive_file:
1431
                logger.info(f"loading weights file {archive_file}")
1432
            else:
1433
                logger.info(f"loading weights file {archive_file} from cache at {resolved_archive_file}")
1434
        else:
thomwolf's avatar
thomwolf committed
1435
            resolved_archive_file = None
1436

1437
1438
1439
1440
1441
        # load pt weights early so that we know which dtype to init the model under
        if from_pt:
            if state_dict is None:
                try:
                    state_dict = torch.load(resolved_archive_file, map_location="cpu")
1442
1443
1444
1445
1446
1447
1448
1449
1450
1451
1452
1453
1454
1455
                except Exception as e:
                    try:
                        with open(resolved_archive_file) as f:
                            if f.read().startswith("version"):
                                raise OSError(
                                    "You seem to have cloned a repository without having git-lfs installed. Please install "
                                    "git-lfs and run `git lfs install` followed by `git lfs pull` in the folder "
                                    "you cloned."
                                )
                            else:
                                raise ValueError from e
                    except (UnicodeDecodeError, ValueError):
                        raise OSError(
                            f"Unable to load weights from pytorch checkpoint file for '{pretrained_model_name_or_path}' "
1456
1457
                            f"at '{resolved_archive_file}'. "
                            "If you tried to load a PyTorch model from a TF 2.0 checkpoint, please set from_tf=True."
1458
                        )
1459
1460
1461
1462
1463
1464
1465
1466
1467
1468
1469
1470
1471
1472
1473
1474
1475

            # set dtype to instantiate the model under:
            # 1. If torch_dtype is not None, we use that dtype
            # 2. If torch_dtype is "auto", we auto-detect dtype from the loaded state_dict, by checking its first
            #    weights entry - we assume all weights are of the same dtype
            # we also may have config.torch_dtype available, but we won't rely on it till v5
            dtype_orig = None
            if torch_dtype is not None:
                if isinstance(torch_dtype, str):
                    if torch_dtype == "auto":
                        torch_dtype = next(iter(state_dict.values())).dtype
                    else:
                        raise ValueError(
                            f"`torch_dtype` can be either a `torch.dtype` or `auto`, but received {torch_dtype}"
                        )
                dtype_orig = cls._set_default_torch_dtype(torch_dtype)

1476
1477
1478
1479
1480
            if low_cpu_mem_usage:
                # save the keys
                loaded_state_dict_keys = [k for k in state_dict.keys()]
                del state_dict  # free CPU memory - will reload again later

1481
1482
        config.name_or_path = pretrained_model_name_or_path

1483
        # Instantiate model.
1484
1485
1486
1487
        if is_deepspeed_zero3_enabled():
            import deepspeed

            logger.info("Detected DeepSpeed ZeRO-3: activating zero.init() for this model")
1488
1489
            # this immediately partitions the model across all gpus, to avoid the overhead in time
            # and memory copying it on CPU or each GPU first
1490
            with deepspeed.zero.Init(config_dict_or_path=deepspeed_config()):
1491
1492
                with no_init_weights(_enable=_fast_init):
                    model = cls(config, *model_args, **model_kwargs)
1493
        else:
1494
1495
            with no_init_weights(_enable=_fast_init):
                model = cls(config, *model_args, **model_kwargs)
1496

1497
1498
1499
1500
1501
        if from_pt:
            # restore default dtype
            if dtype_orig is not None:
                torch.set_default_dtype(dtype_orig)

1502
        if from_tf:
1503
            if resolved_archive_file.endswith(".index"):
1504
1505
1506
1507
1508
                # 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:
1509
                    from .modeling_tf_pytorch_utils import load_tf2_checkpoint_in_pytorch_model
1510

1511
                    model = load_tf2_checkpoint_in_pytorch_model(model, resolved_archive_file, allow_missing_keys=True)
1512
                except ImportError:
1513
1514
1515
1516
                    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."
                    )
1517
                    raise
1518
1519
1520
1521
1522
1523
1524
1525
1526
1527
1528
        elif from_flax:
            try:
                from .modeling_flax_pytorch_utils import load_flax_checkpoint_in_pytorch_model

                model = load_flax_checkpoint_in_pytorch_model(model, resolved_archive_file)
            except ImportError:
                logger.error(
                    "Loading a Flax model in PyTorch, requires both PyTorch and Flax to be installed. Please see "
                    "https://pytorch.org/ and https://flax.readthedocs.io/en/latest/installation.html for installation instructions."
                )
                raise
1529
        elif from_pt:
1530
1531
1532
1533
1534
1535
1536
1537
1538
1539
1540

            if low_cpu_mem_usage:
                cls._load_state_dict_into_model_low_mem(model, loaded_state_dict_keys, resolved_archive_file)
            else:
                model, missing_keys, unexpected_keys, mismatched_keys, error_msgs = cls._load_state_dict_into_model(
                    model,
                    state_dict,
                    pretrained_model_name_or_path,
                    ignore_mismatched_sizes=ignore_mismatched_sizes,
                    _fast_init=_fast_init,
                )
1541

1542
1543
        # make sure token embedding weights are still tied if needed
        model.tie_weights()
1544

1545
        # Set model in evaluation mode to deactivate DropOut modules by default
1546
1547
        model.eval()

thomwolf's avatar
thomwolf committed
1548
        if output_loading_info:
1549
1550
1551
            loading_info = {
                "missing_keys": missing_keys,
                "unexpected_keys": unexpected_keys,
1552
                "mismatched_keys": mismatched_keys,
1553
1554
                "error_msgs": error_msgs,
            }
thomwolf's avatar
thomwolf committed
1555
1556
            return model, loading_info

1557
1558
        return model

1559
    @classmethod
1560
1561
1562
    def _load_state_dict_into_model(
        cls, model, state_dict, pretrained_model_name_or_path, ignore_mismatched_sizes=False, _fast_init=True
    ):
1563
1564
1565
1566
1567
1568
1569
1570
1571
1572
1573
1574
1575
1576
1577
1578
1579

        # Convert old format to new format if needed from a PyTorch state_dict
        old_keys = []
        new_keys = []
        for key in state_dict.keys():
            new_key = None
            if "gamma" in key:
                new_key = key.replace("gamma", "weight")
            if "beta" in key:
                new_key = key.replace("beta", "bias")
            if new_key:
                old_keys.append(key)
                new_keys.append(new_key)
        for old_key, new_key in zip(old_keys, new_keys):
            state_dict[new_key] = state_dict.pop(old_key)

        # Retrieve missing & unexpected_keys
1580
1581
        model_state_dict = model.state_dict()
        expected_keys = list(model_state_dict.keys())
1582
1583
1584
1585
1586
        loaded_keys = list(state_dict.keys())
        prefix = model.base_model_prefix

        has_prefix_module = any(s.startswith(prefix) for s in loaded_keys)
        expects_prefix_module = any(s.startswith(prefix) for s in expected_keys)
Patrick von Platen's avatar
Patrick von Platen committed
1587
1588
1589

        # key re-naming operations are never done on the keys
        # that are loaded, but always on the keys of the newly initialized model
1590
1591
        remove_prefix_from_model = not has_prefix_module and expects_prefix_module
        add_prefix_to_model = has_prefix_module and not expects_prefix_module
1592

1593
        if remove_prefix_from_model:
1594
            expected_keys_not_prefixed = [s for s in expected_keys if not s.startswith(prefix)]
1595
            expected_keys = [".".join(s.split(".")[1:]) if s.startswith(prefix) else s for s in expected_keys]
1596
        elif add_prefix_to_model:
1597
1598
1599
1600
1601
            expected_keys = [".".join([prefix, s]) for s in expected_keys]

        missing_keys = list(set(expected_keys) - set(loaded_keys))
        unexpected_keys = list(set(loaded_keys) - set(expected_keys))

1602
1603
1604
1605
1606
1607
        # Mistmatched keys contains tuples key/shape1/shape2 of weights in the checkpoint that have a shape not
        # matching the weights in the model.
        mismatched_keys = []
        if ignore_mismatched_sizes:
            for checkpoint_key in loaded_keys:
                model_key = checkpoint_key
1608
1609
                if remove_prefix_from_model:
                    # The model key starts with `prefix` but `checkpoint_key` doesn't so we add it.
1610
                    model_key = f"{prefix}.{checkpoint_key}"
1611
1612
                elif add_prefix_to_model:
                    # The model key doesn't start with `prefix` but `checkpoint_key` does so we remove it.
1613
                    model_key = ".".join(checkpoint_key.split(".")[1:])
1614
1615
1616
1617
1618
1619
1620
1621
1622
1623

                if (
                    model_key in model_state_dict
                    and state_dict[checkpoint_key].shape != model_state_dict[model_key].shape
                ):
                    mismatched_keys.append(
                        (checkpoint_key, state_dict[checkpoint_key].shape, model_state_dict[model_key].shape)
                    )
                    del state_dict[checkpoint_key]

1624
1625
1626
1627
1628
1629
1630
1631
1632
1633
        # Some models may have keys that are not in the state by design, removing them before needlessly warning
        # the user.
        if cls._keys_to_ignore_on_load_missing is not None:
            for pat in cls._keys_to_ignore_on_load_missing:
                missing_keys = [k for k in missing_keys if re.search(pat, k) is None]

        if cls._keys_to_ignore_on_load_unexpected is not None:
            for pat in cls._keys_to_ignore_on_load_unexpected:
                unexpected_keys = [k for k in unexpected_keys if re.search(pat, k) is None]

1634
1635
        if _fast_init:
            # retrieve unintialized modules and initialize
1636
            uninitialized_modules = model.retrieve_modules_from_names(
1637
                missing_keys, add_prefix=add_prefix_to_model, remove_prefix=remove_prefix_from_model
1638
            )
1639
            for module in uninitialized_modules:
1640
1641
                model._init_weights(module)

1642
1643
1644
1645
1646
1647
1648
1649
1650
1651
1652
1653
1654
1655
1656
1657
1658
1659
1660
1661
1662
1663
1664
1665
1666
1667
1668
1669
1670
1671
1672
1673
1674
1675
1676
1677
        # copy state_dict so _load_from_state_dict can modify it
        metadata = getattr(state_dict, "_metadata", None)
        state_dict = state_dict.copy()
        if metadata is not None:
            state_dict._metadata = metadata

        error_msgs = []

        # PyTorch's `_load_from_state_dict` does not copy parameters in a module's descendants
        # so we need to apply the function recursively.
        def load(module: nn.Module, prefix=""):
            local_metadata = {} if metadata is None else metadata.get(prefix[:-1], {})
            args = (state_dict, prefix, local_metadata, True, [], [], error_msgs)
            if is_deepspeed_zero3_enabled():
                import deepspeed

                # because zero3 puts placeholders in model params, this context
                # manager gathers (unpartitions) the params of the current layer, then loads from
                # the state dict and then re-partitions them again
                with deepspeed.zero.GatheredParameters(list(module.parameters(recurse=False)), modifier_rank=0):
                    if torch.distributed.get_rank() == 0:
                        module._load_from_state_dict(*args)
            else:
                module._load_from_state_dict(*args)

            for name, child in module._modules.items():
                if child is not None:
                    load(child, prefix + name + ".")

        # Make sure we are able to load base models as well as derived models (with heads)
        start_prefix = ""
        model_to_load = model
        if not hasattr(model, cls.base_model_prefix) and has_prefix_module:
            start_prefix = cls.base_model_prefix + "."
        if hasattr(model, cls.base_model_prefix) and not has_prefix_module:
            model_to_load = getattr(model, cls.base_model_prefix)
1678
1679
1680
1681
1682
            if any(key in expected_keys_not_prefixed for key in loaded_keys):
                raise ValueError(
                    "The state dictionary of the model you are training to load is corrupted. Are you sure it was "
                    "properly saved?"
                )
1683
1684
1685

        load(model_to_load, prefix=start_prefix)

1686
1687
1688
1689
        if len(error_msgs) > 0:
            error_msg = "\n\t".join(error_msgs)
            raise RuntimeError(f"Error(s) in loading state_dict for {model.__class__.__name__}:\n\t{error_msg}")

1690
1691
1692
1693
1694
1695
1696
1697
1698
1699
1700
1701
1702
1703
1704
1705
1706
        if len(unexpected_keys) > 0:
            logger.warning(
                f"Some weights of the model checkpoint at {pretrained_model_name_or_path} were not used when "
                f"initializing {model.__class__.__name__}: {unexpected_keys}\n"
                f"- This IS expected if you are initializing {model.__class__.__name__} from the checkpoint of a model trained on another task "
                f"or with another architecture (e.g. initializing a BertForSequenceClassification model from a BertForPreTraining model).\n"
                f"- This IS NOT expected if you are initializing {model.__class__.__name__} from the checkpoint of a model that you expect "
                f"to be exactly identical (initializing a BertForSequenceClassification model from a BertForSequenceClassification model)."
            )
        else:
            logger.info(f"All model checkpoint weights were used when initializing {model.__class__.__name__}.\n")
        if len(missing_keys) > 0:
            logger.warning(
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at {pretrained_model_name_or_path} "
                f"and are newly initialized: {missing_keys}\n"
                f"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference."
            )
1707
        elif len(mismatched_keys) == 0:
1708
1709
1710
1711
1712
            logger.info(
                f"All the weights of {model.__class__.__name__} were initialized from the model checkpoint at {pretrained_model_name_or_path}.\n"
                f"If your task is similar to the task the model of the checkpoint was trained on, "
                f"you can already use {model.__class__.__name__} for predictions without further training."
            )
1713
1714
1715
1716
1717
1718
1719
1720
1721
1722
1723
1724
        if len(mismatched_keys) > 0:
            mismatched_warning = "\n".join(
                [
                    f"- {key}: found shape {shape1} in the checkpoint and {shape2} in the model instantiated"
                    for key, shape1, shape2 in mismatched_keys
                ]
            )
            logger.warning(
                f"Some weights of {model.__class__.__name__} were not initialized from the model checkpoint at {pretrained_model_name_or_path} "
                f"and are newly initialized because the shapes did not match:\n{mismatched_warning}\n"
                f"You should probably TRAIN this model on a down-stream task to be able to use it for predictions and inference."
            )
1725

1726
        return model, missing_keys, unexpected_keys, mismatched_keys, error_msgs
1727
1728
1729
1730

    def retrieve_modules_from_names(self, names, add_prefix=False, remove_prefix=False):
        module_keys = set([".".join(key.split(".")[:-1]) for key in names])

Patrick von Platen's avatar
Patrick von Platen committed
1731
1732
1733
1734
        # torch.nn.ParameterList is a special case where two parameter keywords
        # are appended to the module name, *e.g.* bert.special_embeddings.0
        module_keys = module_keys.union(set([".".join(key.split(".")[:-2]) for key in names if key[-1].isdigit()]))

1735
1736
1737
1738
1739
1740
        retrieved_modules = []
        # retrieve all modules that has at least one missing weight name
        for name, module in self.named_modules():
            if remove_prefix:
                name = ".".join(name.split(".")[1:]) if name.startswith(self.base_model_prefix) else name
            elif add_prefix:
Patrick von Platen's avatar
Patrick von Platen committed
1741
                name = ".".join([self.base_model_prefix, name]) if len(name) > 0 else self.base_model_prefix
1742
1743
1744
1745
1746
1747

            if name in module_keys:
                retrieved_modules.append(module)

        return retrieved_modules

1748
1749
1750
1751
1752
1753
1754
1755
1756
1757
1758
1759
1760
1761
1762
1763
1764
1765
1766
1767
1768
1769
1770
1771
1772
1773
1774
1775
1776
1777
1778
1779
1780
1781
1782
1783
1784
1785
1786
1787
1788
1789
1790
1791
1792
1793
1794
1795
1796
1797
1798
1799
1800
1801
1802
1803
1804
1805
1806
1807
1808
1809
1810
1811
1812
1813
    @classmethod
    def _load_state_dict_into_model_low_mem(cls, model, loaded_state_dict_keys, resolved_archive_file):
        """
        This is an experimental function that loads the model using ~1.x model size CPU memory

        Before it gets called we do:

        1. save which state_dict keys we have
        2. drop state_dict before model is created, since the latter takes 1x model size memory

        Here then we continue:

        3. switch to the meta device all params/buffers that are going to be replaced from the loaded state_dict
        4. load state_dict 2nd time
        5. replace the params/buffers from the state_dict

        Currently, it doesn't handle missing_keys, unexpected_keys, mismatched_keys. It can't handle deepspeed.
        """

        require_version_core("torch>=1.9")
        if is_deepspeed_zero3_enabled():
            raise ValueError("low_cpu_mem_usage arg cannot be used with DeepSpeed ZeRO-3")

        # a helper util to find the last sub-module and the param/buffer name
        def find_submodule_and_param_name(model, long_key):
            split_key = long_key.split(".")
            submodule = model
            while len(split_key) > 1:
                if hasattr(submodule, split_key[0]):
                    submodule = getattr(submodule, split_key[0])
                    del split_key[0]
                else:
                    submodule = None
                    break
            return submodule, split_key[0]

        # dematerialize param storage for keys that are going to be replaced by state_dict, by
        # putting those on the meta device
        for k in loaded_state_dict_keys:
            submodule, param_name = find_submodule_and_param_name(model, k)
            if submodule is not None:
                # selectively switch to the meta device only those params/buffers that will
                # be next replaced from state_dict. This a complex way to do p.to_("meta")
                # since we have no in-place to_ for tensors.
                new_val = getattr(submodule, param_name)
                if isinstance(new_val, torch.nn.Parameter):
                    # isinstance returns False for Params on meta device, so switch after the check
                    new_val = torch.nn.Parameter(new_val.to("meta"))
                else:
                    new_val = new_val.to("meta")
                setattr(submodule, param_name, new_val)

        # only now can load state_dict
        state_dict = torch.load(resolved_archive_file, map_location="cpu")

        # materialize state_dict entries one by one on CPU
        for k in loaded_state_dict_keys:
            submodule, param_name = find_submodule_and_param_name(model, k)
            if submodule is not None:
                new_val = state_dict[k]
                if isinstance(getattr(submodule, param_name), torch.nn.Parameter):
                    new_val = torch.nn.Parameter(new_val)
                setattr(submodule, param_name, new_val)

        del state_dict

1814
1815
1816
1817
1818
1819
1820
1821
1822
1823
1824
1825
1826
1827
1828
1829
1830
1831
1832
1833
    @classmethod
    def register_for_auto_class(cls, auto_class="AutoModel"):
        """
        Register this class with a given auto class. This should only be used for custom models as the ones in the
        library are already mapped with an auto class.

        Args:
            auto_class (`str` or `type`, *optional*, defaults to `"AutoModel"`):
                The auto class to register this new model with.
        """
        if not isinstance(auto_class, str):
            auto_class = auto_class.__name__

        import transformers.models.auto as auto_module

        if not hasattr(auto_module, auto_class):
            raise ValueError(f"{auto_class} is not a valid auto class.")

        cls._auto_class = auto_class

thomwolf's avatar
thomwolf committed
1834

1835
1836
1837
1838
1839
1840
1841
# To update the docstring, we need to copy the method, otherwise we change the original docstring.
PreTrainedModel.push_to_hub = copy_func(PreTrainedModel.push_to_hub)
PreTrainedModel.push_to_hub.__doc__ = PreTrainedModel.push_to_hub.__doc__.format(
    object="model", object_class="AutoModel", object_files="model checkpoint"
)


thomwolf's avatar
thomwolf committed
1842
class Conv1D(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
1843
1844
1845
1846
1847
1848
    """
    1D-convolutional layer as defined by Radford et al. for OpenAI GPT (and also used in GPT-2).

    Basically works like a linear layer but the weights are transposed.

    Args:
1849
1850
        nf (`int`): The number of output features.
        nx (`int`): The number of input features.
Sylvain Gugger's avatar
Sylvain Gugger committed
1851
1852
    """

thomwolf's avatar
thomwolf committed
1853
    def __init__(self, nf, nx):
Julien Chaumond's avatar
Julien Chaumond committed
1854
        super().__init__()
thomwolf's avatar
thomwolf committed
1855
1856
1857
1858
1859
1860
1861
1862
1863
1864
1865
1866
1867
        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
1868
class PoolerStartLogits(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
1869
1870
    """
    Compute SQuAD start logits from sequence hidden states.
1871

Sylvain Gugger's avatar
Sylvain Gugger committed
1872
    Args:
1873
1874
        config ([`PretrainedConfig`]):
            The config used by the model, will be used to grab the `hidden_size` of the model.
Sylvain Gugger's avatar
Sylvain Gugger committed
1875
1876
1877
    """

    def __init__(self, config: PretrainedConfig):
Julien Chaumond's avatar
Julien Chaumond committed
1878
        super().__init__()
thomwolf's avatar
thomwolf committed
1879
1880
        self.dense = nn.Linear(config.hidden_size, 1)

Sylvain Gugger's avatar
Sylvain Gugger committed
1881
1882
1883
1884
1885
    def forward(
        self, hidden_states: torch.FloatTensor, p_mask: Optional[torch.FloatTensor] = None
    ) -> torch.FloatTensor:
        """
        Args:
1886
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1887
                The final hidden states of the model.
1888
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1889
1890
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.
Sylvain Gugger's avatar
Sylvain Gugger committed
1891
1892

        Returns:
1893
            `torch.FloatTensor`: The start logits for SQuAD.
thomwolf's avatar
thomwolf committed
1894
        """
thomwolf's avatar
thomwolf committed
1895
1896
1897
        x = self.dense(hidden_states).squeeze(-1)

        if p_mask is not None:
Lysandre Debut's avatar
Lysandre Debut committed
1898
            if get_parameter_dtype(self) == torch.float16:
1899
1900
1901
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e30 * p_mask
thomwolf's avatar
thomwolf committed
1902
1903
1904
1905
1906
1907

        return x


class PoolerEndLogits(nn.Module):
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
1908
    Compute SQuAD end logits from sequence hidden states.
1909

Sylvain Gugger's avatar
Sylvain Gugger committed
1910
    Args:
1911
        config ([`PretrainedConfig`]):
Sylvain Gugger's avatar
Sylvain Gugger committed
1912
1913
            The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps`
            to use.
Sylvain Gugger's avatar
Sylvain Gugger committed
1914
1915
1916
    """

    def __init__(self, config: PretrainedConfig):
Julien Chaumond's avatar
Julien Chaumond committed
1917
        super().__init__()
thomwolf's avatar
thomwolf committed
1918
1919
1920
1921
1922
        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)

Sylvain Gugger's avatar
Sylvain Gugger committed
1923
1924
1925
1926
1927
1928
1929
1930
1931
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        start_states: Optional[torch.FloatTensor] = None,
        start_positions: Optional[torch.LongTensor] = None,
        p_mask: Optional[torch.FloatTensor] = None,
    ) -> torch.FloatTensor:
        """
        Args:
1932
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
Sylvain Gugger's avatar
Sylvain Gugger committed
1933
                The final hidden states of the model.
1934
            start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1935
                The hidden states of the first tokens for the labeled span.
1936
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1937
                The position of the first token for the labeled span.
1938
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
1939
1940
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.
Sylvain Gugger's avatar
Sylvain Gugger committed
1941

1942
        <Tip>
Sylvain Gugger's avatar
Sylvain Gugger committed
1943

Stas Bekman's avatar
Stas Bekman committed
1944
1945
        One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides
        `start_states`.
1946
1947

        </Tip>
Sylvain Gugger's avatar
Sylvain Gugger committed
1948
1949

        Returns:
1950
            `torch.FloatTensor`: The end logits for SQuAD.
thomwolf's avatar
thomwolf committed
1951
        """
1952
1953
1954
        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
1955
        if start_positions is not None:
1956
            slen, hsz = hidden_states.shape[-2:]
1957
1958
1959
            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
1960
1961
1962
1963
1964
1965
1966

        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:
Lysandre Debut's avatar
Lysandre Debut committed
1967
            if get_parameter_dtype(self) == torch.float16:
1968
1969
1970
                x = x * (1 - p_mask) - 65500 * p_mask
            else:
                x = x * (1 - p_mask) - 1e30 * p_mask
thomwolf's avatar
thomwolf committed
1971
1972
1973
1974
1975

        return x


class PoolerAnswerClass(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
1976
1977
1978
1979
    """
    Compute SQuAD 2.0 answer class from classification and start tokens hidden states.

    Args:
1980
1981
        config ([`PretrainedConfig`]):
            The config used by the model, will be used to grab the `hidden_size` of the model.
Sylvain Gugger's avatar
Sylvain Gugger committed
1982
    """
1983

thomwolf's avatar
thomwolf committed
1984
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
1985
        super().__init__()
thomwolf's avatar
thomwolf committed
1986
1987
1988
1989
        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)

Sylvain Gugger's avatar
Sylvain Gugger committed
1990
1991
1992
1993
1994
1995
1996
    def forward(
        self,
        hidden_states: torch.FloatTensor,
        start_states: Optional[torch.FloatTensor] = None,
        start_positions: Optional[torch.LongTensor] = None,
        cls_index: Optional[torch.LongTensor] = None,
    ) -> torch.FloatTensor:
1997
1998
        """
        Args:
1999
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
Sylvain Gugger's avatar
Sylvain Gugger committed
2000
                The final hidden states of the model.
2001
            start_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2002
                The hidden states of the first tokens for the labeled span.
2003
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2004
                The position of the first token for the labeled span.
2005
2006
2007
2008
            cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Position of the CLS token for each sentence in the batch. If `None`, takes the last token.

        <Tip>
Sylvain Gugger's avatar
Sylvain Gugger committed
2009

Stas Bekman's avatar
Stas Bekman committed
2010
2011
        One of `start_states` or `start_positions` should be not `None`. If both are set, `start_positions` overrides
        `start_states`.
Sylvain Gugger's avatar
Sylvain Gugger committed
2012

2013
        </Tip>
Sylvain Gugger's avatar
Sylvain Gugger committed
2014
2015

        Returns:
2016
            `torch.FloatTensor`: The SQuAD 2.0 answer class.
thomwolf's avatar
thomwolf committed
2017
        """
Sylvain Gugger's avatar
Sylvain Gugger committed
2018
        # No dependency on end_feature so that we can obtain one single `cls_logits` for each sample.
2019
        hsz = hidden_states.shape[-1]
2020
2021
2022
        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
2023
        if start_positions is not None:
2024
2025
            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
2026
2027

        if cls_index is not None:
2028
2029
            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
2030
        else:
2031
            cls_token_state = hidden_states[:, -1, :]  # shape (bsz, hsz)
thomwolf's avatar
thomwolf committed
2032
2033
2034
2035
2036
2037
2038
2039

        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


2040
2041
2042
@dataclass
class SquadHeadOutput(ModelOutput):
    """
2043
    Base class for outputs of question answering models using a [`~modeling_utils.SQuADHead`].
2044
2045

    Args:
2046
        loss (`torch.FloatTensor` of shape `(1,)`, *optional*, returned if both `start_positions` and `end_positions` are provided):
Sylvain Gugger's avatar
Sylvain Gugger committed
2047
2048
            Classification loss as the sum of start token, end token (and is_impossible if provided) classification
            losses.
2049
        start_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
2050
            Log probabilities for the top config.start_n_top start token possibilities (beam-search).
2051
        start_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
2052
            Indices for the top config.start_n_top start token possibilities (beam-search).
2053
2054
        end_top_log_probs (`torch.FloatTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
            Log probabilities for the top `config.start_n_top * config.end_n_top` end token possibilities
Sylvain Gugger's avatar
Sylvain Gugger committed
2055
            (beam-search).
2056
2057
2058
2059
        end_top_index (`torch.LongTensor` of shape `(batch_size, config.start_n_top * config.end_n_top)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
            Indices for the top `config.start_n_top * config.end_n_top` end token possibilities (beam-search).
        cls_logits (`torch.FloatTensor` of shape `(batch_size,)`, *optional*, returned if `start_positions` or `end_positions` is not provided):
            Log probabilities for the `is_impossible` label of the answers.
2060
2061
2062
2063
2064
2065
2066
2067
2068
2069
2070

    """

    loss: Optional[torch.FloatTensor] = None
    start_top_log_probs: Optional[torch.FloatTensor] = None
    start_top_index: Optional[torch.LongTensor] = None
    end_top_log_probs: Optional[torch.FloatTensor] = None
    end_top_index: Optional[torch.LongTensor] = None
    cls_logits: Optional[torch.FloatTensor] = None


thomwolf's avatar
thomwolf committed
2071
class SQuADHead(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
2072
2073
    r"""
    A SQuAD head inspired by XLNet.
2074

Sylvain Gugger's avatar
Sylvain Gugger committed
2075
    Args:
2076
        config ([`PretrainedConfig`]):
Sylvain Gugger's avatar
Sylvain Gugger committed
2077
2078
            The config used by the model, will be used to grab the `hidden_size` of the model and the `layer_norm_eps`
            to use.
thomwolf's avatar
thomwolf committed
2079
    """
2080

thomwolf's avatar
thomwolf committed
2081
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
2082
        super().__init__()
thomwolf's avatar
thomwolf committed
2083
2084
2085
2086
2087
2088
2089
        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)

Sylvain Gugger's avatar
Sylvain Gugger committed
2090
    @replace_return_docstrings(output_type=SquadHeadOutput, config_class=PretrainedConfig)
2091
    def forward(
2092
        self,
Sylvain Gugger's avatar
Sylvain Gugger committed
2093
2094
2095
2096
2097
2098
        hidden_states: torch.FloatTensor,
        start_positions: Optional[torch.LongTensor] = None,
        end_positions: Optional[torch.LongTensor] = None,
        cls_index: Optional[torch.LongTensor] = None,
        is_impossible: Optional[torch.LongTensor] = None,
        p_mask: Optional[torch.FloatTensor] = None,
2099
        return_dict: bool = False,
Sylvain Gugger's avatar
Sylvain Gugger committed
2100
2101
    ) -> Union[SquadHeadOutput, Tuple[torch.FloatTensor]]:
        """
Lysandre's avatar
Lysandre committed
2102
        Args:
2103
            hidden_states (`torch.FloatTensor` of shape `(batch_size, seq_len, hidden_size)`):
Lysandre's avatar
Lysandre committed
2104
                Final hidden states of the model on the sequence tokens.
2105
            start_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Lysandre's avatar
Lysandre committed
2106
                Positions of the first token for the labeled span.
2107
            end_positions (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Lysandre's avatar
Lysandre committed
2108
                Positions of the last token for the labeled span.
2109
2110
2111
            cls_index (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
                Position of the CLS token for each sentence in the batch. If `None`, takes the last token.
            is_impossible (`torch.LongTensor` of shape `(batch_size,)`, *optional*):
Lysandre's avatar
Lysandre committed
2112
                Whether the question has a possible answer in the paragraph or not.
2113
            p_mask (`torch.FloatTensor` of shape `(batch_size, seq_len)`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2114
2115
                Mask for tokens at invalid position, such as query and special symbols (PAD, SEP, CLS). 1.0 means token
                should be masked.
2116
2117
            return_dict (`bool`, *optional*, defaults to `False`):
                Whether or not to return a [`~file_utils.ModelOutput`] instead of a plain tuple.
Sylvain Gugger's avatar
Sylvain Gugger committed
2118

Lysandre's avatar
Lysandre committed
2119
        Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
2120
        """
thomwolf's avatar
thomwolf committed
2121
        start_logits = self.start_logits(hidden_states, p_mask=p_mask)
thomwolf's avatar
thomwolf committed
2122
2123
2124
2125
2126
2127
2128
2129
2130
2131
2132
2133
2134
2135
2136
2137
2138
2139
2140
2141
2142
2143
2144

        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
2145

2146
            return SquadHeadOutput(loss=total_loss) if return_dict else (total_loss,)
thomwolf's avatar
thomwolf committed
2147
2148
2149
2150

        else:
            # during inference, compute the end logits based on beam search
            bsz, slen, hsz = hidden_states.size()
2151
            start_log_probs = nn.functional.softmax(start_logits, dim=-1)  # shape (bsz, slen)
2152
2153
2154
2155
2156
2157
2158
2159
2160
2161
2162

            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
2163
2164
            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)
2165
            end_log_probs = nn.functional.softmax(end_logits, dim=1)  # shape (bsz, slen, start_n_top)
thomwolf's avatar
thomwolf committed
2166

2167
2168
2169
            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
2170
2171
2172
2173
2174
2175
            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)

2176
            if not return_dict:
2177
2178
2179
2180
2181
2182
2183
2184
2185
                return (start_top_log_probs, start_top_index, end_top_log_probs, end_top_index, cls_logits)
            else:
                return SquadHeadOutput(
                    start_top_log_probs=start_top_log_probs,
                    start_top_index=start_top_index,
                    end_top_log_probs=end_top_log_probs,
                    end_top_index=end_top_index,
                    cls_logits=cls_logits,
                )
thomwolf's avatar
thomwolf committed
2186
2187
2188


class SequenceSummary(nn.Module):
Sylvain Gugger's avatar
Sylvain Gugger committed
2189
2190
2191
2192
    r"""
    Compute a single vector summary of a sequence hidden states.

    Args:
2193
        config ([`PretrainedConfig`]):
Sylvain Gugger's avatar
Sylvain Gugger committed
2194
2195
            The config used by the model. Relevant arguments in the config class of the model are (refer to the actual
            config class of your model for the default values it uses):
Sylvain Gugger's avatar
Sylvain Gugger committed
2196

2197
            - **summary_type** (`str`) -- The method to use to make this summary. Accepted values are:
Sylvain Gugger's avatar
Sylvain Gugger committed
2198

2199
2200
2201
2202
2203
                - `"last"` -- Take the last token hidden state (like XLNet)
                - `"first"` -- Take the first token hidden state (like Bert)
                - `"mean"` -- Take the mean of all tokens hidden states
                - `"cls_index"` -- Supply a Tensor of classification token position (GPT/GPT-2)
                - `"attn"` -- Not implemented now, use multi-head attention
Sylvain Gugger's avatar
Sylvain Gugger committed
2204

2205
            - **summary_use_proj** (`bool`) -- Add a projection after the vector extraction.
Sylvain Gugger's avatar
Sylvain Gugger committed
2206
2207
2208
2209
2210
2211
            - **summary_proj_to_labels** (`bool`) -- If `True`, the projection outputs to `config.num_labels` classes
              (otherwise to `config.hidden_size`).
            - **summary_activation** (`Optional[str]`) -- Set to `"tanh"` to add a tanh activation to the output,
              another string or `None` will add no activation.
            - **summary_first_dropout** (`float`) -- Optional dropout probability before the projection and activation.
            - **summary_last_dropout** (`float`)-- Optional dropout probability after the projection and activation.
thomwolf's avatar
thomwolf committed
2212
    """
2213

2214
    def __init__(self, config: PretrainedConfig):
Julien Chaumond's avatar
Julien Chaumond committed
2215
        super().__init__()
thomwolf's avatar
thomwolf committed
2216

2217
        self.summary_type = getattr(config, "summary_type", "last")
2218
        if self.summary_type == "attn":
thomwolf's avatar
thomwolf committed
2219
2220
2221
2222
2223
            # 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
2224
        self.summary = Identity()
2225
2226
        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:
2227
                num_classes = config.num_labels
thomwolf's avatar
thomwolf committed
2228
2229
2230
2231
            else:
                num_classes = config.hidden_size
            self.summary = nn.Linear(config.hidden_size, num_classes)

2232
        activation_string = getattr(config, "summary_activation", None)
Lysandre's avatar
Lysandre committed
2233
        self.activation: Callable = get_activation(activation_string) if activation_string else Identity()
thomwolf's avatar
thomwolf committed
2234

thomwolf's avatar
thomwolf committed
2235
        self.first_dropout = Identity()
2236
        if hasattr(config, "summary_first_dropout") and config.summary_first_dropout > 0:
2237
2238
            self.first_dropout = nn.Dropout(config.summary_first_dropout)

thomwolf's avatar
thomwolf committed
2239
        self.last_dropout = Identity()
2240
        if hasattr(config, "summary_last_dropout") and config.summary_last_dropout > 0:
2241
            self.last_dropout = nn.Dropout(config.summary_last_dropout)
thomwolf's avatar
thomwolf committed
2242

Sylvain Gugger's avatar
Sylvain Gugger committed
2243
2244
2245
2246
2247
2248
2249
    def forward(
        self, hidden_states: torch.FloatTensor, cls_index: Optional[torch.LongTensor] = None
    ) -> torch.FloatTensor:
        """
        Compute a single vector summary of a sequence hidden states.

        Args:
2250
            hidden_states (`torch.FloatTensor` of shape `[batch_size, seq_len, hidden_size]`):
Sylvain Gugger's avatar
Sylvain Gugger committed
2251
                The hidden states of the last layer.
2252
            cls_index (`torch.LongTensor` of shape `[batch_size]` or `[batch_size, ...]` where ... are optional leading dimensions of `hidden_states`, *optional*):
Sylvain Gugger's avatar
Sylvain Gugger committed
2253
                Used if `summary_type == "cls_index"` and takes the last token of the sequence as classification token.
Sylvain Gugger's avatar
Sylvain Gugger committed
2254
2255

        Returns:
2256
            `torch.FloatTensor`: The summary of the sequence hidden states.
thomwolf's avatar
thomwolf committed
2257
        """
2258
        if self.summary_type == "last":
thomwolf's avatar
thomwolf committed
2259
            output = hidden_states[:, -1]
2260
        elif self.summary_type == "first":
thomwolf's avatar
thomwolf committed
2261
            output = hidden_states[:, 0]
2262
        elif self.summary_type == "mean":
thomwolf's avatar
thomwolf committed
2263
            output = hidden_states.mean(dim=1)
2264
        elif self.summary_type == "cls_index":
thomwolf's avatar
thomwolf committed
2265
            if cls_index is None:
Lysandre's avatar
Lysandre committed
2266
2267
2268
2269
2270
                cls_index = torch.full_like(
                    hidden_states[..., :1, :],
                    hidden_states.shape[-2] - 1,
                    dtype=torch.long,
                )
thomwolf's avatar
thomwolf committed
2271
            else:
thomwolf's avatar
thomwolf committed
2272
                cls_index = cls_index.unsqueeze(-1).unsqueeze(-1)
2273
                cls_index = cls_index.expand((-1,) * (cls_index.dim() - 1) + (hidden_states.size(-1),))
thomwolf's avatar
thomwolf committed
2274
            # shape of cls_index: (bsz, XX, 1, hidden_size) where XX are optional leading dim of hidden_states
2275
2276
            output = hidden_states.gather(-2, cls_index).squeeze(-2)  # shape (bsz, XX, hidden_size)
        elif self.summary_type == "attn":
thomwolf's avatar
thomwolf committed
2277
2278
            raise NotImplementedError

2279
        output = self.first_dropout(output)
thomwolf's avatar
thomwolf committed
2280
2281
        output = self.summary(output)
        output = self.activation(output)
2282
        output = self.last_dropout(output)
thomwolf's avatar
thomwolf committed
2283
2284
2285
2286

        return output


2287
def unwrap_model(model: nn.Module) -> nn.Module:
2288
2289
2290
2291
    """
    Recursively unwraps a model from potential containers (as used in distributed training).

    Args:
2292
        model (`torch.nn.Module`): The model to unwrap.
2293
2294
2295
2296
2297
2298
2299
2300
    """
    # since there could be multiple levels of wrapping, unwrap recursively
    if hasattr(model, "module"):
        return unwrap_model(model.module)
    else:
        return model


2301
def prune_linear_layer(layer: nn.Linear, index: torch.LongTensor, dim: int = 0) -> nn.Linear:
Sylvain Gugger's avatar
Sylvain Gugger committed
2302
2303
2304
2305
2306
2307
    """
    Prune a linear layer to keep only entries in index.

    Used to remove heads.

    Args:
2308
2309
2310
        layer (`torch.nn.Linear`): The layer to prune.
        index (`torch.LongTensor`): The indices to keep in the layer.
        dim (`int`, *optional*, defaults to 0): The dimension on which to keep the indices.
Sylvain Gugger's avatar
Sylvain Gugger committed
2311
2312

    Returns:
2313
        `torch.nn.Linear`: The pruned layer as a new layer with `requires_grad=True`.
2314
2315
2316
2317
2318
2319
2320
2321
2322
2323
2324
2325
2326
2327
2328
2329
2330
2331
2332
2333
2334
    """
    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


Sylvain Gugger's avatar
Sylvain Gugger committed
2335
2336
2337
2338
2339
2340
2341
2342
def prune_conv1d_layer(layer: Conv1D, index: torch.LongTensor, dim: int = 1) -> Conv1D:
    """
    Prune a Conv1D layer to keep only entries in index. A Conv1D work as a Linear layer (see e.g. BERT) but the weights
    are transposed.

    Used to remove heads.

    Args:
2343
2344
2345
        layer ([`~modeling_utils.Conv1D`]): The layer to prune.
        index (`torch.LongTensor`): The indices to keep in the layer.
        dim (`int`, *optional*, defaults to 1): The dimension on which to keep the indices.
Sylvain Gugger's avatar
Sylvain Gugger committed
2346
2347

    Returns:
2348
        [`~modeling_utils.Conv1D`]: The pruned layer as a new layer with `requires_grad=True`.
2349
2350
2351
2352
2353
2354
2355
2356
2357
2358
2359
2360
2361
2362
2363
2364
2365
    """
    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
2366
2367


Sylvain Gugger's avatar
Sylvain Gugger committed
2368
def prune_layer(
2369
2370
    layer: Union[nn.Linear, Conv1D], index: torch.LongTensor, dim: Optional[int] = None
) -> Union[nn.Linear, Conv1D]:
Sylvain Gugger's avatar
Sylvain Gugger committed
2371
2372
2373
2374
2375
2376
    """
    Prune a Conv1D or linear layer to keep only entries in index.

    Used to remove heads.

    Args:
2377
2378
2379
        layer (`Union[torch.nn.Linear, Conv1D]`): The layer to prune.
        index (`torch.LongTensor`): The indices to keep in the layer.
        dim (`int`, *optional*): The dimension on which to keep the indices.
Sylvain Gugger's avatar
Sylvain Gugger committed
2380
2381

    Returns:
Sylvain Gugger's avatar
Sylvain Gugger committed
2382
        `torch.nn.Linear` or [`~modeling_utils.Conv1D`]: The pruned layer as a new layer with `requires_grad=True`.
2383
2384
2385
2386
2387
2388
    """
    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:
2389
        raise ValueError(f"Can't prune layer of class {layer.__class__}")
Patrick von Platen's avatar
Patrick von Platen committed
2390
2391
2392


def apply_chunking_to_forward(
2393
    forward_fn: Callable[..., torch.Tensor], chunk_size: int, chunk_dim: int, *input_tensors
Patrick von Platen's avatar
Patrick von Platen committed
2394
2395
) -> torch.Tensor:
    """
Sylvain Gugger's avatar
Sylvain Gugger committed
2396
2397
    This function chunks the `input_tensors` into smaller input tensor parts of size `chunk_size` over the dimension
    `chunk_dim`. It then applies a layer `forward_fn` to each chunk independently to save memory.
2398

Sylvain Gugger's avatar
Sylvain Gugger committed
2399
2400
    If the `forward_fn` is independent across the `chunk_dim` this function will yield the same result as directly
    applying `forward_fn` to `input_tensors`.
Patrick von Platen's avatar
Patrick von Platen committed
2401
2402

    Args:
2403
        forward_fn (`Callable[..., torch.Tensor]`):
2404
            The forward function of the model.
2405
2406
2407
2408
2409
2410
        chunk_size (`int`):
            The chunk size of a chunked tensor: `num_chunks = len(input_tensors[0]) / chunk_size`.
        chunk_dim (`int`):
            The dimension over which the `input_tensors` should be chunked.
        input_tensors (`Tuple[torch.Tensor]`):
            The input tensors of `forward_fn` which will be chunked
Sylvain Gugger's avatar
Sylvain Gugger committed
2411

Patrick von Platen's avatar
Patrick von Platen committed
2412
    Returns:
2413
        `torch.Tensor`: A tensor with the same shape as the `forward_fn` would have given if applied`.
Patrick von Platen's avatar
Patrick von Platen committed
2414
2415


2416
    Examples:
Patrick von Platen's avatar
Patrick von Platen committed
2417

2418
2419
2420
2421
2422
    ```python
    # rename the usual forward() fn to forward_chunk()
    def forward_chunk(self, hidden_states):
        hidden_states = self.decoder(hidden_states)
        return hidden_states
Patrick von Platen's avatar
Patrick von Platen committed
2423

Sylvain Gugger's avatar
Sylvain Gugger committed
2424

2425
2426
2427
2428
    # implement a chunked forward function
    def forward(self, hidden_states):
        return apply_chunking_to_forward(self.forward_chunk, self.chunk_size_lm_head, self.seq_len_dim, hidden_states)
    ```"""
Patrick von Platen's avatar
Patrick von Platen committed
2429

2430
    assert len(input_tensors) > 0, f"{input_tensors} has to be a tuple/list of tensors"
Patrick von Platen's avatar
Patrick von Platen committed
2431

2432
    # inspect.signature exist since python 3.5 and is a python method -> no problem with backward compatibility
Patrick von Platen's avatar
Patrick von Platen committed
2433
    num_args_in_forward_chunk_fn = len(inspect.signature(forward_fn).parameters)
2434
2435
2436
2437
2438
    if num_args_in_forward_chunk_fn != len(input_tensors):
        raise ValueError(
            f"forward_chunk_fn expects {num_args_in_forward_chunk_fn} arguments, but only {len(input_tensors)} input "
            "tensors are given"
        )
Patrick von Platen's avatar
Patrick von Platen committed
2439
2440

    if chunk_size > 0:
2441
2442
2443
2444
2445
2446
2447
2448
        tensor_shape = input_tensors[0].shape[chunk_dim]
        for input_tensor in input_tensors:
            if input_tensor.shape[chunk_dim] != tensor_shape:
                raise ValueError(
                    f"All input tenors have to be of the same shape: {tensor_shape}, "
                    f"found shape {input_tensor.shape[chunk_dim]}"
                )

2449
2450
2451
2452
2453
        if input_tensors[0].shape[chunk_dim] % chunk_size != 0:
            raise ValueError(
                f"The dimension to be chunked {input_tensors[0].shape[chunk_dim]} has to be a multiple of the chunk "
                f"size {chunk_size}"
            )
Patrick von Platen's avatar
Patrick von Platen committed
2454
2455
2456
2457
2458
2459
2460
2461
2462
2463
2464

        num_chunks = input_tensors[0].shape[chunk_dim] // chunk_size

        # chunk input tensor into tuples
        input_tensors_chunks = tuple(input_tensor.chunk(num_chunks, dim=chunk_dim) for input_tensor in input_tensors)
        # apply forward fn to every tuple
        output_chunks = tuple(forward_fn(*input_tensors_chunk) for input_tensors_chunk in zip(*input_tensors_chunks))
        # concatenate output at same dimension
        return torch.cat(output_chunks, dim=chunk_dim)

    return forward_fn(*input_tensors)