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

thomwolf's avatar
thomwolf committed
18
from __future__ import absolute_import, division, print_function, unicode_literals
thomwolf's avatar
thomwolf committed
19
20
21

import json
import logging
thomwolf's avatar
thomwolf committed
22
23
24
25
import math
import os
import sys
from io import open
thomwolf's avatar
thomwolf committed
26
27
28

import torch
from torch import nn
29
from torch.nn import CrossEntropyLoss, MSELoss
thomwolf's avatar
thomwolf committed
30

31
from .modeling_utils import WEIGHTS_NAME, CONFIG_NAME, PretrainedConfig, PreTrainedModel, prune_linear_layer
thomwolf's avatar
thomwolf committed
32
33
34

logger = logging.getLogger(__name__)

35
BERT_PRETRAINED_MODEL_ARCHIVE_MAP = {
36
37
38
39
40
41
42
43
44
45
    'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-pytorch_model.bin",
    'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-pytorch_model.bin",
    'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-pytorch_model.bin",
    'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-pytorch_model.bin",
    'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-pytorch_model.bin",
    'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-pytorch_model.bin",
    'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-pytorch_model.bin",
    'bert-base-german-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-pytorch_model.bin",
    'bert-large-uncased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-pytorch_model.bin",
    'bert-large-cased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-pytorch_model.bin",
thomwolf's avatar
thomwolf committed
46
47
    'bert-large-uncased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-pytorch_model.bin",
    'bert-large-cased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-pytorch_model.bin",
thomwolf's avatar
thomwolf committed
48
    'bert-base-cased-finetuned-mrpc': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-pytorch_model.bin",
49
}
50

51
BERT_PRETRAINED_CONFIG_ARCHIVE_MAP = {
52
53
54
55
56
57
58
59
60
61
    'bert-base-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-uncased-config.json",
    'bert-large-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-config.json",
    'bert-base-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-config.json",
    'bert-large-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-config.json",
    'bert-base-multilingual-uncased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-uncased-config.json",
    'bert-base-multilingual-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-multilingual-cased-config.json",
    'bert-base-chinese': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-chinese-config.json",
    'bert-base-german-cased': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-german-cased-config.json",
    'bert-large-uncased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-config.json",
    'bert-large-cased-whole-word-masking': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-config.json",
thomwolf's avatar
thomwolf committed
62
63
64
    'bert-large-uncased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-uncased-whole-word-masking-finetuned-squad-config.json",
    'bert-large-cased-whole-word-masking-finetuned-squad': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-large-cased-whole-word-masking-finetuned-squad-config.json",
    'bert-base-cased-finetuned-mrpc': "https://s3.amazonaws.com/models.huggingface.co/bert/bert-base-cased-finetuned-mrpc-config.json",
thomwolf's avatar
thomwolf committed
65
66
}

thomwolf's avatar
thomwolf committed
67

68
def load_tf_weights_in_bert(model, config, tf_checkpoint_path):
69
70
    """ Load tf checkpoints in a pytorch model
    """
71
72
73
74
    try:
        import re
        import numpy as np
        import tensorflow as tf
thomwolf's avatar
thomwolf committed
75
    except ImportError:
76
77
78
        print("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions.")
        raise
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
    tf_path = os.path.abspath(tf_checkpoint_path)
    print("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:
        print("Loading TF weight {} with shape {}".format(name, shape))
        array = tf.train.load_variable(tf_path, name)
        names.append(name)
        arrays.append(array)

    for name, array in zip(names, arrays):
        name = name.split('/')
        # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
        # which are not required for using pretrained model
95
        if any(n in ["adam_v", "adam_m", "global_step"] for n in name):
96
97
98
99
100
101
102
103
104
105
106
107
108
109
            print("Skipping {}".format("/".join(name)))
            continue
        pointer = model
        for m_name in name:
            if re.fullmatch(r'[A-Za-z]+_\d+', m_name):
                l = re.split(r'_(\d+)', m_name)
            else:
                l = [m_name]
            if l[0] == 'kernel' or l[0] == 'gamma':
                pointer = getattr(pointer, 'weight')
            elif l[0] == 'output_bias' or l[0] == 'beta':
                pointer = getattr(pointer, 'bias')
            elif l[0] == 'output_weights':
                pointer = getattr(pointer, 'weight')
thomwolf's avatar
thomwolf committed
110
111
            elif l[0] == 'squad':
                pointer = getattr(pointer, 'classifier')
112
            else:
113
114
115
116
117
                try:
                    pointer = getattr(pointer, l[0])
                except AttributeError:
                    print("Skipping {}".format("/".join(name)))
                    continue
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
            if len(l) >= 2:
                num = int(l[1])
                pointer = pointer[num]
        if m_name[-11:] == '_embeddings':
            pointer = getattr(pointer, 'weight')
        elif m_name == 'kernel':
            array = np.transpose(array)
        try:
            assert pointer.shape == array.shape
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
        print("Initialize PyTorch weight {}".format(name))
        pointer.data = torch.from_numpy(array)
    return model


thomwolf's avatar
thomwolf committed
135
136
137
138
def gelu(x):
    """Implementation of the gelu activation function.
        For information: OpenAI GPT's gelu is slightly different (and gives slightly different results):
        0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))
139
        Also see https://arxiv.org/abs/1606.08415
thomwolf's avatar
thomwolf committed
140
141
142
143
144
145
146
147
148
149
150
    """
    return x * 0.5 * (1.0 + torch.erf(x / math.sqrt(2.0)))


def swish(x):
    return x * torch.sigmoid(x)


ACT2FN = {"gelu": gelu, "relu": torch.nn.functional.relu, "swish": swish}


151
class BertConfig(PretrainedConfig):
152
    r"""
153
        :class:`~pytorch_transformers.BertConfig` is the configuration class to store the configuration of a
154
        `BertModel`.
155

156
        Arguments:
thomwolf's avatar
thomwolf committed
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`.
            hidden_size: Size of the encoder layers and the pooler layer.
            num_hidden_layers: Number of hidden layers in the Transformer encoder.
            num_attention_heads: Number of attention heads for each attention layer in
                the Transformer encoder.
            intermediate_size: The size of the "intermediate" (i.e., feed-forward)
                layer in the Transformer encoder.
            hidden_act: The non-linear activation function (function or string) in the
                encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
            hidden_dropout_prob: The dropout probabilitiy for all fully connected
                layers in the embeddings, encoder, and pooler.
            attention_probs_dropout_prob: The dropout ratio for the attention
                probabilities.
            max_position_embeddings: The maximum sequence length that this model might
                ever be used with. Typically set this to something large just in case
                (e.g., 512 or 1024 or 2048).
            type_vocab_size: The vocabulary size of the `token_type_ids` passed into
                `BertModel`.
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
177
            layer_norm_eps: The epsilon used by LayerNorm.
178
    """
thomwolf's avatar
thomwolf committed
179
    pretrained_config_archive_map = BERT_PRETRAINED_CONFIG_ARCHIVE_MAP
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195

    def __init__(self,
                 vocab_size_or_config_json_file=30522,
                 hidden_size=768,
                 num_hidden_layers=12,
                 num_attention_heads=12,
                 intermediate_size=3072,
                 hidden_act="gelu",
                 hidden_dropout_prob=0.1,
                 attention_probs_dropout_prob=0.1,
                 max_position_embeddings=512,
                 type_vocab_size=2,
                 initializer_range=0.02,
                 layer_norm_eps=1e-12,
                 **kwargs):
        """Constructs BertConfig.
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218

        Arguments:
            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `BertModel`.
            hidden_size: Size of the encoder layers and the pooler layer.
            num_hidden_layers: Number of hidden layers in the Transformer encoder.
            num_attention_heads: Number of attention heads for each attention layer in
                the Transformer encoder.
            intermediate_size: The size of the "intermediate" (i.e., feed-forward)
                layer in the Transformer encoder.
            hidden_act: The non-linear activation function (function or string) in the
                encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
            hidden_dropout_prob: The dropout probabilitiy for all fully connected
                layers in the embeddings, encoder, and pooler.
            attention_probs_dropout_prob: The dropout ratio for the attention
                probabilities.
            max_position_embeddings: The maximum sequence length that this model might
                ever be used with. Typically set this to something large just in case
                (e.g., 512 or 1024 or 2048).
            type_vocab_size: The vocabulary size of the `token_type_ids` passed into
                `BertModel`.
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
            layer_norm_eps: The epsilon used by LayerNorm.
thomwolf's avatar
thomwolf committed
219
        """
thomwolf's avatar
thomwolf committed
220
        super(BertConfig, self).__init__(**kwargs)
thomwolf's avatar
thomwolf committed
221
222
        if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
                        and isinstance(vocab_size_or_config_json_file, unicode)):
