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

18
19
from __future__ import absolute_import, division, print_function, unicode_literals

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

import torch
import torch.nn as nn
thomwolf's avatar
thomwolf committed
30
from torch.nn import CrossEntropyLoss
thomwolf's avatar
thomwolf committed
31
32
from torch.nn.parameter import Parameter

33
from .modeling_utils import (Conv1D, CONFIG_NAME, WEIGHTS_NAME, PretrainedConfig,
34
                             PreTrainedModel, prune_conv1d_layer, SequenceSummary)
thomwolf's avatar
thomwolf committed
35
from .modeling_bert import BertLayerNorm as LayerNorm
thomwolf's avatar
thomwolf committed
36

thomwolf's avatar
thomwolf committed
37
38
logger = logging.getLogger(__name__)

39
40
OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_MAP = {"openai-gpt": "https://s3.amazonaws.com/models.huggingface.co/bert/openai-gpt-pytorch_model.bin"}
OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP = {"openai-gpt": "https://s3.amazonaws.com/models.huggingface.co/bert/openai-gpt-config.json"}
41

42

43
def load_tf_weights_in_openai_gpt(model, config, openai_checkpoint_folder_path):
44
45
    """ Load tf pre-trained weights in a pytorch model (from NumPy arrays here)
    """
46
47
    import re
    import numpy as np
48
49
50
51
52
53

    if '.ckpt' in openai_checkpoint_folder_path:
        openai_checkpoint_folder_path = os.path.dirname(openai_checkpoint_folder_path)

    logger.info("Loading weights from {}".format(openai_checkpoint_folder_path))

54
55
56
57
58
59
60
    names = json.load(open(openai_checkpoint_folder_path + '/parameters_names.json', "r", encoding='utf-8'))
    shapes = json.load(open(openai_checkpoint_folder_path + '/params_shapes.json', "r", encoding='utf-8'))
    offsets = np.cumsum([np.prod(shape) for shape in shapes])
    init_params = [np.load(openai_checkpoint_folder_path + '/params_{}.npy'.format(n)) for n in range(10)]
    init_params = np.split(np.concatenate(init_params, 0), offsets)[:-1]
    init_params = [param.reshape(shape) for param, shape in zip(init_params, shapes)]

thomwolf's avatar
thomwolf committed
61
    # This was used when we had a single embedding matrix for positions and tokens
62
63
    # init_params[0] = np.concatenate([init_params[1], init_params[0]], 0)
    # del init_params[1]
64
65
66
    init_params = [arr.squeeze() for arr in init_params]

    try:
67
68
        assert model.tokens_embed.weight.shape == init_params[1].shape
        assert model.positions_embed.weight.shape == init_params[0].shape
69
    except AssertionError as e:
70
71
        e.args += (model.tokens_embed.weight.shape, init_params[1].shape)
        e.args += (model.positions_embed.weight.shape, init_params[0].shape)
72
73
        raise

74
75
    model.tokens_embed.weight.data = torch.from_numpy(init_params[1])
    model.positions_embed.weight.data = torch.from_numpy(init_params[0])
76
    names.pop(0)
77
78
    # Pop position and token embedding arrays
    init_params.pop(0)
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
    init_params.pop(0)

    for name, array in zip(names, init_params): # names[1:n_transfer], init_params[1:n_transfer]):
        name = name[6:]  # skip "model/"
        assert name[-2:] == ":0"
        name = name[:-2]
        name = name.split('/')
        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] == 'g':
                pointer = getattr(pointer, 'weight')
            elif l[0] == 'b':
                pointer = getattr(pointer, 'bias')
            elif l[0] == 'w':
                pointer = getattr(pointer, 'weight')
            else:
                pointer = getattr(pointer, l[0])
            if len(l) >= 2:
                num = int(l[1])
                pointer = pointer[num]
        try:
            assert pointer.shape == array.shape
        except AssertionError as e:
            e.args += (pointer.shape, array.shape)
            raise
        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
117
118
119
120
121
122
123
124
125

def gelu(x):
    return 0.5 * x * (1 + torch.tanh(math.sqrt(2 / math.pi) * (x + 0.044715 * torch.pow(x, 3))))


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


126
127
ACT_FNS = {"relu": nn.ReLU, "swish": swish, "gelu": gelu}

thomwolf's avatar
thomwolf committed
128

129
class OpenAIGPTConfig(PretrainedConfig):
thomwolf's avatar
thomwolf committed
130
131
    """Configuration class to store the configuration of a `OpenAIGPTModel`.
    """
132
    pretrained_config_archive_map = OPENAI_GPT_PRETRAINED_CONFIG_ARCHIVE_MAP
