





import copy
from typing import Any, Dict, List, Optional, Union

import numpy as np

from ...audio_utils import mel_filter_bank, spectrogram, window_function
from ...feature_extraction_sequence_utils import SequenceFeatureExtractor
from ...feature_extraction_utils import BatchFeature
from ...utils import TensorType, logging


lowercase      :       Union[str, Any]									=  logging.get_logger(__name__)





class  A__    (							__UpperCAmelCase	):




   """simple docstring"""



   __A   :				Tuple         =  ['''input_features''']

   def __init__(					self					,  lowercase=80					,  lowercase=1_6000					,  lowercase=160					,  lowercase=30					,  lowercase=400					,  lowercase=0.0					,  lowercase=False					,  **lowercase					,  )			->						Any:







          '''simple docstring'''

          super().__init__(
              feature_size=lowercase					,  sampling_rate=lowercase					,  padding_value=lowercase					,  return_attention_mask=lowercase					,  **lowercase					,  )
          a__						:	Union[str, Any] =				n_fft
          a__						:	Dict =				hop_length
          a__						:	List[str] =				chunk_length
          a__						:	str =				chunk_length * sampling_rate
          a__						:	List[Any] =				self.n_samples // hop_length
          a__						:	List[str] =				sampling_rate
          a__						:	Any =				mel_filter_bank(
              num_frequency_bins=1 + n_fft // 2					,  num_mel_filters=lowercase					,  min_frequency=0.0					,  max_frequency=80_00.0					,  sampling_rate=lowercase					,  norm='slaney'					,  mel_scale='slaney'					,  )

   def 							__lowercase		(					self					,  lowercase)			->						np.ndarray:







          '''simple docstring'''

          a__						:	Union[str, Any] =				spectrogram(
              lowercase					,  window_function(self.n_fft					,  'hann')					,  frame_length=self.n_fft					,  hop_length=self.hop_length					,  power=2.0					,  mel_filters=self.mel_filters					,  log_mel='log10'					,  )
          a__						:	Optional[Any] =				log_spec[:, :-1]
          a__						:	Optional[int] =				np.maximum(lowercase					,  log_spec.max() - 8.0)
          a__						:	Optional[int] =				(log_spec + 4.0) / 4.0
          return log_spec

   @staticmethod
   # Copied from transformers.models.wav2vec2.feature_extraction_wav2vec2.Wav2Vec2FeatureExtractor.zero_mean_unit_var_norm
   def 							__lowercase		(					lowercase					,  lowercase					,  lowercase = 0.0)			->						List[np.ndarray]:







          '''simple docstring'''

          if attention_mask is not None:
                 a__						:	List[Any] =				np.array(lowercase					,  np.intaa)
                 a__						:	int =				[]

                 for vector, length in zip(lowercase					,  attention_mask.sum(-1)):
                        a__						:	List[str] =				(vector - vector[:length].mean()) / np.sqrt(vector[:length].var() + 1e-7)
                        if length < normed_slice.shape[0]:
                               a__						:	int =				padding_value

                        normed_input_values.append(lowercase)
          else:
                 a__						:	Optional[Any] =				[(x - x.mean()) / np.sqrt(x.var() + 1e-7) for x in input_values]

          return normed_input_values

   def __call__(					self					,  lowercase					,  lowercase = True					,  lowercase = None					,  lowercase = None					,  lowercase = None					,  lowercase = "max_length"					,  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.__class__.__name__} was trained using a'
                            F' sampling rate of {self.sampling_rate}. Please make sure that the provided `raw_speech` input'
                            F' was sampled with {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__						:	List[Any] =				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__						:	Optional[int] =				is_batched_numpy or (
              isinstance(lowercase					,  (list, tuple)) and (isinstance(raw_speech[0]					,  (np.ndarray, tuple, list)))
          )

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

          # always return batch
          if not is_batched:
                 a__						:	List[Any] =				[np.asarray([raw_speech]).T]

          a__						:	Optional[int] =				BatchFeature({'input_features': raw_speech})

          # convert into correct format for padding

          a__						:	Dict =				self.pad(
              lowercase					,  padding=lowercase					,  max_length=max_length if max_length else self.n_samples					,  truncation=lowercase					,  pad_to_multiple_of=lowercase					,  return_attention_mask=return_attention_mask or do_normalize					,  )

          # zero-mean and unit-variance normalization
          if do_normalize:
                 a__						:	List[str] =				self.zero_mean_unit_var_norm(
                     padded_inputs['input_features']					,  attention_mask=padded_inputs['attention_mask']					,  padding_value=self.padding_value					,  )
                 a__						:	Tuple =				np.stack(padded_inputs['input_features']					,  axis=0)

          # make sure list is in array format
          a__						:	Tuple =				padded_inputs.get('input_features').transpose(2					,  0					,  1)

          a__						:	Tuple =				[self._np_extract_fbank_features(lowercase) for waveform in input_features[0]]

          if isinstance(input_features[0]					,  lowercase):
                 a__						:	Optional[Any] =				[np.asarray(lowercase					,  dtype=np.floataa) for feature in input_features]
          else:
                 a__						:	int =				input_features

          if return_attention_mask:
                 # rescale from sample (48000) to feature (3000)
                 a__						:	List[Any] =				padded_inputs['attention_mask'][:, :: self.hop_length]

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

          return padded_inputs

   def 							__lowercase		(					self)			->						Dict[str, Any]:







          '''simple docstring'''

          a__						:	Tuple =				copy.deepcopy(self.__dict__)
          a__						:	Optional[Any] =				self.__class__.__name__
          if "mel_filters" in output:
                 del output["mel_filters"]
          return output









