



"""simple docstring"""






import argparse

import torch

from transformers import (
    SpeechTaConfig,
    SpeechTaFeatureExtractor,
    SpeechTaForSpeechToSpeech,
    SpeechTaForSpeechToText,
    SpeechTaForTextToSpeech,
    SpeechTaProcessor,
    SpeechTaTokenizer,
    logging,
)
from transformers.tokenization_utils import AddedToken


logging.set_verbosity_info()
lowerCAmelCase__      =    logging.get_logger('''transformers.models.speecht5''')

lowerCAmelCase__      =    {
    '''speech_encoder_prenet.layer_norm''': '''speecht5.encoder.prenet.feature_projection.layer_norm''',
    '''speech_encoder_prenet.post_extract_proj''': '''speecht5.encoder.prenet.feature_projection.projection''',
    '''speech_encoder_prenet.pos_conv.0''': '''speecht5.encoder.prenet.pos_conv_embed.conv''',
    '''speech_encoder_prenet.mask_emb''': '''speecht5.encoder.prenet.masked_spec_embed''',
}
lowerCAmelCase__      =    {
    '''text_encoder_prenet.encoder_prenet.0''': '''speecht5.encoder.prenet.embed_tokens''',
    '''text_encoder_prenet.encoder_prenet.1.alpha''': '''speecht5.encoder.prenet.encode_positions.alpha''',
}
lowerCAmelCase__      =    {
    '''speech_decoder_prenet.decoder_prenet.0.0.prenet.0.0''': '''speecht5.decoder.prenet.layers.0''',
    '''speech_decoder_prenet.decoder_prenet.0.0.prenet.1.0''': '''speecht5.decoder.prenet.layers.1''',
    '''speech_decoder_prenet.decoder_prenet.0.1''': '''speecht5.decoder.prenet.final_layer''',
    '''speech_decoder_prenet.decoder_prenet.1.alpha''': '''speecht5.decoder.prenet.encode_positions.alpha''',
    '''speech_decoder_prenet.spkembs_layer.0''': '''speecht5.decoder.prenet.speaker_embeds_layer''',
}
lowerCAmelCase__      =    {
    '''speech_decoder_postnet.feat_out''': '''speech_decoder_postnet.feat_out''',
    '''speech_decoder_postnet.prob_out''': '''speech_decoder_postnet.prob_out''',
    '''speech_decoder_postnet.postnet.postnet.0.0''': '''speech_decoder_postnet.layers.0.conv''',
    '''speech_decoder_postnet.postnet.postnet.0.1''': '''speech_decoder_postnet.layers.0.batch_norm''',
    '''speech_decoder_postnet.postnet.postnet.1.0''': '''speech_decoder_postnet.layers.1.conv''',
    '''speech_decoder_postnet.postnet.postnet.1.1''': '''speech_decoder_postnet.layers.1.batch_norm''',
    '''speech_decoder_postnet.postnet.postnet.2.0''': '''speech_decoder_postnet.layers.2.conv''',
    '''speech_decoder_postnet.postnet.postnet.2.1''': '''speech_decoder_postnet.layers.2.batch_norm''',
    '''speech_decoder_postnet.postnet.postnet.3.0''': '''speech_decoder_postnet.layers.3.conv''',
    '''speech_decoder_postnet.postnet.postnet.3.1''': '''speech_decoder_postnet.layers.3.batch_norm''',
    '''speech_decoder_postnet.postnet.postnet.4.0''': '''speech_decoder_postnet.layers.4.conv''',
    '''speech_decoder_postnet.postnet.postnet.4.1''': '''speech_decoder_postnet.layers.4.batch_norm''',
}
lowerCAmelCase__      =    {
    '''text_decoder_prenet.embed_tokens''': '''speecht5.decoder.prenet.embed_tokens''',
}
lowerCAmelCase__      =    {
    '''text_decoder_postnet.output_projection''': '''text_decoder_postnet.lm_head''',
}
lowerCAmelCase__      =    {
    '''encoder.layers.*.self_attn.k_proj''': '''speecht5.encoder.wrapped_encoder.layers.*.attention.k_proj''',
    '''encoder.layers.*.self_attn.v_proj''': '''speecht5.encoder.wrapped_encoder.layers.*.attention.v_proj''',
    '''encoder.layers.*.self_attn.q_proj''': '''speecht5.encoder.wrapped_encoder.layers.*.attention.q_proj''',
    '''encoder.layers.*.self_attn.out_proj''': '''speecht5.encoder.wrapped_encoder.layers.*.attention.out_proj''',
    '''encoder.layers.*.self_attn_layer_norm''': '''speecht5.encoder.wrapped_encoder.layers.*.layer_norm''',
    '''encoder.layers.*.fc1''': '''speecht5.encoder.wrapped_encoder.layers.*.feed_forward.intermediate_dense''',
    '''encoder.layers.*.fc2''': '''speecht5.encoder.wrapped_encoder.layers.*.feed_forward.output_dense''',
    '''encoder.layers.*.final_layer_norm''': '''speecht5.encoder.wrapped_encoder.layers.*.final_layer_norm''',
    '''encoder.layer_norm''': '''speecht5.encoder.wrapped_encoder.layer_norm''',
    '''encoder.pos_emb.pe_k''': '''speecht5.encoder.wrapped_encoder.embed_positions.pe_k''',
}
lowerCAmelCase__      =    {
    '''decoder.layers.*.self_attn.k_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn.k_proj''',
    '''decoder.layers.*.self_attn.v_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn.v_proj''',
    '''decoder.layers.*.self_attn.q_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn.q_proj''',
    '''decoder.layers.*.self_attn.out_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn.out_proj''',
    '''decoder.layers.*.self_attn_layer_norm''': '''speecht5.decoder.wrapped_decoder.layers.*.self_attn_layer_norm''',
    '''decoder.layers.*.encoder_attn.k_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn.k_proj''',
    '''decoder.layers.*.encoder_attn.v_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn.v_proj''',
    '''decoder.layers.*.encoder_attn.q_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn.q_proj''',
    '''decoder.layers.*.encoder_attn.out_proj''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn.out_proj''',
    '''decoder.layers.*.encoder_attn_layer_norm''': '''speecht5.decoder.wrapped_decoder.layers.*.encoder_attn_layer_norm''',
    '''decoder.layers.*.fc1''': '''speecht5.decoder.wrapped_decoder.layers.*.feed_forward.intermediate_dense''',
    '''decoder.layers.*.fc2''': '''speecht5.decoder.wrapped_decoder.layers.*.feed_forward.output_dense''',
    '''decoder.layers.*.final_layer_norm''': '''speecht5.decoder.wrapped_decoder.layers.*.final_layer_norm''',
}
lowerCAmelCase__      =    {
    **MAPPING_SPEECH_ENCODER_PRENET,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
    **MAPPING_TEXT_DECODER_PRENET,
    **MAPPING_TEXT_DECODER_POSTNET,
}
lowerCAmelCase__      =    {
    **MAPPING_TEXT_ENCODER_PRENET,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
    **MAPPING_SPEECH_DECODER_PRENET,
    **MAPPING_SPEECH_DECODER_POSTNET,
}
lowerCAmelCase__      =    {
    **MAPPING_SPEECH_ENCODER_PRENET,
    **MAPPING_ENCODER,
    **MAPPING_DECODER,
    **MAPPING_SPEECH_DECODER_PRENET,
    **MAPPING_SPEECH_DECODER_POSTNET,
}
lowerCAmelCase__      =    []
lowerCAmelCase__      =    [
    '''encoder.version''',
    '''encoder.layers.*.norm_k.weight''',
    '''encoder.layers.*.norm_k.bias''',
    '''decoder.version''',
    '''decoder.layers.*.norm_k.weight''',
    '''decoder.layers.*.norm_k.bias''',
    '''decoder.pos_emb.pe_k''',
    '''speech_encoder_prenet.embed_positions._float_tensor''',
    '''text_decoder_prenet.embed_positions._float_tensor''',
]
lowerCAmelCase__      =    IGNORE_KEYS + [
    '''encoder.proj''',
    '''text_encoder_prenet.*''',
    '''speech_decoder_prenet.*''',
    '''speech_decoder_postnet.*''',
]
lowerCAmelCase__      =    IGNORE_KEYS + [
    '''encoder.proj''',
    '''speech_encoder_prenet.*''',
    '''text_decoder_prenet.*''',
    '''text_decoder_postnet.*''',
]
lowerCAmelCase__      =    IGNORE_KEYS + [
    '''encoder.proj''',
    '''text_encoder_prenet.*''',
    '''text_decoder_prenet.*''',
    '''text_decoder_postnet.*''',
]






