"""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   UpperCAmelCase					(	UpperCamelCase__       ,    UpperCamelCase__				):






							"""simple docstring"""





							A__	     =							[]
							for part_id in partition_order:
														A__	     =							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   UpperCAmelCase					(	):






							"""simple docstring"""





							A__	     =							pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
							A__	     =							spark.range(100				).repartition(1				)
							A__	     =							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   UpperCAmelCase					(	):






							"""simple docstring"""





							A__	     =							pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
							A__	     =							spark.range(10				).repartition(2				)
							A__	     =							[1, 0]
							A__	     =							_generate_iterable_examples(UpperCamelCase__       ,    UpperCamelCase__				)  # Reverse the partitions.
							A__	     =							_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase__       ,    UpperCamelCase__				)

							for i, (row_id, row_dict) in enumerate(generate_fn()				):
														A__     ,  A__	     =							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   UpperCAmelCase					(	):






							"""simple docstring"""





							A__	     =							pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
							A__	     =							spark.range(10				).repartition(1				)
							A__	     =							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   UpperCAmelCase					(	):






							"""simple docstring"""





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

														A__	     =							SparkExamplesIterable(UpperCamelCase__				).shuffle_data_sources(UpperCamelCase__				)
														assert shuffled_it.n_shards == 3
														for i, (row_id, row_dict) in enumerate(UpperCamelCase__				):
																					A__     ,  A__	     =							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   UpperCAmelCase					(	):






							"""simple docstring"""





							A__	     =							pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
							A__	     =							spark.range(20				).repartition(4				)

							# Partitions 0 and 2
							A__	     =							SparkExamplesIterable(UpperCamelCase__				).shard_data_sources(worker_id=0       ,    num_workers=2				)
							assert shard_it_a.n_shards == 2
							A__	     =							_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase__       ,    [0, 2]				)
							for i, (row_id, row_dict) in enumerate(UpperCamelCase__				):
														A__     ,  A__	     =							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__	     =							SparkExamplesIterable(UpperCamelCase__				).shard_data_sources(worker_id=1       ,    num_workers=2				)
							assert shard_it_a.n_shards == 2
							A__	     =							_get_expected_row_ids_and_row_dicts_for_partition_order(UpperCamelCase__       ,    [1, 3]				)
							for i, (row_id, row_dict) in enumerate(UpperCamelCase__				):
														A__     ,  A__	     =							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   UpperCAmelCase					(	):






							"""simple docstring"""





							A__	     =							pyspark.sql.SparkSession.builder.master('local[*]'				).appName('pyspark'				).getOrCreate()
							A__	     =							spark.range(100				).repartition(1				)
							A__	     =							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"""




import os
import unittest

from transformers.models.cpmant.tokenization_cpmant import VOCAB_FILES_NAMES, CpmAntTokenizer
from transformers.testing_utils import require_jieba, tooslow

from ...test_tokenization_common import TokenizerTesterMixin
@require_jieba
class     UpperCamelCase__(			__A		,						unittest.TestCase ):
						lowerCAmelCase__  :						Optional[int]													=		CpmAntTokenizer
						lowerCAmelCase__  :						List[str]													=		False





						def 				snake_case__	(       self     )					->				Dict:
													super().setUp()

													A__	     =							[
													    '<d>',
													    '</d>',
													    '<s>',
													    '</s>',
													    '</_>',
													    '<unk>',
													    '<pad>',
													    '</n>',
													    '我',
													    '是',
													    'C',
													    'P',
													    'M',
													    'A',
													    'n',
													    't',
													]
													A__	     =							os.path.join(self.tmpdirname		,VOCAB_FILES_NAMES['vocab_file']     )
													with open(self.vocab_file		,'w'		,encoding='utf-8'     ) as vocab_writer:
																				vocab_writer.write(''.join([x + '\n' for x in vocab_tokens]     )     )





						@tooslow
						def 				snake_case__	(       self     )					->				List[Any]:
													A__	     =							CpmAntTokenizer.from_pretrained('openbmb/cpm-ant-10b'     )
													A__	     =							'今天天气真好！'
													A__	     =							['今天', '天气', '真', '好', '！']
													A__	     =							tokenizer.tokenize(__UpperCAmelCase     )
													self.assertListEqual(__UpperCAmelCase		,__UpperCAmelCase     )
													A__	     =							'今天天气真好！'
													A__	     =							[tokenizer.bos_token] + tokens

													A__	     =							[6, 98_02, 1_49_62, 20_82, 8_31, 2_44]
													self.assertListEqual(tokenizer.convert_tokens_to_ids(__UpperCAmelCase     )		,__UpperCAmelCase     )

													A__	     =							tokenizer.decode(__UpperCAmelCase     )
													self.assertEqual(__UpperCAmelCase		,__UpperCAmelCase     )