import argparse
import logging
import os

import datasets
import tensorflow as tf

from transformers import AutoTokenizer


lowerCAmelCase_    =				logging.getLogger(__name__)






def      lowerCamelCase_						(				)		->						Optional[Any]:






				"""simple docstring"""


				snake_case_						:    List[str]				=  argparse.ArgumentParser(
				    description='''Prepare TFRecord shards from pre-tokenized samples of the wikitext dataset.''' )
				parser.add_argument(
				    '''--dataset_name''' ,				type=_UpperCamelCase ,				default='''wikitext''' ,				help='''Name of the training. Explore datasets at: hf.co/datasets.''' ,				)
				parser.add_argument(
				    '''--dataset_config''' ,				type=_UpperCamelCase ,				default='''wikitext-103-raw-v1''' ,				help='''Configuration name of the dataset.''' )
				parser.add_argument(
				    '''--tokenizer_name_or_path''' ,				type=_UpperCamelCase ,				default='''sayakpaul/unigram-tokenizer-wikitext''' ,				help='''Tokenizer identifier. Can be a local filepath or a Hub identifier.''' ,				)
				parser.add_argument(
				    '''--shard_size''' ,				type=_UpperCamelCase ,				default=1_000 ,				help='''Number of entries to go in a single shard.''' ,				)
				parser.add_argument('''--split''' ,				type=_UpperCamelCase ,				default='''train''' ,				choices=['''train''', '''test''', '''validation'''] )
				parser.add_argument(
				    '''--limit''' ,				default=_UpperCamelCase ,				type=_UpperCamelCase ,				help='''Limit the number of shards (used for debugging).''' ,				)
				parser.add_argument(
				    '''--max_length''' ,				type=_UpperCamelCase ,				default=512 ,				help='''Maximum sequence length. For training on TPUs, it helps to have a maximum'''
				    ''' sequence length that is a multiple of 8.''' ,				)
				parser.add_argument(
				    '''--output_dir''' ,				default='''tf-tpu''' ,				type=_UpperCamelCase ,				help='''Output directory where the TFRecord shards will be saved. If the'''
				    ''' path is appended with `gs://` (\'gs://tf-tpu\', for example) then the TFRecord'''
				    ''' shards will be directly saved to a Google Cloud Storage bucket.''' ,				)

				snake_case_						:    List[Any]				=  parser.parse_args()
				return args






def      lowerCamelCase_						(				_UpperCamelCase )		->						Tuple:






				"""simple docstring"""


				def fn(_UpperCamelCase ):
								return tokenizer(examples['''text'''] )

				return fn






def      lowerCamelCase_						(				_UpperCamelCase )		->						Union[str, Any]:






				"""simple docstring"""


				snake_case_						:    Any				=  []
				for i in range(len(tokenized_data['''input_ids'''] ) ):
								snake_case_						:    Any				=  {
								    '''input_ids''': tf.train.Feature(intaa_list=tf.train.IntaaList(value=tokenized_data['''input_ids'''][i] ) ),
								    '''attention_mask''': tf.train.Feature(
								        intaa_list=tf.train.IntaaList(value=tokenized_data['''attention_mask'''][i] ) ),
								}
								snake_case_						:    Optional[int]				=  tf.train.Features(feature=_UpperCamelCase )
								snake_case_						:    Optional[Any]				=  tf.train.Example(features=_UpperCamelCase )
								snake_case_						:    Optional[Any]				=  example.SerializeToString()
								records.append(_UpperCamelCase )
				return records






