
from typing import List, Optional, Union

import numpy as np
import torch
import torchaudio.compliance.kaldi as ta_kaldi

from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import PaddingStrategy, TensorType, logging


_lowercase		:				Tuple      		=logging.get_logger(__name__)

class        snake_case__    (SCREAMING_SNAKE_CASE__		):







						"""simple docstring"""




						__lowerCAmelCase					:Optional[int]      =					["input_features", "attention_mask"]

						def __init__(       self ,       __lowercase=8_0 ,       __lowercase=1_6_0_0_0 ,       __lowercase=8_0 ,       __lowercase=0.0 ,       __lowercase=True ,       __lowercase=True ,       __lowercase=True ,       **__lowercase ,       ) ->				List[str]:
									"""simple docstring"""




									super().__init__(feature_size=__lowercase ,       sampling_rate=__lowercase ,       padding_value=__lowercase ,       **__lowercase  )
									a__       :						List[Any]							=  num_mel_bins
									a__       :						int							=  do_ceptral_normalize
									a__       :						Tuple							=  normalize_means
									a__       :						int							=  normalize_vars
									a__       :						Any							=  True

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       ) ->				np.ndarray:
									"""simple docstring"""




									a__       :						List[Any]							=  waveform * (2**1_5)  # Kaldi compliance: 16-bit signed integers
									a__       :						str							=  torch.from_numpy(__lowercase  ).unsqueeze(0  )
									a__       :						int							=  ta_kaldi.fbank(__lowercase ,       num_mel_bins=self.num_mel_bins ,       sample_frequency=self.sampling_rate  )
									return features.numpy()

						@staticmethod
						def 							SCREAMING_SNAKE_CASE__(       __lowercase ,       __lowercase ,       __lowercase = True ,       __lowercase = True ,       __lowercase = 0.0 ,       ) ->				np.ndarray:
									"""simple docstring"""




									if normalize_means:
												a__       :						Optional[int]							=  x[:input_length].mean(axis=0  )
												a__       :						Dict							=  np.subtract(__lowercase ,       __lowercase  )
									if normalize_vars:
												a__       :						Optional[Any]							=  x[:input_length].std(axis=0  )
												a__       :						str							=  np.divide(__lowercase ,       __lowercase  )

									if input_length < x.shape[0]:
												a__       :						List[str]							=  padding_value

									# make sure array is in float32
									a__       :						int							=  x.astype(np.floataa  )

									return x

						def 							SCREAMING_SNAKE_CASE__(       self ,       __lowercase ,       __lowercase = None  ) ->				List[np.ndarray]:
									"""simple docstring"""




									a__       :						Union[str, Any]							=  attention_mask.sum(-1  ) if attention_mask is not None else [x.shape[0] for x in input_features]
									return [
									    self.utterance_cmvn(__lowercase ,       __lowercase ,       self.normalize_means ,       self.normalize_vars ,       self.padding_value  )
									    for x, n in zip(__lowercase ,       __lowercase  )
									]

						def __call__(       self ,       __lowercase ,       __lowercase = False ,       __lowercase = None ,       __lowercase = False ,       __lowercase = None ,       __lowercase = None ,       __lowercase = None ,       __lowercase = None ,       **__lowercase ,       ) ->				BatchFeature:
									"""simple docstring"""




									if sampling_rate is not None:
												if sampling_rate != self.sampling_rate:
															raise ValueError(
															    F'''The model corresponding to this feature extractor: {self} was trained using a sampling rate of'''
															    F''' {self.sampling_rate}. Please make sure that the provided `raw_speech` input was sampled with'''
															    F''' {self.sampling_rate} and not {sampling_rate}.'''  )
									else:
												logger.warning(
												    """It is strongly recommended to pass the `sampling_rate` argument to this function. """
												    """Failing to do so can result in silent errors that might be hard to debug."""  )

									a__       :						str							=  isinstance(__lowercase ,       np.ndarray  ) and len(raw_speech.shape  ) > 1
									if is_batched_numpy and len(raw_speech.shape  ) > 2:
												raise ValueError(F'''Only mono-channel audio is supported for input to {self}'''  )
									a__       :						int							=  is_batched_numpy or (
									    isinstance(__lowercase ,       (list, tuple)  ) and (isinstance(raw_speech[0] ,       (np.ndarray, tuple, list)  ))
									)

									if is_batched:
												a__       :						str							=  [np.asarray(__lowercase ,       dtype=np.floataa  ) for speech in raw_speech]
									elif not is_batched and not isinstance(__lowercase ,       np.ndarray  ):
												a__       :						Any							=  np.asarray(__lowercase ,       dtype=np.floataa  )
									elif isinstance(__lowercase ,       np.ndarray  ) and raw_speech.dtype is np.dtype(np.floataa  ):
												a__       :						int							=  raw_speech.astype(np.floataa  )

									# always return batch
									if not is_batched:
												a__       :						Dict							=  [raw_speech]

									# extract fbank features
									a__       :						Optional[int]							=  [self._extract_fbank_features(__lowercase  ) for waveform in raw_speech]

									# convert into correct format for padding
									a__       :						Dict							=  BatchFeature({"""input_features""": features}  )

									a__       :						int							=  self.pad(
									    __lowercase ,       padding=__lowercase ,       max_length=__lowercase ,       truncation=__lowercase ,       pad_to_multiple_of=__lowercase ,       return_attention_mask=__lowercase ,       **__lowercase ,       )

									# make sure list is in array format
									a__       :						Optional[int]							=  padded_inputs.get("""input_features"""  )
									if isinstance(input_features[0] ,       __lowercase  ):
												a__       :						List[str]							=  [np.asarray(__lowercase ,       dtype=np.floataa  ) for feature in input_features]

									a__       :						List[Any]							=  padded_inputs.get("""attention_mask"""  )
									if attention_mask is not None:
												a__       :						List[Any]							=  [np.asarray(__lowercase ,       dtype=np.intaa  ) for array in attention_mask]

									# Utterance-level cepstral mean and variance normalization
									if self.do_ceptral_normalize:
												a__       :						Dict							=  (
												    np.array(__lowercase ,       dtype=np.intaa  )
												    if self._get_padding_strategies(__lowercase ,       max_length=__lowercase  ) is not PaddingStrategy.DO_NOT_PAD
												    else None
												)
												a__       :						str							=  self.normalize(
												    padded_inputs["""input_features"""] ,       attention_mask=__lowercase  )

									if return_tensors is not None:
												a__       :						Tuple							=  padded_inputs.convert_to_tensors(__lowercase  )

									return padded_inputs









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
_lowercase		:				Union[str, Any]      		=data_utils.TransfoXLTokenizer
_lowercase		:				Tuple      		=data_utils.TransfoXLCorpus
_lowercase		:				Optional[int]      		=data_utils
_lowercase		:				int      		=data_utils