223
            with open(vocab_size_or_config_json_file, "r", encoding='utf-8') as reader:
thomwolf's avatar
thomwolf committed
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
                json_config = json.loads(reader.read())
            for key, value in json_config.items():
                self.__dict__[key] = value
        elif isinstance(vocab_size_or_config_json_file, int):
            self.vocab_size = vocab_size_or_config_json_file
            self.hidden_size = hidden_size
            self.num_hidden_layers = num_hidden_layers
            self.num_attention_heads = num_attention_heads
            self.hidden_act = hidden_act
            self.intermediate_size = intermediate_size
            self.hidden_dropout_prob = hidden_dropout_prob
            self.attention_probs_dropout_prob = attention_probs_dropout_prob
            self.max_position_embeddings = max_position_embeddings
            self.type_vocab_size = type_vocab_size
            self.initializer_range = initializer_range
239
            self.layer_norm_eps = layer_norm_eps
thomwolf's avatar
thomwolf committed
240
241
242
243
        else:
            raise ValueError("First argument must be either a vocabulary size (int)"
                             "or the path to a pretrained model config file (str)")

244

245

246
247
248
try:
    from apex.normalization.fused_layer_norm import FusedLayerNorm as BertLayerNorm
except ImportError:
249
    logger.info("Better speed can be achieved with apex installed from https://www.github.com/nvidia/apex .")
250
251
252
253
254
255
256
257
258
259
260
261
262
263
    class BertLayerNorm(nn.Module):
        def __init__(self, hidden_size, eps=1e-12):
            """Construct a layernorm module in the TF style (epsilon inside the square root).
            """
            super(BertLayerNorm, self).__init__()
            self.weight = nn.Parameter(torch.ones(hidden_size))
            self.bias = nn.Parameter(torch.zeros(hidden_size))
            self.variance_epsilon = eps

        def forward(self, x):
            u = x.mean(-1, keepdim=True)
            s = (x - u).pow(2).mean(-1, keepdim=True)
            x = (x - u) / torch.sqrt(s + self.variance_epsilon)
            return self.weight * x + self.bias
thomwolf's avatar
thomwolf committed
264
265
266
267
268
269

class BertEmbeddings(nn.Module):
    """Construct the embeddings from word, position and token_type embeddings.
    """
    def __init__(self, config):
        super(BertEmbeddings, self).__init__()
270
        self.word_embeddings = nn.Embedding(config.vocab_size, config.hidden_size, padding_idx=0)
271
272
        self.position_embeddings = nn.Embedding(config.max_position_embeddings, config.hidden_size)
        self.token_type_embeddings = nn.Embedding(config.type_vocab_size, config.hidden_size)
thomwolf's avatar
thomwolf committed
273
274
275

        # self.LayerNorm is not snake-cased to stick with TensorFlow model variable name and be able to load
        # any TensorFlow checkpoint file
276
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
thomwolf's avatar
thomwolf committed
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, input_ids, token_type_ids=None):
        seq_length = input_ids.size(1)
        position_ids = torch.arange(seq_length, dtype=torch.long, device=input_ids.device)
        position_ids = position_ids.unsqueeze(0).expand_as(input_ids)
        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)

        words_embeddings = self.word_embeddings(input_ids)
        position_embeddings = self.position_embeddings(position_ids)
        token_type_embeddings = self.token_type_embeddings(token_type_ids)

        embeddings = words_embeddings + position_embeddings + token_type_embeddings
        embeddings = self.LayerNorm(embeddings)
        embeddings = self.dropout(embeddings)
        return embeddings


class BertSelfAttention(nn.Module):
thomwolf's avatar
thomwolf committed
297
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
298
299
300
301
302
        super(BertSelfAttention, self).__init__()
        if config.hidden_size % config.num_attention_heads != 0:
            raise ValueError(
                "The hidden size (%d) is not a multiple of the number of attention "
                "heads (%d)" % (config.hidden_size, config.num_attention_heads))
thomwolf's avatar
thomwolf committed
303
        self.output_attentions = config.output_attentions
304

thomwolf's avatar
thomwolf committed
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
        self.num_attention_heads = config.num_attention_heads
        self.attention_head_size = int(config.hidden_size / config.num_attention_heads)
        self.all_head_size = self.num_attention_heads * self.attention_head_size

        self.query = nn.Linear(config.hidden_size, self.all_head_size)
        self.key = nn.Linear(config.hidden_size, self.all_head_size)
        self.value = nn.Linear(config.hidden_size, self.all_head_size)

        self.dropout = nn.Dropout(config.attention_probs_dropout_prob)

    def transpose_for_scores(self, x):
        new_x_shape = x.size()[:-1] + (self.num_attention_heads, self.attention_head_size)
        x = x.view(*new_x_shape)
        return x.permute(0, 2, 1, 3)

320
    def forward(self, hidden_states, attention_mask, head_mask=None):
thomwolf's avatar
thomwolf committed
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
        mixed_query_layer = self.query(hidden_states)
        mixed_key_layer = self.key(hidden_states)
        mixed_value_layer = self.value(hidden_states)

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

342
343
344
345
        # Mask heads if we want to
        if head_mask is not None:
            attention_probs = attention_probs * head_mask

thomwolf's avatar
thomwolf committed
346
        context_layer = torch.matmul(attention_probs, value_layer)
347

thomwolf's avatar
thomwolf committed
348
349
350
        context_layer = context_layer.permute(0, 2, 1, 3).contiguous()
        new_context_layer_shape = context_layer.size()[:-2] + (self.all_head_size,)
        context_layer = context_layer.view(*new_context_layer_shape)
351

352
        outputs = (context_layer, attention_probs) if self.output_attentions else (context_layer,)
353
        return outputs
thomwolf's avatar
thomwolf committed
354
355
356
357
358
359


