modeling_transfo_xl.py 45.2 KB
Newer Older
thomwolf's avatar
thomwolf committed
1
# coding=utf-8
thomwolf's avatar
thomwolf committed
2
# Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
thomwolf's avatar
thomwolf committed
3
4
5
6
7
8
9
10
11
12
13
14
15
16
# 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 Transformer XL model.
17
    Adapted from https://github.com/kimiyoung/transformer-xl.
thomwolf's avatar
thomwolf committed
18
19
20
    In particular https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/mem_transformer.py
"""

21

thomwolf's avatar
thomwolf committed
22
import logging
23
24
from dataclasses import dataclass
from typing import List, Optional, Tuple
thomwolf's avatar
thomwolf committed
25
26
27

import torch
import torch.nn as nn
28
import torch.nn.functional as F
thomwolf's avatar
thomwolf committed
29

30
from .configuration_transfo_xl import TransfoXLConfig
31
from .file_utils import ModelOutput, add_code_sample_docstrings, add_start_docstrings, add_start_docstrings_to_callable
32
from .modeling_transfo_xl_utilities import ProjectedAdaptiveLogSoftmax
33
from .modeling_utils import PreTrainedModel
Aymeric Augustin's avatar
Aymeric Augustin committed
34

thomwolf's avatar
thomwolf committed
35
36
37

logger = logging.getLogger(__name__)

38
_CONFIG_FOR_DOC = "TransfoXLConfig"
39
40
_TOKENIZER_FOR_DOC = "TransfoXLTokenizer"

41
42
43
44
TRANSFO_XL_PRETRAINED_MODEL_ARCHIVE_LIST = [
    "transfo-xl-wt103",
    # See all Transformer XL models at https://huggingface.co/models?filter=transfo-xl
]
45

46

47
48
49
50
51
def build_tf_to_pytorch_map(model, config):
    """ A map of modules from TF to PyTorch.
        This time I use a map to keep the PyTorch model as identical to the original PyTorch model as possible.
    """
    tf_to_pt_map = {}
52

53
    if hasattr(model, "transformer"):
54
        # We are loading in a TransfoXLLMHeadModel => we will load also the Adaptive Softmax
55
56
57
58
59
60
61
62
63
        tf_to_pt_map.update(
            {
                "transformer/adaptive_softmax/cutoff_0/cluster_W": model.crit.cluster_weight,
                "transformer/adaptive_softmax/cutoff_0/cluster_b": model.crit.cluster_bias,
            }
        )
        for i, (out_l, proj_l, tie_proj) in enumerate(
            zip(model.crit.out_layers, model.crit.out_projs, config.tie_projs)
        ):
64
65
            layer_str = "transformer/adaptive_softmax/cutoff_%d/" % i
            if config.tie_weight:
66
                tf_to_pt_map.update({layer_str + "b": out_l.bias})
67
68
69
            else:
                raise NotImplementedError
                # I don't think this is implemented in the TF code
70
                tf_to_pt_map.update({layer_str + "lookup_table": out_l.weight, layer_str + "b": out_l.bias})
71
            if not tie_proj:
72
                tf_to_pt_map.update({layer_str + "proj": proj_l})
73
74
75
        # Now load the rest of the transformer
        model = model.transformer

thomwolf's avatar
thomwolf committed
76
    # Embeddings
77
78
    for i, (embed_l, proj_l) in enumerate(zip(model.word_emb.emb_layers, model.word_emb.emb_projs)):
        layer_str = "transformer/adaptive_embed/cutoff_%d/" % i
79
        tf_to_pt_map.update({layer_str + "lookup_table": embed_l.weight, layer_str + "proj_W": proj_l})
80
81
82
83

    # Transformer blocks
    for i, b in enumerate(model.layers):
        layer_str = "transformer/layer_%d/" % i
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
        tf_to_pt_map.update(
            {
                layer_str + "rel_attn/LayerNorm/gamma": b.dec_attn.layer_norm.weight,
                layer_str + "rel_attn/LayerNorm/beta": b.dec_attn.layer_norm.bias,
                layer_str + "rel_attn/o/kernel": b.dec_attn.o_net.weight,
                layer_str + "rel_attn/qkv/kernel": b.dec_attn.qkv_net.weight,
                layer_str + "rel_attn/r/kernel": b.dec_attn.r_net.weight,
                layer_str + "ff/LayerNorm/gamma": b.pos_ff.layer_norm.weight,
                layer_str + "ff/LayerNorm/beta": b.pos_ff.layer_norm.bias,
                layer_str + "ff/layer_1/kernel": b.pos_ff.CoreNet[0].weight,
                layer_str + "ff/layer_1/bias": b.pos_ff.CoreNet[0].bias,
                layer_str + "ff/layer_2/kernel": b.pos_ff.CoreNet[3].weight,
                layer_str + "ff/layer_2/bias": b.pos_ff.CoreNet[3].bias,
            }
        )
99
100
101
102
103
104
105
106
107
108
109

    # Relative positioning biases
    if config.untie_r:
        r_r_list = []
        r_w_list = []
        for b in model.layers:
            r_r_list.append(b.dec_attn.r_r_bias)
            r_w_list.append(b.dec_attn.r_w_bias)
    else:
        r_r_list = [model.r_r_bias]
        r_w_list = [model.r_w_bias]
110
    tf_to_pt_map.update({"transformer/r_r_bias": r_r_list, "transformer/r_w_bias": r_w_list})
111
112
    return tf_to_pt_map

113

114
115
116
def load_tf_weights_in_transfo_xl(model, config, tf_path):
    """ Load tf checkpoints in a pytorch model
    """
117
118
119
    try:
        import numpy as np
        import tensorflow as tf
thomwolf's avatar
thomwolf committed
120
    except ImportError:
121
122
123
124
        logger.error(
            "Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions."
        )
125
        raise
126
127
128
129
130
131
132
    # Build TF to PyTorch weights loading map
    tf_to_pt_map = build_tf_to_pytorch_map(model, config)

    # Load weights from TF model
    init_vars = tf.train.list_variables(tf_path)
    tf_weights = {}
    for name, shape in init_vars:
thomwolf's avatar
thomwolf committed
133
        logger.info("Loading TF weight {} with shape {}".format(name, shape))
134
135
136
137
138
139
140
141
        array = tf.train.load_variable(tf_path, name)
        tf_weights[name] = array

    for name, pointer in tf_to_pt_map.items():
        assert name in tf_weights
        array = tf_weights[name]
        # adam_v and adam_m are variables used in AdamWeightDecayOptimizer to calculated m and v
        # which are not required for using pretrained model
142
        if "kernel" in name or "proj" in name:
143
            array = np.transpose(array)
144
        if ("r_r_bias" in name or "r_w_bias" in name) and len(pointer) > 1:
Julien Chaumond's avatar
Julien Chaumond committed
145
            # Here we will split the TF weights
146
147
148
149
150
151
152
153
            assert len(pointer) == array.shape[0]
            for i, p_i in enumerate(pointer):
                arr_i = array[i, ...]
                try:
                    assert p_i.shape == arr_i.shape
                except AssertionError as e:
                    e.args += (p_i.shape, arr_i.shape)
                    raise
thomwolf's avatar
thomwolf committed
154
                logger.info("Initialize PyTorch weight {} for layer {}".format(name, i))
155
156
157
158
159
160
161
                p_i.data = torch.from_numpy(arr_i)
        else:
            try:
                assert pointer.shape == array.shape
            except AssertionError as e:
                e.args += (pointer.shape, array.shape)
                raise
thomwolf's avatar
thomwolf committed
162
            logger.info("Initialize PyTorch weight {}".format(name))
163
164
            pointer.data = torch.from_numpy(array)
        tf_weights.pop(name, None)
165
166
        tf_weights.pop(name + "/Adam", None)
        tf_weights.pop(name + "/Adam_1", None)
167

168
    logger.info("Weights not copied to PyTorch model: {}".format(", ".join(tf_weights.keys())))
169
170
171
    return model


thomwolf's avatar
thomwolf committed
172
173
class PositionalEmbedding(nn.Module):
    def __init__(self, demb):
Julien Chaumond's avatar
Julien Chaumond committed
174
        super().__init__()
thomwolf's avatar
thomwolf committed
175
176
177
178

        self.demb = demb

        inv_freq = 1 / (10000 ** (torch.arange(0.0, demb, 2.0) / demb))
179
        self.register_buffer("inv_freq", inv_freq)
thomwolf's avatar
thomwolf committed
180
181
182
183
184
185

    def forward(self, pos_seq, bsz=None):
        sinusoid_inp = torch.ger(pos_seq, self.inv_freq)
        pos_emb = torch.cat([sinusoid_inp.sin(), sinusoid_inp.cos()], dim=-1)

        if bsz is not None:
186
            return pos_emb[:, None, :].expand(-1, bsz, -1)
thomwolf's avatar
thomwolf committed
187
        else:
188
            return pos_emb[:, None, :]
thomwolf's avatar
thomwolf committed
189

thomwolf's avatar
thomwolf committed
190

thomwolf's avatar
thomwolf committed
191
class PositionwiseFF(nn.Module):
192
    def __init__(self, d_model, d_inner, dropout, pre_lnorm=False, layer_norm_epsilon=1e-5):
Julien Chaumond's avatar
Julien Chaumond committed
193
        super().__init__()
thomwolf's avatar
thomwolf committed
194
195
196
197
198
199

        self.d_model = d_model
        self.d_inner = d_inner
        self.dropout = dropout

        self.CoreNet = nn.Sequential(
200
201
            nn.Linear(d_model, d_inner),
            nn.ReLU(inplace=True),
thomwolf's avatar
thomwolf committed
202
203
204
205
206
            nn.Dropout(dropout),
            nn.Linear(d_inner, d_model),
            nn.Dropout(dropout),
        )

207
        self.layer_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
208
209
210
211
212

        self.pre_lnorm = pre_lnorm

    def forward(self, inp):
        if self.pre_lnorm:
213
            # layer normalization + positionwise feed-forward
thomwolf's avatar
thomwolf committed
214
215
            core_out = self.CoreNet(self.layer_norm(inp))

216
            # residual connection
thomwolf's avatar
thomwolf committed
217
218
            output = core_out + inp
        else:
219
            # positionwise feed-forward
thomwolf's avatar
thomwolf committed
220
221
            core_out = self.CoreNet(inp)

222
            # residual connection + layer normalization
thomwolf's avatar
thomwolf committed
223
224
225
226
            output = self.layer_norm(inp + core_out)

        return output

thomwolf's avatar
thomwolf committed
227

228
class RelPartialLearnableMultiHeadAttn(nn.Module):
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
    def __init__(
        self,
        n_head,
        d_model,
        d_head,
        dropout,
        dropatt=0,
        tgt_len=None,
        ext_len=None,
        mem_len=None,
        pre_lnorm=False,
        r_r_bias=None,
        r_w_bias=None,
        layer_norm_epsilon=1e-5,
    ):
Julien Chaumond's avatar
Julien Chaumond committed
244
        super().__init__()
thomwolf's avatar
thomwolf committed
245
246
247
248
249
250
251
252
253
254
255
256

        self.n_head = n_head
        self.d_model = d_model
        self.d_head = d_head
        self.dropout = dropout

        self.qkv_net = nn.Linear(d_model, 3 * n_head * d_head, bias=False)

        self.drop = nn.Dropout(dropout)
        self.dropatt = nn.Dropout(dropatt)
        self.o_net = nn.Linear(n_head * d_head, d_model, bias=False)

257
        self.layer_norm = nn.LayerNorm(d_model, eps=layer_norm_epsilon)
thomwolf's avatar
thomwolf committed
258
259
260
261
262

        self.scale = 1 / (d_head ** 0.5)

        self.pre_lnorm = pre_lnorm

263
        if r_r_bias is None or r_w_bias is None:  # Biases are not shared
thomwolf's avatar
thomwolf committed
264
265
            self.r_r_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
            self.r_w_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
thomwolf's avatar
thomwolf committed
266
267
268
269
        else:
            self.r_r_bias = r_r_bias
            self.r_w_bias = r_w_bias

270
        self.r_net = nn.Linear(self.d_model, self.n_head * self.d_head, bias=False)
thomwolf's avatar
thomwolf committed
271

272
    def _rel_shift(self, x):
thomwolf's avatar
thomwolf committed
273
274
        zero_pad_shape = (x.size(0), 1) + x.size()[2:]
        zero_pad = torch.zeros(zero_pad_shape, device=x.device, dtype=x.dtype)
thomwolf's avatar
thomwolf committed
275
276
        x_padded = torch.cat([zero_pad, x], dim=1)

thomwolf's avatar
thomwolf committed
277
278
        x_padded_shape = (x.size(1) + 1, x.size(0)) + x.size()[2:]
        x_padded = x_padded.view(*x_padded_shape)
thomwolf's avatar
thomwolf committed
279
280
281
282
283

        x = x_padded[1:].view_as(x)

        return x

284
    def forward(self, w, r, attn_mask=None, mems=None, head_mask=None, output_attentions=False):
thomwolf's avatar
thomwolf committed
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
        qlen, rlen, bsz = w.size(0), r.size(0), w.size(1)

        if mems is not None:
            cat = torch.cat([mems, w], 0)
            if self.pre_lnorm:
                w_heads = self.qkv_net(self.layer_norm(cat))
            else:
                w_heads = self.qkv_net(cat)
            r_head_k = self.r_net(r)

            w_head_q, w_head_k, w_head_v = torch.chunk(w_heads, 3, dim=-1)
            w_head_q = w_head_q[-qlen:]
        else:
            if self.pre_lnorm:
                w_heads = self.qkv_net(self.layer_norm(w))
            else:
                w_heads = self.qkv_net(w)
            r_head_k = self.r_net(r)

            w_head_q, w_head_k, w_head_v = torch.chunk(w_heads, 3, dim=-1)

        klen = w_head_k.size(0)

308
309
310
        w_head_q = w_head_q.view(qlen, bsz, self.n_head, self.d_head)  # qlen x bsz x n_head x d_head
        w_head_k = w_head_k.view(klen, bsz, self.n_head, self.d_head)  # qlen x bsz x n_head x d_head
        w_head_v = w_head_v.view(klen, bsz, self.n_head, self.d_head)  # qlen x bsz x n_head x d_head
thomwolf's avatar
thomwolf committed
311

312
        r_head_k = r_head_k.view(rlen, self.n_head, self.d_head)  # qlen x n_head x d_head
thomwolf's avatar
thomwolf committed
313

314
        # compute attention score
315
316
        rw_head_q = w_head_q + self.r_w_bias  # qlen x bsz x n_head x d_head
        AC = torch.einsum("ibnd,jbnd->ijbn", (rw_head_q, w_head_k))  # qlen x klen x bsz x n_head
thomwolf's avatar
thomwolf committed
317

thomwolf's avatar
thomwolf committed
318
        rr_head_q = w_head_q + self.r_r_bias
319
        BD = torch.einsum("ibnd,jnd->ijbn", (rr_head_q, r_head_k))  # qlen x klen x bsz x n_head
thomwolf's avatar
thomwolf committed
320
321
322
323
324
325
        BD = self._rel_shift(BD)

        # [qlen x klen x bsz x n_head]
        attn_score = AC + BD
        attn_score.mul_(self.scale)

326
        # compute attention probability
327
        if attn_mask is not None and torch.sum(attn_mask).item():
328
            attn_mask = attn_mask == 1  # Switch to bool
thomwolf's avatar
thomwolf committed
329
            if attn_mask.dim() == 2:
330
                if next(self.parameters()).dtype == torch.float16:
331
332
333
                    attn_score = (
                        attn_score.float().masked_fill(attn_mask[None, :, :, None], -65000).type_as(attn_score)
                    )
334
                else:
335
                    attn_score = attn_score.float().masked_fill(attn_mask[None, :, :, None], -1e30).type_as(attn_score)
thomwolf's avatar
thomwolf committed
336
            elif attn_mask.dim() == 3:
337
                if next(self.parameters()).dtype == torch.float16:
338
                    attn_score = attn_score.float().masked_fill(attn_mask[:, :, :, None], -65000).type_as(attn_score)
339
                else:
340
                    attn_score = attn_score.float().masked_fill(attn_mask[:, :, :, None], -1e30).type_as(attn_score)
thomwolf's avatar
thomwolf committed
341
342
343
344
345

        # [qlen x klen x bsz x n_head]
        attn_prob = F.softmax(attn_score, dim=1)
        attn_prob = self.dropatt(attn_prob)

thomwolf's avatar
thomwolf committed
346
347
348
349
        # Mask heads if we want to
        if head_mask is not None:
            attn_prob = attn_prob * head_mask

350
        # compute attention vector
351
        attn_vec = torch.einsum("ijbn,jbnd->ibnd", (attn_prob, w_head_v))
thomwolf's avatar
thomwolf committed
352
353

        # [qlen x bsz x n_head x d_head]
354
        attn_vec = attn_vec.contiguous().view(attn_vec.size(0), attn_vec.size(1), self.n_head * self.d_head)
thomwolf's avatar
thomwolf committed
355

356
        # linear projection
thomwolf's avatar
thomwolf committed
357
358
359
360
        attn_out = self.o_net(attn_vec)
        attn_out = self.drop(attn_out)

        if self.pre_lnorm:
361
            # residual connection
thomwolf's avatar
thomwolf committed
362
            outputs = [w + attn_out]
thomwolf's avatar
thomwolf committed
363
        else:
364
            # residual connection + layer normalization
thomwolf's avatar
thomwolf committed
365
            outputs = [self.layer_norm(w + attn_out)]
thomwolf's avatar
thomwolf committed
366

367
        if output_attentions:
thomwolf's avatar
thomwolf committed
368
369
370
            outputs.append(attn_prob)

        return outputs
thomwolf's avatar
thomwolf committed
371
372
373


class RelPartialLearnableDecoderLayer(nn.Module):
374
    def __init__(self, n_head, d_model, d_head, d_inner, dropout, layer_norm_epsilon=1e-5, **kwargs):
Julien Chaumond's avatar
Julien Chaumond committed
375
        super().__init__()
thomwolf's avatar
thomwolf committed
376

377
378
379
380
381
382
        self.dec_attn = RelPartialLearnableMultiHeadAttn(
            n_head, d_model, d_head, dropout, layer_norm_epsilon=layer_norm_epsilon, **kwargs
        )
        self.pos_ff = PositionwiseFF(
            d_model, d_inner, dropout, pre_lnorm=kwargs.get("pre_lnorm"), layer_norm_epsilon=layer_norm_epsilon
        )
thomwolf's avatar
thomwolf committed
383

384
    def forward(self, dec_inp, r, dec_attn_mask=None, mems=None, head_mask=None, output_attentions=False):
thomwolf's avatar
thomwolf committed
385

386
387
388
        attn_outputs = self.dec_attn(
            dec_inp, r, attn_mask=dec_attn_mask, mems=mems, head_mask=head_mask, output_attentions=output_attentions,
        )
thomwolf's avatar
thomwolf committed
389
390
391
392
393
        ff_output = self.pos_ff(attn_outputs[0])

        outputs = [ff_output] + attn_outputs[1:]

        return outputs
thomwolf's avatar
thomwolf committed
394
395
396


class AdaptiveEmbedding(nn.Module):
397
    def __init__(self, n_token, d_embed, d_proj, cutoffs, div_val=1, sample_softmax=False):
Julien Chaumond's avatar
Julien Chaumond committed
398
        super().__init__()
thomwolf's avatar
thomwolf committed
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413

        self.n_token = n_token
        self.d_embed = d_embed

        self.cutoffs = cutoffs + [n_token]
        self.div_val = div_val
        self.d_proj = d_proj

        self.emb_scale = d_proj ** 0.5

        self.cutoff_ends = [0] + self.cutoffs

        self.emb_layers = nn.ModuleList()
        self.emb_projs = nn.ParameterList()
        if div_val == 1:
414
            self.emb_layers.append(nn.Embedding(n_token, d_embed, sparse=sample_softmax > 0))
thomwolf's avatar
thomwolf committed
415
            if d_proj != d_embed:
thomwolf's avatar
thomwolf committed
416
                self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_embed)))
thomwolf's avatar
thomwolf committed
417
418
        else:
            for i in range(len(self.cutoffs)):
419
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
thomwolf's avatar
thomwolf committed
420
                d_emb_i = d_embed // (div_val ** i)
421
                self.emb_layers.append(nn.Embedding(r_idx - l_idx, d_emb_i))
thomwolf's avatar
thomwolf committed
422
                self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_emb_i)))
thomwolf's avatar
thomwolf committed
423
424
425
426
427

    def forward(self, inp):
        if self.div_val == 1:
            embed = self.emb_layers[0](inp)
            if self.d_proj != self.d_embed:
428
                embed = F.linear(embed, self.emb_projs[0])
thomwolf's avatar
thomwolf committed
429
430
431
        else:
            param = next(self.parameters())
            inp_flat = inp.view(-1)
432
            emb_flat = torch.zeros([inp_flat.size(0), self.d_proj], dtype=param.dtype, device=param.device)
thomwolf's avatar
thomwolf committed
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
            for i in range(len(self.cutoffs)):
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]

                mask_i = (inp_flat >= l_idx) & (inp_flat < r_idx)
                indices_i = mask_i.nonzero().squeeze()

                if indices_i.numel() == 0:
                    continue

                inp_i = inp_flat.index_select(0, indices_i) - l_idx
                emb_i = self.emb_layers[i](inp_i)
                emb_i = F.linear(emb_i, self.emb_projs[i])

                emb_flat.index_copy_(0, indices_i, emb_i)

thomwolf's avatar
thomwolf committed
448
449
            embed_shape = inp.size() + (self.d_proj,)
            embed = emb_flat.view(embed_shape)
thomwolf's avatar
thomwolf committed
450
451
452
453
454
455

        embed.mul_(self.emb_scale)

        return embed


456
class TransfoXLPreTrainedModel(PreTrainedModel):
457
    """ An abstract class to handle weights initialization and