133
134
135
136
137

    def __init__(
        self,
        vocab_size_or_config_json_file=40478,
        n_special=0,
thomwolf's avatar
thomwolf committed
138
        n_positions=512,
139
140
141
142
143
144
145
146
        n_ctx=512,
        n_embd=768,
        n_layer=12,
        n_head=12,
        afn="gelu",
        resid_pdrop=0.1,
        embd_pdrop=0.1,
        attn_pdrop=0.1,
147
        layer_norm_epsilon=1e-5,
148
        initializer_range=0.02,
thomwolf's avatar
thomwolf committed
149
        predict_special_tokens=True,
thomwolf's avatar
thomwolf committed
150
151
152
        summary_type='token_ids',
        summary_use_proj=True,
        summary_activation=None,
153
        summary_first_dropout=0.1,
thomwolf's avatar
thomwolf committed
154
        **kwargs
155
    ):
thomwolf's avatar
thomwolf committed
156
157
158
159
160
        """Constructs OpenAIGPTConfig.

        Args:
            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `OpenAIGPTModel` or a configuration json file.
            n_special: The number of special tokens to learn during fine-tuning ('[SEP]', '[CLF]', ...)
thomwolf's avatar
thomwolf committed
161
162
            n_positions: Number of positional embeddings.
            n_ctx: Size of the causal mask (usually same as n_positions).
thomwolf's avatar
thomwolf committed
163
164
165
166
167
168
169
170
171
172
173
            n_embd: Dimensionality of the embeddings and hidden states.
            n_layer: Number of hidden layers in the Transformer encoder.
            n_head: Number of attention heads for each attention layer in
                the Transformer encoder.
            afn: The non-linear activation function (function or string) in the
                encoder and pooler. If string, "gelu", "relu" and "swish" are supported.
            resid_pdrop: The dropout probabilitiy for all fully connected
                layers in the embeddings, encoder, and pooler.
            attn_pdrop: The dropout ratio for the attention
                probabilities.
            embd_pdrop: The dropout ratio for the embeddings.
174
            layer_norm_epsilon: epsilon to use in the layer norm layers
thomwolf's avatar
thomwolf committed
175
176
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
177
            predict_special_tokens: should we predict special tokens (when the model has a LM head)
thomwolf's avatar
thomwolf committed
178
        """
thomwolf's avatar
thomwolf committed
179
180
        super(OpenAIGPTConfig, self).__init__(**kwargs)

thomwolf's avatar
thomwolf committed
181
182
        if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
                        and isinstance(vocab_size_or_config_json_file, unicode)):
183
            with open(vocab_size_or_config_json_file, "r", encoding="utf-8") as reader:
thomwolf's avatar
thomwolf committed
184
185
186
187
188
189
190
                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.n_special = n_special
            self.n_ctx = n_ctx
thomwolf's avatar
thomwolf committed
191
            self.n_positions = n_positions
thomwolf's avatar
thomwolf committed
192
193
194
195
196
197
198
            self.n_embd = n_embd
            self.n_layer = n_layer
            self.n_head = n_head
            self.afn = afn
            self.resid_pdrop = resid_pdrop
            self.embd_pdrop = embd_pdrop
            self.attn_pdrop = attn_pdrop
199
            self.layer_norm_epsilon = layer_norm_epsilon
thomwolf's avatar
thomwolf committed
200
            self.initializer_range = initializer_range
201
            self.predict_special_tokens = predict_special_tokens
thomwolf's avatar
thomwolf committed
202
203
204
            self.summary_type = summary_type
            self.summary_use_proj = summary_use_proj
            self.summary_activation = summary_activation
205
            self.summary_first_dropout = summary_first_dropout
thomwolf's avatar
thomwolf committed
206
        else:
207
208
209
210
            raise ValueError(
                "First argument must be either a vocabulary size (int)"
                "or the path to a pretrained model config file (str)"
            )
thomwolf's avatar
thomwolf committed
211
212

    @property
213
214
    def total_tokens_embeddings(self):
        return self.vocab_size + self.n_special
thomwolf's avatar
thomwolf committed
215

thomwolf's avatar
thomwolf committed
216
217
218
219
220
221
222
223
224
225
226
227
    @property
    def hidden_size(self):
        return self.n_embd

    @property
    def num_attention_heads(self):
        return self.n_head

    @property
    def num_hidden_layers(self):
        return self.n_layer

thomwolf's avatar
thomwolf committed
228
229