class BertSelfOutput(nn.Module):
    def __init__(self, config):
        super(BertSelfOutput, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
360
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
thomwolf's avatar
thomwolf committed
361
362
363
364
365
366
367
368
369
370
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class BertAttention(nn.Module):
thomwolf's avatar
thomwolf committed
371
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
372
        super(BertAttention, self).__init__()
thomwolf's avatar
thomwolf committed
373
        self.self = BertSelfAttention(config)
thomwolf's avatar
thomwolf committed
374
375
        self.output = BertSelfOutput(config)

thomwolf's avatar
thomwolf committed
376
    def prune_heads(self, heads):
thomwolf's avatar
thomwolf committed
377
378
        if len(heads) == 0:
            return
thomwolf's avatar
thomwolf committed
379
        mask = torch.ones(self.self.num_attention_heads, self.self.attention_head_size)
thomwolf's avatar
thomwolf committed
380
381
382
383
384
385
386
387
        for head in heads:
            mask[head] = 0
        mask = mask.view(-1).contiguous().eq(1)
        index = torch.arange(len(mask))[mask].long()
        # Prune linear layers
        self.self.query = prune_linear_layer(self.self.query, index)
        self.self.key = prune_linear_layer(self.self.key, index)
        self.self.value = prune_linear_layer(self.self.value, index)
thomwolf's avatar
thomwolf committed
388
        self.output.dense = prune_linear_layer(self.output.dense, index, dim=1)
thomwolf's avatar
thomwolf committed
389
390
391
392
        # Update hyper params
        self.self.num_attention_heads = self.self.num_attention_heads - len(heads)
        self.self.all_head_size = self.self.attention_head_size * self.self.num_attention_heads

393
    def forward(self, input_tensor, attention_mask, head_mask=None):
394
395
        self_outputs = self.self(input_tensor, attention_mask, head_mask)
        attention_output = self.output(self_outputs[0], input_tensor)
396
        outputs = (attention_output,) + self_outputs[1:]  # add attentions if we output them
397
        return outputs
thomwolf's avatar
thomwolf committed
398
399
400
401
402
403


class BertIntermediate(nn.Module):
    def __init__(self, config):
        super(BertIntermediate, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.intermediate_size)
thomwolf's avatar
thomwolf committed
404
405
406
407
        if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
            self.intermediate_act_fn = ACT2FN[config.hidden_act]
        else:
            self.intermediate_act_fn = config.hidden_act
thomwolf's avatar
thomwolf committed
408
409
410
411
412
413
414
415
416
417
418

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.intermediate_act_fn(hidden_states)
        return hidden_states


class BertOutput(nn.Module):
    def __init__(self, config):
        super(BertOutput, self).__init__()
        self.dense = nn.Linear(config.intermediate_size, config.hidden_size)
419
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
thomwolf's avatar
thomwolf committed
420
421
422
423
424
425
426
427
428
429
        self.dropout = nn.Dropout(config.hidden_dropout_prob)

    def forward(self, hidden_states, input_tensor):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.dropout(hidden_states)
        hidden_states = self.LayerNorm(hidden_states + input_tensor)
        return hidden_states


class BertLayer(nn.Module):
thomwolf's avatar
thomwolf committed
430
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
431
        super(BertLayer, self).__init__()
thomwolf's avatar
thomwolf committed
432
        self.attention = BertAttention(config)
thomwolf's avatar
thomwolf committed
433
434
435
        self.intermediate = BertIntermediate(config)
        self.output = BertOutput(config)

436
    def forward(self, hidden_states, attention_mask, head_mask=None):
437
        attention_outputs = self.attention(hidden_states, attention_mask, head_mask)
thomwolf's avatar
thomwolf committed
438
439
        attention_output = attention_outputs[0]
        intermediate_output = self.intermediate(attention_output)
thomwolf's avatar
thomwolf committed
440
        layer_output = self.output(intermediate_output, attention_output)
441
        outputs = (layer_output,) + attention_outputs[1:]  # add attentions if we output them
442
        return outputs
thomwolf's avatar
thomwolf committed
443
444
445


class BertEncoder(nn.Module):
thomwolf's avatar
thomwolf committed
446
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
447
        super(BertEncoder, self).__init__()
thomwolf's avatar
thomwolf committed
448
449
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states
450
        self.layer = nn.ModuleList([BertLayer(config) for _ in range(config.num_hidden_layers)])
thomwolf's avatar
thomwolf committed
451

452
    def forward(self, hidden_states, attention_mask, head_mask=None):
453
454
        all_hidden_states = ()
        all_attentions = ()
455
        for i, layer_module in enumerate(self.layer):
456
            if self.output_hidden_states:
457
                all_hidden_states = all_hidden_states + (hidden_states,)
458
459
460
461

            layer_outputs = layer_module(hidden_states, attention_mask, head_mask[i])
            hidden_states = layer_outputs[0]

thomwolf's avatar
thomwolf committed
462
            if self.output_attentions:
463
                all_attentions = all_attentions + (layer_outputs[1],)
464
465
466

        # Add last layer
        if self.output_hidden_states:
467
            all_hidden_states = all_hidden_states + (hidden_states,)
468

469
        outputs = (hidden_states,)
470
        if self.output_hidden_states:
471
            outputs = outputs + (all_hidden_states,)
thomwolf's avatar
thomwolf committed
472
        if self.output_attentions:
473
            outputs = outputs + (all_attentions,)
474
        return outputs  # outputs, (hidden states), (attentions)
thomwolf's avatar
thomwolf committed
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495


class BertPooler(nn.Module):
    def __init__(self, config):
        super(BertPooler, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
        self.activation = nn.Tanh()

    def forward(self, hidden_states):
        # We "pool" the model by simply taking the hidden state corresponding
        # to the first token.
        first_token_tensor = hidden_states[:, 0]
        pooled_output = self.dense(first_token_tensor)
        pooled_output = self.activation(pooled_output)
        return pooled_output


class BertPredictionHeadTransform(nn.Module):
    def __init__(self, config):
        super(BertPredictionHeadTransform, self).__init__()
        self.dense = nn.Linear(config.hidden_size, config.hidden_size)
thomwolf's avatar
thomwolf committed
496
497
498
499
        if isinstance(config.hidden_act, str) or (sys.version_info[0] == 2 and isinstance(config.hidden_act, unicode)):
            self.transform_act_fn = ACT2FN[config.hidden_act]
        else:
            self.transform_act_fn = config.hidden_act
500
        self.LayerNorm = BertLayerNorm(config.hidden_size, eps=config.layer_norm_eps)
thomwolf's avatar
thomwolf committed
501
502
503
504
505
506
507
508
509
510
511
512

    def forward(self, hidden_states):
        hidden_states = self.dense(hidden_states)
        hidden_states = self.transform_act_fn(hidden_states)
        hidden_states = self.LayerNorm(hidden_states)
        return hidden_states


class BertLMPredictionHead(nn.Module):
    def __init__(self, config, bert_model_embedding_weights):
        super(BertLMPredictionHead, self).__init__()
        self.transform = BertPredictionHeadTransform(config)
513
        self.torchscript = config.torchscript
thomwolf's avatar
thomwolf committed
514
515
516
517
518
519

        # The output weights are the same as the input embeddings, but there is
        # an output-only bias for each token.
        self.decoder = nn.Linear(bert_model_embedding_weights.size(1),
                                 bert_model_embedding_weights.size(0),
                                 bias=False)
520
521
522
523
524
525

        if self.torchscript:
            self.decoder.weight = nn.Parameter(bert_model_embedding_weights.clone())
        else:
            self.decoder.weight = bert_model_embedding_weights

thomwolf's avatar
thomwolf committed
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
        self.bias = nn.Parameter(torch.zeros(bert_model_embedding_weights.size(0)))

    def forward(self, hidden_states):
        hidden_states = self.transform(hidden_states)
        hidden_states = self.decoder(hidden_states) + self.bias
        return hidden_states


class BertOnlyMLMHead(nn.Module):
    def __init__(self, config, bert_model_embedding_weights):
        super(BertOnlyMLMHead, self).__init__()
        self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)

    def forward(self, sequence_output):
        prediction_scores = self.predictions(sequence_output)
        return prediction_scores


class BertOnlyNSPHead(nn.Module):
    def __init__(self, config):
        super(BertOnlyNSPHead, self).__init__()
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, pooled_output):
        seq_relationship_score = self.seq_relationship(pooled_output)
        return seq_relationship_score


