






import math

import torch
from torch import nn

from ..configuration_utils import ConfigMixin, register_to_config
from .attention_processor import Attention
from .embeddings import get_timestep_embedding
from .modeling_utils import ModelMixin




class 				__A     (			a      ,     a ):





				"""simple docstring"""







				@register_to_config
				def __init__( self							,  lowerCamelCase__ = 128							,  lowerCamelCase__ = 256							,  lowerCamelCase__ = 2_000.0							,  lowerCamelCase__ = 768							,  lowerCamelCase__ = 12							,  lowerCamelCase__ = 12							,  lowerCamelCase__ = 64							,  lowerCamelCase__ = 2048							,  lowerCamelCase__ = 0.1							,  ):




									"""simple docstring"""

									super().__init__()

									__UpperCamelCase      :					str       		=nn.Sequential(
									    nn.Linear(lowerCamelCase__							,  d_model * 4							,  bias=lowerCamelCase__    )							,  nn.SiLU()							,  nn.Linear(d_model * 4							,  d_model * 4							,  bias=lowerCamelCase__    )							,  nn.SiLU()							,  )

									__UpperCamelCase      :					Any       		=nn.Embedding(lowerCamelCase__							,  lowerCamelCase__    )
									__UpperCamelCase      :					Optional[int]       		=False

									__UpperCamelCase      :					Dict       		=nn.Linear(lowerCamelCase__							,  lowerCamelCase__							,  bias=lowerCamelCase__    )

									__UpperCamelCase      :					Optional[Any]       		=nn.Dropout(p=lowerCamelCase__    )

									__UpperCamelCase      :					int       		=nn.ModuleList()
									for lyr_num in range(lowerCamelCase__    ):
														# FiLM conditional T5 decoder
														__UpperCamelCase      :					Tuple       		=DecoderLayer(d_model=lowerCamelCase__							,  d_kv=lowerCamelCase__							,  num_heads=lowerCamelCase__							,  d_ff=lowerCamelCase__							,  dropout_rate=lowerCamelCase__    )
														self.decoders.append(lowerCamelCase__    )

									__UpperCamelCase      :					List[str]       		=TaLayerNorm(lowerCamelCase__    )

									__UpperCamelCase      :					Optional[int]       		=nn.Dropout(p=lowerCamelCase__    )
									__UpperCamelCase      :					Optional[int]       		=nn.Linear(lowerCamelCase__							,  lowerCamelCase__							,  bias=lowerCamelCase__    )


				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[int]       		=torch.mul(query_input.unsqueeze(-1    )							,  key_input.unsqueeze(-2    )    )
									return mask.unsqueeze(-3    )







				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase		,				__UpperCamelCase		,				__UpperCamelCase      :					List[str]       		=decoder_input_tokens.shape
									assert decoder_noise_time.shape == (batch,)

									# decoder_noise_time is in [0, 1), so rescale to expected timing range.
									__UpperCamelCase      :					Tuple       		=get_timestep_embedding(
									    decoder_noise_time * self.config.max_decoder_noise_time							,  embedding_dim=self.config.d_model							,  max_period=self.config.max_decoder_noise_time							,  ).to(dtype=self.dtype    )

									__UpperCamelCase      :					Union[str, Any]       		=self.conditioning_emb(lowerCamelCase__    ).unsqueeze(1    )

									assert conditioning_emb.shape == (batch, 1, self.config.d_model * 4)

									__UpperCamelCase      :					Optional[int]       		=decoder_input_tokens.shape[1]

									# If we want to use relative positions for audio context, we can just offset
									# this sequence by the length of encodings_and_masks.
									__UpperCamelCase      :					List[Any]       		=torch.broadcast_to(
									    torch.arange(lowerCamelCase__							,  device=decoder_input_tokens.device    )							,  (batch, seq_length)							,  )

									__UpperCamelCase      :					Any       		=self.position_encoding(lowerCamelCase__    )

									__UpperCamelCase      :					Any       		=self.continuous_inputs_projection(lowerCamelCase__    )
									inputs += position_encodings
									__UpperCamelCase      :					Optional[Any]       		=self.dropout(lowerCamelCase__    )

									# decoder: No padding present.
									__UpperCamelCase      :					Any       		=torch.ones(
									    decoder_input_tokens.shape[:2]							,  device=decoder_input_tokens.device							,  dtype=inputs.dtype    )

									# Translate encoding masks to encoder-decoder masks.
									__UpperCamelCase      :					str       		=[(x, self.encoder_decoder_mask(lowerCamelCase__							,  lowerCamelCase__    )) for x, y in encodings_and_masks]

									# cross attend style: concat encodings
									__UpperCamelCase      :					Any       		=torch.cat([x[0] for x in encodings_and_encdec_masks]							,  dim=1    )
									__UpperCamelCase      :					List[Any]       		=torch.cat([x[1] for x in encodings_and_encdec_masks]							,  dim=-1    )

									for lyr in self.decoders:
														__UpperCamelCase      :					Optional[Any]       		=lyr(
														    lowerCamelCase__							,  conditioning_emb=lowerCamelCase__							,  encoder_hidden_states=lowerCamelCase__							,  encoder_attention_mask=lowerCamelCase__							,  )[0]

									__UpperCamelCase      :					Tuple       		=self.decoder_norm(lowerCamelCase__    )
									__UpperCamelCase      :					Tuple       		=self.post_dropout(lowerCamelCase__    )

									__UpperCamelCase      :					Dict       		=self.spec_out(lowerCamelCase__    )
									return spec_out




