modeling_albert.py 43.6 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

25
from .configuration_albert import AlbertConfig
26
from .file_utils import add_start_docstrings, add_start_docstrings_to_callable
27
28
from .modeling_bert import ACT2FN, BertEmbeddings, BertSelfAttention, prune_linear_layer
from .modeling_utils import PreTrainedModel
29

Aymeric Augustin's avatar
Aymeric Augustin committed
30

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

Lysandre's avatar
Lysandre committed
33
34

ALBERT_PRETRAINED_MODEL_ARCHIVE_MAP = {
Julien Chaumond's avatar
Julien Chaumond committed
35
36
37
38
39
40
41
42
    "albert-base-v1": "https://cdn.huggingface.co/albert-base-v1-pytorch_model.bin",
    "albert-large-v1": "https://cdn.huggingface.co/albert-large-v1-pytorch_model.bin",
    "albert-xlarge-v1": "https://cdn.huggingface.co/albert-xlarge-v1-pytorch_model.bin",
    "albert-xxlarge-v1": "https://cdn.huggingface.co/albert-xxlarge-v1-pytorch_model.bin",
    "albert-base-v2": "https://cdn.huggingface.co/albert-base-v2-pytorch_model.bin",
    "albert-large-v2": "https://cdn.huggingface.co/albert-large-v2-pytorch_model.bin",
    "albert-xlarge-v2": "https://cdn.huggingface.co/albert-xlarge-v2-pytorch_model.bin",
    "albert-xxlarge-v2": "https://cdn.huggingface.co/albert-xxlarge-v2-pytorch_model.bin",
Lysandre's avatar
Lysandre committed
43
44
45
}


Lysandre's avatar
Lysandre committed
46
47
48
49
50
51
52
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:
53
54
55
56
        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
57
58
59
60
61
62
63
64
65
66
67
68
69
        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
70
71
    for name, array in zip(names, arrays):
        print(name)
72

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

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

        # Renaming and simplifying
Lysandre's avatar
Lysandre committed
80
        name = name.replace("ffn_1", "ffn")
Lysandre's avatar
Lysandre committed
81
        name = name.replace("bert/", "albert/")
82
        name = name.replace("attention_1", "attention")
Lysandre's avatar
Lysandre committed
83
        name = name.replace("transform/", "")
84
85
        name = name.replace("LayerNorm_1", "full_layer_layer_norm")
        name = name.replace("LayerNorm", "attention/LayerNorm")
Lysandre's avatar
Lysandre committed
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
        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
104
105
106
        name = name.replace("embeddings/attention", "embeddings")
        name = name.replace("inner_group_", "albert_layers/")
        name = name.replace("group_", "albert_layer_groups/")
107
108
109
110
111

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

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

116
        name = name.split("/")
117
118

        # Ignore the gradients applied by the LAMB/ADAM optimizers.
119
120
121
122
123
124
125
        if (
            "adam_m" in name
            or "adam_v" in name
            or "AdamWeightDecayOptimizer" in name
            or "AdamWeightDecayOptimizer_1" in name
            or "global_step" in name
        ):
126
127
128
            logger.info("Skipping {}".format("/".join(name)))
            continue

Lysandre's avatar
Lysandre committed
129
130
        pointer = model
        for m_name in name:
131
            if re.fullmatch(r"[A-Za-z]+_\d+", m_name):
132
                scope_names = re.split(r"_(\d+)", m_name)
Lysandre's avatar
Lysandre committed
133
            else:
134
                scope_names = [m_name]
Lysandre's avatar
Lysandre committed
135

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

154
155
156
        if m_name[-11:] == "_embeddings":
            pointer = getattr(pointer, "weight")
        elif m_name == "kernel":
Lysandre's avatar
Lysandre committed
157
158
159
160
161
162
            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
163
        print("Initialize PyTorch weight {} from {}".format(name, original_name))
Lysandre's avatar
Lysandre committed
164
165
166
167
168
        pointer.data = torch.from_numpy(array)

    return model


Lysandre's avatar
Lysandre committed
169
class AlbertEmbeddings(BertEmbeddings):
Lysandre's avatar
Lysandre committed
170
171
172
    """
    Construct the embeddings from word, position and token_type embeddings.
    """
