modeling_transfo_xl.py 45.4 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

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

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

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

thomwolf's avatar
thomwolf committed
35

Lysandre Debut's avatar
Lysandre Debut committed
36
logger = logging.get_logger(__name__)
thomwolf's avatar
thomwolf committed
37

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
def build_tf_to_pytorch_map(model, config):
Lysandre's avatar
Lysandre committed
48
49
    """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.
50
51
    """
    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
            layer_str = "transformer/adaptive_softmax/cutoff_%d/" % i
65
            if config.tie_word_embeddings:
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
def load_tf_weights_in_transfo_xl(model, config, tf_path):
Lysandre's avatar
Lysandre committed
115
    """Load tf checkpoints in a pytorch model"""
116
117
118
    try:
        import numpy as np
        import tensorflow as tf
thomwolf's avatar
thomwolf committed
119
    except ImportError:
120
121
122
123
        logger.error(
            "Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
            "https://www.tensorflow.org/install/ for installation instructions."
        )
124
        raise
125
126
127
128
129
130
131
    # 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
132
        logger.info("Loading TF weight {} with shape {}".format(name, shape))
133
134
135
136
137
138
139
140
        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
141
        if "kernel" in name or "proj" in name:
142
            array = np.transpose(array)
143
        if ("r_r_bias" in name or "r_w_bias" in name) and len(pointer) > 1:
Julien Chaumond's avatar
Julien Chaumond committed
144
            # Here we will split the TF weights
145
146
147
148
149
150
151
152
            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
153
                logger.info("Initialize PyTorch weight {} for layer {}".format(name, i))
154
155
156
                p_i.data = torch.from_numpy(arr_i)
        else:
            try:
Teven's avatar
Teven committed
157
158
159
                assert (
                    pointer.shape == array.shape
                ), f"Pointer shape {pointer.shape} and array shape {array.shape} mismatched"
160
161
162
            except AssertionError as e:
                e.args += (pointer.shape, array.shape)
                raise
thomwolf's avatar
thomwolf committed
163
            logger.info("Initialize PyTorch weight {}".format(name))
164
165
            pointer.data = torch.from_numpy(array)
        tf_weights.pop(name, None)
166
167
        tf_weights.pop(name + "/Adam", None)
        tf_weights.pop(name + "/Adam_1", None)
168

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


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

        self.demb = demb

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

    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:
187
            return pos_emb[:, None, :].expand(-1, bsz, -1)
thomwolf's avatar
thomwolf committed
188
        else:
189
            return pos_emb[:, None, :]
thomwolf's avatar
thomwolf committed
190

thomwolf's avatar
thomwolf committed
191

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

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

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

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

        self.pre_lnorm = pre_lnorm

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

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

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

        return output

thomwolf's avatar
thomwolf committed
228

229
class RelPartialLearnableMultiHeadAttn(nn.Module):
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
    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
245
        super().__init__()
thomwolf's avatar
thomwolf committed
246
247
248
249
250
251
252
253
254
255
256
257

        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)

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

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

        self.pre_lnorm = pre_lnorm

264
        if r_r_bias is None or r_w_bias is None:  # Biases are not shared
thomwolf's avatar
thomwolf committed
265
266
            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
267
268
269
270
        else:
            self.r_r_bias = r_r_bias
            self.r_w_bias = r_w_bias

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

273
    def _rel_shift(self, x):
thomwolf's avatar
thomwolf committed
274
275
        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
276
277
        x_padded = torch.cat([zero_pad, x], dim=1)

thomwolf's avatar
thomwolf committed
278
279
        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
280
281
282
283
284

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

        return x

285
    def forward(self, w, r, attn_mask=None, mems=None, head_mask=None, output_attentions=False):
thomwolf's avatar
thomwolf committed
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
        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)

309
310
311
        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
312

313
        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
314

315
        # compute attention score
316
317
        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
318

thomwolf's avatar
thomwolf committed
319
        rr_head_q = w_head_q + self.r_r_bias
320
        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
321
322
323
324
325
326
        BD = self._rel_shift(BD)

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

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

        # [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
347
348
349
350
        # Mask heads if we want to
        if head_mask is not None:
            attn_prob = attn_prob * head_mask

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

        # [qlen x bsz x n_head x d_head]
355
        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
356

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

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

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

        return outputs
thomwolf's avatar
thomwolf committed
372
373
374


class RelPartialLearnableDecoderLayer(nn.Module):
375
    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
376
        super().__init__()
thomwolf's avatar
thomwolf committed
377

378
379
380
381
382
383
        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
384

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

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

        outputs = [ff_output] + attn_outputs[1:]

        return outputs