class BertPreTrainingHeads(nn.Module):
    def __init__(self, config, bert_model_embedding_weights):
        super(BertPreTrainingHeads, self).__init__()
        self.predictions = BertLMPredictionHead(config, bert_model_embedding_weights)
        self.seq_relationship = nn.Linear(config.hidden_size, 2)

    def forward(self, sequence_output, pooled_output):
        prediction_scores = self.predictions(sequence_output)
        seq_relationship_score = self.seq_relationship(pooled_output)
        return prediction_scores, seq_relationship_score


566
class BertPreTrainedModel(PreTrainedModel):
thomwolf's avatar
thomwolf committed
567
568
569
    """ An abstract class to handle weights initialization and
        a simple interface for dowloading and loading pretrained models.
    """
570
    config_class = BertConfig
571
    pretrained_model_archive_map = BERT_PRETRAINED_MODEL_ARCHIVE_MAP
572
573
574
    load_tf_weights = load_tf_weights_in_bert
    base_model_prefix = "bert"

575
576
577
    def __init__(self, *inputs, **kwargs):
        super(BertPreTrainedModel, self).__init__(*inputs, **kwargs)

thomwolf's avatar
thomwolf committed
578
    def init_weights(self, module):
thomwolf's avatar
thomwolf committed
579
580
581
582
583
584
585
        """ 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)
        elif isinstance(module, BertLayerNorm):
Li Dong's avatar
Li Dong committed
586
587
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
thomwolf's avatar
thomwolf committed
588
589
590
591
        if isinstance(module, nn.Linear) and module.bias is not None:
            module.bias.data.zero_()


thomwolf's avatar
thomwolf committed
592
class BertModel(BertPreTrainedModel):
593
594
    r"""BERT model ("Bidirectional Embedding Representations from a Transformer").

595
    :class:`~pytorch_transformers.BertModel` is the basic BERT Transformer model with a layer of summed token, \
596
597
598
599
600
601
602
603
    position and sequence embeddings followed by a series of identical self-attention blocks (12 for BERT-base, 24 \
    for BERT-large). The model is instantiated with the following parameters.

    Arguments:
        config: a BertConfig class instance with the configuration to build a new model
        output_attentions: If True, also output attentions weights computed by the model at each layer. Default: False
        output_hidden_states: If True, also output hidden states computed by the model at each layer. Default: Fals

thomwolf's avatar
thomwolf committed
604

605
    Example::
thomwolf's avatar
thomwolf committed
606

607
608
609
610
        config = modeling.BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
            num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

        model = modeling.BertModel(config=config)
thomwolf's avatar
thomwolf committed
611
612

    """
thomwolf's avatar
thomwolf committed
613
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
614
        super(BertModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
615

thomwolf's avatar
thomwolf committed
616
        self.embeddings = BertEmbeddings(config)
thomwolf's avatar
thomwolf committed
617
        self.encoder = BertEncoder(config)
thomwolf's avatar
thomwolf committed
618
        self.pooler = BertPooler(config)
thomwolf's avatar
thomwolf committed
619

thomwolf's avatar
thomwolf committed
620
        self.apply(self.init_weights)
thomwolf's avatar
thomwolf committed
621

thomwolf's avatar
thomwolf committed
622
    def _prune_heads(self, heads_to_prune):
thomwolf's avatar
thomwolf committed
623
624
        """ Prunes heads of the model.
            heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
thomwolf's avatar
thomwolf committed
625
            See base class PreTrainedModel
thomwolf's avatar
thomwolf committed
626
627
628
629
        """
        for layer, heads in heads_to_prune.items():
            self.encoder.layer[layer].attention.prune_heads(heads)

630
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, head_mask=None):
631
        """
632
        Performs a model forward pass. **Can be called by calling the class directly, once it has been instantiated.**
633
634
635


        Arguments:
636
            input_ids: a ``torch.LongTensor`` of shape [batch_size, sequence_length] with the word token indices in the \
637
638
                vocabulary(see the tokens pre-processing logic in the scripts `run_bert_extract_features.py`, \
                `run_bert_classifier.py` and `run_bert_squad.py`)
639
            token_type_ids: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with the token \
640
641
                types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to \
                a `sentence B` token (see BERT paper for more details).
642
            attention_mask: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with indices \
643
644
645
646
647
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max \
                input sequence length in the current batch. It's the mask that we typically use for attention when \
                a batch has varying length sentences.
            output_all_encoded_layers: boolean which controls the content of the `encoded_layers` output as described \
            below. Default: `True`.
648
            head_mask: an optional ``torch.Tensor`` of shape [num_heads] or [num_layers, num_heads] with indices between 0 \
649
650
651
652
653
654
655
656
657
658
659
            and 1. It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 \
            => head is not masked.


        Returns:
            A tuple composed of (encoded_layers, pooled_output). Encoded layers are controlled by the \
            ``output_all_encoded_layers`` argument.

            If ``output_all_encoded_layers`` is set to True, outputs a list of the full sequences of \
            encoded-hidden-states at the end of each attention \
            block (i.e. 12 full sequences for BERT-base, 24 for BERT-large), each encoded-hidden-state is a\
660
            ``torch.FloatTensor`` of size [batch_size, sequence_length, hidden_size].
661
662
663
664

            If set to False, outputs only the full sequence of hidden-states corresponding \
            to the last attention block of shape [batch_size, sequence_length, hidden_size].

665
            ``pooled_output`` is a ``torch.FloatTensor`` of size [batch_size, hidden_size] which is the output of a \
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
            classifier pretrained on top of the hidden state associated to the first character of the \
            input (`CLS`) to train on the Next-Sentence task (see BERT's paper).

        Example::

            # Already been converted into WordPiece token ids
            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
            input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
            token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])


            all_encoder_layers, pooled_output = model(input_ids, token_type_ids, input_mask)
            # or
            all_encoder_layers, pooled_output = model.forward(input_ids, token_type_ids, input_mask)


        """
thomwolf's avatar
thomwolf committed
683
684
685
686
687
688
689
690
691
692
693
694
695
696
697
698
699
700
701
702
        if attention_mask is None:
            attention_mask = torch.ones_like(input_ids)
        if token_type_ids is None:
            token_type_ids = torch.zeros_like(input_ids)

        # We create a 3D attention mask from a 2D tensor mask.
        # Sizes are [batch_size, 1, 1, to_seq_length]
        # So we can broadcast to [batch_size, num_heads, from_seq_length, to_seq_length]
        # this attention mask is more simple than the triangular masking of causal attention
        # used in OpenAI GPT, we just need to prepare the broadcast dimension here.
        extended_attention_mask = attention_mask.unsqueeze(1).unsqueeze(2)

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

thomwolf's avatar
thomwolf committed
703
        # Prepare head mask if needed
thomwolf's avatar
thomwolf committed
704
        # 1.0 in head_mask indicate we keep the head
thomwolf's avatar
thomwolf committed
705
        # attention_probs has shape bsz x n_heads x N x N
thomwolf's avatar
thomwolf committed
706
707
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads]
        # and head_mask is converted to shape [num_hidden_layers x batch x num_heads x seq_length x seq_length]
thomwolf's avatar
thomwolf committed
708
709
        if head_mask is not None:
            if head_mask.dim() == 1:
710
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
thomwolf's avatar
thomwolf committed
711
                head_mask = head_mask.expand(self.config.num_hidden_layers, -1, -1, -1, -1)
thomwolf's avatar
thomwolf committed
712
            elif head_mask.dim() == 2:
713
                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
thomwolf's avatar
thomwolf committed
714
            head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility
715
716
        else:
            head_mask = [None] * self.config.num_hidden_layers
thomwolf's avatar
thomwolf committed
717

thomwolf's avatar
thomwolf committed
718
        embedding_output = self.embeddings(input_ids, token_type_ids)
719
720
721
722
        encoder_outputs = self.encoder(embedding_output,
                                       extended_attention_mask,
                                       head_mask=head_mask)
        sequence_output = encoder_outputs[0]
thomwolf's avatar
thomwolf committed
723
        pooled_output = self.pooler(sequence_output)
724

725
        outputs = (sequence_output, pooled_output,) + encoder_outputs[1:]  # add hidden_states and attentions if they are here
726
        return outputs  # sequence_output, pooled_output, (hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
727
728


thomwolf's avatar
thomwolf committed
729
class BertForPreTraining(BertPreTrainedModel):
thomwolf's avatar
thomwolf committed
730
731
    """BERT model with pre-training heads.
    This module comprises the BERT model followed by the two pre-training heads:
732

thomwolf's avatar
thomwolf committed
733
        - the masked language modeling head, and
734

thomwolf's avatar
thomwolf committed
735
736
        - the next sentence classification head.

737
    Args:
738
739
        `config`: a BertConfig class instance with the configuration to build a new model
        `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
740
        `output_hidden_states`: If True, also output hidden states computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
741

742
743
744
745
746
747
    Example ::

        config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
            num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

        model = BertForPreTraining(config)
thomwolf's avatar
thomwolf committed
748
    """
thomwolf's avatar
thomwolf committed
749
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
750
        super(BertForPreTraining, self).__init__(config)
751

thomwolf's avatar
thomwolf committed
752
        self.bert = BertModel(config)
thomwolf's avatar
thomwolf committed
753
        self.cls = BertPreTrainingHeads(config, self.bert.embeddings.word_embeddings.weight)
thomwolf's avatar
thomwolf committed
754

thomwolf's avatar
thomwolf committed
755
        self.apply(self.init_weights)
thomwolf's avatar
thomwolf committed
756

757
758
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, masked_lm_labels=None,
                next_sentence_label=None, head_mask=None):
759
        """
760
        Performs a model forward pass. **Can be called by calling the class directly, once it has been instantiated.**
761
762

        Args:
763
            `input_ids`: a ``torch.LongTensor`` of shape [batch_size, sequence_length]
764
765
                with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
                `run_bert_extract_features.py`, `run_bert_classifier.py` and `run_bert_squad.py`)
766
            `token_type_ids`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with the token
767
768
                types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
                a `sentence B` token (see BERT paper for more details).
769
            `attention_mask`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with indices
770
771
772
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
                input sequence length in the current batch. It's the mask that we typically use for attention when
                a batch has varying length sentences.
773
            `masked_lm_labels`: optional masked language modeling labels: ``torch.LongTensor`` of shape [batch_size, sequence_length]
774
775
                with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss
                is only computed for the labels set in [0, ..., vocab_size]
776
            `next_sentence_label`: optional next sentence classification loss: ``torch.LongTensor`` of shape [batch_size]
777
778
                with indices selected in [0, 1].
                0 => next sentence is the continuation, 1 => next sentence is a random sentence.
779
            `head_mask`: an optional ``torch.Tensor`` of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
780
781
782
783
                It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.


        Returns:
784
            Either a ``torch.Tensor`` or ``tuple(torch.Tensor, torch.Tensor)``.
785
786
787
788
789

            if ``masked_lm_labels`` and ``next_sentence_label`` are not ``None``, outputs the total_loss which is the \
             sum of the masked language modeling loss and the next \
            sentence classification loss.

790
791
792
793
            if ``masked_lm_labels`` or ``next_sentence_label`` is ``None``, outputs a tuple made of:

                - the masked language modeling logits of shape [batch_size, sequence_length, vocab_size]

794
795
796
797
798
799
800
801
802
803
804
805
806
807
808
809
810
                - the next sentence classification logits of shape [batch_size, 2].

        Example ::

            # Already been converted into WordPiece token ids
            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
            input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
            token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])

            config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
                num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

            model = BertForPreTraining(config)
            masked_lm_logits_scores, seq_relationship_logits = model(input_ids, token_type_ids, input_mask)
            # or
            masked_lm_logits_scores, seq_relationship_logits = model.forward(input_ids, token_type_ids, input_mask)
        """
811
812
813
        outputs = self.bert(input_ids, token_type_ids, attention_mask, head_mask=head_mask)

        sequence_output, pooled_output = outputs[:2]
thomwolf's avatar
thomwolf committed
814
815
        prediction_scores, seq_relationship_score = self.cls(sequence_output, pooled_output)

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

thomwolf's avatar
thomwolf committed
818
819
        if masked_lm_labels is not None and next_sentence_label is not None:
            loss_fct = CrossEntropyLoss(ignore_index=-1)
820
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
821
            next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
thomwolf's avatar
thomwolf committed
822
            total_loss = masked_lm_loss + next_sentence_loss
823
            outputs = (total_loss,) + outputs
824
825

        return outputs  # (loss), prediction_scores, seq_relationship_score, (hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
826
827


thomwolf's avatar
thomwolf committed
828
class BertForMaskedLM(BertPreTrainedModel):
thomwolf's avatar
thomwolf committed
829
830
831
    """BERT model with the masked language modeling head.
    This module comprises the BERT model followed by the masked language modeling head.

832
    Args:
833
834
        `config`: a BertConfig class instance with the configuration to build a new model
        `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
835
        `output_hidden_states`: If True, also output hidden states computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
836

837
838
839
840
841
842
    Example::

        config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
            num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

        model = BertForMaskedLM(config)
thomwolf's avatar
thomwolf committed
843
    """
thomwolf's avatar
thomwolf committed
844
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
845
        super(BertForMaskedLM, self).__init__(config)
thomwolf's avatar
thomwolf committed
846

thomwolf's avatar
thomwolf committed
847
        self.bert = BertModel(config)
thomwolf's avatar
thomwolf committed
848
        self.cls = BertOnlyMLMHead(config, self.bert.embeddings.word_embeddings.weight)
thomwolf's avatar
thomwolf committed
849

thomwolf's avatar
thomwolf committed
850
        self.apply(self.init_weights)
thomwolf's avatar
thomwolf committed
851

852
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, masked_lm_labels=None, head_mask=None):
853
        """
854
        Performs a model forward pass. **Can be called by calling the class directly, once it has been instantiated.**
855
856

        Args:
857
            `input_ids`: a ``torch.LongTensor`` of shape [batch_size, sequence_length]
858
859
                with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
                `run_bert_extract_features.py`, `run_bert_classifier.py` and `run_bert_squad.py`)
860
            `token_type_ids`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with the token
861
862
                types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
                a `sentence B` token (see BERT paper for more details).
863
            `attention_mask`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with indices
864
865
866
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
                input sequence length in the current batch. It's the mask that we typically use for attention when
                a batch has varying length sentences.
867
            `masked_lm_labels`: masked language modeling labels: ``torch.LongTensor`` of shape [batch_size, sequence_length]
868
869
                with indices selected in [-1, 0, ..., vocab_size]. All labels set to -1 are ignored (masked), the loss
                is only computed for the labels set in [0, ..., vocab_size]
