modeling_gpt2.py 33.3 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-2 model."""

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

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

import torch
import torch.nn as nn
from torch.nn import CrossEntropyLoss
from torch.nn.parameter import Parameter

33
34
from .configuration_gpt2 import GPT2Config
from .file_utils import add_start_docstrings
Aymeric Augustin's avatar
Aymeric Augustin committed
35
36
from .modeling_utils import Conv1D, PreTrainedModel, SequenceSummary, prune_conv1d_layer

thomwolf's avatar
thomwolf committed
37
38
39

logger = logging.getLogger(__name__)

40
41
42
43
44
45
46
47
GPT2_PRETRAINED_MODEL_ARCHIVE_MAP = {
    "gpt2": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-pytorch_model.bin",
    "gpt2-medium": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-medium-pytorch_model.bin",
    "gpt2-large": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-large-pytorch_model.bin",
    "gpt2-xl": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-xl-pytorch_model.bin",
    "distilgpt2": "https://s3.amazonaws.com/models.huggingface.co/bert/distilgpt2-pytorch_model.bin",
}

thomwolf's avatar
thomwolf committed
48

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

    for name, array in zip(names, arrays):
thomwolf's avatar
thomwolf committed
75
        name = name[6:]  # skip "model/"
76
        name = name.split("/")
thomwolf's avatar
thomwolf committed
77
78
        pointer = model
        for m_name in name:
79
80
            if re.fullmatch(r"[A-Za-z]+\d+", m_name):
                l = re.split(r"(\d+)", m_name)
thomwolf's avatar
thomwolf committed
81
82
            else:
                l = [m_name]
83
84
85
86
87
            if l[0] == "w" or l[0] == "g":
                pointer = getattr(pointer, "weight")
            elif l[0] == "b":
                pointer = getattr(pointer, "bias")
            elif l[0] == "wpe" or l[0] == "wte":
thomwolf's avatar
thomwolf committed
88
                pointer = getattr(pointer, l[0])
89
                pointer = getattr(pointer, "weight")
thomwolf's avatar
thomwolf committed
90
91
92
93
94
95
96
97
98
99
            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
thomwolf's avatar
thomwolf committed
100
        logger.info("Initialize PyTorch weight {}".format(name))
thomwolf's avatar
thomwolf committed
101
102
103
104
105
106
107
108
109
        pointer.data = torch.from_numpy(array)
    return model


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


class Attention(nn.Module):
thomwolf's avatar
thomwolf committed
110
    def __init__(self, nx, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
111
        super(Attention, self).__init__()
thomwolf's avatar
thomwolf committed
112
113
        self.output_attentions = config.output_attentions

thomwolf's avatar
thomwolf committed
114
115
116
117
118
119
120
        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]
        assert n_state % config.n_head == 0
        self.register_buffer("bias", torch.tril(torch.ones(n_ctx, n_ctx)).view(1, 1, n_ctx, n_ctx))
        self.n_head = config.n_head
        self.split_size = n_state
        self.scale = scale
121

thomwolf's avatar
thomwolf committed
122
123
        self.c_attn = Conv1D(n_state * 3, nx)
        self.c_proj = Conv1D(n_state, nx)
124
125
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
126
        self.pruned_heads = set()
thomwolf's avatar
thomwolf committed
127

128
    def prune_heads(self, heads):
thomwolf's avatar
thomwolf committed
129
130
        if len(heads) == 0:
            return
131
        mask = torch.ones(self.n_head, self.split_size // self.n_head)
132
        heads = set(heads) - self.pruned_heads  # Convert to set and emove already pruned heads
133
        for head in heads:
134
135
            # Compute how many pruned heads are before the head and move the index accordingly
            head = head - sum(1 if h < head else 0 for h in self.pruned_heads)
136
137
138
            mask[head] = 0
        mask = mask.view(-1).contiguous().eq(1)
        index = torch.arange(len(mask))[mask].long()
139
        index_attn = torch.cat([index, index + self.split_size, index + (2 * self.split_size)])
140

141
142
143
        # 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)
144

145
146
147
        # 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)
148
        self.pruned_heads = self.pruned_heads.union(heads)
149

150
    def _attn(self, q, k, v, attention_mask=None, head_mask=None):
thomwolf's avatar
thomwolf committed
151
152
153
        w = torch.matmul(q, k)
        if self.scale:
            w = w / math.sqrt(v.size(-1))
thomwolf's avatar
thomwolf committed
154
        nd, ns = w.size(-2), w.size(-1)
155
        b = self.bias[:, :, ns - nd : ns, :ns]
156
        w = w * b - 1e4 * (1 - b)
thomwolf's avatar
thomwolf committed
157

158
159
160
161
        if attention_mask is not None:
            # Apply the attention mask
            w = w + attention_mask

thomwolf's avatar
thomwolf committed
162
        w = nn.Softmax(dim=-1)(w)
163
        w = self.attn_dropout(w)
164
165
166
167
168

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

thomwolf's avatar
thomwolf committed
169
        outputs = [torch.matmul(w, v)]
thomwolf's avatar
thomwolf committed
170
        if self.output_attentions:
thomwolf's avatar
thomwolf committed
171
172
            outputs.append(w)
        return outputs
thomwolf's avatar
thomwolf committed
173
174
175
176
177
178
179
180
181
182

    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:
thomwolf's avatar
thomwolf committed
183
            return x.permute(0, 2, 3, 1)  # (batch, head, head_features, seq_length)
thomwolf's avatar
thomwolf committed
184
        else:
thomwolf's avatar
thomwolf committed
185
            return x.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
thomwolf's avatar
thomwolf committed
186

187
    def forward(self, x, layer_past=None, attention_mask=None, head_mask=None):
thomwolf's avatar
thomwolf committed
188
189
190
191
192
        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)
thomwolf's avatar
thomwolf committed
193
        if layer_past is not None:
thomwolf's avatar
thomwolf committed
194
            past_key, past_value = layer_past[0].transpose(-2, -1), layer_past[1]  # transpose back cf below
thomwolf's avatar
thomwolf committed
195
            key = torch.cat((past_key, key), dim=-1)
thomwolf's avatar
thomwolf committed
196
            value = torch.cat((past_value, value), dim=-2)
thomwolf's avatar
thomwolf committed
197
        present = torch.stack((key.transpose(-2, -1), value))  # transpose to have same shapes for stacking
198

199
        attn_outputs = self._attn(query, key, value, attention_mask, head_mask)
thomwolf's avatar
thomwolf committed
200
        a = attn_outputs[0]
201

thomwolf's avatar
thomwolf committed
202
203
        a = self.merge_heads(a)
        a = self.c_proj(a)
204
        a = self.resid_dropout(a)
thomwolf's avatar
thomwolf committed
205
206
207

        outputs = [a, present] + attn_outputs[1:]
        return outputs  # a, present, (attentions)
thomwolf's avatar
thomwolf committed
208
209
210
211
212
213
214
215
216


class MLP(nn.Module):
    def __init__(self, n_state, config):  # in MLP: n_state=3072 (4 * n_embd)
        super(MLP, self).__init__()
        nx = config.n_embd
        self.c_fc = Conv1D(n_state, nx)
        self.c_proj = Conv1D(nx, n_state)
        self.act = gelu
217
        self.dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
218
219
220
221

    def forward(self, x):
        h = self.act(self.c_fc(x))
        h2 = self.c_proj(h)
222
        return self.dropout(h2)
thomwolf's avatar
thomwolf committed
223
224
225


class Block(nn.Module):
thomwolf's avatar
thomwolf committed
226
    def __init__(self, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
227
228
        super(Block, self).__init__()
        nx = config.n_embd
229
        self.ln_1 = nn.LayerNorm(nx, eps=config.layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
230
        self.attn = Attention(nx, n_ctx, config, scale)
231
        self.ln_2 = nn.LayerNorm(nx, eps=config.layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
232
233
        self.mlp = MLP(4 * nx, config)

234
    def forward(self, x, layer_past=None, attention_mask=None, head_mask=None):
235
236
237
        output_attn = self.attn(
            self.ln_1(x), layer_past=layer_past, attention_mask=attention_mask, head_mask=head_mask
        )
thomwolf's avatar
thomwolf committed
238
239
        a = output_attn[0]  # output_attn: a, present, (attentions)

thomwolf's avatar
thomwolf committed
240
        x = x + a
thomwolf's avatar
thomwolf committed
241
        m = self.mlp(self.ln_2(x))
thomwolf's avatar
thomwolf committed
242
        x = x + m
thomwolf's avatar
thomwolf committed
243
244
245

        outputs = [x] + output_attn[1:]
        return outputs  # x, present, (attentions)
thomwolf's avatar
thomwolf committed
246
247


248
class GPT2PreTrainedModel(PreTrainedModel):
thomwolf's avatar
thomwolf committed
249
250
251
    """ An abstract class to handle weights initialization and
        a simple interface for dowloading and loading pretrained models.
    """
252

253
    config_class = GPT2Config
254
    pretrained_model_archive_map = GPT2_PRETRAINED_MODEL_ARCHIVE_MAP
255
256
    load_tf_weights = load_tf_weights_in_gpt2
    base_model_prefix = "transformer"
thomwolf's avatar
thomwolf committed
257

258
259
260
    def __init__(self, *inputs, **kwargs):
        super(GPT2PreTrainedModel, self).__init__(*inputs, **kwargs)

261
    def _init_weights(self, module):
thomwolf's avatar
thomwolf committed
262
263
        """ Initialize the weights.
        """
264
        if isinstance(module, (nn.Linear, nn.Embedding, Conv1D)):
thomwolf's avatar
thomwolf committed
265
266
267
            # 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)
268
269
            if isinstance(module, (nn.Linear, Conv1D)) and module.bias is not None:
                module.bias.data.zero_()
270
        elif isinstance(module, nn.LayerNorm):
thomwolf's avatar
thomwolf committed
271
272
273
274
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


thomwolf's avatar
thomwolf committed
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
GPT2_START_DOCSTRING = r"""    OpenAI GPT-2 model was proposed in
    `Language Models are Unsupervised Multitask Learners`_
    by Alec Radford*, Jeffrey Wu*, Rewon Child, David Luan, Dario Amodei** and Ilya Sutskever**.
    It's a causal (unidirectional) transformer pre-trained using  language modeling on a very large
    corpus of ~40 GB of text data.

    This model is a PyTorch `torch.nn.Module`_ sub-class. Use it as a regular PyTorch Module and
    refer to the PyTorch documentation for all matter related to general usage and behavior.

    .. _`Language Models are Unsupervised Multitask Learners`:
        https://openai.com/blog/better-language-models/

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

    Parameters:
291
        config (:class:`~transformers.GPT2Config`): Model configuration class with all the parameters of the model.
292
            Initializing with a config file does not load the weights associated with the model, only the configuration.
293
            Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights.
thomwolf's avatar
thomwolf committed
294
295
"""

thomwolf's avatar
thomwolf committed
296
GPT2_INPUTS_DOCSTRING = r"""    Inputs:
thomwolf's avatar
thomwolf committed
297
298
        **input_ids**: ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of input sequence tokens in the vocabulary.
thomwolf's avatar
thomwolf committed
299
300
            GPT-2 is a model with absolute position embeddings so it's usually advised to pad the inputs on
            the right rather than the left.
301
302
303
            Indices can be obtained using :class:`transformers.GPT2Tokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
            :func:`transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details.
thomwolf's avatar
thomwolf committed
304
305
306
        **past**:
            list of ``torch.FloatTensor`` (one for each layer):
            that contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
307
308
            (see `past` output below). Can be used to speed up sequential decoding. The token ids which have their past given to this model 
            should not be passed as input ids as they have already been computed.
309
310
311
312
313
314
315
316
317
318
319
        **attention_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length)``:
            Mask to avoid performing attention on padding token indices.
            Mask values selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
        **token_type_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            A parallel sequence of tokens (can be used to indicate various portions of the inputs).
            The embeddings from these tokens will be summed with the respective token embeddings.
            Indices are selected in the vocabulary (unlike BERT which has a specific vocabulary for segment indices).
        **position_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of positions of each input sequence tokens in the position embeddings.
            Selected in the range ``[0, config.max_position_embeddings - 1]``.
