





'''simple docstring'''






from ...configuration_utils import PretrainedConfig
from ...utils import logging
from ...utils.backbone_utils import BackboneConfigMixin, get_aligned_output_features_output_indices


__lowerCamelCase						: int		=   logging.get_logger(__name__)

__lowerCamelCase						: List[str]		=   {
    'shi-labs/dinat-mini-in1k-224': 'https://huggingface.co/shi-labs/dinat-mini-in1k-224/resolve/main/config.json',
    # See all Dinat models at https://huggingface.co/models?filter=dinat
}







class 		UpperCAmelCase   (				lowercase_   ,     lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				"""dinat"""

  lowerCAmelCase_					=				{
      """num_attention_heads""": """num_heads""",
      """num_hidden_layers""": """num_layers""",
  }


  def __init__(				self    :					Any							,	UpperCamelCase__    :					List[Any]=4							,	UpperCamelCase__    :					Any=3							,	UpperCamelCase__    :					Any=64							,	UpperCamelCase__    :					Dict=[3, 4, 6, 5]							,	UpperCamelCase__    :					str=[2, 4, 8, 16]							,	UpperCamelCase__    :					List[Any]=7							,	UpperCamelCase__    :					Dict=[[1, 8, 1], [1, 4, 1, 4], [1, 2, 1, 2, 1, 2], [1, 1, 1, 1, 1]]							,	UpperCamelCase__    :					Optional[int]=3.0							,	UpperCamelCase__    :					Optional[Any]=True							,	UpperCamelCase__    :					Tuple=0.0							,	UpperCamelCase__    :					Tuple=0.0							,	UpperCamelCase__    :					Dict=0.1							,	UpperCamelCase__    :					Any="gelu"							,	UpperCamelCase__    :					Any=0.02							,	UpperCamelCase__    :					Dict=1E-5							,	UpperCamelCase__    :					str=0.0							,	UpperCamelCase__    :					Tuple=None							,	UpperCamelCase__    :					str=None							,	**UpperCamelCase__    :					List[str]							,	)     ->						Any:
    super().__init__(**UpperCamelCase__  )

    _UpperCamelCase       			=patch_size
    _UpperCamelCase       			=num_channels
    _UpperCamelCase       			=embed_dim
    _UpperCamelCase       			=depths
    _UpperCamelCase       			=len(UpperCamelCase__  )
    _UpperCamelCase       			=num_heads
    _UpperCamelCase       			=kernel_size
    _UpperCamelCase       			=dilations
    _UpperCamelCase       			=mlp_ratio
    _UpperCamelCase       			=qkv_bias
    _UpperCamelCase       			=hidden_dropout_prob
    _UpperCamelCase       			=attention_probs_dropout_prob
    _UpperCamelCase       			=drop_path_rate
    _UpperCamelCase       			=hidden_act
    _UpperCamelCase       			=layer_norm_eps
    _UpperCamelCase       			=initializer_range
    # we set the hidden_size attribute in order to make Dinat work with VisionEncoderDecoderModel
    # this indicates the channel dimension after the last stage of the model
    _UpperCamelCase       			=int(embed_dim * 2 ** (len(UpperCamelCase__  ) - 1)  )
    _UpperCamelCase       			=layer_scale_init_value
    _UpperCamelCase       			=['''stem'''] + [F'''stage{idx}''' for idx in range(1							,	len(UpperCamelCase__  ) + 1  )]
    _UpperCamelCase			,     _UpperCamelCase       			=get_aligned_output_features_output_indices(
        out_features=UpperCamelCase__							,	out_indices=UpperCamelCase__							,	stage_names=self.stage_names  )













'''simple docstring'''






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
__lowerCamelCase						: List[Any]		=   data_utils.TransfoXLTokenizer
__lowerCamelCase						: str		=   data_utils.TransfoXLCorpus
__lowerCamelCase						: Optional[int]		=   data_utils
__lowerCamelCase						: int		=   data_utils


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  if transfo_xl_dataset_file:
    # Convert a pre-processed corpus (see original TensorFlow repo)
    with open(__SCREAMING_SNAKE_CASE  ,      '''rb'''   ) as fp:
      _UpperCamelCase       			=pickle.load(__SCREAMING_SNAKE_CASE  ,      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(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

    _UpperCamelCase       			=corpus.__dict__
    corpus_dict_no_vocab.pop('''vocab'''  ,      __SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=pytorch_dump_folder_path + '''/''' + CORPUS_NAME
    print(f'''Save dataset to {pytorch_dataset_dump_path}'''   )
    torch.save(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

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

    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(__SCREAMING_SNAKE_CASE   )
    print(f'''Building PyTorch model from configuration: {config}'''   )
    _UpperCamelCase       			=TransfoXLLMHeadModel(__SCREAMING_SNAKE_CASE   )

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


if __name__ == "__main__":
      __lowerCamelCase						: Any		=   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.',
      )
      __lowerCamelCase						: Optional[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,
      )