173

Lysandre's avatar
Lysandre committed
174
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
175
        super().__init__(config)
Lysandre's avatar
Lysandre committed
176
177
178
179

        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
180
        self.LayerNorm = torch.nn.LayerNorm(config.embedding_size, eps=config.layer_norm_eps)
Lysandre's avatar
Lysandre committed
181
182


Lysandre's avatar
Lysandre committed
183
class AlbertAttention(BertSelfAttention):
Lysandre's avatar
Lysandre committed
184
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
185
        super().__init__(config)
Lysandre's avatar
Lysandre committed
186

187
        self.output_attentions = config.output_attentions
Lysandre's avatar
Lysandre committed
188
        self.num_attention_heads = config.num_attention_heads
189
        self.hidden_size = config.hidden_size
Lysandre's avatar
Lysandre committed
190
        self.attention_head_size = config.hidden_size // config.num_attention_heads
Lysandre's avatar
Lysandre committed
191
192
193
        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
194
195
        self.pruned_heads = set()

Lysandre's avatar
Lysandre committed
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
    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
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
    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
249
250

        # Should find a better way to do this
251
252
253
254
255
        w = (
            self.dense.weight.t()
            .view(self.num_attention_heads, self.attention_head_size, self.hidden_size)
            .to(context_layer.dtype)
        )
256
        b = self.dense.bias.to(context_layer.dtype)
Lysandre's avatar
Lysandre committed
257
258

        projected_context_layer = torch.einsum("bfnd,ndh->bfh", context_layer, w) + b
Lysandre's avatar
Lysandre committed
259
260
        projected_context_layer_dropout = self.dropout(projected_context_layer)
        layernormed_context_layer = self.LayerNorm(input_ids + projected_context_layer_dropout)
261
        return (layernormed_context_layer, attention_probs) if self.output_attentions else (layernormed_context_layer,)
Lysandre's avatar
Lysandre committed
262
263


Lysandre's avatar
Lysandre committed
264
class AlbertLayer(nn.Module):
Lysandre's avatar
Lysandre committed
265
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
266
        super().__init__()
267

Lysandre's avatar
Lysandre committed
268
        self.config = config
Lysandre's avatar
Lysandre committed
269
        self.full_layer_layer_norm = nn.LayerNorm(config.hidden_size, eps=config.layer_norm_eps)
Lysandre's avatar
Lysandre committed
270
        self.attention = AlbertAttention(config)
271
        self.ffn = nn.Linear(config.hidden_size, config.intermediate_size)
Lysandre's avatar
Lysandre committed
272
        self.ffn_output = nn.Linear(config.intermediate_size, config.hidden_size)
273
        self.activation = ACT2FN[config.hidden_act]
Lysandre's avatar
Lysandre committed
274
275

    def forward(self, hidden_states, attention_mask=None, head_mask=None):
Lysandre's avatar
Lysandre committed
276
        attention_output = self.attention(hidden_states, attention_mask, head_mask)
277
        ffn_output = self.ffn(attention_output[0])
278
        ffn_output = self.activation(ffn_output)
279
        ffn_output = self.ffn_output(ffn_output)
280
        hidden_states = self.full_layer_layer_norm(ffn_output + attention_output[0])
Lysandre's avatar
Lysandre committed
281

282
        return (hidden_states,) + attention_output[1:]  # add attentions if we output them
Lysandre's avatar
Lysandre committed
283
284


Lysandre's avatar
Lysandre committed
285
class AlbertLayerGroup(nn.Module):
Lysandre's avatar
Lysandre committed
286
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
287
        super().__init__()
288

289
290
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
Lysandre's avatar
Lysandre committed
291
292
293
        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):
294
295
296
        layer_hidden_states = ()
        layer_attentions = ()

Lysandre's avatar
Lysandre committed
297
298
        for layer_index, albert_layer in enumerate(self.albert_layers):
            layer_output = albert_layer(hidden_states, attention_mask, head_mask[layer_index])
299
300
301
302
303
            hidden_states = layer_output[0]

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

