modeling_albert.py 40.4 KB
Newer Older
Lysandre's avatar
Lysandre committed
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# coding=utf-8
# Copyright 2018 Google AI, Google Brain and the HuggingFace Inc. team.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
"""PyTorch ALBERT model. """

Lysandre's avatar
Lysandre committed
17
import logging
Aymeric Augustin's avatar
Aymeric Augustin committed
18
19
20
import math
import os

Lysandre's avatar
Lysandre committed
21
22
import torch
import torch.nn as nn
Lysandre's avatar
Lysandre committed
23
from torch.nn import CrossEntropyLoss, MSELoss
Aymeric Augustin's avatar
Aymeric Augustin committed
24

Lysandre's avatar
Lysandre committed
25
from transformers.configuration_albert import AlbertConfig
Aymeric Augustin's avatar
Aymeric Augustin committed
26
27
28
from transformers.modeling_bert import ACT2FN, BertEmbeddings, BertSelfAttention, prune_linear_layer
from transformers.modeling_utils import PreTrainedModel

29
from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
30

Aymeric Augustin's avatar
Aymeric Augustin committed
31

Lysandre's avatar
Lysandre committed
32
33
logger = logging.getLogger(__name__)

Lysandre's avatar
Lysandre committed
34
35

ALBERT_PRETRAINED_MODEL_ARCHIVE_MAP = {
36
37
38
39
40
41
42
43
    "albert-base-v1": "https://s3.amazonaws.com/models.huggingface.co/bert/albert-base-pytorch_model.bin",
    "albert-large-v1": "https://s3.amazonaws.com/models.huggingface.co/bert/albert-large-pytorch_model.bin",
    "albert-xlarge-v1": "https://s3.amazonaws.com/models.huggingface.co/bert/albert-xlarge-pytorch_model.bin",
    "albert-xxlarge-v1": "https://s3.amazonaws.com/models.huggingface.co/bert/albert-xxlarge-pytorch_model.bin",
    "albert-base-v2": "https://s3.amazonaws.com/models.huggingface.co/bert/albert-base-v2-pytorch_model.bin",
    "albert-large-v2": "https://s3.amazonaws.com/models.huggingface.co/bert/albert-large-v2-pytorch_model.bin",
    "albert-xlarge-v2": "https://s3.amazonaws.com/models.huggingface.co/bert/albert-xlarge-v2-pytorch_model.bin",
    "albert-xxlarge-v2": "https://s3.amazonaws.com/models.huggingface.co/bert/albert-xxlarge-v2-pytorch_model.bin",
Lysandre's avatar
Lysandre committed
44
45
46
}


Lysandre's avatar
Lysandre committed
47
48
49
50
51
52
53
def load_tf_weights_in_albert(model, config, tf_checkpoint_path):
    """ Load tf checkpoints in a pytorch model."""
    try:
        import re
        import numpy as np
        import tensorflow as tf
    except ImportError:
54
55
56
57
        logger.error(
            "Loading a TensorFlow model in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions."
        )
Lysandre's avatar
Lysandre committed
58
59
60
61
62
63
64
65
66
67
68
69
70
        raise
    tf_path = os.path.abspath(tf_checkpoint_path)
    logger.info("Converting TensorFlow checkpoint from {}".format(tf_path))
    # Load weights from TF model
    init_vars = tf.train.list_variables(tf_path)
    names = []
    arrays = []
    for name, shape in init_vars:
        logger.info("Loading TF weight {} with shape {}".format(name, shape))
        array = tf.train.load_variable(tf_path, name)
        names.append(name)
        arrays.append(array)

Lysandre's avatar
Lysandre committed
71
72
    for name, array in zip(names, arrays):
        print(name)
73

Lysandre's avatar
Lysandre committed
74
    for name, array in zip(names, arrays):
Lysandre's avatar
Lysandre committed
75
        original_name = name
Lysandre's avatar
Lysandre committed
76
77
78
79
80

        # If saved from the TF HUB module
        name = name.replace("module/", "")

        # Renaming and simplifying
Lysandre's avatar
Lysandre committed
81
        name = name.replace("ffn_1", "ffn")
Lysandre's avatar
Lysandre committed
82
        name = name.replace("bert/", "albert/")
83
        name = name.replace("attention_1", "attention")
Lysandre's avatar
Lysandre committed
84
        name = name.replace("transform/", "")
85
86
        name = name.replace("LayerNorm_1", "full_layer_layer_norm")
        name = name.replace("LayerNorm", "attention/LayerNorm")
