






'''simple docstring'''







import json
from typing import List, Optional, Tuple

from tokenizers import pre_tokenizers, processors

from ...tokenization_utils_base import AddedToken, BatchEncoding
from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_roberta import RobertaTokenizer


__lowerCAmelCase       		=      logging.get_logger(__name__)

__lowerCAmelCase       		=      {'vocab_file': 'vocab.json', 'merges_file': 'merges.txt', 'tokenizer_file': 'tokenizer.json'}

__lowerCAmelCase       		=      {
    'vocab_file': {
        'roberta-base': 'https://huggingface.co/roberta-base/resolve/main/vocab.json',
        'roberta-large': 'https://huggingface.co/roberta-large/resolve/main/vocab.json',
        'roberta-large-mnli': 'https://huggingface.co/roberta-large-mnli/resolve/main/vocab.json',
        'distilroberta-base': 'https://huggingface.co/distilroberta-base/resolve/main/vocab.json',
        'roberta-base-openai-detector': 'https://huggingface.co/roberta-base-openai-detector/resolve/main/vocab.json',
        'roberta-large-openai-detector': (
            'https://huggingface.co/roberta-large-openai-detector/resolve/main/vocab.json'
        ),
    },
    'merges_file': {
        'roberta-base': 'https://huggingface.co/roberta-base/resolve/main/merges.txt',
        'roberta-large': 'https://huggingface.co/roberta-large/resolve/main/merges.txt',
        'roberta-large-mnli': 'https://huggingface.co/roberta-large-mnli/resolve/main/merges.txt',
        'distilroberta-base': 'https://huggingface.co/distilroberta-base/resolve/main/merges.txt',
        'roberta-base-openai-detector': 'https://huggingface.co/roberta-base-openai-detector/resolve/main/merges.txt',
        'roberta-large-openai-detector': (
            'https://huggingface.co/roberta-large-openai-detector/resolve/main/merges.txt'
        ),
    },
    'tokenizer_file': {
        'roberta-base': 'https://huggingface.co/roberta-base/resolve/main/tokenizer.json',
        'roberta-large': 'https://huggingface.co/roberta-large/resolve/main/tokenizer.json',
        'roberta-large-mnli': 'https://huggingface.co/roberta-large-mnli/resolve/main/tokenizer.json',
        'distilroberta-base': 'https://huggingface.co/distilroberta-base/resolve/main/tokenizer.json',
        'roberta-base-openai-detector': (
            'https://huggingface.co/roberta-base-openai-detector/resolve/main/tokenizer.json'
        ),
        'roberta-large-openai-detector': (
            'https://huggingface.co/roberta-large-openai-detector/resolve/main/tokenizer.json'
        ),
    },
}

__lowerCAmelCase       		=      {
    'roberta-base': 512,
    'roberta-large': 512,
    'roberta-large-mnli': 512,
    'distilroberta-base': 512,
    'roberta-base-openai-detector': 512,
    'roberta-large-openai-detector': 512,
}