Lysandre's avatar
Lysandre committed
304
305
            if self.output_hidden_states:
                layer_hidden_states = layer_hidden_states + (hidden_states,)
Lysandre's avatar
Lysandre committed
306

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

Lysandre's avatar
Lysandre committed
314

Lysandre's avatar
Lysandre committed
315
316
class AlbertTransformer(nn.Module):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
317
        super().__init__()
318

Lysandre's avatar
Lysandre committed
319
        self.config = config
Lysandre's avatar
Lysandre committed
320
321
322
        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
323
        self.albert_layer_groups = nn.ModuleList([AlbertLayerGroup(config) for _ in range(config.num_hidden_groups)])
Lysandre's avatar
Lysandre committed
324
325
326
327

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

328
329
330
331
332
        all_attentions = ()

        if self.output_hidden_states:
            all_hidden_states = (hidden_states,)

333
334
        for i in range(self.config.num_hidden_layers):
            # Number of layers in a hidden group
Lysandre's avatar
Lysandre committed
335
            layers_per_group = int(self.config.num_hidden_layers / self.config.num_hidden_groups)
336
337
338
339

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

340
341
342
343
344
            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],
            )
345
346
347
            hidden_states = layer_group_output[0]

            if self.output_attentions:
Lysandre's avatar
Lysandre committed
348
                all_attentions = all_attentions + layer_group_output[-1]
349
350
351
352
353
354
355
356
357
358

            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
359

Lysandre's avatar
Lysandre committed
360

361
362
class AlbertPreTrainedModel(PreTrainedModel):
    """ An abstract class to handle weights initialization and
363
        a simple interface for downloading and loading pretrained models.
364
    """
365

366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
    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)


Lysandre's avatar
Lysandre committed
384
ALBERT_START_DOCSTRING = r"""
385

Lysandre's avatar
Lysandre committed
386
387
388
    This model is a PyTorch `torch.nn.Module <https://pytorch.org/docs/stable/nn.html#torch.nn.Module>`_ sub-class.
    Use it as a regular PyTorch Module and refer to the PyTorch documentation for all matter related to general
    usage and behavior.
389

390
    Args:
391
        config (:class:`~transformers.AlbertConfig`): Model configuration class with all the parameters of the model.
392
393
394
395
396
            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"""
397
398
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`):
Lysandre's avatar
Lysandre committed
399
400
            Indices of input sequence tokens in the vocabulary.

401
402
            Indices can be obtained using :class:`transformers.AlbertTokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
403
            :func:`transformers.PreTrainedTokenizer.encode_plus` for details.
Lysandre's avatar
Lysandre committed
404

405
406
            `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`):
407
408
409
            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.
Lysandre's avatar
Lysandre committed
410

411
            `What are attention masks? <../glossary.html#attention-mask>`__
Lysandre's avatar
Lysandre committed
412
        token_type_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
413
414
415
            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
Lysandre's avatar
Lysandre committed
416

417
418
            `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`):
419
420
            Indices of positions of each input sequence tokens in the position embeddings.
            Selected in the range ``[0, config.max_position_embeddings - 1]``.
Lysandre's avatar
Lysandre committed
421

422
423
            `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`):
424
425
            Mask to nullify selected heads of the self-attention modules.
            Mask values selected in ``[0, 1]``:
426
            :obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
427
428
429
430
        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.
431
432
"""

433
434
435
436
437

@add_start_docstrings(
    "The bare ALBERT Model transformer outputting raw hidden-states without any specific head on top.",
    ALBERT_START_DOCSTRING,
)
438
439
440
441
442
443
444
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
445
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
446
        super().__init__(config)
Lysandre's avatar
Lysandre committed
447
448
449
450
451
452
453

        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
454
455
        self.init_weights()

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

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

462
463
464
465
466
    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
467

Lysandre's avatar
Lysandre committed
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
    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)

486
    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
487
488
489
490
491
492
493
494
495
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
    ):
496
497
        r"""
    Return:
Lysandre's avatar
Fixes  
Lysandre committed
498
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
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
533
        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
534
535
536
537
538
539
540
541
542
543
544

        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
545

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

        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)
552
        extended_attention_mask = extended_attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