Lysandre's avatar
Lysandre committed
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
        name = name.replace("transformer/", "")

        # The feed forward layer had an 'intermediate' step which has been abstracted away
        name = name.replace("intermediate/dense/", "")
        name = name.replace("ffn/intermediate/output/dense/", "ffn_output/")

        # ALBERT attention was split between self and output which have been abstracted away
        name = name.replace("/output/", "/")
        name = name.replace("/self/", "/")

        # The pooler is a linear layer
        name = name.replace("pooler/dense", "pooler")

        # The classifier was simplified to predictions from cls/predictions
        name = name.replace("cls/predictions", "predictions")
        name = name.replace("predictions/attention", "predictions")

        # Naming was changed to be more explicit
105
106
107
        name = name.replace("embeddings/attention", "embeddings")
        name = name.replace("inner_group_", "albert_layers/")
        name = name.replace("group_", "albert_layer_groups/")
108
109
110
111
112

        # Classifier
        if len(name.split("/")) == 1 and ("output_bias" in name or "output_weights" in name):
            name = "classifier/" + name

113
        # No ALBERT model currently handles the next sentence prediction task
114
115
116
        if "seq_relationship" in name:
            continue

117
        name = name.split("/")
118
119
120
121
122
123

        # Ignore the gradients applied by the LAMB/ADAM optimizers.
        if "adam_m" in name or "adam_v" in name or "global_step" in name:
            logger.info("Skipping {}".format("/".join(name)))
            continue

Lysandre's avatar
Lysandre committed
124
125
        pointer = model
        for m_name in name:
126
            if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
127
                scope_names = re.split(r"_(\d+)", m_name)
Lysandre's avatar
Lysandre committed
128
            else:
129
                scope_names = [m_name]
Lysandre's avatar
Lysandre committed
130

131
            if scope_names[0] == "kernel" or scope_names[0] == "gamma":
132
                pointer = getattr(pointer, "weight")
133
            elif scope_names[0] == "output_bias" or scope_names[0] == "beta":
134
                pointer = getattr(pointer, "bias")
135
            elif scope_names[0] == "output_weights":
136
                pointer = getattr(pointer, "weight")
137
            elif scope_names[0] == "squad":
138
                pointer = getattr(pointer, "classifier")
Lysandre's avatar
Lysandre committed
139
140
            else:
                try:
141
                    pointer = getattr(pointer, scope_names[0])
Lysandre's avatar
Lysandre committed
142
143
144
                except AttributeError:
                    logger.info("Skipping {}".format("/".join(name)))
                    continue
145
146
            if len(scope_names) >= 2:
                num = int(scope_names[1])
Lysandre's avatar
Lysandre committed
147
148
                pointer = pointer[num]

149
150
151
        if m_name[-11:] == "_embeddings":
            pointer = getattr(pointer, "weight")
        elif m_name == "kernel":
Lysandre's avatar
Lysandre committed
152
153
154
155
156
157
            array = np.transpose(array)
        try:
            assert pointer.shape == array.shape
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
Lysandre's avatar
Lysandre committed
158
        print("Initialize PyTorch weight {} from {}".format(name, original_name))
Lysandre's avatar
Lysandre committed
159
160
161
162
163
        pointer.data = torch.from_numpy(array)

    return model


Lysandre's avatar
Lysandre committed
164
class AlbertEmbeddings(BertEmbeddings):
Lysandre's avatar
Lysandre committed
165
166
167
    """
    Construct the embeddings from word, position and token_type embeddings.
    """
168

Lysandre's avatar
Lysandre committed
169
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
170
        super().__init__(config)
Lysandre's avatar
Lysandre committed
171
172
173
174

        self.word_embeddings = nn.Embedding(config.vocab_size, config.embedding_size, padding_idx=0)
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.embedding_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.embedding_size)
Lysandre's avatar
Lysandre committed
175
        self.LayerNorm = torch.nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
Lysandre's avatar
Lysandre committed
176
177


Lysandre's avatar
Lysandre committed
178
class AlbertAttention(BertSelfAttention):
Lysandre's avatar
Lysandre committed
179
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
180
        super().__init__(config)
Lysandre's avatar
Lysandre committed
181

182
        self.output_attentions = config.output_attentions
Lysandre's avatar
Lysandre committed
183
        self.num_attention_heads = config.num_attention_heads
184
        self.hidden_size = config.hidden_size
Lysandre's avatar
Lysandre committed
185
        self.attention_head_size = config.hidden_size // config.num_attention_heads
Lysandre's avatar
Lysandre committed
186
187
188
        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.LayerNorm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