class 	SCREAMING_SNAKE_CASE					(					SCREAMING_SNAKE_CASE_       ):
			snake_case__  =				VOCAB_FILES_NAMES
			snake_case__  =				PRETRAINED_VOCAB_FILES_MAP
			snake_case__  =				PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
			snake_case__  =				["input_ids", "attention_mask"]
			snake_case__  =				RobertaTokenizer






			def __init__( self      :			Tuple  ,						__SCREAMING_SNAKE_CASE      :			Dict=None  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]=None  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]=None  ,						__SCREAMING_SNAKE_CASE      :			str="replace"  ,						__SCREAMING_SNAKE_CASE      :			Optional[Any]="<s>"  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]="</s>"  ,						__SCREAMING_SNAKE_CASE      :			List[Any]="</s>"  ,						__SCREAMING_SNAKE_CASE      :			List[Any]="<s>"  ,						__SCREAMING_SNAKE_CASE      :			int="<unk>"  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]="<pad>"  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]="<mask>"  ,						__SCREAMING_SNAKE_CASE      :			Optional[int]=False  ,						__SCREAMING_SNAKE_CASE      :			int=True  ,						**__SCREAMING_SNAKE_CASE      :			str  ,						)			->			Optional[Any]:
							super().__init__(
							    __SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE  ,						tokenizer_file=__SCREAMING_SNAKE_CASE  ,						errors=__SCREAMING_SNAKE_CASE  ,						bos_token=__SCREAMING_SNAKE_CASE  ,						eos_token=__SCREAMING_SNAKE_CASE  ,						sep_token=__SCREAMING_SNAKE_CASE  ,						cls_token=__SCREAMING_SNAKE_CASE  ,						unk_token=__SCREAMING_SNAKE_CASE  ,						pad_token=__SCREAMING_SNAKE_CASE  ,						mask_token=__SCREAMING_SNAKE_CASE  ,						add_prefix_space=__SCREAMING_SNAKE_CASE  ,						trim_offsets=__SCREAMING_SNAKE_CASE  ,						**__SCREAMING_SNAKE_CASE  ,						)

							a_		:  Optional[int]											=      json.loads(self.backend_tokenizer.pre_tokenizer.__getstate__()					)
							if pre_tok_state.get('''add_prefix_space'''  ,						__SCREAMING_SNAKE_CASE					) != add_prefix_space:
											a_		:  List[Any]											=      getattr(__SCREAMING_SNAKE_CASE  ,						pre_tok_state.pop('''type'''					)					)
											a_		:  str											=      add_prefix_space
											a_		:  Dict											=      pre_tok_class(**__SCREAMING_SNAKE_CASE					)

							a_		:  Optional[Any]											=      add_prefix_space

							a_		:  str											=      '''post_processor'''
							a_		:  Any											=      getattr(self.backend_tokenizer  ,						__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)
							if tokenizer_component_instance:
											a_		:  Dict											=      json.loads(tokenizer_component_instance.__getstate__()					)

											# The lists 'sep' and 'cls' must be cased in tuples for the object `post_processor_class`
											if "sep" in state:
															a_		:  List[str]											=      tuple(state['''sep''']					)
											if "cls" in state:
															a_		:  List[str]											=      tuple(state['''cls''']					)

											a_		:  Any											=      False

											if state.get('''add_prefix_space'''  ,						__SCREAMING_SNAKE_CASE					) != add_prefix_space:
															a_		:  Union[str, Any]											=      add_prefix_space
															a_		:  Optional[int]											=      True

											if state.get('''trim_offsets'''  ,						__SCREAMING_SNAKE_CASE					) != trim_offsets:
															a_		:  Optional[int]											=      trim_offsets
															a_		:  List[Any]											=      True

											if changes_to_apply:
															a_		:  Optional[Any]											=      getattr(__SCREAMING_SNAKE_CASE  ,						state.pop('''type'''					)					)
															a_		:  str											=      component_class(**__SCREAMING_SNAKE_CASE					)
															setattr(self.backend_tokenizer  ,						__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					)






			@property
			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[Any]					)			->			str:
							if self._mask_token is None:
											if self.verbose:
															logger.error('''Using mask_token, but it is not set yet.'''					)
											return None
							return str(self._mask_token					)






			@mask_token.setter
			def 			SCREAMING_SNAKE_CASE      ( self      :			List[str]  ,						__SCREAMING_SNAKE_CASE      :			int					)			->			Tuple:
							a_		:  Tuple											=      AddedToken(__SCREAMING_SNAKE_CASE  ,						lstrip=__SCREAMING_SNAKE_CASE  ,						rstrip=__SCREAMING_SNAKE_CASE					) if isinstance(__SCREAMING_SNAKE_CASE  ,						__SCREAMING_SNAKE_CASE					) else value
							a_		:  List[str]											=      value






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]  ,						*__SCREAMING_SNAKE_CASE      :			Any  ,						**__SCREAMING_SNAKE_CASE      :			List[str]					)			->			BatchEncoding:
							a_		:  List[str]											=      kwargs.get('''is_split_into_words'''  ,						__SCREAMING_SNAKE_CASE					)
							assert self.add_prefix_space or not is_split_into_words, (
							    f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
							    "to use it with pretokenized inputs."
							)

							return super()._batch_encode_plus(*__SCREAMING_SNAKE_CASE  ,						**__SCREAMING_SNAKE_CASE					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]  ,						*__SCREAMING_SNAKE_CASE      :			Optional[Any]  ,						**__SCREAMING_SNAKE_CASE      :			str					)			->			BatchEncoding:
							a_		:  str											=      kwargs.get('''is_split_into_words'''  ,						__SCREAMING_SNAKE_CASE					)

							assert self.add_prefix_space or not is_split_into_words, (
							    f'You need to instantiate {self.__class__.__name__} with add_prefix_space=True '
							    "to use it with pretokenized inputs."
							)

							return super()._encode_plus(*__SCREAMING_SNAKE_CASE  ,						**__SCREAMING_SNAKE_CASE					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			int  ,						__SCREAMING_SNAKE_CASE      :			str  ,						__SCREAMING_SNAKE_CASE      :			Optional[str] = None					)			->			Tuple[str]:
							a_		:  str											=      self._tokenizer.model.save(__SCREAMING_SNAKE_CASE  ,						name=__SCREAMING_SNAKE_CASE					)
							return tuple(__SCREAMING_SNAKE_CASE					)






			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			int  ,						__SCREAMING_SNAKE_CASE      :			str=None					)			->			List[Any]:
							a_		:  Any											=      [self.bos_token_id] + token_ids_a + [self.eos_token_id]
							if token_ids_a is None:
											return output

							return output + [self.eos_token_id] + token_ids_a + [self.eos_token_id]







			def 			SCREAMING_SNAKE_CASE      ( self      :			Optional[int]  ,						__SCREAMING_SNAKE_CASE      :			List[int]  ,						__SCREAMING_SNAKE_CASE      :			Optional[List[int]] = None					)			->			List[int]:
							a_		:  int											=      [self.sep_token_id]
							a_		:  Dict											=      [self.cls_token_id]

							if token_ids_a is None:
											return len(cls + token_ids_a + sep					) * [0]
							return len(cls + token_ids_a + sep + sep + token_ids_a + sep					) * [0]