def 	snake_case_  (				A_ :		Optional[Any],    A_ :		Tuple,    A_ :		List[Any],    A_ :		Union[str, Any],    A_ :		Optional[int]						):


       '''simple docstring'''





       for attribute in key.split('''.'''						):
              _lowerCamelCase						:    int         =			getattr(A_,    A_						)

       if weight_type is not None:
              _lowerCamelCase						:    Optional[int]         =			getattr(A_,    A_						).shape
       else:
              _lowerCamelCase						:    Tuple         =			hf_pointer.shape

       if hf_shape != value.shape:
              raise ValueError(
                  F'''Shape of hf {key + "." + weight_type if weight_type is not None else ""} is {hf_shape}, but should be'''
                  F''' {value.shape} for {full_name}'''						)

       if weight_type == "weight":
              _lowerCamelCase						:    List[Any]         =			value
       elif weight_type == "weight_g":
              _lowerCamelCase						:    List[str]         =			value
       elif weight_type == "weight_v":
              _lowerCamelCase						:    Tuple         =			value
       elif weight_type == "bias":
              _lowerCamelCase						:    Dict         =			value
       elif weight_type == "running_mean":
              _lowerCamelCase						:    List[Any]         =			value
       elif weight_type == "running_var":
              _lowerCamelCase						:    Dict         =			value
       elif weight_type == "num_batches_tracked":
              _lowerCamelCase						:    Optional[int]         =			value
       else:
              _lowerCamelCase						:    List[str]         =			value

       logger.info(F'''{key + ("." + weight_type if weight_type is not None else "")} was initialized from {full_name}.'''						)