Lysandre's avatar
Lysandre committed
189
190
        self.pruned_heads = set()

Lysandre's avatar
Lysandre committed
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
    def prune_heads(self, heads):
        if len(heads) == 0:
            return
        mask = torch.ones(self.num_attention_heads, self.attention_head_size)
        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
        for head in heads:
            # Compute how many pruned heads are before the head and move the index accordingly
            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
            mask[head] = 0
        mask = mask.view(-1).contiguous().eq(1)
        index = torch.arange(len(mask))[mask].long()

        # Prune linear layers
        self.query = prune_linear_layer(self.query, index)
        self.key = prune_linear_layer(self.key, index)
        self.value = prune_linear_layer(self.value, index)
        self.dense = prune_linear_layer(self.dense, index, dim=1)

        # Update hyper params and store pruned heads
        self.num_attention_heads = self.num_attention_heads - len(heads)
        self.all_head_size = self.attention_head_size * self.num_attention_heads
        self.pruned_heads = self.pruned_heads.union(heads)

Lysandre's avatar
Lysandre committed
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
    def forward(self, input_ids, attention_mask=None, head_mask=None):
        mixed_query_layer = self.query(input_ids)
        mixed_key_layer = self.key(input_ids)
        mixed_value_layer = self.value(input_ids)

        query_layer = self.transpose_for_scores(mixed_query_layer)
        key_layer = self.transpose_for_scores(mixed_key_layer)
        value_layer = self.transpose_for_scores(mixed_value_layer)

        # Take the dot product between "query" and "key" to get the raw attention scores.
        attention_scores = torch.matmul(query_layer, key_layer.transpose(-1, -2))
        attention_scores = attention_scores / math.sqrt(self.attention_head_size)
        if attention_mask is not None:
            # Apply the attention mask is (precomputed for all layers in BertModel forward() function)
            attention_scores = attention_scores + attention_mask

        # Normalize the attention scores to probabilities.
        attention_probs = nn.Softmax(dim=-1)(attention_scores)

        # This is actually dropping out entire tokens to attend to, which might
        # seem a bit unusual, but is taken from the original Transformer paper.
        attention_probs = self.dropout(attention_probs)

        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

        context_layer = torch.matmul(attention_probs, value_layer)

        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
Lysandre's avatar
Lysandre committed
244
245

        # Should find a better way to do this
246
247
248
249
250
        w = (
            self.dense.weight.t()
            .view(self.num_attention_heads, self.attention_head_size, self.hidden_size)
            .to(context_layer.dtype)
        )
251
        b = self.dense.bias.to(context_layer.dtype)
Lysandre's avatar
Lysandre committed
252
253

        projected_context_layer = torch.einsum("bfnd,ndh->bfh", context_layer, w) + b
Lysandre's avatar
Lysandre committed
254
255
        projected_context_layer_dropout = self.dropout(projected_context_layer)
        layernormed_context_layer = self.LayerNorm(input_ids + projected_context_layer_dropout)
256
        return (layernormed_context_layer, attention_probs) if self.output_attentions else (layernormed_context_layer,)
Lysandre's avatar
Lysandre committed
257
258


Lysandre's avatar
Lysandre committed
259
class AlbertLayer(nn.Module):
Lysandre's avatar
Lysandre committed
260
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
261
        super().__init__()
262

Lysandre's avatar
Lysandre committed
263
        self.config = config
Lysandre's avatar
Lysandre committed
264
        self.full_layer_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
Lysandre's avatar
Lysandre committed
265
        self.attention = AlbertAttention(config)
266
        self.ffn = nn.Linear(config.hidden_size, config.intermediate_size)
Lysandre's avatar
Lysandre committed
267
        self.ffn_output = nn.Linear(config.intermediate_size, config.hidden_size)
268
        self.activation = ACT2FN[config.hidden_act]
Lysandre's avatar
Lysandre committed
269
270

    def forward(self, hidden_states, attention_mask=None, head_mask=None):
Lysandre's avatar
Lysandre committed
271
        attention_output = self.attention(hidden_states, attention_mask, head_mask)
272
        ffn_output = self.ffn(attention_output[0])
273
        ffn_output = self.activation(ffn_output)
274
        ffn_output = self.ffn_output(ffn_output)
275
        hidden_states = self.full_layer_layer_norm(ffn_output + attention_output[0])
Lysandre's avatar
Lysandre committed
276

277
        return (hidden_states,) + attention_output[1:]  # add attentions if we output them
Lysandre's avatar
Lysandre committed
278
279