class 				__A     (			nn.Module ):





				"""simple docstring"""







				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__=1E-6    ):




									"""simple docstring"""

									super().__init__()
									__UpperCamelCase      :					Optional[int]       		=nn.ModuleList()

									# cond self attention: layer 0
									self.layer.append(
									    TaLayerSelfAttentionCond(d_model=lowerCamelCase__							,  d_kv=lowerCamelCase__							,  num_heads=lowerCamelCase__							,  dropout_rate=lowerCamelCase__    )    )

									# cross attention: layer 1
									self.layer.append(
									    TaLayerCrossAttention(
									        d_model=lowerCamelCase__							,  d_kv=lowerCamelCase__							,  num_heads=lowerCamelCase__							,  dropout_rate=lowerCamelCase__							,  layer_norm_epsilon=lowerCamelCase__							,  )    )

									# Film Cond MLP + dropout: last layer
									self.layer.append(
									    TaLayerFFCond(d_model=lowerCamelCase__							,  d_ff=lowerCamelCase__							,  dropout_rate=lowerCamelCase__							,  layer_norm_epsilon=lowerCamelCase__    )    )







				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__=None							,  lowerCamelCase__=None							,  lowerCamelCase__=None							,  lowerCamelCase__=None							,  lowerCamelCase__=None							,  ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[Any]       		=self.layer[0](
									    lowerCamelCase__							,  conditioning_emb=lowerCamelCase__							,  attention_mask=lowerCamelCase__							,  )

									if encoder_hidden_states is not None:
														__UpperCamelCase      :					str       		=torch.where(encoder_attention_mask > 0							,  0							,  -1E10    ).to(
														    encoder_hidden_states.dtype    )

														__UpperCamelCase      :					int       		=self.layer[1](
														    lowerCamelCase__							,  key_value_states=lowerCamelCase__							,  attention_mask=lowerCamelCase__							,  )

									# Apply Film Conditional Feed Forward layer
									__UpperCamelCase      :					Union[str, Any]       		=self.layer[-1](lowerCamelCase__							,  lowerCamelCase__    )

									return (hidden_states,)




