'''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_						(					snake_case_       :       Optional[int]						,		snake_case_       :       Optional[Any]      )					->     Optional[Any]:






 '''simple docstring'''

 __lowerCAmelCase											=						[]
 for part_id in partition_order:
  __lowerCAmelCase											=						df.where(f"""SPARK_PARTITION_ID() = {part_id}"""      ).collect()
  for row_idx, row in enumerate(_snake_case      ):
   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_						(					)					->     int:






 '''simple docstring'''

 __lowerCAmelCase											=						pyspark.sql.SparkSession.builder.master("""local[*]"""      ).appName("""pyspark"""      ).getOrCreate()
 __lowerCAmelCase											=						spark.range(1_00      ).repartition(1      )
 __lowerCAmelCase											=						Spark(_snake_case      )
 # 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_						(					)					->     int:






 '''simple docstring'''

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

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






 '''simple docstring'''

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







@require_not_windows
@require_dill_gt_0_3_2
def      UpperCamelCase_						(					)					->     Optional[Any]:






 '''simple docstring'''

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

  __lowerCAmelCase											=						SparkExamplesIterable(_snake_case      ).shuffle_data_sources(_snake_case      )
  assert shuffled_it.n_shards == 3
  for i, (row_id, row_dict) in enumerate(_snake_case      ):
   __lowerCAmelCase							,       __lowerCAmelCase											=						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_						(					)					->     Dict:






 '''simple docstring'''

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

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

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






 '''simple docstring'''

 __lowerCAmelCase											=						pyspark.sql.SparkSession.builder.master("""local[*]"""      ).appName("""pyspark"""      ).getOrCreate()
 __lowerCAmelCase											=						spark.range(1_00      ).repartition(1      )
 __lowerCAmelCase											=						Spark(_snake_case      )
 # 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() == 1_00









import unittest

from transformers import GPTSwaTokenizer
from transformers.testing_utils import get_tests_dir, require_sentencepiece, require_tokenizers, slow

from ...test_tokenization_common import TokenizerTesterMixin


_lowerCAmelCase					:			List[Any]  				=			get_tests_dir("fixtures/test_sentencepiece_with_bytefallback.model")






@require_sentencepiece
@require_tokenizers
class 	__magic_name__    (			lowerCAmelCase_  ,    unittest.TestCase      ):
			SCREAMING_SNAKE_CASE							   =						GPTSwaTokenizer
			SCREAMING_SNAKE_CASE							   =						False
			SCREAMING_SNAKE_CASE							   =						True
			SCREAMING_SNAKE_CASE							   =						False





			def 		__magic_name__    (		self					)   -> Dict:




						'''simple docstring'''
						super().setUp()

						# We have a SentencePiece fixture for testing
						__a        =GPTSwaTokenizer(__snake_case					,  eos_token='<unk>'					,  bos_token='<unk>'					,  pad_token='<unk>'					)

						tokenizer.save_pretrained(self.tmpdirname					)





			def 		__magic_name__    (		self					,  __snake_case					)   -> Tuple:




						'''simple docstring'''
						__a        ='This is a test'
						__a        ='This is a test'
						return input_text, output_text





			def 		__magic_name__    (		self					)   -> Optional[Any]:




						'''simple docstring'''
						__a        ='<s>'
						__a        =1

						self.assertEqual(self.get_tokenizer()._convert_token_to_id(__snake_case					)					,  __snake_case					)
						self.assertEqual(self.get_tokenizer()._convert_id_to_token(__snake_case					)					,  __snake_case					)





			def 		__magic_name__    (		self					)   -> Dict:




						'''simple docstring'''
						__a        =list(self.get_tokenizer().get_vocab().keys()					)

						self.assertEqual(vocab_keys[0]					,  '<unk>'					)
						self.assertEqual(vocab_keys[1]					,  '<s>'					)
						self.assertEqual(vocab_keys[-1]					,  'j'					)
						self.assertEqual(len(__snake_case					)					,  2000					)





			def 		__magic_name__    (		self					)   -> Optional[Any]:




						'''simple docstring'''
						self.assertEqual(self.get_tokenizer().vocab_size					,  2000					)





			def 		__magic_name__    (		self					)   -> Any:




						'''simple docstring'''
						__a        =GPTSwaTokenizer(__snake_case					)

						__a        =tokenizer.tokenize('This is a test'					)
						self.assertListEqual(__snake_case					,  ['▁This', '▁is', '▁a', '▁t', 'est']					)

						self.assertListEqual(tokenizer.convert_tokens_to_ids(__snake_case					)					,  [465, 287, 265, 631, 842]					)

						__a        =tokenizer.tokenize('I was born in 92000, and this is falsé.'					)
						# fmt: off
						self.assertListEqual(
						    __snake_case					,  ['▁I', '▁was', '▁bor', 'n', '▁in', '▁', '<0x39>', '2', '0', '0', '0', ',', '▁and', '▁this', '▁is', '▁f', 'al', 's', '<0xC3>', '<0xA9>', '.']					,  )
						# fmt: on

						__a        =tokenizer.convert_tokens_to_ids(__snake_case					)
						self.assertListEqual(
						    __snake_case					,  [262, 272, 1525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260]					,  )

						__a        =tokenizer.convert_ids_to_tokens(__snake_case					)
						# fmt: off
						self.assertListEqual(
						    __snake_case					,  ['▁I', '▁was', '▁bor', 'n', '▁in', '▁', '<0x39>', '2', '0', '0', '0', ',', '▁and', '▁this', '▁is', '▁f', 'al', 's', '<0xC3>', '<0xA9>', '.']					)
						# fmt: on





			def 		__magic_name__    (		self					)   -> Optional[int]:




						'''simple docstring'''
						__a        =GPTSwaTokenizer(__snake_case					)
						__a        =['This is a test', 'I was born in 92000, and this is falsé.']
						__a        =[
						    [465, 287, 265, 631, 842],
						    [262, 272, 1525, 286, 271, 268, 60, 916, 633, 633, 633, 259, 266, 301, 287, 384, 367, 263, 198, 172, 260],
						]

						# Test that encode_fast returns the same as tokenize + convert_tokens_to_ids
						for text, expected_ids in zip(__snake_case					,  __snake_case					):
									self.assertListEqual(tokenizer.encode_fast(__snake_case					)					,  __snake_case					)

						# Test that decode_fast returns the input text
						for text, token_ids in zip(__snake_case					,  __snake_case					):
									self.assertEqual(tokenizer.decode_fast(__snake_case					)					,  __snake_case					)





			@slow
			def 		__magic_name__    (		self					)   -> List[str]:




						'''simple docstring'''
						__a        =[
						    '<|python|>def fibonacci(n)\n    if n < 0:\n        print(\'Incorrect input\')',
						    'Hey there, how are you doing this fine day?',
						    'This is a text with a trailing spaces followed by a dot     .',
						    'Häj sväjs lillebrör! =)',
						    'Det är inget fel på Mr. Cool',
						]

						# fmt: off
						__a        ={'input_ids': [[6_3423, 5, 6811, 1_4954, 282, 816, 3821, 6_3466, 6_3425, 6_3462, 18, 6_3978, 678, 301, 1320, 6_3423, 6_3455, 6_3458, 18, 6_3982, 4246, 3940, 1901, 4_7789, 5547, 1_8994], [1_9630, 1100, 6_3446, 1342, 633, 544, 4488, 593, 5102, 2416, 6_3495, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1652, 428, 268, 1936, 515, 268, 5_8593, 2_2413, 9106, 546, 268, 3_3213, 6_3979, 698, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [5_5130, 6_3450, 924, 6_3449, 2249, 4062, 1558, 318, 6_3504, 2_1498, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [509, 377, 2827, 2559, 332, 6575, 6_3443, 2_6801, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], 'token_type_ids': [[0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]], 'attention_mask': [[1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0], [1, 1, 1, 1, 1, 1, 1, 1, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0, 0]]}
						# fmt: on
						self.tokenizer_integration_test_util(
						    expected_encoding=__snake_case					,  model_name='AI-Sweden/gpt-sw3-126m'					,  sequences=__snake_case					,  )






