




'''simple docstring'''


from unittest.mock import patch

import pyspark

from datasets.packaged_modules.spark.spark import (
    Spark,
    SparkExamplesIterable,
    _generate_iterable_examples,
)

from ..utils import (
    require_dill_gt_0_3_2,
    require_not_windows,
)



def  _SCREAMING_SNAKE_CASE     (      UpperCamelCase__    :					Dict ,    UpperCamelCase__    :					Tuple				):




     """simple docstring"""


     a_     :							Optional[Any]     						=  []
     for part_id in partition_order:
          a_     :							int     						=  df.where(F"SPARK_PARTITION_ID() = {part_id}"				).collect()
          for row_idx, row in enumerate(UpperCamelCase__				):
               expected_row_ids_and_row_dicts.append((F"{part_id}_{row_idx}", row.asDict())				)
     return expected_row_ids_and_row_dicts



@require_not_windows
@require_dill_gt_0_3_2
def  _SCREAMING_SNAKE_CASE     (      ):




     """simple docstring"""


     a_     :							Union[str, Any]     						=  pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
     a_     :							List[str]     						=  spark.range(100				).repartition(1				)
     a_     :							List[str]     						=  Spark(UpperCamelCase__				)
     # The id ints will be converted to Pyarrow int64s, so each row will be 8 bytes. Setting a max_shard_size of 16 means
     # that each partition can hold 2 rows.
     spark_builder._repartition_df_if_needed(max_shard_size=16				)
     # Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
     assert spark_builder.df.rdd.getNumPartitions() == 50



@require_not_windows
@require_dill_gt_0_3_2
def  _SCREAMING_SNAKE_CASE     (      ):




     """simple docstring"""


     a_     :							Any     						=  pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
     a_     :							str     						=  spark.range(10				).repartition(2				)
     a_     :							int     						=  [1, 0]
     a_     :							int     						=  _generate_iterable_examples(UpperCamelCase__ ,    UpperCamelCase__				)  # Reverse the partitions.
     a_     :							Any     						=  _get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase__ ,    UpperCamelCase__				)

     for i, (row_id, row_dict) in enumerate(generate_fn()				):
          a_						,     a_     :							Any     						=  expected_row_ids_and_row_dicts[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict



@require_not_windows
@require_dill_gt_0_3_2
def  _SCREAMING_SNAKE_CASE     (      ):




     """simple docstring"""


     a_     :							List[Any]     						=  pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
     a_     :							List[str]     						=  spark.range(10				).repartition(1				)
     a_     :							List[Any]     						=  SparkExamplesIterable(UpperCamelCase__				)
     assert it.n_shards == 1
     for i, (row_id, row_dict) in enumerate(UpperCamelCase__				):
          assert row_id == F"0_{i}"
          assert row_dict == {"id": i}



@require_not_windows
@require_dill_gt_0_3_2
def  _SCREAMING_SNAKE_CASE     (      ):




     """simple docstring"""


     a_     :							str     						=  pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
     a_     :							Optional[int]     						=  spark.range(30				).repartition(3				)
     # Mock the generator so that shuffle reverses the partition indices.
     with patch("""numpy.random.Generator"""				) as generator_mock:
          a_     :							Union[str, Any]     						=  lambda UpperCamelCase__				: x.reverse()
          a_     :							Tuple     						=  _get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase__ ,    [2, 1, 0]				)

          a_     :							List[str]     						=  SparkExamplesIterable(UpperCamelCase__				).shuffle_data_sources(UpperCamelCase__				)
          assert shuffled_it.n_shards == 3
          for i, (row_id, row_dict) in enumerate(UpperCamelCase__				):
               a_						,     a_     :							Dict     						=  expected_row_ids_and_row_dicts[i]
               assert row_id == expected_row_id
               assert row_dict == expected_row_dict



@require_not_windows
@require_dill_gt_0_3_2
def  _SCREAMING_SNAKE_CASE     (      ):




     """simple docstring"""


     a_     :							Any     						=  pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
     a_     :							str     						=  spark.range(20				).repartition(4				)

     # Partitions 0 and 2
     a_     :							Any     						=  SparkExamplesIterable(UpperCamelCase__				).shard_data_sources(worker_id=0 ,    num_workers=2				)
     assert shard_it_a.n_shards == 2
     a_     :							int     						=  _get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase__ ,    [0, 2]				)
     for i, (row_id, row_dict) in enumerate(UpperCamelCase__				):
          a_						,     a_     :							Dict     						=  expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict

     # Partitions 1 and 3
     a_     :							List[Any]     						=  SparkExamplesIterable(UpperCamelCase__				).shard_data_sources(worker_id=1 ,    num_workers=2				)
     assert shard_it_a.n_shards == 2
     a_     :							Union[str, Any]     						=  _get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase__ ,    [1, 3]				)
     for i, (row_id, row_dict) in enumerate(UpperCamelCase__				):
          a_						,     a_     :							Tuple     						=  expected_row_ids_and_row_dicts_a[i]
          assert row_id == expected_row_id
          assert row_dict == expected_row_dict