def      lowerCamelCase_						(				_UpperCamelCase )		->						Optional[int]:






				"""simple docstring"""


				snake_case_						:    int				=  datasets.load_dataset(args.dataset_name ,				args.dataset_config ,				split=args.split )

				if args.limit is not None:
								snake_case_						:    Union[str, Any]				=  min(len(_UpperCamelCase ) ,				args.limit )
								snake_case_						:    int				=  dataset.select(range(_UpperCamelCase ) )
								print(f'''Limiting the dataset to {args.limit} entries.''' )

				snake_case_						:    Dict				=  AutoTokenizer.from_pretrained(args.tokenizer_name_or_path )

				# Handle output directory creation.
				# For serializing into a Google Cloud Storage Bucket, one needs to first
				# create a bucket.
				if "gs" not in args.output_dir:
								if not os.path.exists(args.output_dir ):
												os.makedirs(args.output_dir )
								snake_case_						:    str				=  os.path.join(args.output_dir ,				args.split )
								if not os.path.exists(_UpperCamelCase ):
												os.makedirs(_UpperCamelCase )
				else:
								snake_case_						:    Optional[Any]				=  os.path.join(args.output_dir ,				args.split )

				# Tokenize the whole dataset at once.
				snake_case_						:    Optional[Any]				=  tokenize_function(_UpperCamelCase )
				snake_case_						:    List[Any]				=  dataset.map(_UpperCamelCase ,				batched=_UpperCamelCase ,				num_proc=4 ,				remove_columns=['''text'''] )

				# We need to concatenate all our texts together, and then split the result
				# into chunks of a fixed size, which we will call block_size. To do this, we
				# will use the map method again, with the option batched=True. When we use batched=True,
				# the function we pass to map() will be passed multiple inputs at once, allowing us
				# to group them into more or fewer examples than we had in the input.
				# This allows us to create our new fixed-length samples. The advantage of this
				# method is that we don't lose a whole lot of content from the dataset compared to the
				# case where we simply tokenize with a pre-defined max_length.

				def group_texts(_UpperCamelCase ):
								# Concatenate all texts.
								snake_case_						:    Tuple				=  {k: sum(examples[k] ,				[] ) for k in examples.keys()}
								snake_case_						:    List[str]				=  len(concatenated_examples[list(examples.keys() )[0]] )
								# We drop the small remainder, though you could add padding instead if the model supports it
								# In this, as in all things, we advise you to follow your heart 🫀
								snake_case_						:    int				=  (total_length // args.max_length) * args.max_length
								# Split by chunks of max_len.
								snake_case_						:    Union[str, Any]				=  {
								    k: [t[i : i + args.max_length] for i in range(0 ,				_UpperCamelCase ,				args.max_length )]
								    for k, t in concatenated_examples.items()
								}
								return result

				snake_case_						:    int				=  dataset_tokenized.map(_UpperCamelCase ,				batched=_UpperCamelCase ,				batch_size=1_000 ,				num_proc=4 )

				snake_case_						:    str				=  0
				snake_case_						:    Optional[Any]				=  0
				for shard in range(0 ,				len(_UpperCamelCase ) ,				args.shard_size ):
								snake_case_						:    Any				=  grouped_dataset[shard : shard + args.shard_size]
								snake_case_						:    str				=  len(dataset_snapshot['''input_ids'''] )
								snake_case_						:    Union[str, Any]				=  os.path.join(_UpperCamelCase ,				f'''dataset-{shard_count}-{records_containing}.tfrecord''' )
								snake_case_						:    Dict				=  get_serialized_examples(_UpperCamelCase )

								with tf.io.TFRecordWriter(_UpperCamelCase ) as out_file:
												for i in range(len(_UpperCamelCase ) ):
																snake_case_						:    List[str]				=  serialized_examples[i]
																out_file.write(_UpperCamelCase )
												print('''Wrote file {} containing {} records'''.format(_UpperCamelCase ,				_UpperCamelCase ) )

								shard_count += 1
								total_records += records_containing

				with open(f'''split-{args.split}-records-count.txt''' ,				'''w''' ) as f:
								print(f'''Total {args.split} records: {total_records}''' ,				file=_UpperCamelCase )


if __name__ == "__main__":
	lowerCAmelCase_    =				parse_args()
	main(args)