def 	snake_case_  (				A_ :		Dict,    A_ :		Union[str, Any]						):


       '''simple docstring'''





       for key in ignore_keys:
              if key.endswith('''.*'''						):
                     if name.startswith(key[:-1]						):
                            return True
              elif ".*." in key:
                     _lowerCamelCase      ,   _lowerCamelCase						:    str         =			key.split('''.*.'''						)
                     if prefix in name and suffix in name:
                            return True
              elif key in name:
                     return True
       return False






def 	snake_case_  (				A_ :		int,    A_ :		List[str],    A_ :		Dict						):


       '''simple docstring'''





       _lowerCamelCase						:    Union[str, Any]         =			[]

       if task == "s2t":
              _lowerCamelCase						:    int         =			hf_model.speechta.encoder.prenet.feature_encoder
              _lowerCamelCase						:    str         =			MAPPING_S2T
              _lowerCamelCase						:    Optional[int]         =			IGNORE_KEYS_S2T
       elif task == "t2s":
              _lowerCamelCase						:    Any         =			None
              _lowerCamelCase						:    Any         =			MAPPING_T2S
              _lowerCamelCase						:    Dict         =			IGNORE_KEYS_T2S
       elif task == "s2s":
              _lowerCamelCase						:    List[Any]         =			hf_model.speechta.encoder.prenet.feature_encoder
              _lowerCamelCase						:    Any         =			MAPPING_S2S
              _lowerCamelCase						:    str         =			IGNORE_KEYS_S2S
       else:
              raise ValueError(F'''Unsupported task: {task}'''						)

       for name, value in fairseq_dict.items():
              if should_ignore(A_,    A_						):
                     logger.info(F'''{name} was ignored'''						)
                     continue

              _lowerCamelCase						:    List[Any]         =			False
              if "conv_layers" in name:
                     load_conv_layer(
                         A_,    A_,    A_,    A_,    hf_model.config.feat_extract_norm == '''group''',    )
                     _lowerCamelCase						:    Any         =			True
              else:
                     for key, mapped_key in MAPPING.items():
                            # mapped_key = "speecht5." + mapped_key if mapped_key not in TOP_LEVEL_KEYS else mapped_key

                            if "*" in key:
                                   _lowerCamelCase      ,   _lowerCamelCase						:    Tuple         =			key.split('''.*.'''						)
                                   if prefix in name and suffix in name:
                                          _lowerCamelCase						:    Optional[int]         =			suffix

                # if key in name or key.split("w2v_model.")[-1] == name.split(".")[0]:
                            if key in name:
                                   _lowerCamelCase						:    str         =			True
                                   if "*" in mapped_key:
                                          _lowerCamelCase						:    Tuple         =			name.split(A_						)[0].split('''.'''						)[-2]
                                          _lowerCamelCase						:    List[Any]         =			mapped_key.replace('''*''',    A_						)
                                   if "weight_g" in name:
                                          _lowerCamelCase						:    int         =			'''weight_g'''
                                   elif "weight_v" in name:
                                          _lowerCamelCase						:    Any         =			'''weight_v'''
                                   elif "bias" in name:
                                          _lowerCamelCase						:    Dict         =			'''bias'''
                                   elif "weight" in name:
                                          _lowerCamelCase						:    Optional[Any]         =			'''weight'''
                                   elif "running_mean" in name:
                                          _lowerCamelCase						:    Any         =			'''running_mean'''
                                   elif "running_var" in name:
                                          _lowerCamelCase						:    List[str]         =			'''running_var'''
                                   elif "num_batches_tracked" in name:
                                          _lowerCamelCase						:    str         =			'''num_batches_tracked'''
                                   else:
                                          _lowerCamelCase						:    str         =			None
                                   set_recursively(A_,    A_,    A_,    A_,    A_						)
                            continue
              if not is_used:
                     unused_weights.append(A_						)

       logger.warning(F'''Unused weights: {unused_weights}'''						)