@require_not_windows
@require_dill_gt_0_3_2
def  _SCREAMING_SNAKE_CASE     (      ):




     """simple docstring"""


     a_     :							Any     						=  pyspark.sql.SparkSession.builder.master("""local[*]"""				).appName("""pyspark"""				).getOrCreate()
     a_     :							str     						=  spark.range(100				).repartition(1				)
     a_     :							List[str]     						=  Spark(UpperCamelCase__				)
     # Choose a small max_shard_size for maximum partitioning.
     spark_builder._repartition_df_if_needed(max_shard_size=1				)
     # The new number of partitions should not be greater than the number of rows.
     assert spark_builder.df.rdd.getNumPartitions() == 100






'''simple docstring'''


from ...configuration_utils import PretrainedConfig
from ...utils import logging


lowerCAmelCase_      :  str				     =  logging.get_logger(__name__)

lowerCAmelCase_      :  Any				     =  {
    'transfo-xl-wt103': 'https://huggingface.co/transfo-xl-wt103/resolve/main/config.json',
}







class      SCREAMING_SNAKE_CASE    (						snake_case_				):
      __magic_name__  :       str   							=					'''transfo-xl'''
      __magic_name__  :       List[str]   							=					['''mems''']
      __magic_name__  :       Dict   							=					{
          '''n_token''': '''vocab_size''',
          '''hidden_size''': '''d_model''',
          '''num_attention_heads''': '''n_head''',
          '''num_hidden_layers''': '''n_layer''',
      }






      def __init__( self   :  Union[str, Any]						,						lowercase__   :  List[Any]=26_7735						,						lowercase__   :  Optional[Any]=[2_0000, 4_0000, 20_0000]						,						lowercase__   :  Optional[Any]=1024						,						lowercase__   :  str=1024						,						lowercase__   :  Any=16						,						lowercase__   :  int=64						,						lowercase__   :  str=4096						,						lowercase__   :  Union[str, Any]=4						,						lowercase__   :  List[Any]=False						,						lowercase__   :  List[Any]=18						,						lowercase__   :  str=1600						,						lowercase__   :  str=1000						,						lowercase__   :  Any=True						,						lowercase__   :  Optional[Any]=True						,						lowercase__   :  Union[str, Any]=0						,						lowercase__   :  str=-1						,						lowercase__   :  int=True						,						lowercase__   :  str=0.1						,						lowercase__   :  Optional[Any]=0.0						,						lowercase__   :  Tuple=True						,						lowercase__   :  Optional[int]="normal"						,						lowercase__   :  str=0.01						,						lowercase__   :  List[str]=0.01						,						lowercase__   :  Union[str, Any]=0.02						,						lowercase__   :  str=1e-5						,						lowercase__   :  Any=0						,						**lowercase__   :  List[str]						,						):


           '''simple docstring'''

           a_     :							Optional[Any]     						=  vocab_size
           a_     :							Optional[int]     						=  []
           self.cutoffs.extend(lowercase__       )
           if proj_share_all_but_first:
                a_     :							Any     						=  [False] + [True] * len(self.cutoffs       )
           else:
                a_     :							Tuple     						=  [False] + [False] * len(self.cutoffs       )
           a_     :							Tuple     						=  d_model
           a_     :							Optional[int]     						=  d_embed
           a_     :							List[Any]     						=  d_head
           a_     :							List[str]     						=  d_inner
           a_     :							Tuple     						=  div_val
           a_     :							Dict     						=  pre_lnorm
           a_     :							Optional[Any]     						=  n_layer
           a_     :							Dict     						=  n_head
           a_     :							Any     						=  mem_len
           a_     :							Union[str, Any]     						=  same_length
           a_     :							Dict     						=  attn_type
           a_     :							List[str]     						=  clamp_len
           a_     :							str     						=  sample_softmax
           a_     :							Any     						=  adaptive
           a_     :							List[Any]     						=  dropout
           a_     :							str     						=  dropatt
           a_     :							Dict     						=  untie_r
           a_     :							Tuple     						=  init
           a_     :							Optional[int]     						=  init_range
           a_     :							List[Any]     						=  proj_init_std
           a_     :							Optional[int]     						=  init_std
           a_     :							int     						=  layer_norm_epsilon
           super().__init__(eos_token_id=lowercase__						,						**lowercase__       )






      @property
      def 							lowercase_						( self   :  int       ):


           '''simple docstring'''

           logger.info(F"The model {self.model_type} is one of the few models that has no sequence length limit."       )
           return -1






      @max_position_embeddings.setter
      def 							lowercase_						( self   :  Optional[int]						,						lowercase__   :  Tuple       ):


           '''simple docstring'''

           raise NotImplementedError(
               F"The model {self.model_type} is one of the few models that has no sequence length limit."       )