Lysandre's avatar
Lysandre committed
553
        extended_attention_mask = (1.0 - extended_attention_mask) * -10000.0
554
        head_mask = self.get_head_mask(head_mask, self.config.num_hidden_layers)
Lysandre's avatar
Lysandre committed
555

556
557
558
559
        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
560

Lysandre's avatar
Lysandre committed
561
        sequence_output = encoder_outputs[0]
Lysandre's avatar
Lysandre committed
562

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

565
566
567
        outputs = (sequence_output, pooled_output) + encoder_outputs[
            1:
        ]  # add hidden_states and attentions if they are here
Lysandre's avatar
Lysandre committed
568
569
        return outputs

570

Lysandre's avatar
Lysandre committed
571
572
class AlbertMLMHead(nn.Module):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
573
        super().__init__()
Lysandre's avatar
Lysandre committed
574
575
576
577
578
579
580

        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]

581
582
583
        # 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
584
585
586
587
588
589
    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)

Martin Malmsten's avatar
Martin Malmsten committed
590
        prediction_scores = hidden_states
Lysandre's avatar
Lysandre committed
591
592
593

        return prediction_scores

Lysandre's avatar
Lysandre committed
594

595
@add_start_docstrings(
596
    "Albert Model with a `language modeling` head on top.", ALBERT_START_DOCSTRING,
597
)
598
class AlbertForMaskedLM(AlbertPreTrainedModel):
Lysandre's avatar
Lysandre committed
599
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
600
        super().__init__(config)
Lysandre's avatar
Lysandre committed
601

Lysandre's avatar
Lysandre committed
602
        self.albert = AlbertModel(config)
Lysandre's avatar
Lysandre committed
603
        self.predictions = AlbertMLMHead(config)
Lysandre's avatar
Lysandre committed
604

Lysandre's avatar
Lysandre committed
605
606
607
        self.init_weights()
        self.tie_weights()

Lysandre's avatar
Lysandre committed
608
    def tie_weights(self):
609
        self._tie_or_clone_weights(self.predictions.decoder, self.albert.embeddings.word_embeddings)
Lysandre's avatar
Lysandre committed
610

LysandreJik's avatar
LysandreJik committed
611
612
613
    def get_output_embeddings(self):
        return self.predictions.decoder

614
    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
615
616
617
618
619
620
621
622
623
624
    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,
    ):
625
626
627
        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.
Lysandre's avatar
Lysandre committed
628
629
            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
630
631
632
633
634
635
            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