Lysandre's avatar
Lysandre committed
280
class AlbertLayerGroup(nn.Module):
Lysandre's avatar
Lysandre committed
281
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
282
        super().__init__()
283

284
285
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
Lysandre's avatar
Lysandre committed
286
287
288
        self.albert_layers = nn.ModuleList([AlbertLayer(config) for _ in range(config.inner_group_num)])

    def forward(self, hidden_states, attention_mask=None, head_mask=None):
289
290
291
        layer_hidden_states = ()
        layer_attentions = ()

Lysandre's avatar
Lysandre committed
292
293
        for layer_index, albert_layer in enumerate(self.albert_layers):
            layer_output = albert_layer(hidden_states, attention_mask, head_mask[layer_index])
294
295
296
297
298
            hidden_states = layer_output[0]

            if self.output_attentions:
                layer_attentions = layer_attentions + (layer_output[1],)

Lysandre's avatar
Lysandre committed
299
300
            if self.output_hidden_states:
                layer_hidden_states = layer_hidden_states + (hidden_states,)
Lysandre's avatar
Lysandre committed
301

302
303
304
305
306
307
        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (layer_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (layer_attentions,)
        return outputs  # last-layer hidden state, (layer hidden states), (layer attentions)
Lysandre's avatar
Lysandre committed
308

Lysandre's avatar
Lysandre committed
309

Lysandre's avatar
Lysandre committed
310
311
class AlbertTransformer(nn.Module):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
312
        super().__init__()
313

Lysandre's avatar
Lysandre committed
314
        self.config = config
Lysandre's avatar
Lysandre committed
315
316
317
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
        self.embedding_hidden_mapping_in = nn.Linear(config.embedding_size, config.hidden_size)
Lysandre's avatar
Lysandre committed
318
        self.albert_layer_groups = nn.ModuleList([AlbertLayerGroup(config) for _ in range(config.num_hidden_groups)])
Lysandre's avatar
Lysandre committed
319
320
321
322

    def forward(self, hidden_states, attention_mask=None, head_mask=None):
        hidden_states = self.embedding_hidden_mapping_in(hidden_states)

323
324
325
326
327
        all_attentions = ()

        if self.output_hidden_states:
            all_hidden_states = (hidden_states,)

328
329
        for i in range(self.config.num_hidden_layers):
            # Number of layers in a hidden group
Lysandre's avatar
Lysandre committed
330
            layers_per_group = int(self.config.num_hidden_layers / self.config.num_hidden_groups)
331
332
333
334

            # Index of the hidden group
            group_idx = int(i / (self.config.num_hidden_layers / self.config.num_hidden_groups))

335
336
337
338
339
            layer_group_output = self.albert_layer_groups[group_idx](
                hidden_states,
                attention_mask,
                head_mask[group_idx * layers_per_group : (group_idx + 1) * layers_per_group],
            )
340
341
342
            hidden_states = layer_group_output[0]

            if self.output_attentions:
Lysandre's avatar
Lysandre committed
343
                all_attentions = all_attentions + layer_group_output[-1]
344
345
346
347
348
349
350
351
352
353

            if self.output_hidden_states:
                all_hidden_states = all_hidden_states + (hidden_states,)

        outputs = (hidden_states,)
        if self.output_hidden_states:
            outputs = outputs + (all_hidden_states,)
        if self.output_attentions:
            outputs = outputs + (all_attentions,)
        return outputs  # last-layer hidden state, (all hidden states), (all attentions)
Lysandre's avatar
Lysandre committed
354

Lysandre's avatar
Lysandre committed
355

356
357
class AlbertPreTrainedModel(PreTrainedModel):
    """ An abstract class to handle weights initialization and
358
        a simple interface for downloading and loading pretrained models.
359
    """
360

361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
    config_class = AlbertConfig
    pretrained_model_archive_map = ALBERT_PRETRAINED_MODEL_ARCHIVE_MAP
    base_model_prefix = "albert"

    def _init_weights(self, module):
        """ Initialize the weights.
        """
        if isinstance(module, (nn.Linear, nn.Embedding)):
            # Slightly different from the TF version which uses truncated_normal for initialization
            # cf https://github.com/pytorch/pytorch/pull/5617
            module.weight.data.normal_(mean=0.0, std=self.config.initializer_range)
            if isinstance(module, (nn.Linear)) and module.bias is not None:
                module.bias.data.zero_()
        elif isinstance(module, nn.LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


379
ALBERT_START_DOCSTRING = r"""    
380
381
382
383
384
385
386
387
388
    This model is a PyTorch `torch.nn.Module`_ sub-class. Use it as a regular PyTorch Module and
    refer to the PyTorch documentation for all matter related to general usage and behavior.

    .. _`ALBERT: A Lite BERT for Self-supervised Learning of Language Representations`:
        https://arxiv.org/abs/1909.11942

    .. _`torch.nn.Module`:
        https://pytorch.org/docs/stable/nn.html#module

389
    Args:
390
        config (:class:`~transformers.AlbertConfig`): Model configuration class with all the parameters of the model.
391
392
393
394
395
            Initializing with a config file does not load the weights associated with the model, only the configuration.
            Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights.
"""

ALBERT_INPUTS_DOCSTRING = r"""
396
397
398
399
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`):
            Indices of input sequence tokens in the vocabulary. 
            
400
401
            Indices can be obtained using :class:`transformers.AlbertTokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
402
403
404
405
            :func:`transformers.PreTrainedTokenizer.encode_plus` for details.
            
            `What are input IDs? <../glossary.html#input-ids>`__
        attention_mask (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
406
407
408
            Mask to avoid performing attention on padding token indices.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
409
410
411
            
            `What are attention masks? <../glossary.html#attention-mask>`__
        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`): 
412
413
414
            Segment token indices to indicate first and second portions of the inputs.
            Indices are selected in ``[0, 1]``: ``0`` corresponds to a `sentence A` token, ``1``
            corresponds to a `sentence B` token
415
416
417
            
            `What are token type IDs? <../glossary.html#token-type-ids>`_
        position_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
418
419
            Indices of positions of each input sequence tokens in the position embeddings.
            Selected in the range ``[0, config.max_position_embeddings - 1]``.
420
421
422
            
            `What are position IDs? <../glossary.html#position-ids>`_
        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`):
423
424
            Mask to nullify selected heads of the self-attention modules.
            Mask values selected in ``[0, 1]``:
425
            :obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
426
427
428
429
        input_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
            than the model's internal embedding lookup matrix.
430
431
"""

432
433
434
435
436

@add_start_docstrings(
    "The bare ALBERT Model transformer outputting raw hidden-states without any specific head on top.",
    ALBERT_START_DOCSTRING,
)
437
438
439
440
441
442
443
class AlbertModel(AlbertPreTrainedModel):

    config_class = AlbertConfig
    pretrained_model_archive_map = ALBERT_PRETRAINED_MODEL_ARCHIVE_MAP
    load_tf_weights = load_tf_weights_in_albert
    base_model_prefix = "albert"

Lysandre's avatar
Lysandre committed
444
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
445
        super().__init__(config)
Lysandre's avatar
Lysandre committed
446
447
448
449
450
451
452

        self.config = config
        self.embeddings = AlbertEmbeddings(config)
        self.encoder = AlbertTransformer(config)
        self.pooler = nn.Linear(config.hidden_size, config.hidden_size)
        self.pooler_activation = nn.Tanh()

Lysandre's avatar
Lysandre committed
453
454
        self.init_weights()

LysandreJik's avatar
LysandreJik committed
455
456
457
458
459
460
    def get_input_embeddings(self):
        return self.embeddings.word_embeddings

    def set_input_embeddings(self, value):
        self.embeddings.word_embeddings = value

461
462
463
464
465
    def _resize_token_embeddings(self, new_num_tokens):
        old_embeddings = self.embeddings.word_embeddings
        new_embeddings = self._get_resized_embeddings(old_embeddings, new_num_tokens)
        self.embeddings.word_embeddings = new_embeddings
        return self.embeddings.word_embeddings
Lysandre's avatar
Lysandre committed
466

Lysandre's avatar
Lysandre committed
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
    def _prune_heads(self, heads_to_prune):
        """ Prunes heads of the model.
            heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
            ALBERT has a different architecture in that its layers are shared across groups, which then has inner groups.
            If an ALBERT model has 12 hidden layers and 2 hidden groups, with two inner groups, there
            is a total of 4 different layers.

            These layers are flattened: the indices [0,1] correspond to the two inner groups of the first hidden layer,
            while [2,3] correspond to the two inner groups of the second hidden layer.

            Any layer with in index other than [0,1,2,3] will result in an error.
            See base class PreTrainedModel for more information about head pruning
        """
        for layer, heads in heads_to_prune.items():
            group_idx = int(layer / self.config.inner_group_num)
            inner_group_idx = int(layer - group_idx * self.config.inner_group_num)
            self.encoder.albert_layer_groups[group_idx].albert_layers[inner_group_idx].attention.prune_heads(heads)

485
    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
486
487
488
489
490
491
492
493
494
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
    ):
495
496
        r"""
    Return:
Lysandre's avatar
Fixes  
Lysandre committed
497
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
        last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden-states at the output of the last layer of the model.
        pooler_output (:obj:`torch.FloatTensor`: of shape :obj:`(batch_size, hidden_size)`):
            Last layer hidden-state of the first token of the sequence (classification token)
            further processed by a Linear layer and a Tanh activation function. The Linear
            layer weights are trained from the next sentence prediction (classification)
            objective during pre-training.

            This output is usually *not* a good summary
            of the semantic content of the input, you're often better with averaging or pooling
            the sequence of hidden-states for the whole input sequence.
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Example::

        from transformers import AlbertModel, AlbertTokenizer
        import torch

        tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
        model = AlbertModel.from_pretrained('albert-base-v2')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids)
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple

        """
LysandreJik's avatar
LysandreJik committed
533
534
535
536
537
538
539
540
541
542
543

        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_shape = input_ids.size()
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

        device = input_ids.device if input_ids is not None else inputs_embeds.device
Lysandre's avatar
Lysandre committed
544

Lysandre's avatar
Lysandre committed
545
        if attention_mask is None:
LysandreJik's avatar
LysandreJik committed
546
            attention_mask = torch.ones(input_shape, device=device)
Lysandre's avatar
Lysandre committed
547
        if token_type_ids is None:
LysandreJik's avatar
LysandreJik committed
548
            token_type_ids = torch.zeros(input_shape, dtype=torch.long, device=device)
Lysandre's avatar
Lysandre committed
549
550

        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
551
        extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
Lysandre's avatar
Lysandre committed
552
553
554
555
556
557
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
                head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1)
            elif head_mask.dim() == 2:
558
559
560
561
562
563
                head_mask = (
                    head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
                )  # We can specify head_mask for each layer
            head_mask = head_mask.to(
                dtype=next(self.parameters()).dtype
            )  # switch to fload if need + fp16 compatibility
Lysandre's avatar
Lysandre committed
564
565
566
        else:
            head_mask = [None] * self.config.num_hidden_layers

567
568
569
570
        embedding_output = self.embeddings(
            input_ids, position_ids=position_ids, token_type_ids=token_type_ids, inputs_embeds=inputs_embeds
        )
        encoder_outputs = self.encoder(embedding_output, extended_attention_mask, head_mask=head_mask)
Lysandre's avatar
Lysandre committed
571

Lysandre's avatar
Lysandre committed
572
        sequence_output = encoder_outputs[0]
Lysandre's avatar
Lysandre committed
573

Lysandre's avatar
Lysandre committed
574
575
        pooled_output = self.pooler_activation(self.pooler(sequence_output[:, 0]))

576
577
578
        outputs = (sequence_output, pooled_output) + encoder_outputs[
            1:
        ]  # add hidden_states and attentions if they are here
Lysandre's avatar
Lysandre committed
579
580
        return outputs

581

Lysandre's avatar
Lysandre committed
582
583
class AlbertMLMHead(nn.Module):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
584
        super().__init__()