thomwolf's avatar
thomwolf committed
400
401
402


class AdaptiveEmbedding(nn.Module):
403
    def __init__(self, n_token, d_embed, d_proj, cutoffs, div_val=1, sample_softmax=False):
Julien Chaumond's avatar
Julien Chaumond committed
404
        super().__init__()
thomwolf's avatar
thomwolf committed
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419

        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:
420
            self.emb_layers.append(nn.Embedding(n_token, d_embed, sparse=sample_softmax > 0))
thomwolf's avatar
thomwolf committed
421
            if d_proj != d_embed:
thomwolf's avatar
thomwolf committed
422
                self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_embed)))
thomwolf's avatar
thomwolf committed
423
424
        else:
            for i in range(len(self.cutoffs)):
425
                l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i + 1]
thomwolf's avatar
thomwolf committed
426
                d_emb_i = d_embed // (div_val ** i)
427
                self.emb_layers.append(nn.Embedding(r_idx - l_idx, d_emb_i))
thomwolf's avatar
thomwolf committed
428
                self.emb_projs.append(nn.Parameter(torch.FloatTensor(d_proj, d_emb_i)))
thomwolf's avatar
thomwolf committed
429
430
431
432
433

    def forward(self, inp):
        if self.div_val == 1:
            embed = self.emb_layers[0](inp)
            if self.d_proj != self.d_embed:
434
                embed = F.linear(embed, self.emb_projs[0])
thomwolf's avatar
thomwolf committed
435
436
437
        else:
            param = next(self.parameters())
            inp_flat = inp.view(-1)
438
            emb_flat = torch.zeros([inp_flat.size(0), self.d_proj], dtype=param.dtype, device=param.device)
thomwolf's avatar
thomwolf committed
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
            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
454
455
            embed_shape = inp.size() + (self.d_proj,)
            embed = emb_flat.view(embed_shape)
thomwolf's avatar
thomwolf committed
456
457
458
459
460
461

        embed.mul_(self.emb_scale)

        return embed


462
class TransfoXLPreTrainedModel(PreTrainedModel):
Lysandre's avatar
Lysandre committed
463
464
    """An abstract class to handle weights initialization and
    a simple interface for downloading and loading pretrained models.
465
    """
466

467
468
469
470
471
    config_class = TransfoXLConfig
    load_tf_weights = load_tf_weights_in_transfo_xl
    base_model_prefix = "transformer"

    def _init_weight(self, weight):
472
        if self.config.init == "uniform":
473
            nn.init.uniform_(weight, -self.config.init_range, self.config.init_range)
474
        elif self.config.init == "normal":
475
            nn.init.normal_(weight, 0.0, self.config.init_std)
thomwolf's avatar
thomwolf committed
476

477
    def _init_bias(self, bias):
478
479
        nn.init.constant_(bias, 0.0)

480
    def _init_weights(self, m):
Lysandre's avatar
Lysandre committed
481
        """Initialize the weights."""
482
        classname = m.__class__.__name__
483
484
        if classname.find("Linear") != -1:
            if hasattr(m, "weight") and m.weight is not None:
485
                self._init_weight(m.weight)
486
            if hasattr(m, "bias") and m.bias is not None:
487
                self._init_bias(m.bias)
488
489
        elif classname.find("AdaptiveEmbedding") != -1:
            if hasattr(m, "emb_projs"):
490
491
492
                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)
493
494
        elif classname.find("Embedding") != -1:
            if hasattr(m, "weight"):
495
                self._init_weight(m.weight)
496
497
        elif classname.find("ProjectedAdaptiveLogSoftmax") != -1:
            if hasattr(m, "cluster_weight") and m.cluster_weight is not None:
498
                self._init_weight(m.cluster_weight)
499
            if hasattr(m, "cluster_bias") and m.cluster_bias is not None:
500
                self._init_bias(m.cluster_bias)
501
            if hasattr(m, "out_projs"):
502
503
504
                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)
505
506
        elif classname.find("LayerNorm") != -1:
            if hasattr(m, "weight"):
507
                nn.init.normal_(m.weight, 1.0, self.config.init_std)
508
            if hasattr(m, "bias") and m.bias is not None:
509
                self._init_bias(m.bias)
510
        else:
511
            if hasattr(m, "r_emb"):
512
                self._init_weight(m.r_emb)
513
            if hasattr(m, "r_w_bias"):
514
                self._init_weight(m.r_w_bias)
515
            if hasattr(m, "r_r_bias"):
516
                self._init_weight(m.r_r_bias)
517
            if hasattr(m, "r_bias"):