458
        a simple interface for downloading and loading pretrained models.
459
    """
460

461
462
463
464
465
    config_class = TransfoXLConfig
    load_tf_weights = load_tf_weights_in_transfo_xl
    base_model_prefix = "transformer"

    def _init_weight(self, weight):
466
        if self.config.init == "uniform":
467
            nn.init.uniform_(weight, -self.config.init_range, self.config.init_range)
468
        elif self.config.init == "normal":
469
            nn.init.normal_(weight, 0.0, self.config.init_std)
thomwolf's avatar
thomwolf committed
470

471
    def _init_bias(self, bias):
472
473
        nn.init.constant_(bias, 0.0)

474
    def _init_weights(self, m):
475
476
477
        """ Initialize the weights.
        """
        classname = m.__class__.__name__
478
479
        if classname.find("Linear") != -1:
            if hasattr(m, "weight") and m.weight is not None:
480
                self._init_weight(m.weight)
481
            if hasattr(m, "bias") and m.bias is not None:
482
                self._init_bias(m.bias)
483
484
        elif classname.find("AdaptiveEmbedding") != -1:
            if hasattr(m, "emb_projs"):
485
486
487
                for i in range(len(m.emb_projs)):
                    if m.emb_projs[i] is not None:
                        nn.init.normal_(m.emb_projs[i], 0.0, self.config.proj_init_std)
488
489
        elif classname.find("Embedding") != -1:
            if hasattr(m, "weight"):
490
                self._init_weight(m.weight)
491
492
        elif classname.find("ProjectedAdaptiveLogSoftmax") != -1:
            if hasattr(m, "cluster_weight") and m.cluster_weight is not None:
493
                self._init_weight(m.cluster_weight)
494
            if hasattr(m, "cluster_bias") and m.cluster_bias is not None:
495
                self._init_bias(m.cluster_bias)
496
            if hasattr(m, "out_projs"):
497
498
499
                for i in range(len(m.out_projs)):
                    if m.out_projs[i] is not None:
                        nn.init.normal_(m.out_projs[i], 0.0, self.config.proj_init_std)
500
501
        elif classname.find("LayerNorm") != -1:
            if hasattr(m, "weight"):
502
                nn.init.normal_(m.weight, 1.0, self.config.init_std)
503
            if hasattr(m, "bias") and m.bias is not None:
504
                self._init_bias(m.bias)
505
        else:
506
            if hasattr(m, "r_emb"):
507
                self._init_weight(m.r_emb)
508
            if hasattr(m, "r_w_bias"):
509
                self._init_weight(m.r_w_bias)
510
            if hasattr(m, "r_r_bias"):
511
                self._init_weight(m.r_r_bias)
512
            if hasattr(m, "r_bias"):
513
                self._init_bias(m.r_bias)
thomwolf's avatar
thomwolf committed
514

515
516
517
518
519
520
521
522
523
524
525
526
527
528
529
530
531
532
533
534
535
536
537
538
539
540
541
542
543
544
545
546
547
548
549
550
551
552
553
554
555
556
557
558
559
560
561
562
563
564
565
566
567
568
569
570
571
572
573
574
575
576
577
578
579
580
581
582
583
584
585
586
587
588
589
590
591
592
593
    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, layer: Optional[int] = -1):
        """ Resize input token embeddings matrix of the model if new_num_tokens != config.vocab_size.
        Take care of tying weights embeddings afterwards if the model class has a `tie_weights()` method.

        Arguments:

            new_num_tokens: (`optional`) int:
                New number of tokens in the embedding matrix. Increasing the size will add newly initialized vectors at the end. Reducing the size will remove vectors from the end.
                If not provided or None: does nothing and just returns a pointer to the input tokens ``torch.nn.Embeddings`` Module of the model.
            layer: (`optional`) int:
                Layer of the `AdaptiveEmbedding` where the resizing should be done. Per default the last layer will be resized.
                Be aware that when resizing other than the last layer, you have to ensure that the new token(s) in the tokenizer are at the corresponding position.

        Return: ``torch.nn.Embeddings``
            Pointer to the input tokens Embeddings Module of the model
        """
        base_model = getattr(self, self.base_model_prefix, self)  # get the base model if needed

        if new_num_tokens is None:
            return self.get_input_embeddings()

        new_num_tokens_layer, layer = self._get_new_num_tokens_layer(new_num_tokens, layer)
        assert new_num_tokens_layer > 0, "The size of the new embedding layer cannot be 0 or less"
        model_embeds = base_model._resize_token_embeddings(new_num_tokens_layer, layer)

        # Update base model and current model config
        self.config.vocab_size = new_num_tokens
        base_model.vocab_size = new_num_tokens
        base_model.n_token = new_num_tokens

        new_embedding_shapes = self._get_embedding_shapes()
        self._resize_cutoffs(new_num_tokens, new_num_tokens_layer, new_embedding_shapes, layer)

        # Tie weights again if needed
        self.tie_weights()

        return model_embeds

    def _get_new_num_tokens_layer(self, new_num_tokens, layer):
        embeddings = self.get_input_embeddings()
        if layer == -1:
            layer = len(embeddings.emb_layers) - 1
        assert 0 <= layer <= len(embeddings.emb_layers) - 1

        new_num_tokens_layer = (
            new_num_tokens
            - sum([emb.weight.shape[0] for emb in embeddings.emb_layers[:layer]])
            - sum([emb.weight.shape[0] for emb in embeddings.emb_layers[layer + 1 :]])
        )
        return new_num_tokens_layer, layer

    def _get_embedding_shapes(self):
        embeddings = self.get_input_embeddings()
        return [emb.weight.shape[0] for emb in embeddings.emb_layers]

    def _resize_token_embeddings(self, new_num_tokens, layer=-1):
        embeddings = self.get_input_embeddings()
        if new_num_tokens is None:
            return embeddings
        new_embeddings_layer = self._get_resized_embeddings(embeddings.emb_layers[layer], new_num_tokens)
        embeddings.emb_layers[layer] = new_embeddings_layer

        self.set_input_embeddings(embeddings)

        return self.get_input_embeddings()

    def _resize_cutoffs(self, new_num_tokens, new_emb_size, new_embedding_shapes, layer):
        embeddings = self.get_input_embeddings()

        for i in range(layer, len(embeddings.cutoffs)):
            embeddings.cutoffs[i] = sum(new_embedding_shapes[: i + 1])

        embeddings.cutoff_ends = [0] + embeddings.cutoffs
        embeddings.n_token = new_num_tokens

        self.config.cutoffs = embeddings.cutoffs[:-1]

        return embeddings.cutoffs

594

595
596
597
598
599
600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
@dataclass
class TransfoXLModelOutput(ModelOutput):
    """
    Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).

    Args:
        last_hidden_state (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`):
            Sequence of hidden-states at the output of the last layer of the model.
        mems (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
            Contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `mems` 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.
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

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

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

    last_hidden_state: torch.FloatTensor
621
    mems: List[torch.FloatTensor] = None
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
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


@dataclass
class TransfoXLLMHeadModelOutput(ModelOutput):
    """
    Base class for model's outputs that may also contain a past key/values (to speed up sequential decoding).

    Args:
        losses (:obj:`torch.FloatTensor` of shape `(batch_size, sequence_length-1)`, `optional`, returned when ``labels`` is provided)
            Language modeling losses (not reduced).
        prediction_scores (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, config.vocab_size)`):
            Prediction scores of the language modeling head (scores for each vocabulary token after SoftMax).
        mems (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
            Contains pre-computed hidden-states (key and values in the attention blocks).
            Can be used (see `mems` 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.
        hidden_states (:obj:`tuple(torch.FloatTensor)`, `optional`, returned when ``output_hidden_states=True`` is passed or when ``config.output_hidden_states=True``):
            Tuple of :obj:`torch.FloatTensor` (one for the output of the embeddings + one for the output of each layer)
            of shape :obj:`(batch_size, sequence_length, hidden_size)`.

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

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

653
654
655
    losses: Optional[torch.FloatTensor] = None
    prediction_scores: torch.FloatTensor = None
    mems: List[torch.FloatTensor] = None
656
657
658
659
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


Lysandre's avatar
Lysandre committed
660
TRANSFO_XL_START_DOCSTRING = r"""
thomwolf's avatar
thomwolf committed
661

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

thomwolf's avatar
thomwolf committed
666
    Parameters:
667
        config (:class:`~transformers.TransfoXLConfig`): Model configuration class with all the parameters of the model.
668
            Initializing with a config file does not load the weights associated with the model, only the configuration.
669
            Check out the :meth:`~transformers.PreTrainedModel.from_pretrained` method to load the model weights.
thomwolf's avatar
thomwolf committed
670
"""
thomwolf's avatar
thomwolf committed
671

thomwolf's avatar
thomwolf committed
672
TRANSFO_XL_INPUTS_DOCSTRING = r"""
Lysandre's avatar
Lysandre committed
673
674
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`):
Lysandre's avatar
Lysandre committed
675
676
            Indices of input sequence tokens in the vocabulary.

677
678
            Indices can be obtained using :class:`transformers.TransfoXLTokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
679
            :func:`transformers.PreTrainedTokenizer.__call__` for details.
Lysandre's avatar
Lysandre committed
680

Lysandre's avatar
Lysandre committed
681
682
683
684
685
686
            `What are input IDs? <../glossary.html#input-ids>`__
        mems (:obj:`List[torch.FloatTensor]` of length :obj:`config.n_layers`):
            Contains pre-computed hidden-states (key and values in the attention blocks) as computed by the model
            (see `mems` output below). Can be used to speed up sequential decoding. The token ids which have their mems
            given to this model should not be passed as input ids as they have already been computed.
        head_mask (:obj:`torch.FloatTensor` of shape :obj:`(num_heads,)` or :obj:`(num_layers, num_heads)`, `optional`, defaults to :obj:`None`):
thomwolf's avatar
thomwolf committed
687
            Mask to nullify selected heads of the self-attention modules.
thomwolf's avatar
thomwolf committed
688
            Mask values selected in ``[0, 1]``:
Lysandre's avatar
Lysandre committed
689
            :obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
flozi00's avatar
flozi00 committed
690
        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
Lysandre's avatar
Lysandre committed
691
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
692
693
            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.
ZhuBaohe's avatar
ZhuBaohe committed
694
        output_attentions (:obj:`bool`, `optional`, defaults to :obj:`None`):
695
            If set to ``True``, the attentions tensors of all attention layers are returned. See ``attentions`` under returned tensors for more detail.
696
697
        output_hidden_states (:obj:`bool`, `optional`, defaults to :obj:`None`):
            If set to ``True``, the hidden states of all layers are returned. See ``hidden_states`` under returned tensors for more detail.
698
699
700
        return_dict (:obj:`bool`, `optional`, defaults to :obj:`None`):
            If set to ``True``, the model will return a :class:`~transformers.file_utils.ModelOutput` instead of a
            plain tuple.
thomwolf's avatar
thomwolf committed
701
"""
702

703
704
705
706
707

@add_start_docstrings(
    "The bare Bert Model transformer outputting raw hidden-states without any specific head on top.",
    TRANSFO_XL_START_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
708
class TransfoXLModel(TransfoXLPreTrainedModel):
709
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
710
        super().__init__(config)
thomwolf's avatar
thomwolf committed
711

thomwolf's avatar
thomwolf committed
712
        self.n_token = config.vocab_size
713
714
715
716
717
718

        self.d_embed = config.d_embed
        self.d_model = config.d_model
        self.n_head = config.n_head
        self.d_head = config.d_head

719
720
721
        self.word_emb = AdaptiveEmbedding(
            config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val
        )
thomwolf's avatar
thomwolf committed
722

723
        self.drop = nn.Dropout(config.dropout)
thomwolf's avatar
thomwolf committed
724

725
726
727
728
729
730
731
732
733
734
        self.n_layer = config.n_layer

        self.tgt_len = config.tgt_len
        self.mem_len = config.mem_len
        self.ext_len = config.ext_len
        self.max_klen = config.tgt_len + config.ext_len + config.mem_len

        self.attn_type = config.attn_type

        if not config.untie_r:
thomwolf's avatar
thomwolf committed
735
736
            self.r_w_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
            self.r_r_bias = nn.Parameter(torch.FloatTensor(self.n_head, self.d_head))
thomwolf's avatar
thomwolf committed
737

thomwolf's avatar
thomwolf committed
738
        self.layers = nn.ModuleList()
739
        if config.attn_type == 0:  # the default attention
740
            for i in range(config.n_layer):
thomwolf's avatar
thomwolf committed
741
742
                self.layers.append(
                    RelPartialLearnableDecoderLayer(
743
744
745
746
747
748
749
750
751
752
                        config.n_head,
                        config.d_model,
                        config.d_head,
                        config.d_inner,
                        config.dropout,
                        tgt_len=config.tgt_len,
                        ext_len=config.ext_len,
                        mem_len=config.mem_len,
                        dropatt=config.dropatt,
                        pre_lnorm=config.pre_lnorm,
753
                        r_w_bias=None if config.untie_r else self.r_w_bias,
thomwolf's avatar
thomwolf committed
754
                        r_r_bias=None if config.untie_r else self.r_r_bias,
755
756
                        layer_norm_epsilon=config.layer_norm_epsilon,
                    )
thomwolf's avatar
thomwolf committed
757
                )
758
        else:  # learnable embeddings and absolute embeddings are not used in our pretrained checkpoints
759
            raise NotImplementedError  # Removed them to avoid maintaining dead code
thomwolf's avatar
thomwolf committed
760

761
762
        self.same_length = config.same_length
        self.clamp_len = config.clamp_len
thomwolf's avatar
thomwolf committed
763

764
        if self.attn_type == 0:  # default attention
thomwolf's avatar
thomwolf committed
765
            self.pos_emb = PositionalEmbedding(self.d_model)
766
        else:  # learnable embeddings and absolute embeddings
767
            raise NotImplementedError  # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
thomwolf's avatar
thomwolf committed
768

769
        self.init_weights()
thomwolf's avatar
thomwolf committed
770

thomwolf's avatar
thomwolf committed
771
    def get_input_embeddings(self):
thomwolf's avatar
thomwolf committed
772
        return self.word_emb
thomwolf's avatar
thomwolf committed
773

thomwolf's avatar
thomwolf committed
774
    def set_input_embeddings(self, new_embeddings):
775
776
        self.word_emb = new_embeddings

thomwolf's avatar
thomwolf committed
777
778
779
    def backward_compatible(self):
        self.sample_softmax = -1

thomwolf's avatar
thomwolf committed
780
781
782
783
784
    def reset_length(self, tgt_len, ext_len, mem_len):
        self.tgt_len = tgt_len
        self.mem_len = mem_len
        self.ext_len = ext_len

thomwolf's avatar
thomwolf committed
785
786
787
788
    def _prune_heads(self, heads):
        logger.info("Head pruning is not implemented for Transformer-XL model")
        pass

789
    def init_mems(self, bsz):
thomwolf's avatar
thomwolf committed
790
791
792
        if self.mem_len > 0:
            mems = []
            param = next(self.parameters())
793
            for i in range(self.n_layer):
794
                empty = torch.zeros(self.mem_len, bsz, self.config.d_model, dtype=param.dtype, device=param.device)
thomwolf's avatar
thomwolf committed
795
796
797
798
799
800
                mems.append(empty)

            return mems
        else:
            return None

801
    def _update_mems(self, hids, mems, mlen, qlen):
thomwolf's avatar
thomwolf committed
802
        # does not deal with None
803
804
        if mems is None:
            return None
thomwolf's avatar
thomwolf committed
805
806

        # mems is not None
807
        assert len(hids) == len(mems), "len(hids) != len(mems)"
thomwolf's avatar
thomwolf committed
808
809
810
811
812
813
814
815
816
817
818
819
820
821
822
823
824

        # There are `mlen + qlen` steps that can be cached into mems
        # For the next step, the last `ext_len` of the `qlen` tokens
        # will be used as the extended context. Hence, we only cache
        # the tokens from `mlen + qlen - self.ext_len - self.mem_len`
        # to `mlen + qlen - self.ext_len`.
        with torch.no_grad():
            new_mems = []
            end_idx = mlen + max(0, qlen - 0 - self.ext_len)
            beg_idx = max(0, end_idx - self.mem_len)
            for i in range(len(hids)):

                cat = torch.cat([mems[i], hids[i]], dim=0)
                new_mems.append(cat[beg_idx:end_idx].detach())

        return new_mems

Lysandre's avatar
Lysandre committed
825
    @add_start_docstrings_to_callable(TRANSFO_XL_INPUTS_DOCSTRING)
826
827
828
829
830
831
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="transfo-xl-wt103",
        output_type=TransfoXLModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
Joseph Liu's avatar
Joseph Liu committed
832
833
834
835
836
837
838
839
    def forward(
        self,
        input_ids=None,
        mems=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
840
        return_dict=None,
Joseph Liu's avatar
Joseph Liu committed
841
    ):
842
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
Joseph Liu's avatar
Joseph Liu committed
843
844
845
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
846
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
847

848
849
        # the original code for Transformer-XL used shapes [len, bsz] but we want a unified interface in the library
        # so we transpose here from shape [bsz, len] to shape [len, bsz]
850
851
852
853
854
855
856
857
858
859
        if input_ids is not None and inputs_embeds is not None:
            raise ValueError("You cannot specify both input_ids and inputs_embeds at the same time")
        elif input_ids is not None:
            input_ids = input_ids.transpose(0, 1).contiguous()
            qlen, bsz = input_ids.size()
        elif inputs_embeds is not None:
            inputs_embeds = inputs_embeds.transpose(0, 1).contiguous()
            qlen, bsz = inputs_embeds.shape[0], inputs_embeds.shape[1]
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")
860
861

        if mems is None:
862
            mems = self.init_mems(bsz)
thomwolf's avatar
thomwolf committed
863

thomwolf's avatar
thomwolf committed
864
865
866
867
868
869
870
871
872
873
874
        # Prepare head mask if needed
        # 1.0 in head_mask indicate we keep the head
        # attention_probs has shape bsz x n_heads x N x N
        # input head_mask has shape [num_heads] or [num_hidden_layers x num_heads] (a head_mask for each layer)
        # and head_mask is converted to shape [num_hidden_layers x qlen x klen x bsz x n_head]
        if head_mask is not None:
            if head_mask.dim() == 1:
                head_mask = head_mask.unsqueeze(0).unsqueeze(0).unsqueeze(0).unsqueeze(0)
                head_mask = head_mask.expand(self.n_layer, -1, -1, -1, -1)
            elif head_mask.dim() == 2:
                head_mask = head_mask.unsqueeze(1).unsqueeze(1).unsqueeze(1)
875
876
877
            head_mask = head_mask.to(
                dtype=next(self.parameters()).dtype
            )  # switch to fload if need + fp16 compatibility
thomwolf's avatar
thomwolf committed
878
879
880
        else:
            head_mask = [None] * self.n_layer

881
882
883
884
        if inputs_embeds is not None:
            word_emb = inputs_embeds
        else:
            word_emb = self.word_emb(input_ids)
thomwolf's avatar
thomwolf committed
885
886
887
888

        mlen = mems[0].size(0) if mems is not None else 0
        klen = mlen + qlen
        if self.same_length:
thomwolf's avatar
thomwolf committed
889
            all_ones = word_emb.new_ones((qlen, klen), dtype=torch.uint8)
thomwolf's avatar
thomwolf committed
890
891
892
893
894
            mask_len = klen - self.mem_len
            if mask_len > 0:
                mask_shift_len = qlen - mask_len
            else:
                mask_shift_len = qlen
895
            dec_attn_mask = (torch.triu(all_ones, 1 + mlen) + torch.tril(all_ones, -mask_shift_len))[:, :, None]  # -1
thomwolf's avatar
thomwolf committed
896
        else:
897
898
899
            dec_attn_mask = torch.triu(word_emb.new_ones((qlen, klen), dtype=torch.uint8), diagonal=1 + mlen)[
                :, :, None
            ]
thomwolf's avatar
thomwolf committed
900
901

        hids = []
902
        attentions = [] if output_attentions else None
903
904
        if self.attn_type == 0:  # default
            pos_seq = torch.arange(klen - 1, -1, -1.0, device=word_emb.device, dtype=word_emb.dtype)
thomwolf's avatar
thomwolf committed
905
906
907
908
909
910
911
912
            if self.clamp_len > 0:
                pos_seq.clamp_(max=self.clamp_len)
            pos_emb = self.pos_emb(pos_seq)

            core_out = self.drop(word_emb)
            pos_emb = self.drop(pos_emb)

            for i, layer in enumerate(self.layers):
913
                hids.append(core_out)
thomwolf's avatar
thomwolf committed
914
                mems_i = None if mems is None else mems[i]
915
                layer_outputs = layer(
916
917
918
919
920
921
                    core_out,
                    pos_emb,
                    dec_attn_mask=dec_attn_mask,
                    mems=mems_i,
                    head_mask=head_mask[i],
                    output_attentions=output_attentions,
922
                )
thomwolf's avatar
thomwolf committed
923
                core_out = layer_outputs[0]
924
                if output_attentions:
thomwolf's avatar
thomwolf committed
925
                    attentions.append(layer_outputs[1])
926
        else:  # learnable embeddings and absolute embeddings
927
            raise NotImplementedError  # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
thomwolf's avatar
thomwolf committed
928
929
930
931
932

        core_out = self.drop(core_out)

        new_mems = self._update_mems(hids, mems, mlen, qlen)

Joseph Liu's avatar
Joseph Liu committed
933
        if output_hidden_states:
thomwolf's avatar
thomwolf committed
934
935
            # Add last layer and transpose to library standard shape [bsz, len, hidden_dim]
            hids.append(core_out)
936
937
938
            hids = tuple(t.transpose(0, 1).contiguous() for t in hids)
        else:
            hids = None
939
        if output_attentions:
thomwolf's avatar
thomwolf committed
940
            # Transpose to library standard shape [bsz, n_heads, query_seq_len, key_seq_len]
941
942
943
944
            attentions = tuple(t.permute(2, 3, 0, 1).contiguous() for t in attentions)
        # We transpose back here to shape [bsz, len, hidden_dim]
        core_out = core_out.transpose(0, 1).contiguous()

945
        if not return_dict:
946
            return tuple(v for v in [core_out, new_mems, hids, attentions] if v is not None)
947

948
949
950
        return TransfoXLModelOutput(
            last_hidden_state=core_out, mems=new_mems, hidden_states=hids, attentions=attentions,
        )
thomwolf's avatar
thomwolf committed
951
952


953
954
@add_start_docstrings(
    """The Transformer-XL Model with a language modeling head on top
thomwolf's avatar
thomwolf committed
955
    (adaptive softmax with weights tied to the adaptive input embeddings)""",
956
957
    TRANSFO_XL_START_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
958
959
class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
960
        super().__init__(config)
thomwolf's avatar
thomwolf committed
961
962
        self.transformer = TransfoXLModel(config)
        self.sample_softmax = config.sample_softmax
963
964
965
966
967
968
969
970
971

        assert (
            self.sample_softmax <= 0
        ), "Sampling from the softmax is not implemented yet. Please look at issue: #3310: https://github.com/huggingface/transformers/issues/3310"

        self.crit = ProjectedAdaptiveLogSoftmax(
            config.vocab_size, config.d_embed, config.d_model, config.cutoffs, div_val=config.div_val
        )

972
        self.init_weights()
thomwolf's avatar
thomwolf committed
973
974

    def tie_weights(self):
975
976
977
        """
        Run this to be sure output and input (adaptive) softmax weights are tied
        """
978
979
980
981
982
983
984
985
986
987
988
989
990
991
992
993

        if self.config.tie_weight:
            for i in range(len(self.crit.out_layers)):
                self._tie_or_clone_weights(self.crit.out_layers[i], self.transformer.word_emb.emb_layers[i])
        if self.config.tie_projs:
            for i, tie_proj in enumerate(self.config.tie_projs):
                if tie_proj and self.config.div_val == 1 and self.config.d_model != self.config.d_embed:
                    if self.config.torchscript:
                        self.crit.out_projs[i] = nn.Parameter(self.transformer.word_emb.emb_projs[0].clone())
                    else:
                        self.crit.out_projs[i] = self.transformer.word_emb.emb_projs[0]
                elif tie_proj and self.config.div_val != 1:
                    if self.config.torchscript:
                        self.crit.out_projs[i] = nn.Parameter(self.transformer.word_emb.emb_projs[i].clone())
                    else:
                        self.crit.out_projs[i] = self.transformer.word_emb.emb_projs[i]
thomwolf's avatar
thomwolf committed
994
995
996
997

    def reset_length(self, tgt_len, ext_len, mem_len):
        self.transformer.reset_length(tgt_len, ext_len, mem_len)

998
999
    def init_mems(self, bsz):
        return self.transformer.init_mems(bsz)
thomwolf's avatar
thomwolf committed
1000

Lysandre's avatar
Lysandre committed
1001
    @add_start_docstrings_to_callable(TRANSFO_XL_INPUTS_DOCSTRING)
1002
1003
1004
1005
1006
1007
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="transfo-xl-wt103",
        output_type=TransfoXLLMHeadModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
1008
    def forward(
Joseph Liu's avatar
Joseph Liu committed
1009
1010
1011
1012
1013
1014
1015
1016
        self,
        input_ids=None,
        mems=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
1017
        return_dict=None,
1018
    ):
Lysandre's avatar
Lysandre committed
1019
1020
1021
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for language modeling.
1022
            Note that the labels **are shifted** inside the model, i.e. you can set ``labels = input_ids``
Lysandre's avatar
Lysandre committed
1023
1024
            Indices are selected in ``[-100, 0, ..., config.vocab_size]``
            All labels set to ``-100`` are ignored (masked), the loss is only
Lysandre's avatar
Lysandre committed
1025
1026
            computed for labels in ``[0, ..., config.vocab_size]``
        """
1027
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1028
1029
1030
1031
1032
1033
        if input_ids is not None:
            bsz, tgt_len = input_ids.size(0), input_ids.size(1)
        elif inputs_embeds is not None:
            bsz, tgt_len = inputs_embeds.size(0), inputs_embeds.size(1)
        else:
            raise ValueError("You have to specify either input_ids or inputs_embeds")
thomwolf's avatar
thomwolf committed
1034

1035
        transformer_outputs = self.transformer(
Joseph Liu's avatar
Joseph Liu committed
1036
1037
1038
1039
1040
1041
            input_ids,
            mems=mems,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
1042
            return_dict=return_dict,
1043
        )
thomwolf's avatar
thomwolf committed
1044

thomwolf's avatar
thomwolf committed
1045
        last_hidden = transformer_outputs[0]
1046
        pred_hid = last_hidden[:, -tgt_len:]
1047

1048
        softmax_output = self.crit(pred_hid, labels)
1049
1050
1051
        prediction_scores = softmax_output.view(bsz, tgt_len, -1) if labels is None else ()
        loss = softmax_output.view(bsz, tgt_len - 1) if labels is not None else None

1052
        if not return_dict:
Sylvain Gugger's avatar
Sylvain Gugger committed
1053
            output = (prediction_scores,) + transformer_outputs[1:]
1054
1055
1056
1057
1058
1059
1060
1061
1062
            return ((loss,) + output) if loss is not None else output

        return TransfoXLLMHeadModelOutput(
            losses=loss,
            prediction_scores=prediction_scores,
            mems=transformer_outputs.mems,
            hidden_states=transformer_outputs.hidden_states,
            attentions=transformer_outputs.attentions,
        )
R茅mi Louf's avatar
R茅mi Louf committed
1063
1064
1065
1066
1067
1068
1069
1070

    def get_output_embeddings(self):
        """ Double-check if you are using adaptive softmax.
        """
        if self.sample_softmax > 0:
            return self.out_layer
        else:
            return self.crit.out_layers[-1]
1071

1072
    def prepare_inputs_for_generation(self, input_ids, past, **model_kwargs):
1073
        inputs = {}
1074
1075

        # if past is defined in model kwargs then use it for faster decoding
1076
1077
        if past:
            inputs["mems"] = past
1078
1079
1080
            inputs["input_ids"] = input_ids[:, -1].unsqueeze(-1)
        else:
            inputs["input_ids"] = input_ids
1081
1082

        return inputs
1083
1084
1085
1086
1087
1088
1089

    def _resize_cutoffs(self, new_num_tokens, new_emb_size, new_embedding_shapes, layer):
        new_cutoffs = super()._resize_cutoffs(new_num_tokens, new_emb_size, new_embedding_shapes, layer)

        self.crit.cutoffs = new_cutoffs
        self.crit.cutoff_ends = [0] + new_cutoffs
        self.crit.n_token = new_num_tokens