870
            `head_mask`: an optional ``torch.LongTensor`` of shape [num_heads] with indices
871
872
873
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
                input sequence length in the current batch. It's the mask that we typically use for attention when
                a batch has varying length sentences.
874
            `head_mask`: an optional ``torch.Tensor`` of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
875
876
877
                It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.

        Returns:
878
            Masked language modeling loss if ``masked_lm_labels`` is specified, masked language modeling
879
880
881
882
883
884
885
886
887
888
889
890
891
            logits of shape [batch_size, sequence_length, vocab_size] otherwise.

        Example::

            # Already been converted into WordPiece token ids
            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
            input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
            token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])

            masked_lm_logits_scores = model(input_ids, token_type_ids, input_mask)
            # or
            masked_lm_logits_scores = model.forward(input_ids, token_type_ids, input_mask)
        """
thomwolf's avatar
thomwolf committed
892
893
894
        outputs = self.bert(input_ids, token_type_ids, attention_mask, head_mask=head_mask)

        sequence_output = outputs[0]
thomwolf's avatar
thomwolf committed
895
896
        prediction_scores = self.cls(sequence_output)

897
        outputs = (prediction_scores,) + outputs[2:]  # Add hidden states and attention is they are here
thomwolf's avatar
thomwolf committed
898
899
        if masked_lm_labels is not None:
            loss_fct = CrossEntropyLoss(ignore_index=-1)
900
            masked_lm_loss = loss_fct(prediction_scores.view(-1, self.config.vocab_size), masked_lm_labels.view(-1))
901
            outputs = (masked_lm_loss,) + outputs
thomwolf's avatar
thomwolf committed
902
903

        return outputs  # (masked_lm_loss), prediction_scores, (hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
904
905


thomwolf's avatar
thomwolf committed
906
class BertForNextSentencePrediction(BertPreTrainedModel):
thomwolf's avatar
thomwolf committed
907
908
909
    """BERT model with next sentence prediction head.
    This module comprises the BERT model followed by the next sentence classification head.

910
    Args:
911
912
        `config`: a BertConfig class instance with the configuration to build a new model
        `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
913
        `output_hidden_states`: If True, also output hidden states computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
914

915
916
917
918
919
920
    Example::

        config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
            num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

        model = BertForNextSentencePrediction(config)
thomwolf's avatar
thomwolf committed
921
    """
thomwolf's avatar
thomwolf committed
922
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
923
        super(BertForNextSentencePrediction, self).__init__(config)
thomwolf's avatar
thomwolf committed
924

thomwolf's avatar
thomwolf committed
925
        self.bert = BertModel(config)
thomwolf's avatar
thomwolf committed
926
        self.cls = BertOnlyNSPHead(config)
thomwolf's avatar
thomwolf committed
927

thomwolf's avatar
thomwolf committed
928
        self.apply(self.init_weights)
thomwolf's avatar
thomwolf committed
929

930
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, next_sentence_label=None, head_mask=None):
931
        """
932
        Performs a model forward pass. **Can be called by calling the class directly, once it has been instantiated.**
933
934

        Args:
935
            `input_ids`: a ``torch.LongTensor`` of shape [batch_size, sequence_length]
936
937
                with the word token indices in the vocabulary(see the tokens pre-processing logic in the scripts
                `run_bert_extract_features.py`, `run_bert_classifier.py` and `run_bert_squad.py`)
938
            `token_type_ids`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with the token
939
940
                types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
                a `sentence B` token (see BERT paper for more details).
941
            `attention_mask`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with indices
942
943
944
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
                input sequence length in the current batch. It's the mask that we typically use for attention when
                a batch has varying length sentences.
945
            `next_sentence_label`: next sentence classification loss: ``torch.LongTensor`` of shape [batch_size]
946
947
                with indices selected in [0, 1].
                0 => next sentence is the continuation, 1 => next sentence is a random sentence.
948
            `head_mask`: an optional ``torch.Tensor`` of shape [num_heads] or [num_layers, num_heads] with indices between
949
950
951
952
                0 and 1.It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked,
                0.0 => head is not masked.

        Returns:
953
954
955
            If ``next_sentence_label`` is specified, outputs the total_loss which is the sum of the masked language
            modeling loss and the next sentence classification loss. If ``next_sentence_label`` is ``None``, outputs
            the next sentence classification logits of shape [batch_size, 2].
956
957
958
959
960
961
962
963
964
965
966
967
968


        Example::

            # Already been converted into WordPiece token ids
            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
            input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
            token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])

            seq_relationship_logits = model(input_ids, token_type_ids, input_mask)
            # or
            seq_relationship_logits = model.forward(input_ids, token_type_ids, input_mask)
        """
thomwolf's avatar
thomwolf committed
969
970
971
        outputs = self.bert(input_ids, token_type_ids, attention_mask, head_mask=head_mask)
        pooled_output = outputs[1]

972
        seq_relationship_score = self.cls(pooled_output)
thomwolf's avatar
thomwolf committed
973

974
        outputs = (seq_relationship_score,) + outputs[2:]  # add hidden states and attention if they are here
thomwolf's avatar
thomwolf committed
975
976
        if next_sentence_label is not None:
            loss_fct = CrossEntropyLoss(ignore_index=-1)
977
            next_sentence_loss = loss_fct(seq_relationship_score.view(-1, 2), next_sentence_label.view(-1))
978
            outputs = (next_sentence_loss,) + outputs
thomwolf's avatar
thomwolf committed
979
980

        return outputs  # (next_sentence_loss), seq_relationship_score, (hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
981
982


thomwolf's avatar
thomwolf committed
983
class BertForSequenceClassification(BertPreTrainedModel):
thomwolf's avatar
thomwolf committed
984
985
986
987
988
    """BERT model for classification.
    This module is composed of the BERT model with a linear layer on top of
    the pooled output.

    Params:
989
990
        `config`: a BertConfig class instance with the configuration to build a new model
        `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
991
        `output_hidden_states`: If True, also output hidden states computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
992
993
        `num_labels`: the number of classes for the classifier. Default = 2.

994
995
996
997
998
999
1000
1001
    Example::

        config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
            num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

        num_labels = 2

        model = BertForSequenceClassification(config, num_labels)
thomwolf's avatar
thomwolf committed
1002
    """
thomwolf's avatar
thomwolf committed
1003
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
1004
        super(BertForSequenceClassification, self).__init__(config)
thomwolf's avatar
thomwolf committed
1005
        self.num_labels = config.num_labels
thomwolf's avatar
thomwolf committed
1006

thomwolf's avatar
thomwolf committed
1007
        self.bert = BertModel(config)
thomwolf's avatar
thomwolf committed
1008
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
thomwolf's avatar
thomwolf committed
1009
        self.classifier = nn.Linear(config.hidden_size, self.config.num_labels)
thomwolf's avatar
thomwolf committed
1010

thomwolf's avatar
thomwolf committed
1011
        self.apply(self.init_weights)
thomwolf's avatar
thomwolf committed
1012

1013
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, head_mask=None):
1014
        """
1015
        Performs a model forward pass. **Can be called by calling the class directly, once it has been instantiated.**
1016
1017

        Parameters:
1018
            `input_ids`: a ``torch.LongTensor`` of shape [batch_size, sequence_length]
1019
1020
                with the word token indices in the vocabulary. Items in the batch should begin with the special "CLS" token. (see the tokens preprocessing logic in the scripts
                `run_bert_extract_features.py`, `run_bert_classifier.py` and `run_bert_squad.py`)
1021
            `token_type_ids`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with the token