class 				__A     (			nn.Module ):





				"""simple docstring"""







				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__    ):




									"""simple docstring"""

									super().__init__()
									__UpperCamelCase      :					Any       		=TaLayerNorm(lowerCamelCase__    )
									__UpperCamelCase      :					Tuple       		=TaFiLMLayer(in_features=d_model * 4							,  out_features=lowerCamelCase__    )
									__UpperCamelCase      :					Union[str, Any]       		=Attention(query_dim=lowerCamelCase__							,  heads=lowerCamelCase__							,  dim_head=lowerCamelCase__							,  out_bias=lowerCamelCase__							,  scale_qk=lowerCamelCase__    )
									__UpperCamelCase      :					List[str]       		=nn.Dropout(lowerCamelCase__    )







				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__=None							,  lowerCamelCase__=None							,  ):




									"""simple docstring"""

									__UpperCamelCase      :					int       		=self.layer_norm(lowerCamelCase__    )

									if conditioning_emb is not None:
														__UpperCamelCase      :					str       		=self.FiLMLayer(lowerCamelCase__							,  lowerCamelCase__    )

									# Self-attention block
									__UpperCamelCase      :					int       		=self.attention(lowerCamelCase__    )

									__UpperCamelCase      :					Any       		=hidden_states + self.dropout(lowerCamelCase__    )

									return hidden_states




class 				__A     (			nn.Module ):





				"""simple docstring"""







				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__    ):




									"""simple docstring"""

									super().__init__()
									__UpperCamelCase      :					Tuple       		=Attention(query_dim=lowerCamelCase__							,  heads=lowerCamelCase__							,  dim_head=lowerCamelCase__							,  out_bias=lowerCamelCase__							,  scale_qk=lowerCamelCase__    )
									__UpperCamelCase      :					List[str]       		=TaLayerNorm(lowerCamelCase__							,  eps=lowerCamelCase__    )
									__UpperCamelCase      :					Any       		=nn.Dropout(lowerCamelCase__    )







				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__=None							,  lowerCamelCase__=None							,  ):




									"""simple docstring"""

									__UpperCamelCase      :					List[Any]       		=self.layer_norm(lowerCamelCase__    )
									__UpperCamelCase      :					Tuple       		=self.attention(
									    lowerCamelCase__							,  encoder_hidden_states=lowerCamelCase__							,  attention_mask=attention_mask.squeeze(1    )							,  )
									__UpperCamelCase      :					int       		=hidden_states + self.dropout(lowerCamelCase__    )
									return layer_output




class 				__A     (			nn.Module ):





				"""simple docstring"""







				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__    ):




									"""simple docstring"""

									super().__init__()
									__UpperCamelCase      :					Any       		=TaDenseGatedActDense(d_model=lowerCamelCase__							,  d_ff=lowerCamelCase__							,  dropout_rate=lowerCamelCase__    )
									__UpperCamelCase      :					Tuple       		=TaFiLMLayer(in_features=d_model * 4							,  out_features=lowerCamelCase__    )
									__UpperCamelCase      :					Dict       		=TaLayerNorm(lowerCamelCase__							,  eps=lowerCamelCase__    )
									__UpperCamelCase      :					Union[str, Any]       		=nn.Dropout(lowerCamelCase__    )







				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__=None    ):




									"""simple docstring"""

									__UpperCamelCase      :					List[str]       		=self.layer_norm(lowerCamelCase__    )
									if conditioning_emb is not None:
														__UpperCamelCase      :					List[str]       		=self.film(lowerCamelCase__							,  lowerCamelCase__    )

									__UpperCamelCase      :					Union[str, Any]       		=self.DenseReluDense(lowerCamelCase__    )
									__UpperCamelCase      :					Any       		=hidden_states + self.dropout(lowerCamelCase__    )
									return hidden_states




