content: # coding=utf-8
  # Copyright 2018 Google AI, Google Brain and Carnegie Mellon University Authors and the HuggingFace Inc. team.
  # 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.
      Adapted from https://github.com/kimiyoung/transformer-xl.
      In particular https://github.com/kimiyoung/transformer-xl/blob/master/pytorch/mem_transformer.py
  """
  
  from __future__ import absolute_import, division, print_function, unicode_literals
  
  import os
  import copy
  import json
  import math
  import logging
  import tarfile
  import tempfile
  import shutil
  import collections
  import sys
  from io import open
  
  import torch
  import torch.nn as nn
  import torch.nn.functional as F
  from torch.nn import CrossEntropyLoss
  from torch.nn.parameter import Parameter
  
  from .modeling import BertLayerNorm as LayerNorm
  from .modeling_transfo_xl_utilities import ProjectedAdaptiveLogSoftmax, sample_logits
  from .file_utils import cached_path, CONFIG_NAME, WEIGHTS_NAME
  
  logger = logging.getLogger(__name__)
  
  PRETRAINED_MODEL_ARCHIVE_MAP = {
      'transfo-xl-wt103': "https://s3.amazonaws.com/models.huggingface.co/bert/transfo-xl-wt103-pytorch_model.bin",
  }
  PRETRAINED_CONFIG_ARCHIVE_MAP = {
      'transfo-xl-wt103': "https://s3.amazonaws.com/models.huggingface.co/bert/transfo-xl-wt103-config.json",
  }
  
  TF_WEIGHTS_NAME = 'model.ckpt'
  
  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 = {}
  
      if hasattr(model, 'transformer'):
          # We are loading in a TransfoXLLMHeadModel => we will load also the Adaptive Softmax
          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)):
              layer_str = "transformer/adaptive_softmax/cutoff_%d/" % i
              if config.tie_weight:
                  tf_to_pt_map.update({
                      layer_str + 'b': out_l.bias})
              else:
                  raise NotImplementedError
                  # I don't think this is implemented in the TF code
                  tf_to_pt_map.update({
                      layer_str + 'lookup_table': out_l.weight,
                      layer_str + 'b': out_l.bias})
              if not tie_proj:
                  tf_to_pt_map.update({
                      layer_str + 'proj': proj_l
                      })
          # Now load the rest of the transformer
          model = model.transformer
  
      # Embeddings
      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
          tf_to_pt_map.update({
              layer_str + 'lookup_table': embed_l.weight,
              layer_str + 'proj_W': proj_l
              })
  
      # Transformer blocks
      for i, b in enumerate(model.layers):
          layer_str = "transformer/layer_%d/" % i
          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,
          })
  
      # 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]
      tf_to_pt_map.update({
          'transformer/r_r_bias': r_r_list,
          'transformer/r_w_bias': r_w_list})
      return tf_to_pt_map
  
  def load_tf_weights_in_transfo_xl(model, config, tf_path):
      """ Load tf checkpoints in a pytorch model
      """
      try:
          import numpy as np
          import tensorflow as tf
      except ImportError:
          print("Loading a TensorFlow models in PyTorch, requires TensorFlow to be installed. Please see "
              "https://www.tensorflow.org/install/ for installation instructions.")
          raise
      # 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:
          print("Loading TF weight {} with shape {}".format(name, shape))
          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
          if 'kernel' in name or 'proj' in name:
              array = np.transpose(array)
          if ('r_r_bias' in name or 'r_w_bias' in name) and len(pointer) > 1:
              # Here we will split the TF weigths
              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
                  print("Initialize PyTorch weight {} for layer {}".format(name, i))
                  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
              print("Initialize PyTorch weight {}".format(name))
              pointer.data = torch.from_numpy(array)
          tf_weights.pop(name, None)
          tf_weights.pop(name + '/Adam', None)
          tf_weights.pop(name + '/Adam_1', None)
  
      print("Weights not copied to PyTorch model: {}".format(', '.join(tf_weights.keys())))
      return model
  
  
  class TransfoXLConfig(object):
      """Configuration class to store the configuration of a `TransfoXLModel`.
      """
      def __init__(self,
                   vocab_size_or_config_json_file=267735,
                   cutoffs=[20000, 40000, 200000],
                   d_model=1024,
                   d_embed=1024,
                   n_head=16,
                   d_head=64,
                   d_inner=4096,
                   div_val=4,
                   pre_lnorm=False,
                   n_layer=18,
                   tgt_len=128,
                   ext_len=0,
                   mem_len=1600,
                   clamp_len=1000,
                   same_length=True,
                   proj_share_all_but_first=True,
                   attn_type=0,
                   sample_softmax=-1,
                   adaptive=True,
                   tie_weight=True,
                   dropout=0.1,
                   dropatt=0.0,
                   untie_r=True,
                   init="normal",
                   init_range=0.01,
                   proj_init_std=0.01,
                   init_std=0.02):
          """Constructs TransfoXLConfig.
  
          Args:
              vocab_size_or_config_json_file: Vocabulary size of `inputs_ids` in `TransfoXLModel` or a configuration json file.
              cutoffs: cutoffs for the adaptive softmax
              d_model: Dimensionality of the model's hidden states.
              d_embed: Dimensionality of the embeddings
              d_head: Dimensionality of the model's heads.
              div_val: divident value for adapative input and softmax
              pre_lnorm: apply LayerNorm to the input instead of the output
              d_inner: Inner dimension in FF
              n_layer: Number of hidden layers in the Transformer encoder.
              n_head: Number of attention heads for each attention layer in
                  the Transformer encoder.
              tgt_len: number of tokens to predict
              ext_len: length of the extended context
              mem_len: length of the retained previous heads
              same_length: use the same attn length for all tokens
              proj_share_all_but_first: True to share all but first projs, False not to share.
              attn_type: attention type. 0 for Transformer-XL, 1 for Shaw et al, 2 for Vaswani et al, 3 for Al Rfou et al.
              clamp_len: use the same pos embeddings after clamp_len
              sample_softmax: number of samples in sampled softmax
              adaptive: use adaptive softmax
              tie_weight: tie the word embedding and softmax weights
              dropout: The dropout probabilitiy for all fully connected
                  layers in the embeddings, encoder, and pooler.
              dropatt: The dropout ratio for the attention probabilities.
              untie_r: untie relative position biases           
              embd_pdrop: The dropout ratio for the embeddings.
              init: parameter initializer to use
              init_range: parameters initialized by U(-init_range, init_range).
              proj_init_std: parameters initialized by N(0, init_std)
              init_std: parameters initialized by N(0, init_std)
          """
          if isinstance(vocab_size_or_config_json_file, str) or (sys.version_info[0] == 2
                          and isinstance(vocab_size_or_config_json_file, unicode)):
              with open(vocab_size_or_config_json_file, "r", encoding='utf-8') as reader:
                  json_config = json.loads(reader.read())
              for key, value in json_config.items():
                  self.__dict__[key] = value
          elif isinstance(vocab_size_or_config_json_file, int):
              self.n_token = vocab_size_or_config_json_file
              self.cutoffs = []
              self.cutoffs.extend(cutoffs)
              self.tie_weight = tie_weight
              if proj_share_all_but_first:
                  self.tie_projs = [False] + [True] * len(self.cutoffs)
              else:
                  self.tie_projs = [False] + [False] * len(self.cutoffs)
              self.d_model = d_model
              self.d_embed = d_embed
              self.d_head = d_head
              self.d_inner = d_inner
              self.div_val = div_val
              self.pre_lnorm = pre_lnorm
              self.n_layer = n_layer
              self.n_head = n_head
              self.tgt_len = tgt_len
              self.ext_len = ext_len
              self.mem_len = mem_len
              self.same_length = same_length
              self.attn_type = attn_type
              self.clamp_len = clamp_len
              self.sample_softmax = sample_softmax
              self.adaptive = adaptive
              self.dropout = dropout
              self.dropatt = dropatt
              self.untie_r = untie_r
              self.init = init
              self.init_range = init_range
              self.proj_init_std = proj_init_std
              self.init_std = init_std
          else:
              raise ValueError("First argument must be either a vocabulary size (int)"
                               "or the path to a pretrained model config file (str)")
  
      @classmethod
      def from_dict(cls, json_object):
          """Constructs a `TransfoXLConfig` from a Python dictionary of parameters."""
          config = TransfoXLConfig(vocab_size_or_config_json_file=-1)
          for key, value in json_object.items():
              config.__dict__[key] = value
          return config
  
      @classmethod
      def from_json_file(cls, json_file):
          """Constructs a `TransfoXLConfig` from a json file of parameters."""
          with open(json_file, "r", encoding='utf-8') as reader:
              text = reader.read()
          return cls.from_dict(json.loads(text))
  
      def __repr__(self):
          return str(self.to_json_string())
  
      def to_dict(self):
          """Serializes this instance to a Python dictionary."""
          output = copy.deepcopy(self.__dict__)
          return output
  
      def to_json_string(self):
          """Serializes this instance to a JSON string."""
          return json.dumps(self.to_dict(), indent=2, sort_keys=True) + "\n"
  
      def to_json_file(self, json_file_path):
          """ Save this instance to a json file."""
          with open(json_file_path, "w", encoding='utf-8') as writer:
              writer.write(self.to_json_string())
  
  
  class PositionalEmbedding(nn.Module):
      def __init__(self, demb):
          super(PositionalEmbedding, self).__init__()
  
          self.demb = demb
  
          inv_freq = 1 / (10000 ** (torch.arange(0.0, demb, 2.0) / demb))
          self.register_buffer('inv_freq', inv_freq)
  
      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:
              return pos_emb[:,None,:].expand(-1, bsz, -1)
          else:
              return pos_emb[:,None,:]
  
  
  class PositionwiseFF(nn.Module):
      def __init__(self, d_model, d_inner, dropout, pre_lnorm=False):
          super(PositionwiseFF, self).__init__()
  
          self.d_model = d_model
          self.d_inner = d_inner
          self.dropout = dropout
  
          self.CoreNet = nn.Sequential(
              nn.Linear(d_model, d_inner), nn.ReLU(inplace=True),
              nn.Dropout(dropout),
              nn.Linear(d_inner, d_model),
              nn.Dropout(dropout),
          )
  
          self.layer_norm = LayerNorm(d_model)
  
          self.pre_lnorm = pre_lnorm
  
      def forward(self, inp):
          if self.pre_lnorm:
              ##### layer normalization + positionwise feed-forward
              core_out = self.CoreNet(self.layer_norm(inp))
  
              ##### residual connection
              output = core_out + inp
          else:
              ##### positionwise feed-forward
              core_out = self.CoreNet(inp)
  
              ##### residual connection + layer normalization
              output = self.layer_norm(inp + core_out)
  
          return output
  
  class MultiHeadAttn(nn.Module):
      def __init__(self, n_head, d_model, d_head, dropout, dropatt=0, 
                   pre_lnorm=False, r_r_bias=None, r_w_bias=None):
          super(MultiHeadAttn, self).__init__()
  
          self.n_head = n_head
          self.d_model = d_model
          self.d_head = d_head
          self.dropout = dropout
  
          self.q_net = nn.Linear(d_model, n_head * d_head, bias=False)
          self.kv_net = nn.Linear(d_model, 2 * 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)
  
          self.layer_norm = LayerNorm(d_model)
  
          self.scale = 1 / (d_head ** 0.5)
  
          self.pre_lnorm = pre_lnorm
  
          if r_r_bias is None or r_w_bias is None: # Biases are not shared
              self.r_r_bias = nn.Parameter(torch.Tensor(self.n_head, self.d_head))
              self.r_w_bias = nn.Parameter(torch.Tensor(self.n_head, self.d_head))
          else:
              self.r_r_bias = r_r_bias
              self.r_w_bias = r_w_bias
  
      def forward(self, h, attn_mask=None, mems=None):
          ##### multihead attention
          # [hlen x bsz x n_head x d_head]
  
          if mems is not None:
              c = torch.cat([mems, h], 0)
          else:
              c = h
  
          if self.pre_lnorm:
              ##### layer normalization
              c = self.layer_norm(c)
  
          head_q = self.q_net(h)
          head_k, head_v = torch.chunk(self.kv_net(c), 2, -1)
  
          head_q = head_q.view(h.size(0), h.size(1), self.n_head, self.d_head)
          head_k = head_k.view(c.size(0), c.size(1), self.n_head, self.d_head)
          head_v = head_v.view(c.size(0), c.size(1), self.n_head, self.d_head)
  
          # [qlen x klen x bsz x n_head]
          attn_score = torch.einsum('ibnd,jbnd->ijbn', (head_q, head_k))
          attn_score.mul_(self.scale)
          if attn_mask is not None and attn_mask.any().item():
              if attn_mask.dim() == 2:
                  attn_score.masked_fill_(attn_mask[None,:,:,None], -float('inf'))
              elif attn_mask.dim() == 3:
                  attn_score.masked_fill_(attn_mask[:,:,:,None], -float('inf'))
  
          # [qlen x klen x bsz x n_head]
          attn_prob = F.softmax(attn_score, dim=1)
          attn_prob = self.dropatt(attn_prob)
  
          # [qlen x klen x bsz x n_head] + [klen x bsz x n_head x d_head] -> [qlen x bsz x n_head x d_head]
          attn_vec = torch.einsum('ijbn,jbnd->ibnd', (attn_prob, head_v))
          attn_vec = attn_vec.contiguous().view(
              attn_vec.size(0), attn_vec.size(1), self.n_head * self.d_head)
  
          ##### linear projection
          attn_out = self.o_net(attn_vec)
          attn_out = self.drop(attn_out)
  
          if self.pre_lnorm:
              ##### residual connection
              output = h + attn_out
          else:
              ##### residual connection + layer normalization
              output = self.layer_norm(h + attn_out)
  
          return output
  
  class RelMultiHeadAttn(nn.Module):
      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):
          super(RelMultiHeadAttn, self).__init__()
  
          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)
  
          self.layer_norm = LayerNorm(d_model)
  
          self.scale = 1 / (d_head ** 0.5)
  
          self.pre_lnorm = pre_lnorm
  
          if r_r_bias is None or r_w_bias is None: # Biases are not shared
              self.r_r_bias = nn.Parameter(torch.Tensor(self.n_head, self.d_head))
              self.r_w_bias = nn.Parameter(torch.Tensor(self.n_head, self.d_head))
          else:
              self.r_r_bias = r_r_bias
              self.r_w_bias = r_w_bias
  
      def _parallelogram_mask(self, h, w, left=False):
          mask = torch.ones((h, w)).byte()
          m = min(h, w)
          mask[:m,:m] = torch.triu(mask[:m,:m])
          mask[-m:,-m:] = torch.tril(mask[-m:,-m:])
  
          if left:
              return mask
          else:
              return mask.flip(0)
  
      def _shift(self, x, qlen, klen, mask, left=False):
          if qlen > 1:
              zero_pad = torch.zeros((x.size(0), qlen-1, x.size(2), x.size(3)),
                                      device=x.device, dtype=x.dtype)
          else:
              zero_pad = torch.zeros(0, device=x.device, dtype=x.dtype)
  
          if left:
              mask = mask.flip(1)
              x_padded = torch.cat([zero_pad, x], dim=1).expand(qlen, -1, -1, -1)
          else:
              x_padded = torch.cat([x, zero_pad], dim=1).expand(qlen, -1, -1, -1)
  
          x = x_padded.masked_select(mask[:,:,None,None]) \
                      .view(qlen, klen, x.size(2), x.size(3))
  
          return x
  
      def _rel_shift(self, x, zero_triu=False):
          zero_pad_shape = (x.size(0), 1) + x.size()[2:]
          zero_pad = torch.zeros(zero_pad_shape, device=x.device, dtype=x.dtype)
          x_padded = torch.cat([zero_pad, x], dim=1)
  
          x_padded_shape = (x.size(1) + 1, x.size(0)) + x.size()[2:]
          x_padded = x_padded.view(*x_padded_shape)
  
          x = x_padded[1:].view_as(x)
  
          if zero_triu:
              ones = torch.ones((x.size(0), x.size(1)))
              x = x * torch.tril(ones, x.size(1) - x.size(0))[:,:,None,None]
  
          return x
  
      def forward(self, w, r, attn_mask=None, mems=None):
          raise NotImplementedError
  
  class RelPartialLearnableMultiHeadAttn(RelMultiHeadAttn):
      def __init__(self, *args, **kwargs):
          super(RelPartialLearnableMultiHeadAttn, self).__init__(*args, **kwargs)
  
          self.r_net = nn.Linear(self.d_model, self.n_head * self.d_head, bias=False)
  
      def forward(self, w, r, attn_mask=None, mems=None):
          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)
  
          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
  
          r_head_k = r_head_k.view(rlen, self.n_head, self.d_head)                # qlen x n_head x d_head
  
          #### compute attention score
          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
  
          rr_head_q = w_head_q + self.r_r_bias
          BD = torch.einsum('ibnd,jnd->ijbn', (rr_head_q, r_head_k))              # qlen x klen x bsz x n_head
          BD = self._rel_shift(BD)
  
          # [qlen x klen x bsz x n_head]
          attn_score = AC + BD
          attn_score.mul_(self.scale)
  
          #### compute attention probability
          if attn_mask is not None and attn_mask.any().item():
              if attn_mask.dim() == 2:
                  attn_score = attn_score.float().masked_fill(
                      attn_mask[None,:,:,None], -1e30).type_as(attn_score)
              elif attn_mask.dim() == 3:
                  attn_score = attn_score.float().masked_fill(
                      attn_mask[:,:,:,None], -1e30).type_as(attn_score)
  
          # [qlen x klen x bsz x n_head]
          attn_prob = F.softmax(attn_score, dim=1)
          attn_prob = self.dropatt(attn_prob)
  
          #### compute attention vector
          attn_vec = torch.einsum('ijbn,jbnd->ibnd', (attn_prob, w_head_v))
  
          # [qlen x bsz x n_head x d_head]
          attn_vec = attn_vec.contiguous().view(
              attn_vec.size(0), attn_vec.size(1), self.n_head * self.d_head)
  
          ##### linear projection
          attn_out = self.o_net(attn_vec)
          attn_out = self.drop(attn_out)
  
          if self.pre_lnorm:
              ##### residual connection
              output = w + attn_out
          else:
              ##### residual connection + layer normalization
              output = self.layer_norm(w + attn_out)
  
          return output
  
  class RelLearnableMultiHeadAttn(RelMultiHeadAttn):
      def __init__(self, *args, **kwargs):
          super(RelLearnableMultiHeadAttn, self).__init__(*args, **kwargs)
  
      def forward(self, w, r_emb, r_w_bias, r_bias, attn_mask=None, mems=None):
          # r_emb: [klen, n_head, d_head], used for term B
          # r_w_bias: [n_head, d_head], used for term C
          # r_bias: [klen, n_head], used for term D
  
          qlen, bsz = w.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)
              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)
              w_head_q, w_head_k, w_head_v = torch.chunk(w_heads, 3, dim=-1)
  
          klen = w_head_k.size(0)
  
          w_head_q = w_head_q.view(qlen, bsz, self.n_head, self.d_head)
          w_head_k = w_head_k.view(klen, bsz, self.n_head, self.d_head)
          w_head_v = w_head_v.view(klen, bsz, self.n_head, self.d_head)
  
          if klen > r_emb.size(0):
              r_emb_pad = r_emb[0:1].expand(klen-r_emb.size(0), -1, -1)
              r_emb = torch.cat([r_emb_pad, r_emb], 0)
              r_bias_pad = r_bias[0:1].expand(klen-r_bias.size(0), -1)
              r_bias = torch.cat([r_bias_pad, r_bias], 0)
          else:
              r_emb = r_emb[-klen:]
              r_bias = r_bias[-klen:]
  
          #### compute attention score
          rw_head_q = w_head_q + r_w_bias[None]                                   # 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
          B_ = torch.einsum('ibnd,jnd->ijbn', (w_head_q, r_emb))                  # qlen x klen x bsz x n_head
          D_ = r_bias[None, :, None]                                              # 1    x klen x 1   x n_head
          BD = self._rel_shift(B_ + D_)
  
          # [qlen x klen x bsz x n_head]
          attn_score = AC + BD
          attn_score.mul_(self.scale)
  
          #### compute attention probability
          if attn_mask is not None and attn_mask.any().item():
              if attn_mask.dim() == 2:
                  attn_score.masked_fill_(attn_mask[None,:,:,None], -float('inf'))
              elif attn_mask.dim() == 3:
                  attn_score.masked_fill_(attn_mask[:,:,:,None], -float('inf'))
  
          # [qlen x klen x bsz x n_head]
          attn_prob = F.softmax(attn_score, dim=1)
          attn_prob = self.dropatt(attn_prob)
  
          #### compute attention vector
          attn_vec = torch.einsum('ijbn,jbnd->ibnd', (attn_prob, w_head_v))
  
          # [qlen x bsz x n_head x d_head]
          attn_vec = attn_vec.contiguous().view(
              attn_vec.size(0), attn_vec.size(1), self.n_head * self.d_head)
  
          ##### linear projection
          attn_out = self.o_net(attn_vec)
          attn_out = self.drop(attn_out)
  
          if self.pre_lnorm:
              ##### residual connection
              output = w + attn_out
          else:
              ##### residual connection + layer normalization
              output = self.layer_norm(w + attn_out)
  
          return output
  
  class DecoderLayer(nn.Module):
      def __init__(self, n_head, d_model, d_head, d_inner, dropout, **kwargs):
          super(DecoderLayer, self).__init__()
  
          self.dec_attn = MultiHeadAttn(n_head, d_model, d_head, dropout, **kwargs)
          self.pos_ff = PositionwiseFF(d_model, d_inner, dropout, 
                                       pre_lnorm=kwargs.get('pre_lnorm'))
  
      def forward(self, dec_inp, dec_attn_mask=None, mems=None):
  
          output = self.dec_attn(dec_inp, attn_mask=dec_attn_mask,
                                 mems=mems)
          output = self.pos_ff(output)
  
          return output
  
  class RelLearnableDecoderLayer(nn.Module):
      def __init__(self, n_head, d_model, d_head, d_inner, dropout,
                   **kwargs):
          super(RelLearnableDecoderLayer, self).__init__()
  
          self.dec_attn = RelLearnableMultiHeadAttn(n_head, d_model, d_head, dropout,
                                           **kwargs)
          self.pos_ff = PositionwiseFF(d_model, d_inner, dropout, 
                                       pre_lnorm=kwargs.get('pre_lnorm'))
  
      def forward(self, dec_inp, r_emb, r_w_bias, r_bias, dec_attn_mask=None, mems=None):
  
          output = self.dec_attn(dec_inp, r_emb, r_w_bias, r_bias,
                                 attn_mask=dec_attn_mask,
                                 mems=mems)
          output = self.pos_ff(output)
  
          return output
  
  class RelPartialLearnableDecoderLayer(nn.Module):
      def __init__(self, n_head, d_model, d_head, d_inner, dropout,
                   **kwargs):
          super(RelPartialLearnableDecoderLayer, self).__init__()
  
          self.dec_attn = RelPartialLearnableMultiHeadAttn(n_head, d_model,
                              d_head, dropout, **kwargs)
          self.pos_ff = PositionwiseFF(d_model, d_inner, dropout, 
                                       pre_lnorm=kwargs.get('pre_lnorm'))
  
      def forward(self, dec_inp, r, dec_attn_mask=None, mems=None):
  
          output = self.dec_attn(dec_inp, r,
                                 attn_mask=dec_attn_mask,
                                 mems=mems)
          output = self.pos_ff(output)
  
          return output
  
  
  class AdaptiveEmbedding(nn.Module):
      def __init__(self, n_token, d_embed, d_proj, cutoffs, div_val=1, 
                   sample_softmax=False):
          super(AdaptiveEmbedding, self).__init__()
  
          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:
              self.emb_layers.append(
                  nn.Embedding(n_token, d_embed, sparse=sample_softmax>0)
              )
              if d_proj != d_embed:
                  self.emb_projs.append(nn.Parameter(torch.Tensor(d_proj, d_embed)))
          else:
              for i in range(len(self.cutoffs)):
                  l_idx, r_idx = self.cutoff_ends[i], self.cutoff_ends[i+1]
                  d_emb_i = d_embed // (div_val ** i)
                  self.emb_layers.append(nn.Embedding(r_idx-l_idx, d_emb_i))
                  self.emb_projs.append(nn.Parameter(torch.Tensor(d_proj, d_emb_i)))
  
      def forward(self, inp):
          if self.div_val == 1:
              embed = self.emb_layers[0](inp)
              if self.d_proj != self.d_embed:
                  embed  = F.linear(embed, self.emb_projs[0])
          else:
              param = next(self.parameters())
              inp_flat = inp.view(-1)
              emb_flat = torch.zeros([inp_flat.size(0), self.d_proj], 
                  dtype=param.dtype, device=param.device)
              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)
  
              embed_shape = inp.size() + (self.d_proj,)
              embed = emb_flat.view(embed_shape)
  
          embed.mul_(self.emb_scale)
  
          return embed
  
  
  class TransfoXLPreTrainedModel(nn.Module):
      """ An abstract class to handle weights initialization and
          a simple interface for dowloading and loading pretrained models.
      """
      def __init__(self, config, *inputs, **kwargs):
          super(TransfoXLPreTrainedModel, self).__init__()
          if not isinstance(config, TransfoXLConfig):
              raise ValueError(
                  "Parameter config in `{}(config)` should be an instance of class `TransfoXLConfig`. "
                  "To create a model from a pretrained model use "
                  "`model = {}.from_pretrained(PRETRAINED_MODEL_NAME)`".format(
                      self.__class__.__name__, self.__class__.__name__
                  ))
          self.config = config
  
      def init_weight(self, weight):
          if self.config.init == 'uniform':
  