thomwolf's avatar
thomwolf committed
320
        **head_mask**: (`optional`) ``torch.FloatTensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:
thomwolf's avatar
thomwolf committed
321
            Mask to nullify selected heads of the self-attention modules.
thomwolf's avatar
thomwolf committed
322
            Mask values selected in ``[0, 1]``:
thomwolf's avatar
thomwolf committed
323
            ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.
Julien Chaumond's avatar
Julien Chaumond committed
324
325
        **inputs_embeds**: (`optional`) ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, embedding_dim)``:
            Optionally, instead of passing ``input_ids`` you can choose to directly pass an embedded representation.
326
327
            This is useful if you want more control over how to convert `input_ids` indices into associated vectors
            than the model's internal embedding lookup matrix.
thomwolf's avatar
thomwolf committed
328
329
"""

330
331
332
333
334
335

@add_start_docstrings(
    "The bare GPT2 Model transformer outputting raw hidden-states without any specific head on top.",
    GPT2_START_DOCSTRING,
    GPT2_INPUTS_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
336
class GPT2Model(GPT2PreTrainedModel):
337
    r"""
thomwolf's avatar
thomwolf committed
338
339
340
341
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **last_hidden_state**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, hidden_size)``
            Sequence of hidden-states at the last layer of the model.
        **past**:
342
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(2, batch_size, num_heads, sequence_length, embed_size_per_head)``:
thomwolf's avatar
thomwolf committed
343
            that contains pre-computed hidden-states (key and values in the attention blocks).
344
345
            Can be used (see `past` input) to speed up sequential decoding. The token ids which have their past given to this model 
            should not be passed as input ids as they have already been computed.
thomwolf's avatar
thomwolf committed
346
347
348
349
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
thomwolf's avatar
thomwolf committed
350
351
352
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
thomwolf's avatar
thomwolf committed
353
354
355

    Examples::

356
357
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        model = GPT2Model.from_pretrained('gpt2')
358
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
359
360
        outputs = model(input_ids)
        last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple
thomwolf's avatar
thomwolf committed
361
362

    """
363

thomwolf's avatar
thomwolf committed
364
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
365
        super(GPT2Model, self).__init__(config)
thomwolf's avatar
thomwolf committed
366
367
        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions
368
        self.output_past = config.output_past
thomwolf's avatar
thomwolf committed
369

thomwolf's avatar
thomwolf committed
370
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
thomwolf's avatar
thomwolf committed
371
        self.wpe = nn.Embedding(config.n_positions, config.n_embd)
372
        self.drop = nn.Dropout(config.embd_pdrop)
373
        self.h = nn.ModuleList([Block(config.n_ctx, config, scale=True) for _ in range(config.n_layer)])
374
        self.ln_f = nn.LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
375

376
        self.init_weights()
thomwolf's avatar
thomwolf committed
377

thomwolf's avatar
thomwolf committed
378
    def get_input_embeddings(self):
thomwolf's avatar
thomwolf committed
379
        return self.wte
thomwolf's avatar
thomwolf committed
380

thomwolf's avatar
thomwolf committed
381
    def set_input_embeddings(self, new_embeddings):
382
383
        self.wte = new_embeddings

thomwolf's avatar
thomwolf committed
384
    def _prune_heads(self, heads_to_prune):
385
386
387
388
389
390
        """ 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)

391
392
393
394
395
396
397
398
399
400
    def forward(
        self,
        input_ids=None,
        past=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
    ):