class 				__A     (			nn.Module ):





				"""simple docstring"""







				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__							,  lowerCamelCase__    ):




									"""simple docstring"""

									super().__init__()
									__UpperCamelCase      :					List[str]       		=nn.Linear(lowerCamelCase__							,  lowerCamelCase__							,  bias=lowerCamelCase__    )
									__UpperCamelCase      :					Optional[int]       		=nn.Linear(lowerCamelCase__							,  lowerCamelCase__							,  bias=lowerCamelCase__    )
									__UpperCamelCase      :					List[Any]       		=nn.Linear(lowerCamelCase__							,  lowerCamelCase__							,  bias=lowerCamelCase__    )
									__UpperCamelCase      :					List[Any]       		=nn.Dropout(lowerCamelCase__    )
									__UpperCamelCase      :					Any       		=NewGELUActivation()







				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase      :					Any       		=self.act(self.wi_a(lowerCamelCase__    )    )
									__UpperCamelCase      :					Union[str, Any]       		=self.wi_a(lowerCamelCase__    )
									__UpperCamelCase      :					Any       		=hidden_gelu * hidden_linear
									__UpperCamelCase      :					List[str]       		=self.dropout(lowerCamelCase__    )

									__UpperCamelCase      :					Optional[Any]       		=self.wo(lowerCamelCase__    )
									return hidden_states




class 				__A     (			nn.Module ):





				"""simple docstring"""







				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__=1E-6    ):




									"""simple docstring"""

									super().__init__()
									__UpperCamelCase      :					Any       		=nn.Parameter(torch.ones(lowerCamelCase__    )    )
									__UpperCamelCase      :					str       		=eps







				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase      :					Any       		=hidden_states.to(torch.floataa    ).pow(2    ).mean(-1							,  keepdim=lowerCamelCase__    )
									__UpperCamelCase      :					Dict       		=hidden_states * torch.rsqrt(variance + self.variance_epsilon    )

									# convert into half-precision if necessary
									if self.weight.dtype in [torch.floataa, torch.bfloataa]:
														__UpperCamelCase      :					List[str]       		=hidden_states.to(self.weight.dtype    )

									return self.weight * hidden_states




class 				__A     (			nn.Module ):





				"""simple docstring"""







				def       __lowercase							( self							,  lowerCamelCase__    ):




									"""simple docstring"""

									return 0.5 * input * (1.0 + torch.tanh(math.sqrt(2.0 / math.pi    ) * (input + 0.044_715 * torch.pow(lowerCamelCase__							,  3.0    ))    ))




class 				__A     (			nn.Module ):





				"""simple docstring"""







				def __init__( self							,  lowerCamelCase__							,  lowerCamelCase__    ):




									"""simple docstring"""

									super().__init__()
									__UpperCamelCase      :					Optional[int]       		=nn.Linear(lowerCamelCase__							,  out_features * 2							,  bias=lowerCamelCase__    )







				def       __lowercase							( self							,  lowerCamelCase__							,  lowerCamelCase__    ):




									"""simple docstring"""

									__UpperCamelCase      :					Optional[int]       		=self.scale_bias(lowerCamelCase__    )
									__UpperCamelCase		,				__UpperCamelCase      :					Optional[int]       		=torch.chunk(lowerCamelCase__							,  2							,  -1    )
									__UpperCamelCase      :					int       		=x * (1 + scale) + shift
									return x







import argparse
import os
import pickle
import sys

import torch

from transformers import TransfoXLConfig, TransfoXLLMHeadModel, load_tf_weights_in_transfo_xl
from transformers.models.transfo_xl import tokenization_transfo_xl as data_utils
from transformers.models.transfo_xl.tokenization_transfo_xl import CORPUS_NAME, VOCAB_FILES_NAMES
from transformers.utils import CONFIG_NAME, WEIGHTS_NAME, logging


logging.set_verbosity_info()

# We do this to be able to load python 2 datasets pickles
# See e.g. https://stackoverflow.com/questions/2121874/python-pickling-after-changing-a-modules-directory/2121918#2121918
snake_case_ : List[Any]    		=       data_utils.TransfoXLTokenizer
snake_case_ : int    		=       data_utils.TransfoXLCorpus
snake_case_ : List[Any]    		=       data_utils
snake_case_ : int    		=       data_utils


