




'''simple docstring'''



import json
import os
import unittest

from transformers import AutoTokenizer, GPTaTokenizer, GPTaTokenizerFast
from transformers.models.gpta.tokenization_gpta import VOCAB_FILES_NAMES
from transformers.testing_utils import require_tokenizers

from ...test_tokenization_common import TokenizerTesterMixin





@require_tokenizers
class      lowerCAmelCase (			__lowerCAmelCase					, unittest.TestCase):
       __lowercase	:   Any										=      GPTaTokenizer
       __lowercase	:   List[str]										=      GPTaTokenizerFast
       __lowercase	:   Optional[int]										=      True
       __lowercase	:   Any										=      {'''add_prefix_space''': True}
       __lowercase	:   List[Any]										=      False


       def    lowerCAmelCase   ( self	)   ->	Any:



            '''simple docstring'''




            super().setUp()

            # Adapted from Sennrich et al. 2015 and https://github.com/rsennrich/subword-nmt
            __snake_case            =	[
                '''l''',
                '''o''',
                '''w''',
                '''e''',
                '''r''',
                '''s''',
                '''t''',
                '''i''',
                '''d''',
                '''n''',
                '''\u0120''',
                '''\u0120l''',
                '''\u0120n''',
                '''\u0120lo''',
                '''\u0120low''',
                '''er''',
                '''\u0120lowest''',
                '''\u0120newer''',
                '''\u0120wider''',
                '''<unk>''',
                '''<|endoftext|>''',
            ]
            __snake_case            =	dict(zip(__SCREAMING_SNAKE_CASE     ,	range(len(__SCREAMING_SNAKE_CASE	)	)	)	)
            __snake_case            =	['''#version: 0.2''', '''\u0120 l''', '''\u0120l o''', '''\u0120lo w''', '''e r''', '''''']
            __snake_case            =	{'''unk_token''': '''<unk>'''}

            __snake_case            =	os.path.join(self.tmpdirname     ,	VOCAB_FILES_NAMES['''vocab_file''']	)
            __snake_case            =	os.path.join(self.tmpdirname     ,	VOCAB_FILES_NAMES['''merges_file''']	)
            with open(self.vocab_file     ,	'''w'''     ,	encoding='''utf-8'''	) as fp:
                 fp.write(json.dumps(__SCREAMING_SNAKE_CASE	) + '''\n'''	)
            with open(self.merges_file     ,	'''w'''     ,	encoding='''utf-8'''	) as fp:
                 fp.write('''\n'''.join(__SCREAMING_SNAKE_CASE	)	)


       def    lowerCAmelCase   ( self     ,	**__SCREAMING_SNAKE_CASE	)   ->	Dict:



            '''simple docstring'''




            kwargs.update(self.special_tokens_map	)
            return GPTaTokenizer.from_pretrained(self.tmpdirname     ,	**__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self     ,	**__SCREAMING_SNAKE_CASE	)   ->	Union[str, Any]:



            '''simple docstring'''




            kwargs.update(self.special_tokens_map	)
            return GPTaTokenizerFast.from_pretrained(self.tmpdirname     ,	**__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE	)   ->	Optional[Any]:



            '''simple docstring'''




            __snake_case            =	'''lower newer'''
            __snake_case            =	'''lower newer'''
            return input_text, output_text


       def    lowerCAmelCase   ( self	)   ->	Union[str, Any]:



            '''simple docstring'''




            __snake_case            =	GPTaTokenizer(self.vocab_file     ,	self.merges_file     ,	**self.special_tokens_map	)
            __snake_case            =	'''lower newer'''
            __snake_case            =	['''\u0120low''', '''er''', '''\u0120''', '''n''', '''e''', '''w''', '''er''']
            __snake_case            =	tokenizer.tokenize(__SCREAMING_SNAKE_CASE     ,	add_prefix_space=__SCREAMING_SNAKE_CASE	)
            self.assertListEqual(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)

            __snake_case            =	tokens + [tokenizer.unk_token]
            __snake_case            =	[14, 15, 10, 9, 3, 2, 15, 19]
            self.assertListEqual(tokenizer.convert_tokens_to_ids(__SCREAMING_SNAKE_CASE	)     ,	__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self	)   ->	List[Any]:



            '''simple docstring'''




            if not self.test_rust_tokenizer:
                 return

            __snake_case            =	self.get_tokenizer()
            __snake_case            =	self.get_rust_tokenizer(add_prefix_space=__SCREAMING_SNAKE_CASE	)

            __snake_case            =	'''lower newer'''

            # Testing tokenization
            __snake_case            =	tokenizer.tokenize(__SCREAMING_SNAKE_CASE     ,	add_prefix_space=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	rust_tokenizer.tokenize(__SCREAMING_SNAKE_CASE	)
            self.assertListEqual(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)

            # Testing conversion to ids without special tokens
            __snake_case            =	tokenizer.encode(__SCREAMING_SNAKE_CASE     ,	add_special_tokens=__SCREAMING_SNAKE_CASE     ,	add_prefix_space=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	rust_tokenizer.encode(__SCREAMING_SNAKE_CASE     ,	add_special_tokens=__SCREAMING_SNAKE_CASE	)
            self.assertListEqual(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)

            # Testing conversion to ids with special tokens
            __snake_case            =	self.get_rust_tokenizer(add_prefix_space=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	tokenizer.encode(__SCREAMING_SNAKE_CASE     ,	add_prefix_space=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	rust_tokenizer.encode(__SCREAMING_SNAKE_CASE	)
            self.assertListEqual(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)

            # Testing the unknown token
            __snake_case            =	tokens + [rust_tokenizer.unk_token]
            __snake_case            =	[14, 15, 10, 9, 3, 2, 15, 19]
            self.assertListEqual(rust_tokenizer.convert_tokens_to_ids(__SCREAMING_SNAKE_CASE	)     ,	__SCREAMING_SNAKE_CASE	)


       def    lowerCAmelCase   ( self     ,	*__SCREAMING_SNAKE_CASE     ,	**__SCREAMING_SNAKE_CASE	)   ->	Dict:



            '''simple docstring'''




            pass


       def    lowerCAmelCase   ( self     ,	__SCREAMING_SNAKE_CASE=15	)   ->	str:



            '''simple docstring'''




            for tokenizer, pretrained_name, kwargs in self.tokenizers_list:
                 with self.subTest(F'''{tokenizer.__class__.__name__} ({pretrained_name})'''	):
                      __snake_case            =	self.rust_tokenizer_class.from_pretrained(__SCREAMING_SNAKE_CASE     ,	**__SCREAMING_SNAKE_CASE	)

                      # Simple input
                      __snake_case            =	'''This is a simple input'''
                      __snake_case            =	['''This is a simple input 1''', '''This is a simple input 2''']
                      __snake_case            =	('''This is a simple input''', '''This is a pair''')
                      __snake_case            =	[
                          ('''This is a simple input 1''', '''This is a simple input 2'''),
                          ('''This is a simple pair 1''', '''This is a simple pair 2'''),
                      ]

                      # Simple input tests
                      self.assertRaises(__SCREAMING_SNAKE_CASE     ,	tokenizer_r.encode     ,	__SCREAMING_SNAKE_CASE     ,	max_length=__SCREAMING_SNAKE_CASE     ,	padding='''max_length'''	)

                      # Simple input
                      self.assertRaises(__SCREAMING_SNAKE_CASE     ,	tokenizer_r.encode_plus     ,	__SCREAMING_SNAKE_CASE     ,	max_length=__SCREAMING_SNAKE_CASE     ,	padding='''max_length'''	)

                      # Simple input
                      self.assertRaises(
                          __SCREAMING_SNAKE_CASE     ,	tokenizer_r.batch_encode_plus     ,	__SCREAMING_SNAKE_CASE     ,	max_length=__SCREAMING_SNAKE_CASE     ,	padding='''max_length'''     ,	)

                      # Pair input
                      self.assertRaises(__SCREAMING_SNAKE_CASE     ,	tokenizer_r.encode     ,	__SCREAMING_SNAKE_CASE     ,	max_length=__SCREAMING_SNAKE_CASE     ,	padding='''max_length'''	)

                      # Pair input
                      self.assertRaises(__SCREAMING_SNAKE_CASE     ,	tokenizer_r.encode_plus     ,	__SCREAMING_SNAKE_CASE     ,	max_length=__SCREAMING_SNAKE_CASE     ,	padding='''max_length'''	)

                      # Pair input
                      self.assertRaises(
                          __SCREAMING_SNAKE_CASE     ,	tokenizer_r.batch_encode_plus     ,	__SCREAMING_SNAKE_CASE     ,	max_length=__SCREAMING_SNAKE_CASE     ,	padding='''max_length'''     ,	)


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



            '''simple docstring'''




            __snake_case            =	GPTaTokenizer.from_pretrained(self.tmpdirname     ,	pad_token='''<pad>'''	)

            # Simple input
            __snake_case            =	'''This is a simple input'''
            __snake_case            =	['''This is a simple input looooooooong''', '''This is a simple input''']
            __snake_case            =	('''This is a simple input''', '''This is a pair''')
            __snake_case            =	[
                ('''This is a simple input loooooong''', '''This is a simple input'''),
                ('''This is a simple pair loooooong''', '''This is a simple pair'''),
            ]

            __snake_case            =	tokenizer.pad_token_id

            __snake_case            =	tokenizer(__SCREAMING_SNAKE_CASE     ,	padding='''max_length'''     ,	max_length=30     ,	return_tensors='''np'''	)
            __snake_case            =	tokenizer(__SCREAMING_SNAKE_CASE     ,	padding=__SCREAMING_SNAKE_CASE     ,	truncate=__SCREAMING_SNAKE_CASE     ,	return_tensors='''np'''	)
            __snake_case            =	tokenizer(*__SCREAMING_SNAKE_CASE     ,	padding='''max_length'''     ,	max_length=60     ,	return_tensors='''np'''	)
            __snake_case            =	tokenizer(__SCREAMING_SNAKE_CASE     ,	padding=__SCREAMING_SNAKE_CASE     ,	truncate=__SCREAMING_SNAKE_CASE     ,	return_tensors='''np'''	)

            # s
            # test single string max_length padding
            self.assertEqual(out_s['''input_ids'''].shape[-1]     ,	30	)
            self.assertTrue(pad_token_id in out_s['''input_ids''']	)
            self.assertTrue(0 in out_s['''attention_mask''']	)

            # s2
            # test automatic padding
            self.assertEqual(out_sa['''input_ids'''].shape[-1]     ,	33	)
            # long slice doesn't have padding
            self.assertFalse(pad_token_id in out_sa['''input_ids'''][0]	)
            self.assertFalse(0 in out_sa['''attention_mask'''][0]	)
            # short slice does have padding
            self.assertTrue(pad_token_id in out_sa['''input_ids'''][1]	)
            self.assertTrue(0 in out_sa['''attention_mask'''][1]	)

            # p
            # test single pair max_length padding
            self.assertEqual(out_p['''input_ids'''].shape[-1]     ,	60	)
            self.assertTrue(pad_token_id in out_p['''input_ids''']	)
            self.assertTrue(0 in out_p['''attention_mask''']	)

            # p2
            # test automatic padding pair
            self.assertEqual(out_pa['''input_ids'''].shape[-1]     ,	52	)
            # long slice pair doesn't have padding
            self.assertFalse(pad_token_id in out_pa['''input_ids'''][0]	)
            self.assertFalse(0 in out_pa['''attention_mask'''][0]	)
            # short slice pair does have padding
            self.assertTrue(pad_token_id in out_pa['''input_ids'''][1]	)
            self.assertTrue(0 in out_pa['''attention_mask'''][1]	)


       def    lowerCAmelCase   ( self	)   ->	List[Any]:



            '''simple docstring'''




            __snake_case            =	'''$$$'''
            __snake_case            =	GPTaTokenizer.from_pretrained(self.tmpdirname     ,	bos_token=__SCREAMING_SNAKE_CASE     ,	add_bos_token=__SCREAMING_SNAKE_CASE	)

            __snake_case            =	'''This is a simple input'''
            __snake_case            =	['''This is a simple input 1''', '''This is a simple input 2''']

            __snake_case            =	tokenizer.bos_token_id

            __snake_case            =	tokenizer(__SCREAMING_SNAKE_CASE	)
            __snake_case            =	tokenizer(__SCREAMING_SNAKE_CASE	)

            self.assertEqual(out_s.input_ids[0]     ,	__SCREAMING_SNAKE_CASE	)
            self.assertTrue(all(o[0] == bos_token_id for o in out_sa.input_ids	)	)

            __snake_case            =	tokenizer.decode(out_s.input_ids	)
            __snake_case            =	tokenizer.batch_decode(out_sa.input_ids	)

            self.assertEqual(decode_s.split()[0]     ,	__SCREAMING_SNAKE_CASE	)
            self.assertTrue(all(d.split()[0] == bos_token for d in decode_sa	)	)


       def    lowerCAmelCase   ( self	)   ->	Dict:



            '''simple docstring'''




            pass




       def    lowerCAmelCase   ( self	)   ->	List[str]:



            '''simple docstring'''




            __snake_case            =	[self.get_tokenizer(do_lower_case=__SCREAMING_SNAKE_CASE     ,	add_bos_token=__SCREAMING_SNAKE_CASE	)]
            for tokenizer in tokenizers:
                 with self.subTest(F'''{tokenizer.__class__.__name__}'''	):
                      __snake_case            =	'''Encode this.'''
                      __snake_case            =	'''This one too please.'''
                      __snake_case            =	tokenizer.encode(__SCREAMING_SNAKE_CASE     ,	add_special_tokens=__SCREAMING_SNAKE_CASE	)
                      encoded_sequence += tokenizer.encode(__SCREAMING_SNAKE_CASE     ,	add_special_tokens=__SCREAMING_SNAKE_CASE	)
                      __snake_case            =	tokenizer.encode_plus(
                          __SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE     ,	add_special_tokens=__SCREAMING_SNAKE_CASE     ,	return_special_tokens_mask=__SCREAMING_SNAKE_CASE     ,	)
                      __snake_case            =	encoded_sequence_dict['''input_ids''']
                      __snake_case            =	encoded_sequence_dict['''special_tokens_mask''']
                      self.assertEqual(len(__SCREAMING_SNAKE_CASE	)     ,	len(__SCREAMING_SNAKE_CASE	)	)

                      __snake_case            =	[
                          (x if not special_tokens_mask[i] else None) for i, x in enumerate(__SCREAMING_SNAKE_CASE	)
                      ]
                      __snake_case            =	[x for x in filtered_sequence if x is not None]
                      self.assertEqual(__SCREAMING_SNAKE_CASE     ,	__SCREAMING_SNAKE_CASE	)





@require_tokenizers
class      lowerCAmelCase (			unittest.TestCase):


       def    lowerCAmelCase   ( self	)   ->	str:



            '''simple docstring'''




            __snake_case            =	AutoTokenizer.from_pretrained('''facebook/opt-350m'''     ,	from_slow=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	'''A photo of a cat'''

            __snake_case            =	tokenizer.encode(
                __SCREAMING_SNAKE_CASE     ,	)
            self.assertEqual(__SCREAMING_SNAKE_CASE     ,	[2, 250, 1345, 9, 10, 4758]	)
            tokenizer.save_pretrained('''test_opt'''	)

            __snake_case            =	AutoTokenizer.from_pretrained('''./test_opt'''	)
            __snake_case            =	tokenizer.encode(
                __SCREAMING_SNAKE_CASE     ,	)
            self.assertEqual(__SCREAMING_SNAKE_CASE     ,	[2, 250, 1345, 9, 10, 4758]	)


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



            '''simple docstring'''




            __snake_case            =	AutoTokenizer.from_pretrained('''facebook/opt-350m'''     ,	use_slow=__SCREAMING_SNAKE_CASE	)
            __snake_case            =	'''A photo of a cat'''

            __snake_case            =	tokenizer.encode(
                __SCREAMING_SNAKE_CASE     ,	)
            # Same as above
            self.assertEqual(__SCREAMING_SNAKE_CASE     ,	[2, 250, 1345, 9, 10, 4758]	)




       @unittest.skip('''This test is failing because of a bug in the fast tokenizer'''	)
       def    lowerCAmelCase   ( self	)   ->	str:



            '''simple docstring'''




            __snake_case            =	AutoTokenizer.from_pretrained('''facebook/opt-350m'''     ,	from_slow=__SCREAMING_SNAKE_CASE	)

            __snake_case            =	'''bos'''
            __snake_case            =	tokenizer.get_vocab()['''bos''']

            __snake_case            =	'''A photo of a cat'''
            __snake_case            =	tokenizer.encode(
                __SCREAMING_SNAKE_CASE     ,	)
            # We changed the bos token
            self.assertEqual(__SCREAMING_SNAKE_CASE     ,	[3_1957, 250, 1345, 9, 10, 4758]	)
            tokenizer.save_pretrained('''./tok'''	)
            __snake_case            =	AutoTokenizer.from_pretrained('''./tok'''	)
            self.assertTrue(tokenizer.is_fast	)
            __snake_case            =	tokenizer.encode(
                __SCREAMING_SNAKE_CASE     ,	)
            self.assertEqual(__SCREAMING_SNAKE_CASE     ,	[3_1957, 250, 1345, 9, 10, 4758]	)












'''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			(_lowerCamelCase				:  Union[str, Any] ,						_lowerCamelCase				:  List[Any]						)-> Optional[Any]:




     '''simple docstring'''





     __snake_case            =	[]
     for part_id in partition_order:
          __snake_case            =	df.where(f'''SPARK_PARTITION_ID() = {part_id}'''						).collect()
          for row_idx, row in enumerate(_lowerCamelCase						):
               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			()-> Any:




     '''simple docstring'''





     __snake_case            =	pyspark.sql.SparkSession.builder.master('''local[*]'''						).appName('''pyspark'''						).getOrCreate()
     __snake_case            =	spark.range(1_00						).repartition(1						)
     __snake_case            =	Spark(_lowerCamelCase						)
     # 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			()-> Tuple:




     '''simple docstring'''





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

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




     '''simple docstring'''





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







@require_not_windows
@require_dill_gt_0_3_2
def       _UpperCamelCase			()-> Union[str, Any]:




     '''simple docstring'''





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

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




     '''simple docstring'''





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

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

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




     '''simple docstring'''





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