518
                self._init_bias(m.r_bias)
thomwolf's avatar
thomwolf committed
519

520
    def resize_token_embeddings(self, new_num_tokens: Optional[int] = None, layer: Optional[int] = -1):
Lysandre's avatar
Lysandre committed
521
        """Resize input token embeddings matrix of the model if new_num_tokens != config.vocab_size.
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
594
595
596
597
598
        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

599

600
601
602
603
604
605
606
607
608
609
610
611
612
613
614
615
616
617
618
619
620
621
622
623
624
625
@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
626
    mems: List[torch.FloatTensor] = None
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
    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.
    """

658
659
660
    losses: Optional[torch.FloatTensor] = None
    prediction_scores: torch.FloatTensor = None
    mems: List[torch.FloatTensor] = None
661
662
663
664
    hidden_states: Optional[Tuple[torch.FloatTensor]] = None
    attentions: Optional[Tuple[torch.FloatTensor]] = None


Lysandre's avatar
Lysandre committed
665
TRANSFO_XL_START_DOCSTRING = r"""
thomwolf's avatar
thomwolf committed
666

Lysandre's avatar
Lysandre committed
667
668
    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
669
    usage and behavior.
670

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

thomwolf's avatar
thomwolf committed
677
TRANSFO_XL_INPUTS_DOCSTRING = r"""
Lysandre's avatar
Lysandre committed
678
679
    Args:
        input_ids (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`):
Lysandre's avatar
Lysandre committed
680
681
            Indices of input sequence tokens in the vocabulary.

682
683
            Indices can be obtained using :class:`transformers.TransfoXLTokenizer`.
            See :func:`transformers.PreTrainedTokenizer.encode` and
684
            :func:`transformers.PreTrainedTokenizer.__call__` for details.
Lysandre's avatar
Lysandre committed
685

Lysandre's avatar
Lysandre committed
686
687
688
689
690
691
            `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
692
            Mask to nullify selected heads of the self-attention modules.
thomwolf's avatar
thomwolf committed
693
            Mask values selected in ``[0, 1]``:
Lysandre's avatar
Lysandre committed
694
            :obj:`1` indicates the head is **not masked**, :obj:`0` indicates the head is **masked**.
flozi00's avatar
flozi00 committed
695
        inputs_embeds (:obj:`torch.FloatTensor` of shape :obj:`(batch_size, sequence_length, hidden_size)`, `optional`, defaults to :obj:`None`):
Lysandre's avatar
Lysandre committed
696
            Optionally, instead of passing :obj:`input_ids` you can choose to directly pass an embedded representation.
697
698
            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
699
        output_attentions (:obj:`bool`, `optional`, defaults to :obj:`None`):
700
            If set to ``True``, the attentions tensors of all attention layers are returned. See ``attentions`` under returned tensors for more detail.
701
702
        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.
703
704
705
        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
706
"""
707

708
709
710
711
712

@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
713
class TransfoXLModel(TransfoXLPreTrainedModel):
714
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
715
        super().__init__(config)
thomwolf's avatar
thomwolf committed
716

thomwolf's avatar
thomwolf committed
717
        self.n_token = config.vocab_size
718
719
720
721
722
723

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

724
725
726
        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
727

728
        self.drop = nn.Dropout(config.dropout)
thomwolf's avatar
thomwolf committed
729

730
731
732
733
734
735
736
737
738
739
        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
740
741
            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
742

thomwolf's avatar
thomwolf committed
743
        self.layers = nn.ModuleList()
744
        if config.attn_type == 0:  # the default attention
745
            for i in range(config.n_layer):
thomwolf's avatar
thomwolf committed
746
747
                self.layers.append(
                    RelPartialLearnableDecoderLayer(
748
749
750
751
752
753
754
755
756
757
                        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,
758
                        r_w_bias=None if config.untie_r else self.r_w_bias,
thomwolf's avatar
thomwolf committed
759
                        r_r_bias=None if config.untie_r else self.r_r_bias,
760
761
                        layer_norm_epsilon=config.layer_norm_epsilon,
                    )
thomwolf's avatar
thomwolf committed
762
                )
763
        else:  # learnable embeddings and absolute embeddings are not used in our pretrained checkpoints
764
            raise NotImplementedError  # Removed them to avoid maintaining dead code
thomwolf's avatar
thomwolf committed
765

766
767
        self.same_length = config.same_length
        self.clamp_len = config.clamp_len
thomwolf's avatar
thomwolf committed
768

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

774
        self.init_weights()
thomwolf's avatar
thomwolf committed
775

thomwolf's avatar
thomwolf committed
776
    def get_input_embeddings(self):
thomwolf's avatar
thomwolf committed
777
        return self.word_emb
thomwolf's avatar
thomwolf committed
778

thomwolf's avatar
thomwolf committed
779
    def set_input_embeddings(self, new_embeddings):
780
781
        self.word_emb = new_embeddings

thomwolf's avatar
thomwolf committed
782
783
784
    def backward_compatible(self):
        self.sample_softmax = -1

thomwolf's avatar
thomwolf committed
785
786
787
788
789
    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
790
791
792
793
    def _prune_heads(self, heads):
        logger.info("Head pruning is not implemented for Transformer-XL model")
        pass

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

            return mems
        else:
            return None

806
    def _update_mems(self, hids, mems, mlen, qlen):
thomwolf's avatar
thomwolf committed
807
        # does not deal with None
808
809
        if mems is None:
            return None
thomwolf's avatar
thomwolf committed
810
811

        # mems is not None
812
        assert len(hids) == len(mems), "len(hids) != len(mems)"
thomwolf's avatar
thomwolf committed
813
814
815
816
817
818
819
820
821
822
823
824
825
826
827
828
829

        # 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
830
    @add_start_docstrings_to_callable(TRANSFO_XL_INPUTS_DOCSTRING)
831
832
833
834
835
836
    @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
837
838
839
840
841
842
843
844
    def forward(
        self,
        input_ids=None,
        mems=None,
        head_mask=None,
        inputs_embeds=None,
        output_attentions=None,
        output_hidden_states=None,
845
        return_dict=None,
Joseph Liu's avatar
Joseph Liu committed
846
    ):
847
        output_attentions = output_attentions if output_attentions is not None else self.config.output_attentions
Joseph Liu's avatar
Joseph Liu committed
848
849
850
        output_hidden_states = (
            output_hidden_states if output_hidden_states is not None else self.config.output_hidden_states
        )
851
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
852

853
854
        # 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]
855
856
857
858
859
860
861
862
863
864
        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")
865
866

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

thomwolf's avatar
thomwolf committed
869
870
871
872
873
874
875
876
877
878
879
        # 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)
880
881
882
            head_mask = head_mask.to(
                dtype=next(self.parameters()).dtype
            )  # switch to fload if need + fp16 compatibility
thomwolf's avatar
thomwolf committed
883
884
885
        else:
            head_mask = [None] * self.n_layer

886
887
888
889
        if inputs_embeds is not None:
            word_emb = inputs_embeds
        else:
            word_emb = self.word_emb(input_ids)
thomwolf's avatar
thomwolf committed
890
891
892
893

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

        hids = []
907
        attentions = [] if output_attentions else None
908
909
        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
910
911
912
913
914
915
916
917
            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):
918
                hids.append(core_out)
thomwolf's avatar
thomwolf committed
919
                mems_i = None if mems is None else mems[i]
920
                layer_outputs = layer(
921
922
923
924
925
926
                    core_out,
                    pos_emb,
                    dec_attn_mask=dec_attn_mask,
                    mems=mems_i,
                    head_mask=head_mask[i],
                    output_attentions=output_attentions,
927
                )
thomwolf's avatar
thomwolf committed
928
                core_out = layer_outputs[0]
929
                if output_attentions:
thomwolf's avatar
thomwolf committed
930
                    attentions.append(layer_outputs[1])
931
        else:  # learnable embeddings and absolute embeddings
932
            raise NotImplementedError  # Removed these to avoid maintaining dead code - They are not used in our pretrained checkpoint
thomwolf's avatar
thomwolf committed
933
934
935
936
937

        core_out = self.drop(core_out)

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

Joseph Liu's avatar
Joseph Liu committed
938
        if output_hidden_states:
thomwolf's avatar
thomwolf committed
939
940
            # Add last layer and transpose to library standard shape [bsz, len, hidden_dim]
            hids.append(core_out)
941
942
943
            hids = tuple(t.transpose(0, 1).contiguous() for t in hids)
        else:
            hids = None
944
        if output_attentions:
thomwolf's avatar
thomwolf committed
945
            # Transpose to library standard shape [bsz, n_heads, query_seq_len, key_seq_len]
946
947
948
949
            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()

950
        if not return_dict:
951
            return tuple(v for v in [core_out, new_mems, hids, attentions] if v is not None)
952

953
        return TransfoXLModelOutput(
Lysandre's avatar
Lysandre committed
954
955
956
957
            last_hidden_state=core_out,
            mems=new_mems,
            hidden_states=hids,
            attentions=attentions,
958
        )
thomwolf's avatar
thomwolf committed
959
960


961
962
@add_start_docstrings(
    """The Transformer-XL Model with a language modeling head on top
thomwolf's avatar
thomwolf committed
963
    (adaptive softmax with weights tied to the adaptive input embeddings)""",
964
965
    TRANSFO_XL_START_DOCSTRING,
)
thomwolf's avatar
thomwolf committed
966
967
class TransfoXLLMHeadModel(TransfoXLPreTrainedModel):
    def __init__(self, config):
Julien Chaumond's avatar
Julien Chaumond committed
968
        super().__init__(config)
thomwolf's avatar
thomwolf committed
969
970
        self.transformer = TransfoXLModel(config)
        self.sample_softmax = config.sample_softmax
971
972
973
974
975
976
977
978
979

        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
        )

980
        self.init_weights()
thomwolf's avatar
thomwolf committed
981
982

    def tie_weights(self):
983
984
985
        """
        Run this to be sure output and input (adaptive) softmax weights are tied
        """
986

987
        if self.config.tie_word_embeddings:
988
989
990
991
992
993
994
995
996
997
998
999
1000
1001
            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
1002
1003
1004
1005

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

1006
1007
    def init_mems(self, bsz):
        return self.transformer.init_mems(bsz)
thomwolf's avatar
thomwolf committed
1008

Lysandre's avatar
Lysandre committed
1009
    @add_start_docstrings_to_callable(TRANSFO_XL_INPUTS_DOCSTRING)
1010
1011
1012
1013
1014
1015
    @add_code_sample_docstrings(
        tokenizer_class=_TOKENIZER_FOR_DOC,
        checkpoint="transfo-xl-wt103",
        output_type=TransfoXLLMHeadModelOutput,
        config_class=_CONFIG_FOR_DOC,
    )
1016
    def forward(
Joseph Liu's avatar
Joseph Liu committed
1017
1018
1019
1020
1021
1022
1023
1024
        self,
        input_ids=None,
        mems=None,
        head_mask=None,
        inputs_embeds=None,
        labels=None,
        output_attentions=None,
        output_hidden_states=None,
1025
        return_dict=None,
1026
    ):
Lysandre's avatar
Lysandre committed
1027
1028
1029
        r"""
        labels (:obj:`torch.LongTensor` of shape :obj:`(batch_size, sequence_length)`, `optional`, defaults to :obj:`None`):
            Labels for language modeling.
1030
            Note that the labels **are shifted** inside the model, i.e. you can set ``labels = input_ids``
Lysandre's avatar
Lysandre committed
1031
1032
            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
1033
1034
            computed for labels in ``[0, ..., config.vocab_size]``
        """
1035
        return_dict = return_dict if return_dict is not None else self.config.use_return_dict
1036
1037
1038
1039
1040
1041
        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
1042

1043
        transformer_outputs = self.transformer(
Joseph Liu's avatar
Joseph Liu committed
1044
1045
1046
1047
1048
1049
            input_ids,
            mems=mems,
            head_mask=head_mask,
            inputs_embeds=inputs_embeds,
            output_attentions=output_attentions,
            output_hidden_states=output_hidden_states,
1050
            return_dict=return_dict,
1051
        )
thomwolf's avatar
thomwolf committed
1052

thomwolf's avatar
thomwolf committed
1053
        last_hidden = transformer_outputs[0]
1054
        pred_hid = last_hidden[:, -tgt_len:]
1055

1056
        softmax_output = self.crit(pred_hid, labels)
1057
1058
1059
        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

1060
        if not return_dict:
Sylvain Gugger's avatar
Sylvain Gugger committed
1061
            output = (prediction_scores,) + transformer_outputs[1:]
1062
1063
1064
1065
1066
1067
1068
1069
1070
            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
1071
1072

    def get_output_embeddings(self):
Lysandre's avatar
Lysandre committed
1073
        """Double-check if you are using adaptive softmax."""
R茅mi Louf's avatar
R茅mi Louf committed
1074
1075
1076
1077
        if self.sample_softmax > 0:
            return self.out_layer
        else:
            return self.crit.out_layers[-1]
1078

1079
    def prepare_inputs_for_generation(self, input_ids, past, **model_kwargs):
1080
        inputs = {}
1081
1082

        # if past is defined in model kwargs then use it for faster decoding
1083
1084
        if past:
            inputs["mems"] = past
1085
1086
1087
            inputs["input_ids"] = input_ids[:, -1].unsqueeze(-1)
        else:
            inputs["input_ids"] = input_ids
1088
1089

        return inputs
1090
1091
1092
1093
1094
1095
1096

    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