def 	snake_case_  (				A_ :		Union[str, Any],    A_ :		Union[str, Any],    A_ :		List[Any],    A_ :		Tuple,    A_ :		List[str]						):


       '''simple docstring'''





       _lowerCamelCase						:    str         =			full_name.split('''conv_layers.'''						)[-1]
       _lowerCamelCase						:    str         =			name.split('''.'''						)
       _lowerCamelCase						:    List[Any]         =			int(items[0]						)
       _lowerCamelCase						:    Tuple         =			int(items[1]						)

       if type_id == 0:
              if "bias" in name:
                     if value.shape != feature_extractor.conv_layers[layer_id].conv.bias.data.shape:
                            raise ValueError(
                                F'''{full_name} has size {value.shape}, but'''
                                F''' {feature_extractor.conv_layers[layer_id].conv.bias.data.shape} was found.'''						)
                     _lowerCamelCase						:    Dict         =			value
                     logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''						)
              elif "weight" in name:
                     if value.shape != feature_extractor.conv_layers[layer_id].conv.weight.data.shape:
                            raise ValueError(
                                F'''{full_name} has size {value.shape}, but'''
                                F''' {feature_extractor.conv_layers[layer_id].conv.weight.data.shape} was found.'''						)
                     _lowerCamelCase						:    List[Any]         =			value
                     logger.info(F'''Feat extract conv layer {layer_id} was initialized from {full_name}.'''						)
       elif (type_id == 2 and not use_group_norm) or (type_id == 2 and layer_id == 0 and use_group_norm):
              if "bias" in name:
                     if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape:
                            raise ValueError(
                                F'''{full_name} has size {value.shape}, but'''
                                F''' {feature_extractor.conv_layers[layer_id].layer_norm.bias.data.shape} was found.'''						)
                     _lowerCamelCase						:    Any         =			value
                     logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''						)
              elif "weight" in name:
                     if value.shape != feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape:
                            raise ValueError(
                                F'''{full_name} has size {value.shape}, but'''
                                F''' {feature_extractor.conv_layers[layer_id].layer_norm.weight.data.shape} was found.'''						)
                     _lowerCamelCase						:    Union[str, Any]         =			value
                     logger.info(F'''Feat extract layer norm weight of layer {layer_id} was initialized from {full_name}.'''						)
       else:
              unused_weights.append(A_						)






@torch.no_grad()
def 	snake_case_  (				A_ :		Optional[Any],    A_ :		List[str],    A_ :		str,    A_ :		Union[str, Any]=None,    A_ :		Dict=None,    A_ :		Optional[Any]=None,    ):


       '''simple docstring'''





       if config_path is not None:
              _lowerCamelCase						:    Optional[int]         =			SpeechTaConfig.from_pretrained(A_						)
       else:
              _lowerCamelCase						:    int         =			SpeechTaConfig()

       if task == "s2t":
              _lowerCamelCase						:    Tuple         =			config.max_text_positions
              _lowerCamelCase						:    str         =			SpeechTaForSpeechToText(A_						)
       elif task == "t2s":
              _lowerCamelCase						:    Tuple         =			18_76
              _lowerCamelCase						:    Any         =			6_00
              _lowerCamelCase						:    Any         =			config.max_speech_positions
              _lowerCamelCase						:    List[Any]         =			SpeechTaForTextToSpeech(A_						)
       elif task == "s2s":
              _lowerCamelCase						:    Optional[int]         =			18_76
              _lowerCamelCase						:    Optional[int]         =			config.max_speech_positions
              _lowerCamelCase						:    Any         =			SpeechTaForSpeechToSpeech(A_						)
       else:
              raise ValueError(F'''Unknown task name: {task}'''						)

       if vocab_path:
              _lowerCamelCase						:    Optional[Any]         =			SpeechTaTokenizer(A_,    model_max_length=config.max_text_positions						)

              # Mask token behaves like a normal word, i.e. include the space before it
              _lowerCamelCase						:    Any         =			AddedToken('''<mask>''',    lstrip=A_,    rstrip=A_						)
              _lowerCamelCase						:    List[str]         =			mask_token
              tokenizer.add_special_tokens({'''mask_token''': mask_token}						)
              tokenizer.add_tokens(['''<ctc_blank>''']						)

       _lowerCamelCase						:    Union[str, Any]         =			SpeechTaFeatureExtractor()
       _lowerCamelCase						:    Union[str, Any]         =			SpeechTaProcessor(tokenizer=A_,    feature_extractor=A_						)
       processor.save_pretrained(A_						)

       _lowerCamelCase						:    List[str]         =			torch.load(A_						)
       recursively_load_weights(fairseq_checkpoint['''model'''],    A_,    A_						)

       model.save_pretrained(A_						)

       if repo_id:
              print('''Pushing to the hub...'''						)
              processor.push_to_hub(A_						)
              model.push_to_hub(A_						)