Julien Chaumond's avatar
Julien Chaumond committed
401
402
403
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
404
405
406
407
408
409
410
            input_shape = input_ids.size()
            input_ids = input_ids.view(-1, input_shape[-1])
        elif inputs_embeds is not None:
            input_shape = inputs_embeds.size()[:-1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")

411
412
413
414
415
        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, input_shape[-1])
        if position_ids is not None:
            position_ids = position_ids.view(-1, input_shape[-1])

thomwolf's avatar
thomwolf committed
416
        if past is None:
thomwolf's avatar
thomwolf committed
417
            past_length = 0
thomwolf's avatar
thomwolf committed
418
            past = [None] * len(self.h)
thomwolf's avatar
thomwolf committed
419
        else:
thomwolf's avatar
thomwolf committed
420
            past_length = past[0][0].size(-2)
thomwolf's avatar
thomwolf committed
421
        if position_ids is None:
422
423
424
            device = input_ids.device if input_ids is not None else inputs_embeds.device
            position_ids = torch.arange(past_length, input_shape[-1] + past_length, dtype=torch.long, device=device)
            position_ids = position_ids.unsqueeze(0).view(-1, input_shape[-1])
thomwolf's avatar
thomwolf committed
425

426
427
        # Attention mask.
        if attention_mask is not None:
428
            attention_mask = attention_mask.view(-1, input_shape[-1])
429
430
431
432
433
434
435
436
437
438
439
440
            # 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.
            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.
441
            attention_mask = attention_mask.to(dtype=next(self.parameters()).dtype)  # fp16 compatibility
442
443
            attention_mask = (1.0 - attention_mask) * -10000.0

444
        # Prepare head mask if needed
thomwolf's avatar
thomwolf committed
445
        # 1.0 in head_mask indicate we keep the head
446
        # attention_probs has shape bsz x n_heads x N x N
447
        # head_mask has shape n_layer x batch x n_heads x N x N
448
449
        if head_mask is not None:
            if head_mask.dim() == 1:
450
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
thomwolf's avatar
thomwolf committed
451
                head_mask = head_mask.expand(self.config.n_layer, -1, -1, -1, -1)
452
            elif head_mask.dim() == 2:
453
454
455
456
457
458
                head_mask = (
                    head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)
                )  # We can specify head_mask for each layer
            head_mask = head_mask.to(
                dtype=next(self.parameters()).dtype
            )  # switch to fload if need + fp16 compatibility
459
460
        else:
            head_mask = [None] * self.config.n_layer
461

462
463
        if inputs_embeds is None:
            inputs_embeds = self.wte(input_ids)
thomwolf's avatar
thomwolf committed
464
465
466
467
468
469
        position_embeds = self.wpe(position_ids)
        if token_type_ids is not None:
            token_type_embeds = self.wte(token_type_ids)
        else:
            token_type_embeds = 0
        hidden_states = inputs_embeds + position_embeds + token_type_embeds
470
471
        hidden_states = self.drop(hidden_states)

472
473
        output_shape = input_shape + (hidden_states.size(-1),)

474
        presents = ()
thomwolf's avatar
thomwolf committed
475
        all_attentions = []
476
        all_hidden_states = ()
477
        for i, (block, layer_past) in enumerate(zip(self.h, past)):
thomwolf's avatar
thomwolf committed
478
            if self.output_hidden_states:
479
                all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)
thomwolf's avatar
thomwolf committed
480

481
482
483
            outputs = block(
                hidden_states, layer_past=layer_past, attention_mask=attention_mask, head_mask=head_mask[i]
            )
484

thomwolf's avatar
thomwolf committed
485
            hidden_states, present = outputs[:2]
486
487
            if self.output_past:
                presents = presents + (present,)
thomwolf's avatar
thomwolf committed
488
489
490
491

            if self.output_attentions:
                all_attentions.append(outputs[2])

thomwolf's avatar
thomwolf committed
492
        hidden_states = self.ln_f(hidden_states)
493