class Attention(nn.Module):
thomwolf's avatar
thomwolf committed
230
    def __init__(self, nx, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
231
232
233
        super(Attention, self).__init__()
        n_state = nx  # in Attention: n_state=768 (nx=n_embd)
        # [switch nx => n_state from Block to Attention to keep identical to TF implem]
234
        assert n_state % config.n_head == 0
thomwolf's avatar
thomwolf committed
235
        self.register_buffer("bias", torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
236
        self.n_head = config.n_head
thomwolf's avatar
thomwolf committed
237
238
        self.split_size = n_state
        self.scale = scale
239

thomwolf's avatar
thomwolf committed
240
        self.output_attentions = config.output_attentions
241

242
243
        self.c_attn = Conv1D(n_state * 3, nx)
        self.c_proj = Conv1D(n_state, nx)
244
245
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
246

247
    def prune_heads(self, heads):
thomwolf's avatar
thomwolf committed
248
249
        if len(heads) == 0:
            return
250
251
252
253
254
255
256
257
258
259
260
261
262
263
        mask = torch.ones(self.n_head, self.split_size // self.n_head)
        for head in heads:
            mask[head] = 0
        mask = mask.view(-1).contiguous().eq(1)
        index = torch.arange(len(mask))[mask].long()
        index_attn = torch.cat([index, index + self.split_size, index + (2*self.split_size)])
        # Prune conv1d layers
        self.c_attn = prune_conv1d_layer(self.c_attn, index_attn, dim=1)
        self.c_proj = prune_conv1d_layer(self.c_proj, index, dim=0)
        # Update hyper params
        self.split_size = (self.split_size // self.n_head) * (self.n_head - len(heads))
        self.n_head = self.n_head - len(heads)

    def _attn(self, q, k, v, head_mask=None):
thomwolf's avatar
thomwolf committed
264
265
266
        w = torch.matmul(q, k)
        if self.scale:
            w = w / math.sqrt(v.size(-1))
thomwolf's avatar
thomwolf committed
267
        # w = w * self.bias + -1e9 * (1 - self.bias)  # TF implem method: mask_attn_weights
thomwolf's avatar
thomwolf committed
268
        # XD: self.b may be larger than w, so we need to crop it
thomwolf's avatar
thomwolf committed
269
        b = self.bias[:, :, : w.size(-2), : w.size(-1)]
thomwolf's avatar
thomwolf committed
270
271
        w = w * b + -1e9 * (1 - b)

thomwolf's avatar
thomwolf committed
272
273
        w = nn.Softmax(dim=-1)(w)
        w = self.attn_dropout(w)
274
275
276
277
278

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

thomwolf's avatar
thomwolf committed
279
        outputs = [torch.matmul(w, v)]
thomwolf's avatar
thomwolf committed
280
        if self.output_attentions:
thomwolf's avatar
thomwolf committed
281
282
            outputs.append(w)
        return outputs
thomwolf's avatar
thomwolf committed
283
284
285
286
287
288
289
290
291
292
293
294
295
296

    def merge_heads(self, x):
        x = x.permute(0, 2, 1, 3).contiguous()
        new_x_shape = x.size()[:-2] + (x.size(-2) * x.size(-1),)
        return x.view(*new_x_shape)  # in Tensorflow implem: fct merge_states

    def split_heads(self, x, k=False):
        new_x_shape = x.size()[:-1] + (self.n_head, x.size(-1) // self.n_head)
        x = x.view(*new_x_shape)  # in Tensorflow implem: fct split_states
        if k:
            return x.permute(0, 2, 3, 1)
        else:
            return x.permute(0, 2, 1, 3)

297
    def forward(self, x, head_mask=None):
thomwolf's avatar
thomwolf committed
298
299
300
301
302
        x = self.c_attn(x)
        query, key, value = x.split(self.split_size, dim=2)
        query = self.split_heads(query)
        key = self.split_heads(key, k=True)
        value = self.split_heads(value)
303

thomwolf's avatar
thomwolf committed
304
305
        attn_outputs = self._attn(query, key, value, head_mask)
        a = attn_outputs[0]
306

thomwolf's avatar
thomwolf committed
307
308
309
        a = self.merge_heads(a)
        a = self.c_proj(a)
        a = self.resid_dropout(a)
thomwolf's avatar
thomwolf committed
310
311
312

        outputs = [a] + attn_outputs[1:]
        return outputs  # a, (attentions)
thomwolf's avatar
thomwolf committed
313
314
315


class MLP(nn.Module):
316
    def __init__(self, n_state, config):  # in MLP: n_state=3072 (4 * n_embd)
thomwolf's avatar
thomwolf committed
317
        super(MLP, self).__init__()
318
        nx = config.n_embd
319
320
        self.c_fc = Conv1D(n_state, nx)
        self.c_proj = Conv1D(nx, n_state)
321
322
        self.act = ACT_FNS[config.afn]
        self.dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
323
324
325
326
327
328
329
330

    def forward(self, x):
        h = self.act(self.c_fc(x))
        h2 = self.c_proj(h)
        return self.dropout(h2)


class Block(nn.Module):
thomwolf's avatar
thomwolf committed
331
    def __init__(self, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
332
        super(Block, self).__init__()
333
        nx = config.n_embd
thomwolf's avatar
thomwolf committed
334
        self.attn = Attention(nx, n_ctx, config, scale)
335
        self.ln_1 = LayerNorm(nx, eps=config.layer_norm_epsilon)
336
        self.mlp = MLP(4 * nx, config)
337
        self.ln_2 = LayerNorm(nx, eps=config.layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
338

339
    def forward(self, x, head_mask=None):
thomwolf's avatar
thomwolf committed
340
341
342
        attn_outputs = self.attn(x, head_mask=head_mask)
        a = attn_outputs[0]

thomwolf's avatar
thomwolf committed
343
344
345
        n = self.ln_1(x + a)
        m = self.mlp(n)
        h = self.ln_2(n + m)
thomwolf's avatar
thomwolf committed
346
347
348

        outputs = [h] + attn_outputs[1:]
        return outputs
thomwolf's avatar
thomwolf committed
349
350


thomwolf's avatar
thomwolf committed
351
class OpenAIGPTLMHead(nn.Module):
thomwolf's avatar
thomwolf committed
352
353
    """ Language Model Head for the transformer """

354
    def __init__(self, model_embeddings_weights, config):
thomwolf's avatar
thomwolf committed
355
        super(OpenAIGPTLMHead, self).__init__()
356
        self.n_embd = config.n_embd
357
358
        self.vocab_size = config.vocab_size
        self.predict_special_tokens = config.predict_special_tokens
359
        self.torchscript = config.torchscript
thomwolf's avatar
thomwolf committed
360
361
        embed_shape = model_embeddings_weights.shape
        self.decoder = nn.Linear(embed_shape[1], embed_shape[0], bias=False)
thomwolf's avatar
thomwolf committed
362
363
        self.set_embeddings_weights(model_embeddings_weights)

364
365
    def set_embeddings_weights(self, model_embeddings_weights, predict_special_tokens=True):
        self.predict_special_tokens = predict_special_tokens
366
367
368
369
370

        if self.torchscript:
            self.decoder.weight = nn.Parameter(model_embeddings_weights.clone())
        else:
            self.decoder.weight = model_embeddings_weights  # Tied weights
thomwolf's avatar
thomwolf committed
371

thomwolf's avatar
thomwolf committed
372
373
    def forward(self, hidden_state):
        lm_logits = self.decoder(hidden_state)
374
375
        if not self.predict_special_tokens:
            lm_logits = lm_logits[..., :self.vocab_size]
thomwolf's avatar
thomwolf committed
376
377
378
        return lm_logits


379
class OpenAIGPTPreTrainedModel(PreTrainedModel):
thomwolf's avatar
thomwolf committed
380
381
382
    """ An abstract class to handle weights initialization and
        a simple interface for dowloading and loading pretrained models.
    """
383
    config_class = OpenAIGPTConfig
384
    pretrained_model_archive_map = OPENAI_GPT_PRETRAINED_MODEL_ARCHIVE_MAP
385
386
    load_tf_weights = load_tf_weights_in_openai_gpt
    base_model_prefix = "transformer"
387

388
389
390
    def __init__(self, *inputs, **kwargs):
        super(OpenAIGPTPreTrainedModel, self).__init__(*inputs, **kwargs)

thomwolf's avatar
thomwolf committed
391
392
393
    def init_weights(self, module):
        """ Initialize the weights.
        """
394
        if isinstance(module, (nn.Linear, nn.Embedding, Conv1D)):
thomwolf's avatar
thomwolf committed
395
396
397
            # 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)
398
399
            if isinstance(module, (nn.Linear, Conv1D)) and module.bias is not None:
                module.bias.data.zero_()
thomwolf's avatar
thomwolf committed
400
401
402
        elif isinstance(module, LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)
thomwolf's avatar
thomwolf committed
403

thomwolf's avatar
thomwolf committed
404
    @classmethod
405
    def from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs):
thomwolf's avatar
thomwolf committed
406
407
408
409
410
        """
        Instantiate a OpenAIGPTPreTrainedModel from a pre-trained model file or a pytorch state dict.
        Download and cache the pre-trained model file if needed.

        Params:
thomwolf's avatar
thomwolf committed
411
            pretrained_model_name_or_path: either:
thomwolf's avatar
thomwolf committed
412
413
                - a str with the name of a pre-trained model to load selected in the list of:
                - a path or url to a pretrained model archive containing:
414
                    . `config.json` a configuration file for the model
thomwolf's avatar
thomwolf committed
415
                    . `pytorch_model.bin` a PyTorch dump of a OpenAIGPTModel instance
416
                - a path or url to a pretrained model archive containing:
417
                    . `config.json` a configuration file for the model
418
419
                    . a series of NumPy files containing OpenAI TensorFlow trained weights
            from_tf: should we load the weights from a locally saved TensorFlow checkpoint
thomwolf's avatar
thomwolf committed
420
421
            cache_dir: an optional path to a folder in which the pre-trained models will be cached.
            state_dict: an optional state dictionnary (collections.OrderedDict object) to use instead of pre-trained models
422
            *inputs, **kwargs: additional input for the specific OpenAI-GPT class
thomwolf's avatar
thomwolf committed
423
        """
424
425
426
427
        num_special_tokens = kwargs.get('num_special_tokens', None)
        kwargs.pop('num_special_tokens', None)

        model = PreTrainedModel.from_pretrained(cls, pretrained_model_name_or_path, *inputs, **kwargs)
428

thomwolf's avatar
thomwolf committed
429
        # Add additional embeddings for special tokens if needed
430
        # This step also make sure we are still sharing the output and input embeddings after loading weights
431
        model.set_num_special_tokens(num_special_tokens)
thomwolf's avatar
thomwolf committed
432
        return model
thomwolf's avatar
thomwolf committed
433
434


thomwolf's avatar
thomwolf committed
435
class OpenAIGPTModel(OpenAIGPTPreTrainedModel):
436
437
    """OpenAI GPT model ("Improving Language Understanding by Generative Pre-Training").

438
439
440
441
442
443
    OpenAI GPT use a single embedding matrix to store the word and special embeddings.
    Special tokens embeddings are additional tokens that are not pre-trained: [SEP], [CLS]...
    Special tokens need to be trained during the fine-tuning if you use them.
    The number of special embeddings can be controled using the `set_num_special_tokens(num_special_tokens)` function.

    The embeddings are ordered as follow in the token embeddings matrice:
444
445
446
447
448
        [0,                                                         ----------------------
         ...                                                        -> word embeddings
         config.vocab_size - 1,                                     ______________________
         config.vocab_size,
         ...                                                        -> special embeddings
449
         config.vocab_size + config.n_special - 1]                  ______________________
450

451
452
    where total_tokens_embeddings can be obtained as config.total_tokens_embeddings and is:
        total_tokens_embeddings = config.vocab_size + config.n_special
453
454
455
    You should use the associate indices to index the embeddings.

    Params:
456
457
458
459
        `config`: a OpenAIGPTConfig 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
        `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
            This can be used to compute head importance metrics. Default: False
460
461
462

    Inputs:
        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] (or more generally [d_1, ..., d_n, sequence_length]
463
            were d_1 ... d_n are arbitrary dimensions) with the word BPE token indices selected in the range [0, total_tokens_embeddings[
464
        `position_ids`: an optional torch.LongTensor with the same shape as input_ids
465
            with the position indices (selected in the range [0, config.n_positions - 1[.
466
        `token_type_ids`: an optional torch.LongTensor with the same shape as input_ids
467
468
469
470
            You can use it to add a third type of embedding to each input token in the sequence
            (the previous two being the word and position embeddings).
            The input, position and token_type embeddings are summed inside the Transformer before the first
            self-attention block.
471
472
        `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 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.
473
474

    Outputs:
475
476
        `hidden_states`: a list of all the encoded-hidden-states in the model (length of the list: number of layers + 1 for the output of the embeddings)
            as torch.FloatTensor of size [batch_size, sequence_length, hidden_size]
477
478
479
480
481
482
483
484
485
486
487
488
489
            (or more generally [d_1, ..., d_n, hidden_size] were d_1 ... d_n are the dimension of input_ids)

    Example usage:
    ```python
    # Already been converted into BPE token ids
    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])

    config = modeling_openai.OpenAIGPTConfig()

    model = modeling_openai.OpenAIGPTModel(config)
    hidden_states = model(input_ids)
    ```
    """
490

thomwolf's avatar
thomwolf committed
491
    def __init__(self, config):
492
        super(OpenAIGPTModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
493
494
495
        self.output_attentions = config.output_attentions
        self.output_hidden_states = config.output_hidden_states

thomwolf's avatar
thomwolf committed
496
        self.tokens_embed = nn.Embedding(config.total_tokens_embeddings, config.n_embd)
497
        self.positions_embed = nn.Embedding(config.n_positions, config.n_embd)
498
        self.drop = nn.Dropout(config.embd_pdrop)
499
        self.h = nn.ModuleList([Block(config.n_ctx, config, scale=True) for _ in range(config.n_layer)])
thomwolf's avatar
thomwolf committed
500

thomwolf's avatar
thomwolf committed
501
502
        self.apply(self.init_weights)

503
    def set_num_special_tokens(self, num_special_tokens=None):
504
        " Update input embeddings with new embedding matrice if needed "
505
        if num_special_tokens is None or self.config.n_special == num_special_tokens:
506
            return
thomwolf's avatar
thomwolf committed
507
508
        # Update config
        self.config.n_special = num_special_tokens
thomwolf's avatar
thomwolf committed
509
        # Build new embeddings and initialize all new embeddings (in particular the special tokens)
510
        old_embed = self.tokens_embed
511
        self.tokens_embed = nn.Embedding(self.config.total_tokens_embeddings, self.config.n_embd)
thomwolf's avatar
thomwolf committed
512
        self.tokens_embed.to(old_embed.weight.device)
513
        self.init_weights(self.tokens_embed)
thomwolf's avatar
thomwolf committed
514
515
        # Copy word embeddings from the previous weights
        self.tokens_embed.weight.data[:self.config.vocab_size, :] = old_embed.weight.data[:self.config.vocab_size, :]
thomwolf's avatar
thomwolf committed
516

thomwolf's avatar
thomwolf committed
517
    def _prune_heads(self, heads_to_prune):
518
519
520
521
522
523
524
        """ Prunes heads of the model.
            heads_to_prune: dict of {layer_num: list of heads to prune in this layer}
        """
        for layer, heads in heads_to_prune.items():
            self.h[layer].attn.prune_heads(heads)

    def forward(self, input_ids, position_ids=None, token_type_ids=None, head_mask=None):
thomwolf's avatar
thomwolf committed
525
        if position_ids is None:
526
527
528
529
530
            # This was used when we had a single embedding matrice from position and token embeddings
            # start = self.config.vocab_size + self.config.n_special
            # end = start + input_ids.size(-1)
            # position_ids = torch.arange(start, end, dtype=torch.long, device=input_ids.device)
            position_ids = torch.arange(input_ids.size(-1), dtype=torch.long, device=input_ids.device)
thomwolf's avatar
thomwolf committed
531
532
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)

533
        # Prepare head mask if needed
thomwolf's avatar
thomwolf committed
534
        # 1.0 in head_mask indicate we keep the head
535
        # attention_probs has shape bsz x n_heads x N x N
536
        # head_mask has shape n_layer x batch x n_heads x N x N
537
538
        if head_mask is not None:
            if head_mask.dim() == 1:
539
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
thomwolf's avatar
thomwolf committed
540
                head_mask = head_mask.expand(self.config.n_layer, -1, -1, -1, -1)
541
            elif head_mask.dim() == 2:
542
                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
543
            head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility
544
545
        else:
            head_mask = [None] * self.config.n_layer
546

thomwolf's avatar
thomwolf committed
547
548
549
550
        input_shape = input_ids.size()
        input_ids = input_ids.view(-1, input_ids.size(-1))
        position_ids = position_ids.view(-1, position_ids.size(-1))

551
552
        inputs_embeds = self.tokens_embed(input_ids)
        position_embeds = self.positions_embed(position_ids)
thomwolf's avatar
thomwolf committed
553
554
        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1))
555
            token_type_embeds = self.tokens_embed(token_type_ids)
thomwolf's avatar
thomwolf committed
556
557
558
        else:
            token_type_embeds = 0
        hidden_states = inputs_embeds + position_embeds + token_type_embeds
559
560
        hidden_states = self.drop(hidden_states)

561
562
        output_shape = input_shape + (hidden_states.size(-1),)

563
564
        all_attentions = ()
        all_hidden_states = ()
565
        for i, block in enumerate(self.h):
thomwolf's avatar
thomwolf committed
566
            if self.output_hidden_states:
567
                all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)
thomwolf's avatar
thomwolf committed
568

569
            outputs = block(hidden_states, head_mask[i])
thomwolf's avatar
thomwolf committed
570
            hidden_states = outputs[0]
thomwolf's avatar
thomwolf committed
571
            if self.output_attentions:
572
                all_attentions = all_attentions + (outputs[1],)
thomwolf's avatar
thomwolf committed
573
574
575

        # Add last layer
        if self.output_hidden_states:
576
            all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)
577

578
        outputs = (hidden_states.view(*output_shape),)
thomwolf's avatar
thomwolf committed
579
        if self.output_hidden_states:
580
            outputs = outputs + (all_hidden_states,)
thomwolf's avatar
thomwolf committed
581
        if self.output_attentions:
582
            outputs = outputs + (all_attentions,)
thomwolf's avatar
thomwolf committed
583
        return outputs  # last hidden state, (all hidden states), (all attentions)
thomwolf's avatar
thomwolf committed
584

585

thomwolf's avatar
thomwolf committed
586
class OpenAIGPTLMHeadModel(OpenAIGPTPreTrainedModel):
587
588
    """OpenAI GPT model with a Language Modeling head ("Improving Language Understanding by Generative Pre-Training").

589
590
591
592
593
594
    OpenAI GPT use a single embedding matrix to store the word and special embeddings.
    Special tokens embeddings are additional tokens that are not pre-trained: [SEP], [CLS]...
    Special tokens need to be trained during the fine-tuning if you use them.
    The number of special embeddings can be controled using the `set_num_special_tokens(num_special_tokens)` function.

    The embeddings are ordered as follow in the token embeddings matrice:
595
596
597
598
599
        [0,                                                         ----------------------
         ...                                                        -> word embeddings
         config.vocab_size - 1,                                     ______________________
         config.vocab_size,
         ...                                                        -> special embeddings
600
         config.vocab_size + config.n_special - 1]                  ______________________
601

602
603
604
    where total_tokens_embeddings can be obtained as config.total_tokens_embeddings and is:
        total_tokens_embeddings = config.vocab_size + config.n_special
    You should use the associate indices to index the embeddings.
605
606

    Params:
607
608
609
610
        `config`: a OpenAIGPTConfig 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
        `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
            This can be used to compute head importance metrics. Default: False
611
612
613

    Inputs:
        `input_ids`: a torch.LongTensor of shape [batch_size, sequence_length] (or more generally [d_1, ..., d_n, sequence_length]
614
            were d_1 ... d_n are arbitrary dimensions) with the word BPE token indices selected in the range [0, total_tokens_embeddings[
615
        `position_ids`: an optional torch.LongTensor with the same shape as input_ids
616
            with the position indices (selected in the range [0, config.n_positions - 1[.
617
        `token_type_ids`: an optional torch.LongTensor with the same shape as input_ids
618
619
620
621
            You can use it to add a third type of embedding to each input token in the sequence
            (the previous two being the word and position embeddings).
            The input, position and token_type embeddings are summed inside the Transformer before the first
            self-attention block.
622
623
624
        `lm_labels`: optional language modeling labels: torch.LongTensor of shape [batch_size, sequence_length]
            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]
625
626
        `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 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.
627
628
629
630
631

    Outputs:
        if `lm_labels` is not `None`:
            Outputs the language modeling loss.
        else:
632
633
            `lm_logits`: the language modeling logits as a torch.FloatTensor of size [batch_size, sequence_length, total_tokens_embeddings]
                (or more generally [d_1, ..., d_n, total_tokens_embeddings] were d_1 ... d_n are the dimension of input_ids)
634
635
636
637
638
639
640
641
642
643
644
645

    Example usage:
    ```python
    # Already been converted into BPE token ids
    input_ids = torch.LongTensor([[31, 51, 99], [15, 5, 0]])

    config = modeling_openai.OpenAIGPTConfig()

    model = modeling_openai.OpenAIGPTLMHeadModel(config)
    lm_logits = model(input_ids)
    ```
    """
646

thomwolf's avatar
thomwolf committed
647
    def __init__(self, config):
648
        super(OpenAIGPTLMHeadModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
649
        self.transformer = OpenAIGPTModel(config)
650
        self.lm_head = OpenAIGPTLMHead(self.transformer.tokens_embed.weight, config)
thomwolf's avatar
thomwolf committed
651
652
        self.apply(self.init_weights)

653
    def set_num_special_tokens(self, num_special_tokens, predict_special_tokens=True):
654
655
656
        """ Update input and output embeddings with new embedding matrice
            Make sure we are sharing the embeddings
        """
657
        self.config.predict_special_tokens = self.transformer.config.predict_special_tokens = predict_special_tokens
thomwolf's avatar
thomwolf committed
658
        self.transformer.set_num_special_tokens(num_special_tokens)
659
        self.lm_head.set_embeddings_weights(self.transformer.tokens_embed.weight, predict_special_tokens=predict_special_tokens)
thomwolf's avatar
thomwolf committed
660

661
    def forward(self, input_ids, position_ids=None, token_type_ids=None, lm_labels=None, head_mask=None):
thomwolf's avatar
thomwolf committed
662
663
        transformer_outputs = self.transformer(input_ids, position_ids, token_type_ids, head_mask)
        hidden_states = transformer_outputs[0]
thomwolf's avatar
thomwolf committed
664
        lm_logits = self.lm_head(hidden_states)
thomwolf's avatar
thomwolf committed
665

666
        outputs = (lm_logits,) + transformer_outputs[1:]
thomwolf's avatar
thomwolf committed
667
        if lm_labels is not None:
668
            # Shift so that tokens < n predict n
thomwolf's avatar
thomwolf committed
669
670
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = lm_labels[..., 1:].contiguous()
Catalin Voss's avatar
Catalin Voss committed
671
            # Flatten the tokens
thomwolf's avatar
thomwolf committed
672
            loss_fct = CrossEntropyLoss(ignore_index=-1)
673
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)),
674
                            shift_labels.view(-1))