636
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`)
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
            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

Lysandre's avatar
Lysandre committed
655
656
        tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
        model = AlbertForMaskedLM.from_pretrained('albert-base-v2')
657
658
659
660
661
        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
662
663
664
665
666
667
        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,
668
            inputs_embeds=inputs_embeds,
LysandreJik's avatar
LysandreJik committed
669
        )
670
        sequence_outputs = outputs[0]
Lysandre's avatar
Lysandre committed
671
672

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

674
675
        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
676
            loss_fct = CrossEntropyLoss()
677
678
            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
679

680
        return outputs
Lysandre's avatar
Lysandre committed
681
682


683
684
@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
685
    the pooled output) e.g. for GLUE tasks. """,
686
687
    ALBERT_START_DOCSTRING,
)
Lysandre's avatar
Lysandre committed
688
689
class AlbertForSequenceClassification(AlbertPreTrainedModel):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
690
        super().__init__(config)
Lysandre's avatar
Lysandre committed
691
692
693
        self.num_labels = config.num_labels

        self.albert = AlbertModel(config)
694
        self.dropout = nn.Dropout(config.classifier_dropout_prob)
Lysandre's avatar
Lysandre committed
695
696
697
698
        self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)

        self.init_weights()

699
    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
700
701
702
703
704
705
706
707
708
709
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
    ):
710
711
712
713
714
715
716
717
        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
718
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
719
720
721
722
723
724
725
726
727
728
729
730
731
732
733
734
735
736
737
738
739
740
741
742
743
744
745
746
747
        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
748

LysandreJik's avatar
LysandreJik committed
749
750
751
752
753
754
        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,
755
            inputs_embeds=inputs_embeds,
LysandreJik's avatar
LysandreJik committed
756
        )
Lysandre's avatar
Lysandre committed
757
758
759
760
761
762
763
764
765
766
767
768
769
770
771
772
773
774

        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
775
776
777
        return outputs  # (loss), logits, (hidden_states), (attentions)


778
779
780
781
782
783
784
785
786
787
788
789
790
791
792
793
794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
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
858
859
860
861
862
863
864
865
866
867
868
869
870
871
872
873
874
@add_start_docstrings(
    """Albert Model with a token classification head on top (a linear layer on top of
    the hidden-states output) e.g. for Named-Entity-Recognition (NER) tasks. """,
    ALBERT_START_DOCSTRING,
)
class AlbertForTokenClassification(AlbertPreTrainedModel):
    def __init__(self, config):
        super().__init__(config)
        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()

    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
    def forward(
        self,
        input_ids=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
    ):
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for computing the token classification loss.
            Indices should be in ``[0, ..., config.num_labels - 1]``.

    Returns:
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
        loss (:obj:`torch.FloatTensor` of shape :obj:`(1,)`, `optional`, returned when ``labels`` is provided) :
            Classification loss.
        scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.num_labels)`)
            Classification scores (before SoftMax).
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

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

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

    Examples::

        from transformers import AlbertTokenizer, AlbertForTokenClassification
        import torch

        tokenizer = AlbertTokenizer.from_pretrained('albert-base-v2')
        model = AlbertForTokenClassification.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
        labels = torch.tensor([1] * input_ids.size(1)).unsqueeze(0)  # Batch size 1
        outputs = model(input_ids, labels=labels)

        loss, scores = outputs[:2]

        """

        outputs = self.albert(
            input_ids,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )

        sequence_output = outputs[0]

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

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

        if labels is not None:
            loss_fct = CrossEntropyLoss()
            # Only keep active parts of the loss
            if attention_mask is not None:
                active_loss = attention_mask.view(-1) == 1
                active_logits = logits.view(-1, self.num_labels)[active_loss]
                active_labels = labels.view(-1)[active_loss]
                loss = loss_fct(active_logits, active_labels)
            else:
                loss = loss_fct(logits.view(-1, self.num_labels), labels.view(-1))
            outputs = (loss,) + outputs

        return outputs  # (loss), logits, (hidden_states), (attentions)


875
876
@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
877
    the hidden-states output to compute `span start logits` and `span end logits`). """,
878
879
    ALBERT_START_DOCSTRING,
)
Lysandre's avatar
Lysandre committed
880
881
class AlbertForQuestionAnswering(AlbertPreTrainedModel):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
882
        super().__init__(config)
Lysandre's avatar
Lysandre committed
883
884
885
886
887
888
889
        self.num_labels = config.num_labels

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

        self.init_weights()

890
    @add_start_docstrings_to_callable(ALBERT_INPUTS_DOCSTRING)
891
892
893
894
895
896
897
898
899
900
901
    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,
    ):
902
903
904
905
906
907
908
909
910
911
912
        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
913
        :obj:`tuple(torch.FloatTensor)` comprising various elements depending on the configuration (:class:`~transformers.AlbertConfig`) and inputs:
914
915
916
917
918
919
920
921
922
923
924
925
926
927
928
929
930
931
        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.

Lysandre's avatar
Lysandre committed
932
    Examples::
933

Lysandre's avatar
Lysandre committed
934
        # The checkpoint albert-base-v2 is not fine-tuned for question answering. Please see the
935
        # examples/question-answering/run_squad.py example to see how to fine-tune a model to a question answering task.
936

Lysandre's avatar
Lysandre committed
937
938
939
940
941
942
943
944
        from transformers import AlbertTokenizer, AlbertForQuestionAnswering
        import torch

        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)
945
946

        """
LysandreJik's avatar
LysandreJik committed
947
948
949
950
951
952
953

        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,
954
            inputs_embeds=inputs_embeds,
LysandreJik's avatar
LysandreJik committed
955
        )
Lysandre's avatar
Lysandre committed
956
957
958
959
960
961
962
963
964
965
966
967
968
969
970
971
972
973
974
975
976
977
978
979
980
981
982

        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)