1022
1023
                types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
                a `sentence B` token (see BERT paper for more details).
1024
            `attention_mask`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with indices
1025
1026
1027
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
                input sequence length in the current batch. It's the mask that we typically use for attention when
                a batch has varying length sentences.
1028
            `labels`: labels for the classification output: ``torch.LongTensor`` of shape [batch_size]
1029
                with indices selected in [0, ..., num_labels].
1030
            `head_mask`: an optional ``torch.Tensor`` of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
1031
1032
1033
                It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.

        Returns:
1034
1035
            If ``labels`` is not ``None``, outputs the CrossEntropy classification loss of the output with the labels.
            If ``labels`` is ``None``, outputs the classification logits of shape [batch_size, num_labels].
1036
1037
1038
1039
1040
1041
1042
1043
1044
1045
1046
1047

        Example::

            # Already been converted into WordPiece token ids
            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
            input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
            token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])

            logits = model(input_ids, token_type_ids, input_mask)
            # or
            logits = model.forward(input_ids, token_type_ids, input_mask)
        """
thomwolf's avatar
thomwolf committed
1048
1049
1050
        outputs = self.bert(input_ids, token_type_ids, attention_mask, head_mask=head_mask)
        pooled_output = outputs[1]

thomwolf's avatar
thomwolf committed
1051
1052
1053
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)

1054
        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
thomwolf's avatar
thomwolf committed
1055

thomwolf's avatar
thomwolf committed
1056
        if labels is not None:
1057
1058
1059
1060
1061
1062
1063
            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))
1064
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
1065
1066

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


thomwolf's avatar
thomwolf committed
1069
class BertForMultipleChoice(BertPreTrainedModel):
1070
    """BERT model for multiple choice tasks.
1071
    This module is composed of the BERT model with a linear layer on top of the pooled output.
1072

1073
    Parameters:
1074
1075
        `config`: a BertConfig class instance with the configuration to build a new model
        `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
1076
        `output_hidden_states`: If True, also output hidden states computed by the model at each layer. Default: False
1077

1078
1079
1080
1081
1082
1083
1084
1085
1086
1087
1088
    Example::

        # Already been converted into WordPiece token ids
        input_ids = torch.LongTensor([[[31, 51, 99], [15, 5, 0]], [[12, 16, 42], [14, 28, 57]]])
        input_mask = torch.LongTensor([[[1, 1, 1], [1, 1, 0]],[[1,1,0], [1, 0, 0]]])
        token_type_ids = torch.LongTensor([[[0, 0, 1], [0, 1, 0]],[[0, 1, 1], [0, 0, 1]]])
        config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
            num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

        model = BertForMultipleChoice(config)
        logits = model(input_ids, token_type_ids, input_mask)
1089
    """
thomwolf's avatar
thomwolf committed
1090
    def __init__(self, config):
1091
        super(BertForMultipleChoice, self).__init__(config)
thomwolf's avatar
thomwolf committed
1092

thomwolf's avatar
thomwolf committed
1093
        self.bert = BertModel(config)
1094
1095
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
        self.classifier = nn.Linear(config.hidden_size, 1)
thomwolf's avatar
thomwolf committed
1096

thomwolf's avatar
thomwolf committed
1097
        self.apply(self.init_weights)
1098

1099
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, head_mask=None):
1100
        """
1101
        Performs a model forward pass. **Can be called by calling the class directly, once it has been instantiated.**
1102
1103

        Parameters:
1104
            `input_ids`: a ``torch.LongTensor`` of shape [batch_size, num_choices, sequence_length]
1105
1106
                with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
                `run_bert_extract_features.py`, `run_bert_classifier.py` and `run_bert_squad.py`)
1107
            `token_type_ids`: an optional ``torch.LongTensor`` of shape [batch_size, num_choices, sequence_length]
1108
1109
                with the token types indices selected in [0, 1]. Type 0 corresponds to a `sentence A`
                and type 1 corresponds to a `sentence B` token (see BERT paper for more details).
1110
            `attention_mask`: an optional ``torch.LongTensor`` of shape [batch_size, num_choices, sequence_length] with indices
1111
1112
1113
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
                input sequence length in the current batch. It's the mask that we typically use for attention when
                a batch has varying length sentences.
1114
            `labels`: labels for the classification output: ``torch.LongTensor`` of shape [batch_size]
1115
                with indices selected in [0, ..., num_choices].
1116
            `head_mask`: an optional ``torch.Tensor`` of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
1117
1118
1119
                It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.

        Returns:
1120
1121
            If ``labels`` is not ``None``, outputs the CrossEntropy classification loss of the output with the labels.
            If ``labels`` is ``None``, outputs the classification logits of shape [batch_size, num_labels].
1122
1123
1124
1125
1126
1127
1128
1129
1130
1131
1132
1133
1134

        Example::

            # Already been converted into WordPiece token ids
            input_ids = torch.LongTensor([[[31, 51, 99], [15, 5, 0]], [[12, 16, 42], [14, 28, 57]]])
            input_mask = torch.LongTensor([[[1, 1, 1], [1, 1, 0]],[[1,1,0], [1, 0, 0]]])
            token_type_ids = torch.LongTensor([[[0, 0, 1], [0, 1, 0]],[[0, 1, 1], [0, 0, 1]]])
            config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
                num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

            model = BertForMultipleChoice(config)
            logits = model(input_ids, token_type_ids, input_mask)
        """
thomwolf's avatar
thomwolf committed
1135
1136
1137
        """ Input shapes should be [bsz, num choices, seq length] """
        num_choices = input_ids.shape[1]

1138
        flat_input_ids = input_ids.view(-1, input_ids.size(-1))
1139
1140
        flat_token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1)) if token_type_ids is not None else None
        flat_attention_mask = attention_mask.view(-1, attention_mask.size(-1)) if attention_mask is not None else None
thomwolf's avatar
thomwolf committed
1141
1142
1143
        outputs = self.bert(flat_input_ids, flat_token_type_ids, flat_attention_mask, head_mask=head_mask)
        pooled_output = outputs[1]

1144
1145
        pooled_output = self.dropout(pooled_output)
        logits = self.classifier(pooled_output)
thomwolf's avatar
thomwolf committed
1146
        reshaped_logits = logits.view(-1, num_choices)
1147

1148
        outputs = (reshaped_logits,) + outputs[2:]  # add hidden states and attention if they are here
thomwolf's avatar
thomwolf committed
1149

1150
1151
1152
        if labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(reshaped_logits, labels)
1153
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
1154
1155

        return outputs  # (loss), reshaped_logits, (hidden_states), (attentions)
1156
1157


thomwolf's avatar
thomwolf committed
1158
class BertForTokenClassification(BertPreTrainedModel):
1159
1160
1161
1162
    """BERT model for token-level classification.
    This module is composed of the BERT model with a linear layer on top of
    the full hidden state of the last layer.

1163
    Parameters:
1164
1165
        `config`: a BertConfig class instance with the configuration to build a new model
        `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
1166
        `output_hidden_states`: If True, also output hidden states computed by the model at each layer. Default: False
1167
1168
        `num_labels`: the number of classes for the classifier. Default = 2.

1169
1170
1171
1172
1173
1174
1175
1176
    Example::

        config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
            num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

        num_labels = 2

        model = BertForTokenClassification(config, num_labels)
