modeling_gpt2.py 35.5 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
from .modeling_utils import (Conv1D, CONFIG_NAME, WEIGHTS_NAME, PretrainedConfig,
34
35
                             PreTrainedModel, prune_conv1d_layer, SequenceSummary,
                             add_start_docstrings)
thomwolf's avatar
thomwolf committed
36
from .modeling_bert import BertLayerNorm as LayerNorm
thomwolf's avatar
thomwolf committed
37
38
39

logger = logging.getLogger(__name__)

40
GPT2_PRETRAINED_MODEL_ARCHIVE_MAP = {"gpt2": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-pytorch_model.bin",
41
                                     "gpt2-medium": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-medium-pytorch_model.bin"}
42
GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP = {"gpt2": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-config.json",
43
                                      "gpt2-medium": "https://s3.amazonaws.com/models.huggingface.co/bert/gpt2-medium-config.json"}
thomwolf's avatar
thomwolf committed
44

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

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


103
class GPT2Config(PretrainedConfig):
thomwolf's avatar
thomwolf committed
104
    """Configuration class to store the configuration of a `GPT2Model`.
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121

    Args:
        vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `GPT2Model` or a configuration json file.
        n_positions: Number of positional embeddings.
        n_ctx: Size of the causal mask (usually same as n_positions).
        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.
        layer_norm_epsilon: epsilon to use in the layer norm layers
        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.
        initializer_range: The sttdev of the truncated_normal_initializer for
            initializing all weight matrices.
thomwolf's avatar
thomwolf committed
122
    """
123
    pretrained_config_archive_map = GPT2_PRETRAINED_CONFIG_ARCHIVE_MAP
thomwolf's avatar
thomwolf committed
124
125
126

    def __init__(
        self,
thomwolf's avatar
thomwolf committed
127
        vocab_size_or_config_json_file=50257,
thomwolf's avatar
thomwolf committed
128
129
130
131
132
        n_positions=1024,
        n_ctx=1024,
        n_embd=768,
        n_layer=12,
        n_head=12,
133
134
135
        resid_pdrop=0.1,
        embd_pdrop=0.1,
        attn_pdrop=0.1,
thomwolf's avatar
thomwolf committed
136
137
        layer_norm_epsilon=1e-5,
        initializer_range=0.02,
thomwolf's avatar
thomwolf committed
138
139

        num_labels=1,
thomwolf's avatar
thomwolf committed
140
141
142
        summary_type='token_ids',
        summary_use_proj=True,
        summary_activation=None,
thomwolf's avatar
thomwolf committed
143
        summary_proj_to_labels=True,
144
        summary_first_dropout=0.1,
thomwolf's avatar
thomwolf committed
145
        **kwargs
thomwolf's avatar
thomwolf committed
146
147
148
149
150
151
152
153
154
155
156
157
    ):
        """Constructs GPT2Config.

        Args:
            vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `GPT2Model` or a configuration json file.
            n_positions: Number of positional embeddings.
            n_ctx: Size of the causal mask (usually same as n_positions).
            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.
            layer_norm_epsilon: epsilon to use in the layer norm layers
158
159
160
161
162
            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.
thomwolf's avatar
thomwolf committed
163
164
165
            initializer_range: The sttdev of the truncated_normal_initializer for
                initializing all weight matrices.
        """
thomwolf's avatar
thomwolf committed
166
167
        super(GPT2Config, self).__init__(**kwargs)

thomwolf's avatar
thomwolf committed
168
169
170
171
172
173
174
175
176
177
178
179
180
        if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
                        and isinstance(vocab_size_or_config_json_file, unicode)):
            with open(vocab_size_or_config_json_file, "r", encoding="utf-8") as reader:
                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_ctx = n_ctx
            self.n_positions = n_positions
            self.n_embd = n_embd
            self.n_layer = n_layer
            self.n_head = n_head
181
182
183
            self.resid_pdrop = resid_pdrop
            self.embd_pdrop = embd_pdrop
            self.attn_pdrop = attn_pdrop