675
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
676
677

        return outputs  # (loss), lm_logits, (all hidden states), (all attentions)
thomwolf's avatar
thomwolf committed
678

679

thomwolf's avatar
thomwolf committed
680
class OpenAIGPTDoubleHeadsModel(OpenAIGPTPreTrainedModel):
thomwolf's avatar
thomwolf committed
681
    """OpenAI GPT model with a Language Modeling and a Multiple Choice head ("Improving Language Understanding by Generative Pre-Training").
682

683
684
685
686
687
688
    OpenAI GPT use a single embedding matrix to store the word and special embeddings.
    Special tokens embeddings are additional tokens that are not pre-trained: [SEP], [CLS]...
    Special tokens need to be trained during the fine-tuning if you use them.
    The number of special embeddings can be controled using the `set_num_special_tokens(num_special_tokens)` function.

    The embeddings are ordered as follow in the token embeddings matrice:
689
690
691
692
693
        [0,                                                         ----------------------
         ...                                                        -> word embeddings
         config.vocab_size - 1,                                     ______________________
         config.vocab_size,
         ...                                                        -> special embeddings
694
         config.vocab_size + config.n_special - 1]                  ______________________
695

696
697
698
    where total_tokens_embeddings can be obtained as config.total_tokens_embeddings and is:
        total_tokens_embeddings = config.vocab_size + config.n_special
    You should use the associate indices to index the embeddings.
699
700

    Params:
701
702
703
704
        `config`: a OpenAIGPTConfig 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
        `keep_multihead_output`: If True, saves output of the multi-head attention module with its gradient.
            This can be used to compute head importance metrics. Default: False
705
706

    Inputs:
thomwolf's avatar
thomwolf committed
707
708
709
710
        `input_ids`: a torch.LongTensor of shape [batch_size, num_choices, sequence_length] with the BPE token
            indices selected in the range [0, total_tokens_embeddings[
        `mc_token_ids`: a torch.LongTensor of shape [batch_size, num_choices] with the index of the token from
            which we should take the hidden state to feed the multiple choice classifier (usually last token of the sequence)
711
        `position_ids`: an optional torch.LongTensor with the same shape as input_ids
712
            with the position indices (selected in the range [0, config.n_positions - 1[.
713
        `token_type_ids`: an optional torch.LongTensor with the same shape as input_ids
714
715
716
717
            You can use it to add a third type of embedding to each input token in the sequence
            (the previous two being the word and position embeddings).
            The input, position and token_type embeddings are summed inside the Transformer before the first
            self-attention block.
718
        `lm_labels`: optional language modeling labels: torch.LongTensor of shape [batch_size, num_choices, sequence_length]
719
720
            with indices selected in [-1, 0, ..., total_tokens_embeddings]. All labels set to -1 are ignored (masked), the loss
            is only computed for the labels set in [0, ..., total_tokens_embeddings]
721
722
        `multiple_choice_labels`: optional multiple choice labels: torch.LongTensor of shape [batch_size]
            with indices selected in [0, ..., num_choices].
723
724
        `head_mask`: an optional torch.Tensor of shape [num_heads] or [num_layers, num_heads] with indices between 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.
725
726
727
728
729

    Outputs:
        if `lm_labels` and `multiple_choice_labels` are not `None`:
            Outputs a tuple of losses with the language modeling loss and the multiple choice loss.
        else: a tuple with
730
            `lm_logits`: the language modeling logits as a torch.FloatTensor of size [batch_size, num_choices, sequence_length, total_tokens_embeddings]
731
732
733
734
735
            `multiple_choice_logits`: the multiple choice logits as a torch.FloatTensor of size [batch_size, num_choices]

    Example usage:
    ```python
    # Already been converted into BPE token ids
thomwolf's avatar
thomwolf committed
736
737
    input_ids = torch.LongTensor([[[31, 51, 99], [15, 5, 0]]])  # (bsz, number of choice, seq length)
    mc_token_ids = torch.LongTensor([[2], [1]]) # (bsz, number of choice)
738
739
740

    config = modeling_openai.OpenAIGPTConfig()

VictorSanh's avatar
VictorSanh committed
741
    model = modeling_openai.OpenAIGPTDoubleHeadsModel(config)
thomwolf's avatar
thomwolf committed
742
    lm_logits, multiple_choice_logits = model(input_ids, mc_token_ids)
743
744
    ```
    """