1177
    """
thomwolf's avatar
thomwolf committed
1178
    def __init__(self, config):
1179
        super(BertForTokenClassification, self).__init__(config)
thomwolf's avatar
thomwolf committed
1180
        self.num_labels = config.num_labels
thomwolf's avatar
thomwolf committed
1181

thomwolf's avatar
thomwolf committed
1182
        self.bert = BertModel(config)
1183
        self.dropout = nn.Dropout(config.hidden_dropout_prob)
thomwolf's avatar
thomwolf committed
1184
        self.classifier = nn.Linear(config.hidden_size, config.num_labels)
thomwolf's avatar
thomwolf committed
1185

thomwolf's avatar
thomwolf committed
1186
        self.apply(self.init_weights)
1187

1188
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, labels=None, head_mask=None):
1189
        """
1190
        Performs a model forward pass. **Can be called by calling the class directly, once it has been instantiated.**
1191
1192

        Parameters:
1193
            `input_ids`: a ``torch.LongTensor`` of shape [batch_size, sequence_length]
1194
1195
                with the word token indices in the vocabulary(see the tokens pre-processing logic in the scripts
                `run_bert_extract_features.py`, `run_bert_classifier.py` and `run_bert_squad.py`)
1196
            `token_type_ids`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with the token
1197
1198
                types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
                a `sentence B` token (see BERT paper for more details).
1199
            `attention_mask`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with indices
1200
1201
1202
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
                input sequence length in the current batch. It's the mask that we typically use for attention when
                a batch has varying length sentences.
1203
            `labels`: labels for the classification output: ``torch.LongTensor`` of shape [batch_size, sequence_length]
1204
                with indices selected in [0, ..., num_labels].
1205
            `head_mask`: an optional ``torch.Tensor`` of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
1206
1207
1208
                It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.

        Returns:
1209
1210
            If ``labels`` is not ``None``, outputs the CrossEntropy classification loss of the output with the labels.
            If ``labels`` is ``None``, outputs the classification logits of shape [batch_size, sequence_length, num_labels].
1211
1212
1213
1214
1215
1216
1217
1218
1219
1220
1221
1222

        Example::

            # Already been converted into WordPiece token ids
            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
            input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
            token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])

            logits = model(input_ids, token_type_ids, input_mask)
            # or
            logits = model.forward(input_ids, token_type_ids, input_mask)
        """
thomwolf's avatar
thomwolf committed
1223
1224
1225
        outputs = self.bert(input_ids, token_type_ids, attention_mask, head_mask=head_mask)
        sequence_output = outputs[0]

1226
1227
        sequence_output = self.dropout(sequence_output)
        logits = self.classifier(sequence_output)
1228

1229
        outputs = (logits,) + outputs[2:]  # add hidden states and attention if they are here
1230
1231
        if labels is not None:
            loss_fct = CrossEntropyLoss()
1232
1233
1234
1235
1236
1237
1238
1239
            # 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))
1240
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
1241
1242

        return outputs  # (loss), logits, (hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
1243
1244


thomwolf's avatar
thomwolf committed
1245
class BertForQuestionAnswering(BertPreTrainedModel):
thomwolf's avatar
thomwolf committed
1246
1247
1248
1249
    """BERT model for Question Answering (span extraction).
    This module is composed of the BERT model with a linear layer on top of
    the sequence output that computes start_logits and end_logits

1250
    Parameters:
1251
1252
        `config`: a BertConfig class instance with the configuration to build a new model
        `output_attentions`: If True, also output attentions weights computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
1253
        `output_hidden_states`: If True, also output hidden states computed by the model at each layer. Default: False
thomwolf's avatar
thomwolf committed
1254

1255
1256
1257
1258
1259
1260
    Example::

        config = BertConfig(vocab_size_or_config_json_file=32000, hidden_size=768,
            num_hidden_layers=12, num_attention_heads=12, intermediate_size=3072)

        model = BertForQuestionAnswering(config)
thomwolf's avatar
thomwolf committed
1261
    """
thomwolf's avatar
thomwolf committed
1262
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
1263
        super(BertForQuestionAnswering, self).__init__(config)
thomwolf's avatar
thomwolf committed
1264
1265
1266
1267
        self.num_labels = config.num_labels

        self.bert = BertModel(config)
        self.qa_outputs = nn.Linear(config.hidden_size, config.num_labels)
thomwolf's avatar
thomwolf committed
1268

thomwolf's avatar
thomwolf committed
1269
        self.apply(self.init_weights)
thomwolf's avatar
thomwolf committed
1270

thomwolf's avatar
thomwolf committed
1271
1272
    def forward(self, input_ids, token_type_ids=None, attention_mask=None, start_positions=None,
                end_positions=None, head_mask=None):
1273
        """
1274
1275
        Performs a model forward pass. **Can be called by calling the class directly, once it has been instantiated.**

1276
        Parameters:
1277
            `input_ids`: a ``torch.LongTensor`` of shape [batch_size, sequence_length]
1278
1279
                with the word token indices in the vocabulary(see the tokens preprocessing logic in the scripts
                `run_bert_extract_features.py`, `run_bert_classifier.py` and `run_bert_squad.py`)
1280
            `token_type_ids`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with the token
1281
1282
                types indices selected in [0, 1]. Type 0 corresponds to a `sentence A` and type 1 corresponds to
                a `sentence B` token (see BERT paper for more details).
1283
            `attention_mask`: an optional ``torch.LongTensor`` of shape [batch_size, sequence_length] with indices
1284
1285
1286
                selected in [0, 1]. It's a mask to be used if the input sequence length is smaller than the max
                input sequence length in the current batch. It's the mask that we typically use for attention when
                a batch has varying length sentences.
1287
            `start_positions`: position of the first token for the labeled span: ``torch.LongTensor`` of shape [batch_size].
1288
1289
                Positions are clamped to the length of the sequence and position outside of the sequence are not taken
                into account for computing the loss.
1290
            `end_positions`: position of the last token for the labeled span: ``torch.LongTensor`` of shape [batch_size].
1291
1292
                Positions are clamped to the length of the sequence and position outside of the sequence are not taken
                into account for computing the loss.
1293
            `head_mask`: an optional ``torch.Tensor`` of shape [num_heads] or [num_layers, num_heads] with indices between 0 and 1.
1294
1295
1296
                It's a mask to be used to nullify some heads of the transformer. 1.0 => head is fully masked, 0.0 => head is not masked.

        Returns:
1297
            If ``start_positions`` and ``end_positions`` are not ``None``, outputs the total_loss which is the sum of the
1298
            CrossEntropy loss for the start and end token positions.
1299
            If ``start_positions`` or ``end_positions`` is ``None``, outputs a tuple of start_logits, end_logits which are the
1300
1301
1302
1303
1304
1305
1306
1307
1308
1309
1310
            logits respectively for the start and end position tokens of shape [batch_size, sequence_length].

        Example::

            # Already been converted into WordPiece token ids
            input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])
            input_mask = torch.LongTensor([[1, 1, 1], [1, 1, 0]])
            token_type_ids = torch.LongTensor([[0, 0, 1], [0, 1, 0]])

            start_logits, end_logits = model(input_ids, token_type_ids, input_mask)
        """
thomwolf's avatar
thomwolf committed
1311
1312
1313
        outputs = self.bert(input_ids, token_type_ids, attention_mask, head_mask=head_mask)
        sequence_output = outputs[0]

thomwolf's avatar
thomwolf committed
1314
1315
1316
1317
1318
        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)

1319
        outputs = (start_logits, end_logits,) + outputs[2:]
thomwolf's avatar
thomwolf committed
1320
1321
1322
1323
1324
1325
1326
1327
1328
1329
1330
1331
1332
1333
1334
        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
1335
            outputs = (total_loss,) + outputs
thomwolf's avatar
thomwolf committed
1336
1337

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