'''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
__UpperCAmelCase												=data_utils.TransfoXLTokenizer
__UpperCAmelCase												=data_utils.TransfoXLCorpus
__UpperCAmelCase												=data_utils
__UpperCAmelCase												=data_utils






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

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

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

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

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


if __name__ == "__main__":
  __UpperCAmelCase												=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.",
  )
  __UpperCAmelCase												=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,
  )

'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_flax_available,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase												={
    "configuration_whisper": ["WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP", "WhisperConfig", "WhisperOnnxConfig"],
    "feature_extraction_whisper": ["WhisperFeatureExtractor"],
    "processing_whisper": ["WhisperProcessor"],
    "tokenization_whisper": ["WhisperTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=["WhisperTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST",
      "WhisperForConditionalGeneration",
      "WhisperModel",
      "WhisperPreTrainedModel",
      "WhisperForAudioClassification",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFWhisperForConditionalGeneration",
      "TFWhisperModel",
      "TFWhisperPreTrainedModel",
  ]

try:
  if not is_flax_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase												=[
      "FlaxWhisperForConditionalGeneration",
      "FlaxWhisperModel",
      "FlaxWhisperPreTrainedModel",
      "FlaxWhisperForAudioClassification",
  ]


if TYPE_CHECKING:
  from .configuration_whisper import WHISPER_PRETRAINED_CONFIG_ARCHIVE_MAP, WhisperConfig, WhisperOnnxConfig
  from .feature_extraction_whisper import WhisperFeatureExtractor
  from .processing_whisper import WhisperProcessor
  from .tokenization_whisper import WhisperTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_whisper_fast import WhisperTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_whisper import (
        WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
        WhisperForAudioClassification,
        WhisperForConditionalGeneration,
        WhisperModel,
        WhisperPreTrainedModel,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_whisper import (
        TF_WHISPER_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFWhisperForConditionalGeneration,
        TFWhisperModel,
        TFWhisperPreTrainedModel,
    )

  try:
    if not is_flax_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_flax_whisper import (
        FlaxWhisperForAudioClassification,
        FlaxWhisperForConditionalGeneration,
        FlaxWhisperModel,
        FlaxWhisperPreTrainedModel,
    )

else:
  import sys

  __UpperCAmelCase												=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