def    A						(__A   :    Dict   ,				__A   :    List[Any]   ,				__A   :    Union[str, Any]   ,				__A   :    Tuple       )   ->      Union[str, Any]:







     """simple docstring"""


     if transfo_xl_dataset_file:
          # Convert a pre-processed corpus (see original TensorFlow repo)
          with open(__A   ,				'''rb'''       ) as fp:
               UpperCAmelCase_          =							pickle.load(__A   ,				encoding='''latin1'''       )
          # Save vocabulary and dataset cache as Dictionaries (should be better than pickles for the long-term)
          UpperCAmelCase_          =							pytorch_dump_folder_path + '''/''' + VOCAB_FILES_NAMES['''pretrained_vocab_file''']
          print(F"""Save vocabulary to {pytorch_vocab_dump_path}"""       )
          UpperCAmelCase_          =							corpus.vocab.__dict__
          torch.save(__A   ,				__A       )

          UpperCAmelCase_          =							corpus.__dict__
          corpus_dict_no_vocab.pop('''vocab'''   ,				__A       )
          UpperCAmelCase_          =							pytorch_dump_folder_path + '''/''' + CORPUS_NAME
          print(F"""Save dataset to {pytorch_dataset_dump_path}"""       )
          torch.save(__A   ,				__A       )

     if tf_checkpoint_path:
          # Convert a pre-trained TensorFlow model
          UpperCAmelCase_          =							os.path.abspath(__A       )
          UpperCAmelCase_          =							os.path.abspath(__A       )

          print(F"""Converting Transformer XL checkpoint from {tf_path} with config at {config_path}."""       )
          # Initialise PyTorch model
          if transfo_xl_config_file == "":
               UpperCAmelCase_          =							TransfoXLConfig()
          else:
               UpperCAmelCase_          =							TransfoXLConfig.from_json_file(__A       )
          print(F"""Building PyTorch model from configuration: {config}"""       )
          UpperCAmelCase_          =							TransfoXLLMHeadModel(__A       )

          UpperCAmelCase_          =							load_tf_weights_in_transfo_xl(__A   ,				__A   ,				__A       )
          # Save pytorch-model
          UpperCAmelCase_          =							os.path.join(__A   ,				__A       )
          UpperCAmelCase_          =							os.path.join(__A   ,				__A       )
          print(F"""Save PyTorch model to {os.path.abspath(__A       )}"""       )
          torch.save(model.state_dict()   ,				__A       )
          print(F"""Save configuration file to {os.path.abspath(__A       )}"""       )
          with open(__A   ,				'''w'''   ,				encoding='''utf-8'''       ) as f:
               f.write(config.to_json_string()       )


if __name__ == "__main__":
    snake_case_ : List[str]    		=       argparse.ArgumentParser()
    parser.add_argument(
        "--pytorch_dump_folder_path",
        default=None,
        type=str,
        required=True,
        help="Path to the folder to store the PyTorch model or dataset/vocab.",
    )
    parser.add_argument(
        "--tf_checkpoint_path",
        default="",
        type=str,
        help="An optional path to a TensorFlow checkpoint path to be converted.",
    )
    parser.add_argument(
        "--transfo_xl_config_file",
        default="",
        type=str,
        help=(
            "An optional config json file corresponding to the pre-trained BERT model. \n"
            "This specifies the model architecture."
        ),
    )
    parser.add_argument(
        "--transfo_xl_dataset_file",
        default="",
        type=str,
        help="An optional dataset file to be converted in a vocabulary.",
    )
    snake_case_ : int    		=       parser.parse_args()
    convert_transfo_xl_checkpoint_to_pytorch(
        args.tf_checkpoint_path,
        args.transfo_xl_config_file,
        args.pytorch_dump_folder_path,
        args.transfo_xl_dataset_file,
    )