def      lowerCAmelCase_					(  _lowercase    :       Optional[Any]						,       _lowercase    :       List[Any]						,       _lowercase    :       Optional[Any]						,       _lowercase    :       Optional[Any])							->	List[Any]:







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

						a__       :						Optional[int]							=  corpus.__dict__
						corpus_dict_no_vocab.pop("""vocab"""						,       _lowercase)
						a__       :						Tuple							=  pytorch_dump_folder_path + """/""" + CORPUS_NAME
						print(F'''Save dataset to {pytorch_dataset_dump_path}''')
						torch.save(_lowercase						,       _lowercase)

			if tf_checkpoint_path:
						# Convert a pre-trained TensorFlow model
						a__       :						Optional[Any]							=  os.path.abspath(_lowercase)
						a__       :						str							=  os.path.abspath(_lowercase)

						print(F'''Converting Transformer XL checkpoint from {tf_path} with config at {config_path}.''')
						# Initialise PyTorch model
						if transfo_xl_config_file == "":
									a__       :						List[str]							=  TransfoXLConfig()
						else:
									a__       :						Any							=  TransfoXLConfig.from_json_file(_lowercase)
						print(F'''Building PyTorch model from configuration: {config}''')
						a__       :						Any							=  TransfoXLLMHeadModel(_lowercase)

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


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