Lysandre's avatar
Lysandre committed
585
586
587
588
589
590
591

        self.LayerNorm = nn.LayerNorm(config.embedding_size)
        self.bias = nn.Parameter(torch.zeros(config.vocab_size))
        self.dense = nn.Linear(config.hidden_size, config.embedding_size)
        self.decoder = nn.Linear(config.embedding_size, config.vocab_size)
        self.activation = ACT2FN[config.hidden_act]

592
593
594
        # Need a link between the two variables so that the bias is correctly resized with `resize_token_embeddings`
        self.decoder.bias = self.bias

Lysandre's avatar
Lysandre committed
595
596
597
598
599
600
601
602
603
604
    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.activation(hidden_states)
        hidden_states = self.LayerNorm(hidden_states)
        hidden_states = self.decoder(hidden_states)

        prediction_scores = hidden_states + self.bias

        return prediction_scores

Lysandre's avatar
Lysandre committed
605

606
@add_start_docstrings(
607
    "Albert Model with a `language modeling` head on top.", ALBERT_START_DOCSTRING,
608
)
609
class AlbertForMaskedLM(AlbertPreTrainedModel):
Lysandre's avatar
Lysandre committed
610
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
611
        super().__init__(config)
Lysandre's avatar
Lysandre committed
612

Lysandre's avatar
Lysandre committed
613
        self.albert = AlbertModel(config)