745

thomwolf's avatar
thomwolf committed
746
    def __init__(self, config):
747
        super(OpenAIGPTDoubleHeadsModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
748

thomwolf's avatar
thomwolf committed
749
        self.transformer = OpenAIGPTModel(config)
750
        self.lm_head = OpenAIGPTLMHead(self.transformer.tokens_embed.weight, config)
thomwolf's avatar
thomwolf committed
751
752
        self.multiple_choice_head = SequenceSummary(config)

thomwolf's avatar
thomwolf committed
753
        self.apply(self.init_weights)
thomwolf's avatar
thomwolf committed
754

755
    def set_num_special_tokens(self, num_special_tokens, predict_special_tokens=True):
756
757
758
        """ Update input and output embeddings with new embedding matrice
            Make sure we are sharing the embeddings
        """
759
        self.config.predict_special_tokens = self.transformer.config.predict_special_tokens = predict_special_tokens
thomwolf's avatar
thomwolf committed
760
        self.transformer.set_num_special_tokens(num_special_tokens)
761
        self.lm_head.set_embeddings_weights(self.transformer.tokens_embed.weight, predict_special_tokens=predict_special_tokens)
thomwolf's avatar
thomwolf committed
762

thomwolf's avatar
thomwolf committed
763
    def forward(self, input_ids, mc_token_ids=None, lm_labels=None, mc_labels=None, token_type_ids=None,
764
                position_ids=None, head_mask=None):
thomwolf's avatar
thomwolf committed
765
766
        transformer_outputs = self.transformer(input_ids, position_ids, token_type_ids, head_mask)
        hidden_states = transformer_outputs[0]
767

thomwolf's avatar
thomwolf committed
768
        lm_logits = self.lm_head(hidden_states)
thomwolf's avatar
thomwolf committed
769
        mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids).squeeze(-1)
thomwolf's avatar
thomwolf committed
770

771
        outputs = (lm_logits, mc_logits) + transformer_outputs[1:]
thomwolf's avatar
thomwolf committed
772
773
774
775
        if mc_labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)),
                            mc_labels.view(-1))
776
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
777
        if lm_labels is not None:
thomwolf's avatar
thomwolf committed
778
779
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = lm_labels[..., 1:].contiguous()
thomwolf's avatar
thomwolf committed
780
            loss_fct = CrossEntropyLoss(ignore_index=-1)
thomwolf's avatar
thomwolf committed
781
782
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)),
                            shift_labels.view(-1))
783
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
784
785

        return outputs  # (lm loss), (mc loss), lm logits, mc logits, (all hidden_states), (attentions)