if __name__ == "__main__":
   lowerCAmelCase__      =    argparse.ArgumentParser()
   parser.add_argument(
       '''--task''',
       default='''s2t''',
       type=str,
       help='''Type of the SpeechT5 model you\'d like to convert. Should be one of \'s2t\', \'t2s\', \'s2s\'.''',
   )
   parser.add_argument('''--checkpoint_path''', required=True, default=None, type=str, help='''Path to fairseq checkpoint''')
   parser.add_argument('''--vocab_path''', default=None, type=str, help='''Path to SentencePiece model''')
   parser.add_argument('''--config_path''', default=None, type=str, help='''Path to hf config.json of model to convert''')
   parser.add_argument(
       '''--pytorch_dump_folder_path''', required=True, default=None, type=str, help='''Path to the output PyTorch model.'''
   )
   parser.add_argument(
       '''--push_to_hub''', default=None, type=str, help='''Where to upload the converted model on the 🤗 hub.'''
   )

   lowerCAmelCase__      =    parser.parse_args()
   convert_speechta_checkpoint(
       args.task,
       args.checkpoint_path,
       args.pytorch_dump_folder_path,
       args.config_path,
       args.vocab_path,
       args.push_to_hub,
   )







"""simple docstring"""






import os
import posixpath
import uuid
from dataclasses import dataclass
from typing import TYPE_CHECKING, Iterable, List, Optional, Tuple, Union

import numpy as np
import pyarrow as pa

import datasets
from datasets.arrow_writer import ArrowWriter, ParquetWriter
from datasets.config import MAX_SHARD_SIZE
from datasets.filesystems import (
    is_remote_filesystem,
    rename,
)
from datasets.iterable_dataset import _BaseExamplesIterable
from datasets.utils.py_utils import convert_file_size_to_int


lowerCAmelCase__      =    datasets.utils.logging.get_logger(__name__)

if TYPE_CHECKING:
   import pyspark






@dataclass
class 	__snake_case	(     datasets.BuilderConfig):
  snake_case__    :       Optional[datasets.Features]        =						None






def 	snake_case_  (				A_ :		"pyspark.sql.DataFrame",    A_ :		List[int],    ):


       '''simple docstring'''





       import pyspark

       def generate_fn():
              _lowerCamelCase						:    int         =			df.select('''*''',    pyspark.sql.functions.spark_partition_id().alias('''part_id'''						)						)
              for partition_id in partition_order:
                     _lowerCamelCase						:    Any         =			df_with_partition_id.select('''*'''						).where(F'''part_id = {partition_id}'''						).drop('''part_id'''						)
                     _lowerCamelCase						:    Optional[int]         =			partition_df.collect()
                     _lowerCamelCase						:    List[str]         =			0
                     for row in rows:
                            yield F'''{partition_id}_{row_id}''', row.asDict()
                            row_id += 1

       return generate_fn