'''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   (							__A							:   List[str]						,		__A							:   List[Any]      ):
				a_		:  Any											=      []
				for part_id in partition_order:
								a_		:  str											=      df.where(f'SPARK_PARTITION_ID() = {part_id}'      ).collect()
								for row_idx, row in enumerate(__A      ):
												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   (							):
				a_		:  List[str]											=      pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				a_		:  Union[str, Any]											=      spark.range(1_00      ).repartition(1      )
				a_		:  Any											=      Spark(__A      )
				# 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   (							):
				a_		:  List[Any]											=      pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				a_		:  int											=      spark.range(10      ).repartition(2      )
				a_		:  Tuple											=      [1, 0]
				a_		:  List[str]											=      _generate_iterable_examples(__A						,		__A      )  # Reverse the partitions.
				a_		:  int											=      _get_expected_row_ids_and_row_dicts_for_partition_order(__A						,		__A      )

				for i, (row_id, row_dict) in enumerate(generate_fn()      ):
								a_				,			a_		:  List[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 		_UpperCAmelCase   (							):
				a_		:  int											=      pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				a_		:  str											=      spark.range(10      ).repartition(1      )
				a_		:  Tuple											=      SparkExamplesIterable(__A      )
				assert it.n_shards == 1
				for i, (row_id, row_dict) in enumerate(__A      ):
								assert row_id == f'0_{i}'
								assert row_dict == {"id": i}





@require_not_windows
@require_dill_gt_0_3_2
def 		_UpperCAmelCase   (							):
				a_		:  Tuple											=      pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				a_		:  str											=      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 __A      : x.reverse()
								a_		:  Any											=      _get_expected_row_ids_and_row_dicts_for_partition_order(__A						,		[2, 1, 0]      )

								a_		:  str											=      SparkExamplesIterable(__A      ).shuffle_data_sources(__A      )
								assert shuffled_it.n_shards == 3
								for i, (row_id, row_dict) in enumerate(__A      ):
												a_				,			a_		:  Optional[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 		_UpperCAmelCase   (							):
				a_		:  int											=      pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				a_		:  List[str]											=      spark.range(20      ).repartition(4      )

				# Partitions 0 and 2
				a_		:  Dict											=      SparkExamplesIterable(__A      ).shard_data_sources(worker_id=0						,		num_workers=2      )
				assert shard_it_a.n_shards == 2
				a_		:  Optional[Any]											=      _get_expected_row_ids_and_row_dicts_for_partition_order(__A						,		[0, 2]      )
				for i, (row_id, row_dict) in enumerate(__A      ):
								a_				,			a_		:  Tuple											=      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(__A      ).shard_data_sources(worker_id=1						,		num_workers=2      )
				assert shard_it_a.n_shards == 2
				a_		:  Optional[int]											=      _get_expected_row_ids_and_row_dicts_for_partition_order(__A						,		[1, 3]      )
				for i, (row_id, row_dict) in enumerate(__A      ):
								a_				,			a_		:  Any											=      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   (							):
				a_		:  Any											=      pyspark.sql.SparkSession.builder.master('''local[*]'''      ).appName('''pyspark'''      ).getOrCreate()
				a_		:  List[Any]											=      spark.range(1_00      ).repartition(1      )
				a_		:  Optional[Any]											=      Spark(__A      )
				# 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