thomwolf's avatar
thomwolf committed
184
185
            self.layer_norm_epsilon = layer_norm_epsilon
            self.initializer_range = initializer_range
thomwolf's avatar
thomwolf committed
186
187

            self.num_labels = num_labels
thomwolf's avatar
thomwolf committed
188
189
190
            self.summary_type = summary_type
            self.summary_use_proj = summary_use_proj
            self.summary_activation = summary_activation
191
            self.summary_first_dropout = summary_first_dropout
thomwolf's avatar
thomwolf committed
192
            self.summary_proj_to_labels = summary_proj_to_labels
thomwolf's avatar
thomwolf committed
193
194
195
196
197
198
        else:
            raise ValueError(
                "First argument must be either a vocabulary size (int)"
                "or the path to a pretrained model config file (str)"
            )

199
200
201
202
    @property
    def max_position_embeddings(self):
        return self.n_positions

thomwolf's avatar
thomwolf committed
203
204
205
206
207
208
209
210
211
212
213
214
215
    @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
216
217

class Attention(nn.Module):
thomwolf's avatar
thomwolf committed
218
    def __init__(self, nx, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
219
        super(Attention, self).__init__()
thomwolf's avatar
thomwolf committed
220
221
        self.output_attentions = config.output_attentions

thomwolf's avatar
thomwolf committed
222
223
224
225
226
227
228
        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
229

thomwolf's avatar
thomwolf committed
230
231
        self.c_attn = Conv1D(n_state * 3, nx)
        self.c_proj = Conv1D(n_state, nx)
232
233
        self.attn_dropout = nn.Dropout(config.attn_pdrop)
        self.resid_dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
234

235
    def prune_heads(self, heads):
thomwolf's avatar
thomwolf committed
236
237
        if len(heads) == 0:
            return
238
239
240
241
242
243
244
245
246
247
248
249
250
251
        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
252
253
254
        w = torch.matmul(q, k)
        if self.scale:
            w = w / math.sqrt(v.size(-1))
thomwolf's avatar
thomwolf committed
255
256
        nd, ns = w.size(-2), w.size(-1)
        b = self.bias[:, :, ns-nd:ns, :ns]
257
        w = w * b - 1e4 * (1 - b)
thomwolf's avatar
thomwolf committed
258
259

        w = nn.Softmax(dim=-1)(w)
260
        w = self.attn_dropout(w)
261
262
263
264
265

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

thomwolf's avatar
thomwolf committed
266
        outputs = [torch.matmul(w, v)]
thomwolf's avatar
thomwolf committed
267
        if self.output_attentions:
thomwolf's avatar
thomwolf committed
268
269
            outputs.append(w)
        return outputs
thomwolf's avatar
thomwolf committed
270
271
272
273
274
275
276
277
278
279

    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
280
            return x.permute(0, 2, 3, 1)  # (batch, head, head_features, seq_length)
thomwolf's avatar
thomwolf committed
281
        else:
thomwolf's avatar
thomwolf committed
282
            return x.permute(0, 2, 1, 3)  # (batch, head, seq_length, head_features)
thomwolf's avatar
thomwolf committed
283

284
    def forward(self, x, layer_past=None, head_mask=None):
thomwolf's avatar
thomwolf committed
285
286
287
288
289
        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
290
        if layer_past is not None:
thomwolf's avatar
thomwolf committed
291
            past_key, past_value = layer_past[0].transpose(-2, -1), layer_past[1]  # transpose back cf below
thomwolf's avatar
thomwolf committed
292
            key = torch.cat((past_key, key), dim=-1)
thomwolf's avatar
thomwolf committed
293
            value = torch.cat((past_value, value), dim=-2)
thomwolf's avatar
thomwolf committed
294
        present = torch.stack((key.transpose(-2, -1), value))  # transpose to have same shapes for stacking
295

thomwolf's avatar
thomwolf committed
296
297
        attn_outputs = self._attn(query, key, value, head_mask)
        a = attn_outputs[0]
298

thomwolf's avatar
thomwolf committed
299
300
        a = self.merge_heads(a)
        a = self.c_proj(a)
301
        a = self.resid_dropout(a)
thomwolf's avatar
thomwolf committed
302
303
304

        outputs = [a, present] + attn_outputs[1:]
        return outputs  # a, present, (attentions)
thomwolf's avatar
thomwolf committed
305
306
307
308
309
310
311
312
313


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
314
        self.dropout = nn.Dropout(config.resid_pdrop)
thomwolf's avatar
thomwolf committed
315
316
317
318

    def forward(self, x):
        h = self.act(self.c_fc(x))
        h2 = self.c_proj(h)
319
        return self.dropout(h2)
thomwolf's avatar
thomwolf committed
320
321
322


class Block(nn.Module):
thomwolf's avatar
thomwolf committed
323
    def __init__(self, n_ctx, config, scale=False):
thomwolf's avatar
thomwolf committed
324
325
326
        super(Block, self).__init__()
        nx = config.n_embd
        self.ln_1 = LayerNorm(nx, eps=config.layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
327
        self.attn = Attention(nx, n_ctx, config, scale)
thomwolf's avatar
thomwolf committed
328
329
330
        self.ln_2 = LayerNorm(nx, eps=config.layer_norm_epsilon)
        self.mlp = MLP(4 * nx, config)

331
332
    def forward(self, x, layer_past=None, head_mask=None):
        output_attn = self.attn(self.ln_1(x), layer_past=layer_past, head_mask=head_mask)
thomwolf's avatar
thomwolf committed
333
334
        a = output_attn[0]  # output_attn: a, present, (attentions)

thomwolf's avatar
thomwolf committed
335
        x = x + a
thomwolf's avatar
thomwolf committed
336
        m = self.mlp(self.ln_2(x))
thomwolf's avatar
thomwolf committed
337
        x = x + m
thomwolf's avatar
thomwolf committed
338
339
340

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


343
class GPT2PreTrainedModel(PreTrainedModel):
thomwolf's avatar
thomwolf committed
344
345
346
    """ An abstract class to handle weights initialization and
        a simple interface for dowloading and loading pretrained models.
    """
347
    config_class = GPT2Config
348
    pretrained_model_archive_map = GPT2_PRETRAINED_MODEL_ARCHIVE_MAP
349
350
    load_tf_weights = load_tf_weights_in_gpt2
    base_model_prefix = "transformer"
thomwolf's avatar
thomwolf committed
351

352
353
354
    def __init__(self, *inputs, **kwargs):
        super(GPT2PreTrainedModel, self).__init__(*inputs, **kwargs)

thomwolf's avatar
thomwolf committed
355
356
357
    def init_weights(self, module):
        """ Initialize the weights.
        """
358
        if isinstance(module, (nn.Linear, nn.Embedding, Conv1D)):
thomwolf's avatar
thomwolf committed
359
360
361
            # 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)
362
363
            if isinstance(module, (nn.Linear, Conv1D)) and module.bias is not None:
                module.bias.data.zero_()
thomwolf's avatar
thomwolf committed
364
365
366
367
368
        elif isinstance(module, LayerNorm):
            module.bias.data.zero_()
            module.weight.data.fill_(1.0)


thomwolf's avatar
thomwolf committed
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
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:
        config (:class:`~pytorch_transformers.BertConfig`): Model configuration class with all the parameters of the model.
"""

GPT2_INPUTS_DOCTRING = r"""    Inputs:
        **input_ids**: ``torch.LongTensor`` of shape ``(batch_size, sequence_length)``:
            Indices of input sequence tokens in the vocabulary.
            Indices can be obtained using :class:`pytorch_transformers.BPT2Tokenizer`.
            See :func:`pytorch_transformers.PreTrainedTokenizer.encode` and
            :func:`pytorch_transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details.
        **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[``.
        **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).
        **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
            (see `past` output below). Can be used to speed up sequential decoding.
        **attention_mask**: (`optional`) ``torch.Tensor`` of shape ``(batch_size, sequence_length)``:
            Mask to avoid performing attention on padding token indices.
            Mask indices selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
        **head_mask**: (`optional`) ``torch.Tensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:
            Mask to nullify selected heads of the self-attention modules.
            Mask indices selected in ``[0, 1]``:
            ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.
"""

@add_start_docstrings("The bare GPT2 Model transformer outputing raw hidden-states without any specific head on top.",
                      GPT2_START_DOCSTRING, GPT2_INPUTS_DOCTRING)
thomwolf's avatar
thomwolf committed
417
class GPT2Model(GPT2PreTrainedModel):
418
    r"""
thomwolf's avatar
thomwolf committed
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
    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**:
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            that contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding.
        **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.
        **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.

    Examples::

        >>> config = GPT2Config.from_pretrained('gpt2')
        >>> tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        >>> model = GPT2Model(config)
        >>> input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        >>> outputs = model(input_ids)
        >>> last_hidden_states = outputs[0]  # The last hidden-state is the first element of the output tuple
thomwolf's avatar
thomwolf committed
442
443

    """
thomwolf's avatar
thomwolf committed
444
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
445
        super(GPT2Model, self).__init__(config)
thomwolf's avatar
thomwolf committed
446
447
448
        self.output_hidden_states = config.output_hidden_states
        self.output_attentions = config.output_attentions

thomwolf's avatar
thomwolf committed
449
        self.wte = nn.Embedding(config.vocab_size, config.n_embd)
thomwolf's avatar
thomwolf committed
450
        self.wpe = nn.Embedding(config.n_positions, config.n_embd)
451
        self.drop = nn.Dropout(config.embd_pdrop)
452
        self.h = nn.ModuleList([Block(config.n_ctx, config, scale=True) for _ in range(config.n_layer)])
thomwolf's avatar
thomwolf committed
453
        self.ln_f = LayerNorm(config.n_embd, eps=config.layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
454
455
456

        self.apply(self.init_weights)

thomwolf's avatar
thomwolf committed
457
458
    def _resize_token_embeddings(self, new_num_tokens):
        self.wte = self._get_resized_embeddings(self.wte, new_num_tokens)
thomwolf's avatar
thomwolf committed
459
        return self.wte
thomwolf's avatar
thomwolf committed
460

thomwolf's avatar
thomwolf committed
461
    def _prune_heads(self, heads_to_prune):
462
463
464
465
466
467
468
        """ 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, past=None, head_mask=None):
thomwolf's avatar
thomwolf committed
469
        if past is None:
thomwolf's avatar
thomwolf committed
470
            past_length = 0
thomwolf's avatar
thomwolf committed
471
            past = [None] * len(self.h)
thomwolf's avatar
thomwolf committed
472
        else:
thomwolf's avatar
thomwolf committed
473
            past_length = past[0][0].size(-2)
thomwolf's avatar
thomwolf committed
474
475
476
477
        if position_ids is None:
            position_ids = torch.arange(past_length, input_ids.size(-1) + past_length, dtype=torch.long, device=input_ids.device)
            position_ids = position_ids.unsqueeze(0).expand_as(input_ids)

478
        # Prepare head mask if needed
thomwolf's avatar
thomwolf committed
479
        # 1.0 in head_mask indicate we keep the head
480
        # attention_probs has shape bsz x n_heads x N x N
481
        # head_mask has shape n_layer x batch x n_heads x N x N
482
483
        if head_mask is not None:
            if head_mask.dim() == 1:
484
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(-1).unsqueeze(-1)
thomwolf's avatar
thomwolf committed
485
                head_mask = head_mask.expand(self.config.n_layer, -1, -1, -1, -1)
486
            elif head_mask.dim() == 2:
487
                head_mask = head_mask.unsqueeze(1).unsqueeze(-1).unsqueeze(-1)  # We can specify head_mask for each layer
488
            head_mask = head_mask.to(dtype=next(self.parameters()).dtype) # switch to fload if need + fp16 compatibility
489
490
        else:
            head_mask = [None] * self.config.n_layer
491

thomwolf's avatar
thomwolf committed
492
493
494
495
496
497
498
499
500
501
502
503
        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))

        inputs_embeds = self.wte(input_ids)
        position_embeds = self.wpe(position_ids)
        if token_type_ids is not None:
            token_type_ids = token_type_ids.view(-1, token_type_ids.size(-1))
            token_type_embeds = self.wte(token_type_ids)
        else:
            token_type_embeds = 0
        hidden_states = inputs_embeds + position_embeds + token_type_embeds
504
505
        hidden_states = self.drop(hidden_states)

506
507
        output_shape = input_shape + (hidden_states.size(-1),)

508
        presents = ()
thomwolf's avatar
thomwolf committed
509
        all_attentions = []
510
        all_hidden_states = ()
511
        for i, (block, layer_past) in enumerate(zip(self.h, past)):
thomwolf's avatar
thomwolf committed
512
            if self.output_hidden_states:
513
                all_hidden_states = all_hidden_states + (hidden_states.view(*output_shape),)
thomwolf's avatar
thomwolf committed
514

515
            outputs = block(hidden_states, layer_past, head_mask[i])
thomwolf's avatar
thomwolf committed
516
            hidden_states, present = outputs[:2]
517
            presents = presents + (present,)
thomwolf's avatar
thomwolf committed
518
519
520
521

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

thomwolf's avatar
thomwolf committed
522
        hidden_states = self.ln_f(hidden_states)
523

thomwolf's avatar
thomwolf committed
524
525
526
        hidden_states = hidden_states.view(*output_shape)
        # Add last hidden state
        if self.output_hidden_states:
527
            all_hidden_states = all_hidden_states + (hidden_states,)
thomwolf's avatar
thomwolf committed
528

529
        outputs = (hidden_states, presents)
thomwolf's avatar
thomwolf committed
530
        if self.output_hidden_states:
531
            outputs = outputs + (all_hidden_states,)
thomwolf's avatar
thomwolf committed
532
        if self.output_attentions:
thomwolf's avatar
thomwolf committed
533
534
            # 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:]
535
            all_attentions = tuple(t.view(*attention_output_shape) for t in all_attentions)
536
            outputs = outputs + (all_attentions,)
thomwolf's avatar
thomwolf committed
537
        return outputs  # last hidden state, presents, (all hidden_states), (attentions)
thomwolf's avatar
thomwolf committed
538
539


thomwolf's avatar
thomwolf committed
540
541
@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_DOCTRING)
thomwolf's avatar
thomwolf committed
542
class GPT2LMHeadModel(GPT2PreTrainedModel):
543
    r"""
thomwolf's avatar
thomwolf committed
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
        **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]``
            All labels set to ``-1`` are ignored (masked), the loss is only
            computed for labels in ``[0, ..., config.vocab_size]``

    Outputs: `Tuple` comprising various elements depending on the configuration (config) and inputs:
        **loss**: (`optional`, returned when ``lm_labels`` is provided) ``torch.FloatTensor`` of shape ``(1,)``:
            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**:
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            that contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding.
        **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.
        **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.

    Examples::

        >>> config = GPT2Config.from_pretrained('gpt2')
        >>> tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        >>> model = GPT2LMHeadModel(config)
        >>> input_ids = torch.tensor(tokenizer.encode("Hello, my dog is cute")).unsqueeze(0)  # Batch size 1
        >>> outputs = model(input_ids, lm_labels=input_ids)
        >>> loss, logits = outputs[:2]
thomwolf's avatar
thomwolf committed
576
577

    """
thomwolf's avatar
thomwolf committed
578
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
579
        super(GPT2LMHeadModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
580
        self.transformer = GPT2Model(config)
thomwolf's avatar
thomwolf committed
581
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
thomwolf's avatar
thomwolf committed
582

thomwolf's avatar
thomwolf committed
583
584
        self.apply(self.init_weights)
        self.tie_weights()
585

thomwolf's avatar
thomwolf committed
586
587
588
    def tie_weights(self):
        """ Make sure we are sharing the input and output embeddings.
            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
thomwolf's avatar
thomwolf committed
589
        """
thomwolf's avatar
thomwolf committed
590
591
        self._tie_or_clone_weights(self.lm_head,
                                   self.transformer.wte)
thomwolf's avatar
thomwolf committed
592

593
    def forward(self, input_ids, position_ids=None, token_type_ids=None, lm_labels=None, past=None, head_mask=None):
thomwolf's avatar
thomwolf committed
594
595
        transformer_outputs = self.transformer(input_ids, position_ids, token_type_ids, past, head_mask)
        hidden_states = transformer_outputs[0]
596

thomwolf's avatar
thomwolf committed
597
        lm_logits = self.lm_head(hidden_states)
thomwolf's avatar
thomwolf committed
598

599
        outputs = (lm_logits,) + transformer_outputs[1:]
thomwolf's avatar
thomwolf committed
600
        if lm_labels is not None:
601
            # Shift so that tokens < n predict n
602
603
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = lm_labels[..., 1:].contiguous()
Catalin Voss's avatar
Catalin Voss committed
604
            # Flatten the tokens
thomwolf's avatar
thomwolf committed
605
            loss_fct = CrossEntropyLoss(ignore_index=-1)
606
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)),
607
                            shift_labels.view(-1))
608
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
609
610

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


thomwolf's avatar
thomwolf committed
613
614
615
616
617
@add_start_docstrings("""The GPT2 Model transformer with a language modeling and a multiple-choice classification
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,
the classification head takes as input the input of a specified classification token index in the intput sequence).
""", GPT2_START_DOCSTRING)
thomwolf's avatar
thomwolf committed
618
class GPT2DoubleHeadsModel(GPT2PreTrainedModel):
619
    r"""    Inputs:
thomwolf's avatar
thomwolf committed
620
621
622
623
624
625
626
627
628
629
630
631
632
633
634
635
636
637
638
639
640
641
642
643
644
645
646
647
648
649
650
651
652
653
654
655
656
657
        **input_ids**: ``torch.LongTensor`` of shape ``(batch_size, num_choices, sequence_length)``:
            Indices of input sequence tokens in the vocabulary.
            The second dimension of the input (`num_choices`) indicates the number of choices to score.
            Indices can be obtained using :class:`pytorch_transformers.BPT2Tokenizer`.
            See :func:`pytorch_transformers.PreTrainedTokenizer.encode` and
            :func:`pytorch_transformers.PreTrainedTokenizer.convert_tokens_to_ids` for details.
        **mc_token_ids**: ``torch.LongTensor`` of shape ``(batch_size, num_choices)``:
            Index of the classification token in each input sequence.
            Selected in the range ``[0, input_ids.size(-1) - 1[``.
        **position_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, num_choices, sequence_length)``:
            Indices of positions of each input sequence tokens in the position embeddings.
            Selected in the range ``[0, config.max_position_embeddings - 1[``.
        **token_type_ids**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size, num_choices, 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).
        **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
            (see `past` output below). Can be used to speed up sequential decoding.
        **attention_mask**: (`optional`) ``torch.Tensor`` of shape ``(batch_size, num_choices, sequence_length)``:
            Mask to avoid performing attention on padding token indices.
            Mask indices selected in ``[0, 1]``:
            ``1`` for tokens that are NOT MASKED, ``0`` for MASKED tokens.
        **head_mask**: (`optional`) ``torch.Tensor`` of shape ``(num_heads,)`` or ``(num_layers, num_heads)``:
            Mask to nullify selected heads of the self-attention modules.
            Mask indices selected in ``[0, 1]``:
            ``1`` indicates the head is **not masked**, ``0`` indicates the head is **masked**.
        **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]``
            All labels set to ``-1`` are ignored (masked), the loss is only
            computed for labels in ``[0, ..., config.vocab_size]``
        **multiple_choice_labels**: (`optional`) ``torch.LongTensor`` of shape ``(batch_size)``:
            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
658

thomwolf's avatar
thomwolf committed
659
660
            `multiple_choice_labels`: optional multiple choice labels: ``torch.LongTensor`` of shape [batch_size]
                with indices selected in [0, ..., num_choices].
thomwolf's avatar
thomwolf committed
661

thomwolf's avatar
thomwolf committed
662
663
664
665
666
667
668
669
670
671
672
673
674
675
676
677
678
679
680
681
682
683
684
685
686
687
688
689
690
691
692
    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**:
            list of ``torch.FloatTensor`` (one for each layer) of shape ``(batch_size, num_heads, sequence_length, sequence_length)``:
            that contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `past` input) to speed up sequential decoding.
        **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.
        **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.

    Examples::

        >>> config = GPT2Config.from_pretrained('gpt2')
        >>> tokenizer = GPT2Tokenizer.from_pretrained('gpt2')
        >>> model = GPT2DoubleHeadsModel(config)
        >>> choices = ["Hello, my dog is cute [CLS]", "Hello, my cat is cute [CLS]"]  # Assume you've added [CLS] to the vocabulary
        >>> input_ids = torch.tensor([tokenizer.encode(s) for s in choices]).unsqueeze(0)  # Batch size 1, 2 choices
        >>> mc_token_ids = torch.tensor([-1, -1]).unsqueeze(0)  # Batch size 1
        >>> outputs = model(input_ids, mc_token_ids)
        >>> lm_prediction_scores, mc_prediction_scores = outputs[:2]
thomwolf's avatar
thomwolf committed
693
694

    """
thomwolf's avatar
thomwolf committed
695
    def __init__(self, config):
thomwolf's avatar
thomwolf committed
696
        super(GPT2DoubleHeadsModel, self).__init__(config)
thomwolf's avatar
thomwolf committed
697
        self.transformer = GPT2Model(config)
thomwolf's avatar
thomwolf committed
698
        self.lm_head = nn.Linear(config.n_embd, config.vocab_size, bias=False)
thomwolf's avatar
thomwolf committed
699
        self.multiple_choice_head = SequenceSummary(config)
thomwolf's avatar
thomwolf committed
700

thomwolf's avatar
thomwolf committed
701
702
        self.apply(self.init_weights)

thomwolf's avatar
thomwolf committed
703
704
705
    def tie_weights(self):
        """ Make sure we are sharing the input and output embeddings.
            Export to TorchScript can't handle parameter sharing so we are cloning them instead.
706
        """
thomwolf's avatar
thomwolf committed
707
708
        self._tie_or_clone_weights(self.lm_head,
                                   self.transformer.wte)
thomwolf's avatar
thomwolf committed
709

thomwolf's avatar
thomwolf committed
710
    def forward(self, input_ids, mc_token_ids=None, lm_labels=None, mc_labels=None, token_type_ids=None,
711
                position_ids=None, past=None, head_mask=None):
thomwolf's avatar
thomwolf committed
712
713
        transformer_outputs = self.transformer(input_ids, position_ids, token_type_ids, past, head_mask)
        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
721
722
        if mc_labels is not None:
            loss_fct = CrossEntropyLoss()
            loss = loss_fct(mc_logits.view(-1, mc_logits.size(-1)),
                            mc_labels.view(-1))
723
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
724
        if lm_labels is not None:
725
726
            shift_logits = lm_logits[..., :-1, :].contiguous()
            shift_labels = lm_labels[..., 1:].contiguous()
thomwolf's avatar
thomwolf committed
727
            loss_fct = CrossEntropyLoss(ignore_index=-1)
thomwolf's avatar
thomwolf committed
728
729
            loss = loss_fct(shift_logits.view(-1, shift_logits.size(-1)),
                            shift_labels.view(-1))
730
            outputs = (loss,) + outputs
thomwolf's avatar
thomwolf committed
731
732

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