class 	__snake_case	(     _BaseExamplesIterable):




  def __init__(			self					:						Tuple						, __lowerCAmelCase					:						"pyspark.sql.DataFrame"						, __lowerCAmelCase					:						Optional[int]=None						, ):






         """simple docstring"""






         _lowerCamelCase						:    Dict         =			df
         _lowerCamelCase						:    Union[str, Any]         =			partition_order or range(self.df.rdd.getNumPartitions()		)
         _lowerCamelCase						:    Dict         =			_generate_iterable_examples(self.df						, self.partition_order		)




  def __iter__(			self					:						List[Any]		):






         """simple docstring"""






         yield from self.generate_examples_fn()




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						np.random.Generator		):






         """simple docstring"""






         _lowerCamelCase						:    Optional[Any]         =			list(range(self.df.rdd.getNumPartitions()		)		)
         generator.shuffle(__lowerCAmelCase		)
         return SparkExamplesIterable(self.df						, partition_order=__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Optional[Any]						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int		):






         """simple docstring"""






         _lowerCamelCase						:    List[Any]         =			self.split_shard_indices_by_worker(__lowerCAmelCase						, __lowerCAmelCase		)
         return SparkExamplesIterable(self.df						, partition_order=__lowerCAmelCase		)




  @property
  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]		):






         """simple docstring"""






         return len(self.partition_order		)