thomwolf's avatar
thomwolf committed
494
495
496
        hidden_states = hidden_states.view(*output_shape)
        # Add last hidden state
        if self.output_hidden_states:
497
            all_hidden_states = all_hidden_states + (hidden_states,)
thomwolf's avatar
thomwolf committed
498

499
500
501
        outputs = (hidden_states,)
        if self.output_past:
            outputs = outputs + (presents,)
thomwolf's avatar
thomwolf committed
502
        if self.output_hidden_states:
503
            outputs = outputs + (all_hidden_states,)
thomwolf's avatar
thomwolf committed
504
        if self.output_attentions:
thomwolf's avatar
thomwolf committed
505
506
            # let the number of heads free (-1) so we can extract attention even after head pruning
            attention_output_shape = input_shape[:-1] + (-1,) + all_attentions[0].shape[-2:]
507
            all_attentions = tuple(t.view(*attention_output_shape) for t in all_attentions)
508
            outputs = outputs + (all_attentions,)
509
        return outputs  # last hidden state, (presents), (all hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
510
511


512
513
514
515
516
517
@add_start_docstrings(
    """The GPT2 Model transformer with a language modeling head on top
(linear layer with weights tied to the input embeddings). """,
    GPT2_START_DOCSTRING,
    GPT2_INPUTS_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
518
class GPT2LMHeadModel(GPT2PreTrainedModel):
519
    r"""
thomwolf's avatar
thomwolf committed
520
        **labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
thomwolf's avatar
thomwolf committed
521
522
523
            Labels for language modeling.
            Note that the labels **are shifted** inside the model, i.e. you can set ``lm_labels = input_ids``
            Indices are selected in ``[-1, 0, ..., config.vocab_size]``
LysandreJik's avatar
LysandreJik committed
524
            All labels set to ``-100`` are ignored (masked), the loss is only
thomwolf's avatar
thomwolf committed
525
526
527
            computed for labels in ``[0, ..., config.vocab_size]``

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
thomwolf's avatar
thomwolf committed
528
        **loss**: (`optional`, returned when ``labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
thomwolf's avatar
thomwolf committed
529
530
531
532
            Language modeling loss.
        **prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **past**:
533
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(2, batch_size, num_heads, sequence_length, embed_size_per_head)``:
thomwolf's avatar
thomwolf committed
534
            that contains pre-computed hidden-states (key and values in the attention blocks).
535
536
            Can be used (see `past` input) to speed up sequential decoding. The token ids which have their past given to this model 
            should not be passed as input ids as they have already been computed.
thomwolf's avatar
thomwolf committed
537
538
539
540
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
thomwolf's avatar
thomwolf committed
541
542
543
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
thomwolf's avatar
thomwolf committed
544
545
546

    Examples::

thomwolf's avatar
thomwolf committed
547
        import torch
548
        from transformers import GPT2Tokenizer, GPT2LMHeadModel
thomwolf's avatar
thomwolf committed
549

550
551
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        model = GPT2LMHeadModel.from_pretrained('gpt2')
thomwolf's avatar
thomwolf committed
552

553
        input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute", add_special_tokens=True)).unsqueeze(0)  # Batch size 1
554
555
        outputs = model(input_ids, labels=input_ids)
        loss, logits = outputs[:2]
thomwolf's avatar
thomwolf committed
556
557

    """
558

thomwolf's avatar
thomwolf committed
559
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
560
        super(GPT2LMHeadModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
561
        self.transformer = GPT2Model(config)
thomwolf's avatar
thomwolf committed
562
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
thomwolf's avatar
thomwolf committed
563

564
        self.init_weights()
565

thomwolf's avatar
thomwolf committed
566
    def get_output_embeddings(self):
567
        return self.lm_head
thomwolf's avatar
thomwolf committed
568

569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
    def forward(
        self,
        input_ids=None,
        past=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
    ):
        transformer_outputs = self.transformer(
            input_ids,
            past=past,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )
thomwolf's avatar
thomwolf committed
589
        hidden_states = transformer_outputs[0]
590

thomwolf's avatar
thomwolf committed
591
        lm_logits = self.lm_head(hidden_states)
thomwolf's avatar
thomwolf committed
592

593
        outputs = (lm_logits,) + transformer_outputs[1:]
thomwolf's avatar
thomwolf committed
594
        if labels is not None:
595
            # Shift so that tokens < n predict n
596
            shift_logits = lm_logits[..., :-1, :].contiguous()
thomwolf's avatar
thomwolf committed
597
            shift_labels = labels[..., 1:].contiguous()
Catalin Voss's avatar
Catalin Voss committed
598
            # Flatten the tokens
LysandreJik's avatar
LysandreJik committed
599
            loss_fct = CrossEntropyLoss()
600
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
601
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
602
603

        return outputs  # (loss), lm_logits, presents, (all hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
604
605


606
607
@add_start_docstrings(
    """The GPT2 Model transformer with a language modeling and a multiple-choice classification
thomwolf's avatar
thomwolf committed
608
609
head on top e.g. for RocStories/SWAG tasks. The two heads are two linear layers.
The language modeling head has its weights tied to the input embeddings,
Julien Chaumond's avatar
Julien Chaumond committed
610
the classification head takes as input the input of a specified classification token index in the input sequence).
611
612
613
614
""",
    GPT2_START_DOCSTRING,
    GPT2_INPUTS_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
615
class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
616
    r"""
thomwolf's avatar
thomwolf committed
617
        **mc_token_ids**: (`optional`, default to index of the last token of the input) ``torch.LongTensor`` of shape ``(batch_size, num_choices)``:
thomwolf's avatar
thomwolf committed
618
619
620
621
622
623
            Index of the classification token in each input sequence.
            Selected in the range ``[0, input_ids.size(-1) - 1[``.
        **lm_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Labels for language modeling.
            Note that the labels **are shifted** inside the model, i.e. you can set ``lm_labels = input_ids``
            Indices are selected in ``[-1, 0, ..., config.vocab_size]``
LysandreJik's avatar
LysandreJik committed
624
            All labels set to ``-100`` are ignored (masked), the loss is only
thomwolf's avatar
thomwolf committed
625
            computed for labels in ``[0, ..., config.vocab_size]``
626
        **mc_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size)``:
thomwolf's avatar
thomwolf committed
627
628
629
            Labels for computing the multiple choice classification loss.
            Indices should be in ``[0, ..., num_choices]`` where `num_choices` is the size of the second dimension
            of the input tensors. (see `input_ids` above)
thomwolf's avatar
thomwolf committed
630

thomwolf's avatar
thomwolf committed
631
632
633
634
635
636
637
638
639
640
    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **lm_loss**: (`optional`, returned when ``lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Language modeling loss.
        **mc_loss**: (`optional`, returned when ``multiple_choice_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            Multiple choice classification loss.
        **lm_prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, num_choices, sequence_length, config.vocab_size)``
            Prediction scores of the language modeling head (scores for each vocabulary token before SoftMax).
        **mc_prediction_scores**: ``torch.FloatTensor`` of shape ``(batch_size, num_choices)``
            Prediction scores of the multiplechoice classification head (scores for each choice before SoftMax).
        **past**:
641
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(2, batch_size, num_heads, sequence_length, embed_size_per_head)``:
thomwolf's avatar
thomwolf committed
642
            that contains pre-computed hidden-states (key and values in the attention blocks).
643
644
            Can be used (see `past` input) to speed up sequential decoding. The token ids which have their past given to this model 
            should not be passed as input ids as they have already been computed.
thomwolf's avatar
thomwolf committed
645
646
647
648
        **hidden_states**: (`optional`, returned when ``config.output_hidden_states=True``)
            list of ``torch.FloatTensor`` (one for the output of each layer + the output of the embeddings)
            of shape ``(batch_size, sequence_length, hidden_size)``:
            Hidden-states of the model at the output of each layer plus the initial embedding outputs.
thomwolf's avatar
thomwolf committed
649
650
651
        **attentions**: (`optional`, returned when ``config.output_attentions=True``)
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            Attentions weights after the attention softmax, used to compute the weighted average in the self-attention heads.
thomwolf's avatar
thomwolf committed
652
653
654

    Examples::

655
        import torch
656
        from transformers import GPT2Tokenizer, GPT2DoubleHeadsModel
657
        
658
659
        tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        model = GPT2DoubleHeadsModel.from_pretrained('gpt2')
thomwolf's avatar
thomwolf committed
660
661
662
663
664
665
        
        # Add a [CLS] to the vocabulary (we should train it also!)
        tokenizer.add_special_tokens({'cls_token': '[CLS]'})
        model.resize_token_embeddings(len(tokenizer))  # Update the model embeddings with the new vocabulary size
        print(tokenizer.cls_token_id, len(tokenizer))  # The newly token the last token of the vocabulary
        
thomwolf's avatar
thomwolf committed
666
        choices = ["Hello, my dog is cute [CLS]", "Hello, my cat is cute [CLS]"]
thomwolf's avatar
thomwolf committed
667
668
669
670
671
672
673
        encoded_choices = [tokenizer.encode(s) for s in choices]
        cls_token_location = [tokens.index(tokenizer.cls_token_id) for tokens in encoded_choices]

        input_ids = torch.tensor(encoded_choices).unsqueeze(0)  # Batch size: 1, number of choices: 2
        mc_token_ids = torch.tensor([cls_token_location])  # Batch size: 1

        outputs = model(input_ids, mc_token_ids=mc_token_ids)
674
        lm_prediction_scores, mc_prediction_scores = outputs[:2]
thomwolf's avatar
thomwolf committed
675
676

    """
677

thomwolf's avatar
thomwolf committed
678
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
679
        super(GPT2DoubleHeadsModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
680
        config.num_labels = 1
thomwolf's avatar
thomwolf committed
681
        self.transformer = GPT2Model(config)
thomwolf's avatar
thomwolf committed
682
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
thomwolf's avatar
thomwolf committed
683
        self.multiple_choice_head = SequenceSummary(config)
thomwolf's avatar
thomwolf committed
684

685
        self.init_weights()
thomwolf's avatar
thomwolf committed
686

thomwolf's avatar
thomwolf committed
687
    def get_output_embeddings(self):
688
        return self.lm_head
thomwolf's avatar
thomwolf committed
689

690
691
692
693
694
695
696
697
698
699
700
701
702
703
704
705
706
707
708
709
710
711
    def forward(
        self,
        input_ids=None,
        past=None,
        attention_mask=None,
        token_type_ids=None,
        position_ids=None,
        head_mask=None,
        inputs_embeds=None,
        mc_token_ids=None,
        lm_labels=None,
        mc_labels=None,
    ):
        transformer_outputs = self.transformer(
            input_ids,
            past=past,
            attention_mask=attention_mask,
            token_type_ids=token_type_ids,
            position_ids=position_ids,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
        )
712

thomwolf's avatar
thomwolf committed
713
        hidden_states = transformer_outputs[0]
714

thomwolf's avatar
thomwolf committed
715
        lm_logits = self.lm_head(hidden_states)
thomwolf's avatar
thomwolf committed
716
        mc_logits = self.multiple_choice_head(hidden_states, mc_token_ids).squeeze(-1)
thomwolf's avatar
thomwolf committed
717

718
        outputs = (lm_logits, mc_logits) + transformer_outputs[1:]
thomwolf's avatar
thomwolf committed
719
720
        if mc_labels is not None:
            loss_fct = CrossEntropyLoss()
721
            loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)), mc_labels.view(-1))
722
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
723
        if lm_labels is not None:
724
725
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = lm_labels[..., 1:].contiguous()
LysandreJik's avatar
LysandreJik committed
726
            loss_fct = CrossEntropyLoss()
727
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)), shift_labels.view(-1))
728
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
729
730

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