





'''simple docstring'''






from collections import OrderedDict
from typing import TYPE_CHECKING, Any, List, Mapping, Optional, Union

from ...configuration_utils import PretrainedConfig
from ...onnx import OnnxConfig
from ...utils import TensorType, logging


if TYPE_CHECKING:
      from ...onnx.config import PatchingSpec
      from ...tokenization_utils_base import PreTrainedTokenizerBase


__lowerCamelCase						: int		=   logging.get_logger(__name__)

__lowerCamelCase						: int		=   {
    'allenai/longformer-base-4096': 'https://huggingface.co/allenai/longformer-base-4096/resolve/main/config.json',
    'allenai/longformer-large-4096': 'https://huggingface.co/allenai/longformer-large-4096/resolve/main/config.json',
    'allenai/longformer-large-4096-finetuned-triviaqa': (
        'https://huggingface.co/allenai/longformer-large-4096-finetuned-triviaqa/resolve/main/config.json'
    ),
    'allenai/longformer-base-4096-extra.pos.embd.only': (
        'https://huggingface.co/allenai/longformer-base-4096-extra.pos.embd.only/resolve/main/config.json'
    ),
    'allenai/longformer-large-4096-extra.pos.embd.only': (
        'https://huggingface.co/allenai/longformer-large-4096-extra.pos.embd.only/resolve/main/config.json'
    ),
}







class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  lowerCAmelCase_					=				"""longformer"""


  def __init__(				self    :					Union[str, Any]							,	UpperCamelCase__    :					Union[List[int], int] = 512							,	UpperCamelCase__    :					int = 2							,	UpperCamelCase__    :					int = 1							,	UpperCamelCase__    :					int = 0							,	UpperCamelCase__    :					int = 2							,	UpperCamelCase__    :					int = 3_0522							,	UpperCamelCase__    :					int = 768							,	UpperCamelCase__    :					int = 12							,	UpperCamelCase__    :					int = 12							,	UpperCamelCase__    :					int = 3072							,	UpperCamelCase__    :					str = "gelu"							,	UpperCamelCase__    :					float = 0.1							,	UpperCamelCase__    :					float = 0.1							,	UpperCamelCase__    :					int = 512							,	UpperCamelCase__    :					int = 2							,	UpperCamelCase__    :					float = 0.02							,	UpperCamelCase__    :					float = 1E-12							,	UpperCamelCase__    :					bool = False							,	**UpperCamelCase__    :					Union[str, Any]							,	)     ->						List[Any]:
    super().__init__(pad_token_id=UpperCamelCase__							,	**UpperCamelCase__  )

    _UpperCamelCase       			=attention_window
    _UpperCamelCase       			=sep_token_id
    _UpperCamelCase       			=bos_token_id
    _UpperCamelCase       			=eos_token_id
    _UpperCamelCase       			=vocab_size
    _UpperCamelCase       			=hidden_size
    _UpperCamelCase       			=num_hidden_layers
    _UpperCamelCase       			=num_attention_heads
    _UpperCamelCase       			=hidden_act
    _UpperCamelCase       			=intermediate_size
    _UpperCamelCase       			=hidden_dropout_prob
    _UpperCamelCase       			=attention_probs_dropout_prob
    _UpperCamelCase       			=max_position_embeddings
    _UpperCamelCase       			=type_vocab_size
    _UpperCamelCase       			=initializer_range
    _UpperCamelCase       			=layer_norm_eps
    _UpperCamelCase       			=onnx_export







class 		UpperCAmelCase   (				lowercase_):


  """simple docstring"""







  def __init__(				self    :					int							,	UpperCamelCase__    :					"PretrainedConfig"							,	UpperCamelCase__    :					str = "default"							,	UpperCamelCase__    :					"List[PatchingSpec]" = None  )     ->						List[str]:
    super().__init__(UpperCamelCase__							,	UpperCamelCase__							,	UpperCamelCase__  )
    _UpperCamelCase       			=True


  @property
  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Mapping[str, Mapping[int, str]]:
    if self.task == "multiple-choice":
      _UpperCamelCase       			={0: '''batch''', 1: '''choice''', 2: '''sequence'''}
    else:
      _UpperCamelCase       			={0: '''batch''', 1: '''sequence'''}
    return OrderedDict(
        [
            ('''input_ids''', dynamic_axis),
            ('''attention_mask''', dynamic_axis),
            ('''global_attention_mask''', dynamic_axis),
        ]  )


  @property
  def   UpperCamelCase__      (				self    :					List[str]  )     ->						Mapping[str, Mapping[int, str]]:
    _UpperCamelCase       			=super().outputs
    if self.task == "default":
      _UpperCamelCase       			={0: '''batch'''}
    return outputs


  @property
  def   UpperCamelCase__      (				self    :					List[str]  )     ->						float:
    return 1E-4


  @property
  def   UpperCamelCase__      (				self    :					Dict  )     ->						int:
    # needs to be >= 14 to support tril operator
    return max(super().default_onnx_opset							,	14  )




  def   UpperCamelCase__      (				self    :					int							,	UpperCamelCase__    :					"PreTrainedTokenizerBase"							,	UpperCamelCase__    :					int = -1							,	UpperCamelCase__    :					int = -1							,	UpperCamelCase__    :					bool = False							,	UpperCamelCase__    :					Optional[TensorType] = None							,	)     ->						Mapping[str, Any]:
    _UpperCamelCase       			=super().generate_dummy_inputs(
        preprocessor=UpperCamelCase__							,	batch_size=UpperCamelCase__							,	seq_length=UpperCamelCase__							,	is_pair=UpperCamelCase__							,	framework=UpperCamelCase__  )
    import torch

    # for some reason, replacing this code by inputs["global_attention_mask"] = torch.randint(2, inputs["input_ids"].shape, dtype=torch.int64)
    # makes the export fail randomly
    _UpperCamelCase       			=torch.zeros_like(inputs['''input_ids''']  )
    # make every second token global
    _UpperCamelCase       			=1

    return inputs








'''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,
      )