class 	__snake_case	(     datasets.DatasetBuilder):
  snake_case__    :       List[Any]        =						SparkConfig




  def __init__(			self					:						Union[str, Any]						, __lowerCAmelCase					:						"pyspark.sql.DataFrame"						, __lowerCAmelCase					:						str = None						, __lowerCAmelCase					:						str = None						, **__lowerCAmelCase					:						int						, ):






         """simple docstring"""






         import pyspark

         _lowerCamelCase						:    Optional[int]         =			pyspark.sql.SparkSession.builder.getOrCreate()
         _lowerCamelCase						:    int         =			df
         _lowerCamelCase						:    Any         =			working_dir

         super().__init__(
             cache_dir=__lowerCAmelCase						, config_name=str(self.df.semanticHash()		)						, **__lowerCAmelCase						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						List[Any]		):






         """simple docstring"""






         def create_cache_and_write_probe(__lowerCAmelCase					:						Optional[int]		):
                # makedirs with exist_ok will recursively create the directory. It will not throw an error if directories
                # already exist.
                os.makedirs(self._cache_dir						, exist_ok=__lowerCAmelCase		)
                _lowerCamelCase						:    Optional[int]         =			os.path.join(self._cache_dir						, '''fs_test''' + uuid.uuida().hex		)
                # Opening the file in append mode will create a new file unless it already exists, in which case it will not
                # change the file contents.
                open(__lowerCAmelCase						, '''a'''		)
                return [probe_file]

         if self._spark.conf.get('''spark.master'''						, ''''''		).startswith('''local'''		):
                return

         # If the cluster is multi-node, make sure that the user provided a cache_dir and that it is on an NFS
         # accessible to the driver.
         # TODO: Stream batches to the driver using ArrowCollectSerializer instead of throwing an error.
         if self._cache_dir:
                _lowerCamelCase						:    Optional[Any]         =			(
                    self._spark.sparkContext.parallelize(range(1		)						, 1		).mapPartitions(__lowerCAmelCase		).collect()
                )
                if os.path.isfile(probe[0]		):
                       return

         raise ValueError(
             '''When using Dataset.from_spark on a multi-node cluster, the driver and all workers should be able to access cache_dir'''		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any		):






         """simple docstring"""






         return datasets.DatasetInfo(features=self.config.features		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						int						, __lowerCAmelCase					:						datasets.download.download_manager.DownloadManager		):






         """simple docstring"""






         return [datasets.SplitGenerator(name=datasets.Split.TRAIN		)]




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						Union[str, Any]		):






         """simple docstring"""






         import pyspark

         def get_arrow_batch_size(__lowerCAmelCase					:						Dict		):
                for batch in it:
                       yield pa.RecordBatch.from_pydict({'''batch_bytes''': [batch.nbytes]}		)

         _lowerCamelCase						:    Any         =			self.df.count()
         _lowerCamelCase						:    Union[str, Any]         =			df_num_rows if df_num_rows <= 1_0_0 else 1_0_0
         # Approximate the size of each row (in Arrow format) by averaging over a max-100-row sample.
         _lowerCamelCase						:    List[Any]         =			(
             self.df.limit(__lowerCAmelCase		)
             .repartition(1		)
             .mapInArrow(__lowerCAmelCase						, '''batch_bytes: long'''		)
             .agg(pyspark.sql.functions.sum('''batch_bytes'''		).alias('''sample_bytes'''		)		)
             .collect()[0]
             .sample_bytes
             / sample_num_rows
         )
         _lowerCamelCase						:    Dict         =			approx_bytes_per_row * df_num_rows
         if approx_total_size > max_shard_size:
                # Make sure there is at least one row per partition.
                _lowerCamelCase						:    List[str]         =			min(__lowerCAmelCase						, int(approx_total_size / max_shard_size		)		)
                _lowerCamelCase						:    Optional[int]         =			self.df.repartition(__lowerCAmelCase		)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Union[str, Any]						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						str						, __lowerCAmelCase					:						int						, ):






         """simple docstring"""






         import pyspark

         _lowerCamelCase						:    Optional[Any]         =			ParquetWriter if file_format == '''parquet''' else ArrowWriter
         _lowerCamelCase						:    List[Any]         =			os.path.join(self._working_dir						, os.path.basename(__lowerCAmelCase		)		) if self._working_dir else fpath
         _lowerCamelCase						:    Dict         =			file_format == '''parquet'''

         # Define these so that we don't reference self in write_arrow, which will result in a pickling error due to
         # pickling the SparkContext.
         _lowerCamelCase						:    str         =			self.config.features
         _lowerCamelCase						:    Dict         =			self._writer_batch_size
         _lowerCamelCase						:    List[str]         =			self._fs.storage_options

         def write_arrow(__lowerCAmelCase					:						List[str]		):
                # Within the same SparkContext, no two task attempts will share the same attempt ID.
                _lowerCamelCase						:    List[str]         =			pyspark.TaskContext().taskAttemptId()
                _lowerCamelCase						:    Any         =			next(__lowerCAmelCase						, __lowerCAmelCase		)
                if first_batch is None:
                       # Some partitions might not receive any data.
                       return pa.RecordBatch.from_arrays(
                           [[task_id], [0], [0]]						, names=['''task_id''', '''num_examples''', '''num_bytes''']						, )
                _lowerCamelCase						:    List[Any]         =			0
                _lowerCamelCase						:    Optional[int]         =			writer_class(
                    features=__lowerCAmelCase						, path=working_fpath.replace('''SSSSS'''						, f'''{shard_id:05d}'''		).replace('''TTTTT'''						, f'''{task_id:05d}'''		)						, writer_batch_size=__lowerCAmelCase						, storage_options=__lowerCAmelCase						, embed_local_files=__lowerCAmelCase						, )
                _lowerCamelCase						:    int         =			pa.Table.from_batches([first_batch]		)
                writer.write_table(__lowerCAmelCase		)
                for batch in it:
                       if max_shard_size is not None and writer._num_bytes >= max_shard_size:
                              _lowerCamelCase      ,   _lowerCamelCase						:    Any         =			writer.finalize()
                              writer.close()
                              yield pa.RecordBatch.from_arrays(
                                  [[task_id], [num_examples], [num_bytes]]						, names=['''task_id''', '''num_examples''', '''num_bytes''']						, )
                              shard_id += 1
                              _lowerCamelCase						:    Optional[int]         =			writer_class(
                                  features=writer._features						, path=working_fpath.replace('''SSSSS'''						, f'''{shard_id:05d}'''		).replace('''TTTTT'''						, f'''{task_id:05d}'''		)						, writer_batch_size=__lowerCAmelCase						, storage_options=__lowerCAmelCase						, embed_local_files=__lowerCAmelCase						, )
                       _lowerCamelCase						:    Optional[int]         =			pa.Table.from_batches([batch]		)
                       writer.write_table(__lowerCAmelCase		)

                if writer._num_bytes > 0:
                       _lowerCamelCase      ,   _lowerCamelCase						:    Optional[int]         =			writer.finalize()
                       writer.close()
                       yield pa.RecordBatch.from_arrays(
                           [[task_id], [num_examples], [num_bytes]]						, names=['''task_id''', '''num_examples''', '''num_bytes''']						, )

                if working_fpath != fpath:
                       for file in os.listdir(os.path.dirname(__lowerCAmelCase		)		):
                              _lowerCamelCase						:    Optional[Any]         =			os.path.join(os.path.dirname(__lowerCAmelCase		)						, os.path.basename(__lowerCAmelCase		)		)
                              shutil.move(__lowerCAmelCase						, __lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			(
             self.df.mapInArrow(__lowerCAmelCase						, '''task_id: long, num_examples: long, num_bytes: long'''		)
             .groupBy('''task_id'''		)
             .agg(
                 pyspark.sql.functions.sum('''num_examples'''		).alias('''total_num_examples'''		)						, pyspark.sql.functions.sum('''num_bytes'''		).alias('''total_num_bytes'''		)						, pyspark.sql.functions.count('''num_bytes'''		).alias('''num_shards'''		)						, pyspark.sql.functions.collect_list('''num_examples'''		).alias('''shard_lengths'''		)						, )
             .collect()
         )
         for row in stats:
                yield row.task_id, (row.total_num_examples, row.total_num_bytes, row.num_shards, row.shard_lengths)




  def 						SCREAMING_SNAKE_CASE    (			self					:						Tuple						, __lowerCAmelCase					:						"datasets.SplitGenerator"						, __lowerCAmelCase					:						str = "arrow"						, __lowerCAmelCase					:						Optional[Union[str, int]] = None						, __lowerCAmelCase					:						Optional[int] = None						, **__lowerCAmelCase					:						Tuple						, ):






         """simple docstring"""






         self._validate_cache_dir()

         _lowerCamelCase						:    str         =			convert_file_size_to_int(max_shard_size or MAX_SHARD_SIZE		)
         self._repartition_df_if_needed(__lowerCAmelCase		)
         _lowerCamelCase						:    str         =			not is_remote_filesystem(self._fs		)
         _lowerCamelCase						:    Tuple         =			os.path.join if is_local else posixpath.join

         _lowerCamelCase						:    int         =			'''-TTTTT-SSSSS-of-NNNNN'''
         _lowerCamelCase						:    Tuple         =			f'''{self.name}-{split_generator.name}{SUFFIX}.{file_format}'''
         _lowerCamelCase						:    List[Any]         =			path_join(self._output_dir						, __lowerCAmelCase		)

         _lowerCamelCase						:    List[Any]         =			0
         _lowerCamelCase						:    Any         =			0
         _lowerCamelCase						:    str         =			0
         _lowerCamelCase						:    int         =			[]
         _lowerCamelCase						:    List[str]         =			[]

         for task_id, content in self._prepare_split_single(__lowerCAmelCase						, __lowerCAmelCase						, __lowerCAmelCase		):
                (
                    (
                    _lowerCamelCase
                )      ,   (
                    _lowerCamelCase
                )      ,   (
                    _lowerCamelCase
                )      ,   (
                    _lowerCamelCase
                )      ,   
                )						:    str         =			content
                if num_bytes > 0:
                       total_num_examples += num_examples
                       total_num_bytes += num_bytes
                       total_shards += num_shards
                       task_id_and_num_shards.append((task_id, num_shards)		)
                       all_shard_lengths.extend(__lowerCAmelCase		)

         _lowerCamelCase						:    int         =			total_num_examples
         _lowerCamelCase						:    str         =			total_num_bytes

         # should rename everything at the end
         logger.debug(f'''Renaming {total_shards} shards.'''		)
         if total_shards > 1:
                _lowerCamelCase						:    Optional[Any]         =			all_shard_lengths

                # Define fs outside of _rename_shard so that we don't reference self in the function, which will result in a
                # pickling error due to pickling the SparkContext.
                _lowerCamelCase						:    str         =			self._fs

                # use the -SSSSS-of-NNNNN pattern
                def _rename_shard(
                    __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, __lowerCAmelCase					:						int						, ):
                       rename(
                           __lowerCAmelCase						, fpath.replace('''SSSSS'''						, f'''{shard_id:05d}'''		).replace('''TTTTT'''						, f'''{task_id:05d}'''		)						, fpath.replace('''TTTTT-SSSSS'''						, f'''{global_shard_id:05d}'''		).replace('''NNNNN'''						, f'''{total_shards:05d}'''		)						, )

                _lowerCamelCase						:    Union[str, Any]         =			[]
                _lowerCamelCase						:    Any         =			0
                for i in range(len(__lowerCAmelCase		)		):
                       _lowerCamelCase      ,   _lowerCamelCase						:    Dict         =			task_id_and_num_shards[i]
                       for shard_id in range(__lowerCAmelCase		):
                              args.append([task_id, shard_id, global_shard_id]		)
                              global_shard_id += 1
                self._spark.sparkContext.parallelize(__lowerCAmelCase						, len(__lowerCAmelCase		)		).map(lambda __lowerCAmelCase		: _rename_shard(*__lowerCAmelCase		)		).collect()
         else:
                # don't use any pattern
                _lowerCamelCase						:    Any         =			0
                _lowerCamelCase						:    List[str]         =			task_id_and_num_shards[0][0]
                self._rename(
                    fpath.replace('''SSSSS'''						, f'''{shard_id:05d}'''		).replace('''TTTTT'''						, f'''{task_id:05d}'''		)						, fpath.replace(__lowerCAmelCase						, ''''''		)						, )




  def 						SCREAMING_SNAKE_CASE    (			self					:						Any						, __lowerCAmelCase					:						"datasets.SplitGenerator"						, ):






         """simple docstring"""






         return SparkExamplesIterable(self.df		)