Lysandre's avatar
Lysandre committed
614
        self.predictions = AlbertMLMHead(config)
Lysandre's avatar
Lysandre committed
615

Lysandre's avatar
Lysandre committed
616
617
618
        self.init_weights()
        self.tie_weights()

Lysandre's avatar
Lysandre committed
619
    def tie_weights(self):
620
        self._tie_or_clone_weights(self.predictions.decoder, self.albert.embeddings.word_embeddings)
Lysandre's avatar
Lysandre committed
621

LysandreJik's avatar
LysandreJik committed
622
623
624
    def get_output_embeddings(self):
        return self.predictions.decoder

625
    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
626
627
628
629
630
631
632
633
634
635
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        masked_lm_labels=None,
    ):
636
637
638
639
640
641
642
643
644
645
646
        r"""
        masked_lm_labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the masked language modeling loss.
            Indices should be in ``[-100, 0, ..., config.vocab_size]`` (see ``input_ids`` docstring)
            Tokens with indices set to ``-100`` are ignored (masked), the loss is only computed for the tokens with
            labels in ``[0, ..., config.vocab_size]``

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
        loss (`optional`, returned when ``masked_lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Masked language modeling loss.
Lysandre's avatar
Lysandre committed
647
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
648
649
650
651
652
653
654
655
656
657
658
659
660
661
662
663
664
665
666
667
668
669
670
671
672
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

    Example::

        from transformers import AlbertTokenizer, AlbertForMaskedLM
        import torch

        tokenizer = BertTokenizer.from_pretrained('albert-base-v2')
        model = BertForMaskedLM.from_pretrained('albert-base-v2')
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, masked_lm_labels=input_ids)
        loss, prediction_scores = outputs[:2]

        """
LysandreJik's avatar
LysandreJik committed
673
674
675
676
677
678
        outputs = self.albert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
679
            inputs_embeds=inputs_embeds,
LysandreJik's avatar
LysandreJik committed
680
        )
681
        sequence_outputs = outputs[0]
Lysandre's avatar
Lysandre committed
682
683

        prediction_scores = self.predictions(sequence_outputs)
Lysandre's avatar
Lysandre committed
684

685
686
        outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention if they are here
        if masked_lm_labels is not None:
LysandreJik's avatar
LysandreJik committed
687
            loss_fct = CrossEntropyLoss()
688
689
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
            outputs = (masked_lm_loss,) + outputs
Lysandre's avatar
Lysandre committed
690

691
        return outputs
Lysandre's avatar
Lysandre committed
692
693


694
695
@add_start_docstrings(
    """Albert Model transformer with a sequence classification/regression head on top (a linear layer on top of
Lysandre's avatar
Lysandre committed
696
    the pooled output) e.g. for GLUE tasks. """,
697
698
    ALBERT_START_DOCSTRING,
)
Lysandre's avatar
Lysandre committed
699
700
class AlbertForSequenceClassification(AlbertPreTrainedModel):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
701
        super().__init__(config)
Lysandre's avatar
Lysandre committed
702
703
704
705
706
707
708
709
        self.num_labels = config.num_labels

        self.albert = AlbertModel(config)
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)

        self.init_weights()

710
    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
711
712
713
714
715
716
717
718
719
720
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
    ):
721
722
723
724
725
726
727
728
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for computing the sequence classification/regression loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.
            If ``config.num_labels == 1`` a regression loss is computed (Mean-Square loss),
            If ``config.num_labels > 1`` a classification loss is computed (Cross-Entropy).

    Returns:
Lysandre's avatar
Fixes  
Lysandre committed
729
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
748
749
750
751
752
753
754
755
756
757
758
        loss: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Classification (or regression if config.num_labels==1) loss.
        logits ``torch.FloatTensor`` of shape ``(batch_size, config.num_labels)``
            Classification (or regression if config.num_labels==1) scores (before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

        Examples::

            from transformers import AlbertTokenizer, AlbertForSequenceClassification
            import torch

            tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
            model = AlbertForSequenceClassification.from_pretrained('albert-base-v2')
            input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
            labels = torch.tensor([1]).unsqueeze(0)  # Batch size 1
            outputs = model(input_ids, labels=labels)
            loss, logits = outputs[:2]

        """
Lysandre's avatar
Lysandre committed
759

LysandreJik's avatar
LysandreJik committed
760
761
762
763
764
765
        outputs = self.albert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
766
            inputs_embeds=inputs_embeds,
LysandreJik's avatar
LysandreJik committed
767
        )
Lysandre's avatar
Lysandre committed
768
769
770
771
772
773
774
775
776
777
778
779
780
781
782
783
784
785

        pooled_output = outputs[1]

        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here

        if labels is not None:
            if self.num_labels == 1:
                #  We are doing regression
                loss_fct = MSELoss()
                loss = loss_fct(logits.view(-1), labels.view(-1))
            else:
                loss_fct = CrossEntropyLoss()
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            outputs = (loss,) + outputs

Lysandre's avatar
Lysandre committed
786
787
788
        return outputs  # (loss), logits, (hidden_states), (attentions)


789
790
@add_start_docstrings(
    """Albert Model with a span classification head on top for extractive question-answering tasks like SQuAD (a linear layers on top of
Lysandre's avatar
Lysandre committed
791
    the hidden-states output to compute `span start logits` and `span end logits`). """,
792
793
    ALBERT_START_DOCSTRING,
)
Lysandre's avatar
Lysandre committed
794
795
class AlbertForQuestionAnswering(AlbertPreTrainedModel):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
796
        super().__init__(config)
Lysandre's avatar
Lysandre committed
797
798
799
800
801
802
803
        self.num_labels = config.num_labels

        self.albert = AlbertModel(config)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)

        self.init_weights()

804
    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
805
806
807
808
809
810
811
812
813
814
815
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        start_positions=None,
        end_positions=None,
    ):
816
817
818
819
820
821
822
823
824
825
826
        r"""
        start_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for position (index) of the start of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.
        end_positions (:obj:`torch.LongTensor` of shape :obj:`(batch_size,)`, `optional`, defaults to :obj:`None`):
            Labels for position (index) of the end of the labelled span for computing the token classification loss.
            Positions are clamped to the length of the sequence (`sequence_length`).
            Position outside of the sequence are not taken into account for computing the loss.

    Returns:
Lysandre's avatar
Fixes  
Lysandre committed
827
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
828
829
830
831
832
833
834
835
836
837
838
839
840
841
842
843
844
845
846
847
848
849
850
851
852
853
854
855
856
857
        loss: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Total span extraction loss is the sum of a Cross-Entropy for the start and end positions.
        start_scores ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-start scores (before SoftMax).
        end_scores: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length,)``
            Span-end scores (before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
        attentions (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_attentions=True``):
            Tuple of :obj:`torch.FloatTensor` (one for each layer) of shape
            :obj:`(batch_size, num_heads, sequence_length, sequence_length)`.

            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention
            heads.

        Examples::

            # The checkpoint albert-base-v2 is not fine-tuned for question answering. Please see the
            # examples/run_squad.py example to see how to fine-tune a model to a question answering task.

            tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
            model = AlbertForQuestionAnswering.from_pretrained('albert-base-v2')
            question, text = "Who was Jim Henson?", "Jim Henson was a nice puppet"
            input_dict = tokenizer.encode_plus(question, text, return_tensors='pt')
            start_scores, end_scores = model(**input_dict)

        """
LysandreJik's avatar
LysandreJik committed
858
859
860
861
862
863
864

        outputs = self.albert(
            input_ids=input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
865
            inputs_embeds=inputs_embeds,
LysandreJik's avatar
LysandreJik committed
866
        )
Lysandre's avatar
Lysandre committed
867
868
869
870
871
872
873
874
875
876
877
878
879
880
881
882
883
884
885
886
887
888
889
890
891
892
893

        sequence_output = outputs[0]

        logits = self.qa_outputs(sequence_output)
        start_logits, end_logits = logits.split(1, dim=-1)
        start_logits = start_logits.squeeze(-1)
        end_logits = end_logits.squeeze(-1)

        outputs = (start_logits, end_logits,) + outputs[2:]
        if start_positions is not None and end_positions is not None:
            # If we are on multi-GPU, split add a dimension
            if len(start_positions.size()) > 1:
                start_positions = start_positions.squeeze(-1)
            if len(end_positions.size()) > 1:
                end_positions = end_positions.squeeze(-1)
            # sometimes the start/end positions are outside our model inputs, we ignore these terms
            ignored_index = start_logits.size(1)
            start_positions.clamp_(0, ignored_index)
            end_positions.clamp_(0, ignored_index)

            loss_fct = CrossEntropyLoss(ignore_index=ignored_index)
            start_loss = loss_fct(start_logits, start_positions)
            end_loss = loss_fct(end_logits, end_positions)
            total_loss = (start_loss + end_loss) / 2
            outputs = (total_loss,) + outputs

        return outputs  # (loss), start_logits, end_logits, (hidden_states), (attentions)