






'''simple docstring'''





import json
import os
from functools import lru_cache
from typing import List, Optional, Tuple

import regex as re

from ...tokenization_utils import AddedToken, PreTrainedTokenizer
from ...utils import logging


_UpperCamelCase						:						Any  =					logging.get_logger(__name__)


_UpperCamelCase						:						Union[str, Any]  =					{'vocab_file': 'vocab.json', 'merges_file': 'merges.txt'}

# See all BART models at https://huggingface.co/models?filter=bart
_UpperCamelCase						:						Tuple  =					{
    'vocab_file': {
        'facebook/bart-base': 'https://huggingface.co/facebook/bart-base/resolve/main/vocab.json',
        'facebook/bart-large': 'https://huggingface.co/facebook/bart-large/resolve/main/vocab.json',
        'facebook/bart-large-mnli': 'https://huggingface.co/facebook/bart-large-mnli/resolve/main/vocab.json',
        'facebook/bart-large-cnn': 'https://huggingface.co/facebook/bart-large-cnn/resolve/main/vocab.json',
        'facebook/bart-large-xsum': 'https://huggingface.co/facebook/bart-large-xsum/resolve/main/vocab.json',
        'yjernite/bart_eli5': 'https://huggingface.co/yjernite/bart_eli5/resolve/main/vocab.json',
    },
    'merges_file': {
        'facebook/bart-base': 'https://huggingface.co/facebook/bart-base/resolve/main/merges.txt',
        'facebook/bart-large': 'https://huggingface.co/facebook/bart-large/resolve/main/merges.txt',
        'facebook/bart-large-mnli': 'https://huggingface.co/facebook/bart-large-mnli/resolve/main/merges.txt',
        'facebook/bart-large-cnn': 'https://huggingface.co/facebook/bart-large-cnn/resolve/main/merges.txt',
        'facebook/bart-large-xsum': 'https://huggingface.co/facebook/bart-large-xsum/resolve/main/merges.txt',
        'yjernite/bart_eli5': 'https://huggingface.co/yjernite/bart_eli5/resolve/main/merges.txt',
    },
}

_UpperCamelCase						:						List[str]  =					{
    'facebook/bart-base': 1_024,
    'facebook/bart-large': 1_024,
    'facebook/bart-large-mnli': 1_024,
    'facebook/bart-large-cnn': 1_024,
    'facebook/bart-large-xsum': 1_024,
    'yjernite/bart_eli5': 1_024,
}
@lru_cache()
def __UpperCAmelCase   (					)    ->   Optional[Any]:
			UpperCAmelCase_      :       Optional[Any]          =       (
			    list(range(ord('''!'''		)			,       ord('''~'''		) + 1		)		) + list(range(ord('''¡'''		)			,       ord('''¬'''		) + 1		)		) + list(range(ord('''®'''		)			,       ord('''ÿ'''		) + 1		)		)
			)
			UpperCAmelCase_      :       Union[str, Any]          =       bs[:]
			UpperCAmelCase_      :       Optional[Any]          =       0
			for b in range(2**8		):
						if b not in bs:
									bs.append(A		)
									cs.append(2**8 + n		)
									n += 1
			UpperCAmelCase_      :       List[Any]          =       [chr(A		) for n in cs]
			return dict(zip(A			,       A		)		)
def __UpperCAmelCase   (					A					:			Optional[Any]		)    ->   Union[str, Any]:
			UpperCAmelCase_      :       Dict          =       set()
			UpperCAmelCase_      :       Optional[int]          =       word[0]
			for char in word[1:]:
						pairs.add((prev_char, char)		)
						UpperCAmelCase_      :       Any          =       char
			return pairs







class  snake_case__      (					UpperCamelCase):
					a_        =					VOCAB_FILES_NAMES
					a_        =					PRETRAINED_VOCAB_FILES_MAP
					a_        =					PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
					a_        =					["input_ids", "attention_mask"]




					def __init__(  self  :		List[str]				, _A  :		List[Any]				, _A  :		Any				, _A  :		int="replace"				, _A  :		List[str]="<s>"				, _A  :		Optional[int]="</s>"				, _A  :		Any="</s>"				, _A  :		Dict="<s>"				, _A  :		Any="<unk>"				, _A  :		int="<pad>"				, _A  :		Optional[int]="<mask>"				, _A  :		str=False				, **_A  :		int				, )    ->							int:
								UpperCAmelCase_      :       Optional[int]          =       AddedToken(_A				, lstrip=_A				, rstrip=_A      ) if isinstance(_A				, _A      ) else bos_token
								UpperCAmelCase_      :       List[str]          =       AddedToken(_A				, lstrip=_A				, rstrip=_A      ) if isinstance(_A				, _A      ) else eos_token
								UpperCAmelCase_      :       Union[str, Any]          =       AddedToken(_A				, lstrip=_A				, rstrip=_A      ) if isinstance(_A				, _A      ) else sep_token
								UpperCAmelCase_      :       Dict          =       AddedToken(_A				, lstrip=_A				, rstrip=_A      ) if isinstance(_A				, _A      ) else cls_token
								UpperCAmelCase_      :       Tuple          =       AddedToken(_A				, lstrip=_A				, rstrip=_A      ) if isinstance(_A				, _A      ) else unk_token
								UpperCAmelCase_      :       Union[str, Any]          =       AddedToken(_A				, lstrip=_A				, rstrip=_A      ) if isinstance(_A				, _A      ) else pad_token

								# Mask token behave like a normal word, i.e. include the space before it
								UpperCAmelCase_      :       Optional[int]          =       AddedToken(_A				, lstrip=_A				, rstrip=_A      ) if isinstance(_A				, _A      ) else mask_token

								super().__init__(
								    errors=_A				, bos_token=_A				, eos_token=_A				, unk_token=_A				, sep_token=_A				, cls_token=_A				, pad_token=_A				, mask_token=_A				, add_prefix_space=_A				, **_A				, )

								with open(_A				, encoding='''utf-8'''      ) as vocab_handle:
											UpperCAmelCase_      :       Optional[Any]          =       json.load(_A      )
								UpperCAmelCase_      :       Any          =       {v: k for k, v in self.encoder.items()}
								UpperCAmelCase_      :       str          =       errors  # how to handle errors in decoding
								UpperCAmelCase_      :       Any          =       bytes_to_unicode()
								UpperCAmelCase_      :       int          =       {v: k for k, v in self.byte_encoder.items()}
								with open(_A				, encoding='''utf-8'''      ) as merges_handle:
											UpperCAmelCase_      :       List[Any]          =       merges_handle.read().split('''\n'''      )[1:-1]
								UpperCAmelCase_      :       List[Any]          =       [tuple(merge.split()      ) for merge in bpe_merges]
								UpperCAmelCase_      :       List[Any]          =       dict(zip(_A				, range(len(_A      )      )      )      )
								UpperCAmelCase_      :       str          =       {}
								UpperCAmelCase_      :       Union[str, Any]          =       add_prefix_space

								# Should have added re.IGNORECASE so BPE merges can happen for capitalized versions of contractions
								UpperCAmelCase_      :       Union[str, Any]          =       re.compile(r'''\'s|\'t|\'re|\'ve|\'m|\'ll|\'d| ?\p{L}+| ?\p{N}+| ?[^\s\p{L}\p{N}]+|\s+(?!\S)|\s+'''      )




					@property
					def 		A      (  self  :		Dict      )    ->							str:
								return len(self.encoder      )




					def 		A      (  self  :		Union[str, Any]      )    ->							Tuple:
								return dict(self.encoder				, **self.added_tokens_encoder      )




					def 		A      (  self  :		Tuple				, _A  :		Union[str, Any]      )    ->							Union[str, Any]:
								if token in self.cache:
											return self.cache[token]
								UpperCAmelCase_      :       Union[str, Any]          =       tuple(_A      )
								UpperCAmelCase_      :       List[Any]          =       get_pairs(_A      )

								if not pairs:
											return token

								while True:
											UpperCAmelCase_      :       Optional[Any]          =       min(_A				, key=lambda _A      : self.bpe_ranks.get(_A				, float('''inf'''      )      )      )
											if bigram not in self.bpe_ranks:
														break
											UpperCAmelCase_    ,      UpperCAmelCase_      :       Optional[Any]          =       bigram
											UpperCAmelCase_      :       Optional[int]          =       []
											UpperCAmelCase_      :       Optional[int]          =       0
											while i < len(_A      ):
														try:
																	UpperCAmelCase_      :       Tuple          =       word.index(_A				, _A      )
														except ValueError:
																	new_word.extend(word[i:]      )
																	break
														else:
																	new_word.extend(word[i:j]      )
																	UpperCAmelCase_      :       Optional[int]          =       j

														if word[i] == first and i < len(_A      ) - 1 and word[i + 1] == second:
																	new_word.append(first + second      )
																	i += 2
														else:
																	new_word.append(word[i]      )
																	i += 1
											UpperCAmelCase_      :       int          =       tuple(_A      )
											UpperCAmelCase_      :       str          =       new_word
											if len(_A      ) == 1:
														break
											else:
														UpperCAmelCase_      :       int          =       get_pairs(_A      )
								UpperCAmelCase_      :       Union[str, Any]          =       ''' '''.join(_A      )
								UpperCAmelCase_      :       Optional[int]          =       word
								return word




					def 		A      (  self  :		Dict				, _A  :		List[Any]      )    ->							Dict:
								UpperCAmelCase_      :       Optional[int]          =       []
								for token in re.findall(self.pat				, _A      ):
											UpperCAmelCase_      :       Union[str, Any]          =       ''''''.join(
											    self.byte_encoder[b] for b in token.encode('''utf-8'''      )      )  # Maps all our bytes to unicode strings, avoiding control tokens of the BPE (spaces in our case)
											bpe_tokens.extend(bpe_token for bpe_token in self.bpe(_A      ).split(''' '''      )      )
								return bpe_tokens




					def 		A      (  self  :		Dict				, _A  :		Union[str, Any]      )    ->							Tuple:
								return self.encoder.get(_A				, self.encoder.get(self.unk_token      )      )




					def 		A      (  self  :		List[Any]				, _A  :		Optional[Any]      )    ->							str:
								return self.decoder.get(_A      )




					def 		A      (  self  :		int				, _A  :		int      )    ->							List[Any]:
								UpperCAmelCase_      :       Optional[int]          =       ''''''.join(_A      )
								UpperCAmelCase_      :       Union[str, Any]          =       bytearray([self.byte_decoder[c] for c in text]      ).decode('''utf-8'''				, errors=self.errors      )
								return text




					def 		A      (  self  :		Union[str, Any]				, _A  :		str				, _A  :		Optional[str] = None      )    ->							Tuple[str]:
								if not os.path.isdir(_A      ):
											logger.error(F"Vocabulary path ({save_directory}) should be a directory"      )
											return
								UpperCAmelCase_      :       List[str]          =       os.path.join(
								    _A				, (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''vocab_file''']      )
								UpperCAmelCase_      :       Tuple          =       os.path.join(
								    _A				, (filename_prefix + '''-''' if filename_prefix else '''''') + VOCAB_FILES_NAMES['''merges_file''']      )

								with open(_A				, '''w'''				, encoding='''utf-8'''      ) as f:
											f.write(json.dumps(self.encoder				, indent=2				, sort_keys=_A				, ensure_ascii=_A      ) + '''\n'''      )

								UpperCAmelCase_      :       Any          =       0
								with open(_A				, '''w'''				, encoding='''utf-8'''      ) as writer:
											writer.write('''#version: 0.2\n'''      )
											for bpe_tokens, token_index in sorted(self.bpe_ranks.items()				, key=lambda _A      : kv[1]      ):
														if index != token_index:
																	logger.warning(
																	    F"Saving vocabulary to {merge_file}: BPE merge indices are not consecutive."
																	    ''' Please check that the tokenizer is not corrupted!'''      )
																	UpperCAmelCase_      :       Tuple          =       token_index
														writer.write(''' '''.join(_A      ) + '''\n'''      )
														index += 1

								return vocab_file, merge_file




					def 		A      (  self  :		int				, _A  :		List[int]				, _A  :		Optional[List[int]] = None      )    ->							List[int]:
								if token_ids_a is None:
											return [self.cls_token_id] + token_ids_a + [self.sep_token_id]
								UpperCAmelCase_      :       str          =       [self.cls_token_id]
								UpperCAmelCase_      :       List[Any]          =       [self.sep_token_id]
								return cls + token_ids_a + sep + sep + token_ids_a + sep




					def 		A      (  self  :		Tuple				, _A  :		List[int]				, _A  :		Optional[List[int]] = None				, _A  :		bool = False      )    ->							List[int]:
								if already_has_special_tokens:
											return super().get_special_tokens_mask(
											    token_ids_a=_A				, token_ids_a=_A				, already_has_special_tokens=_A      )

								if token_ids_a is None:
											return [1] + ([0] * len(_A      )) + [1]
								return [1] + ([0] * len(_A      )) + [1, 1] + ([0] * len(_A      )) + [1]




					def 		A      (  self  :		Tuple				, _A  :		List[int]				, _A  :		Optional[List[int]] = None      )    ->							List[int]:
								UpperCAmelCase_      :       Optional[int]          =       [self.sep_token_id]
								UpperCAmelCase_      :       int          =       [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]





					def 		A      (  self  :		List[str]				, _A  :		Dict				, _A  :		int=False				, **_A  :		int      )    ->							Optional[int]:
								UpperCAmelCase_      :       Union[str, Any]          =       kwargs.pop('''add_prefix_space'''				, self.add_prefix_space      )
								if (is_split_into_words or add_prefix_space) and (len(_A      ) > 0 and not text[0].isspace()):
											UpperCAmelCase_      :       int          =       ''' ''' + text
								return (text, kwargs)










'''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					:			Optional[int]			,       A					:			Optional[int]		)    ->   str:
			UpperCAmelCase_      :       List[Any]          =       []
			for part_id in partition_order:
						UpperCAmelCase_      :       Any          =       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   (					)    ->   Any:
			UpperCAmelCase_      :       Any          =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
			UpperCAmelCase_      :       List[str]          =       spark.range(1_0_0		).repartition(1		)
			UpperCAmelCase_      :       List[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=1_6		)
			# Given that the dataframe has 100 rows and each partition has 2 rows, we expect 50 partitions.
			assert spark_builder.df.rdd.getNumPartitions() == 5_0
@require_not_windows
@require_dill_gt_0_3_2
def __UpperCAmelCase   (					)    ->   str:
			UpperCAmelCase_      :       List[str]          =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
			UpperCAmelCase_      :       Optional[Any]          =       spark.range(1_0		).repartition(2		)
			UpperCAmelCase_      :       int          =       [1, 0]
			UpperCAmelCase_      :       str          =       _generate_iterable_examples(A			,       A		)  # Reverse the partitions.
			UpperCAmelCase_      :       Optional[Any]          =       _get_expected_row_ids_and_row_dicts_for_partition_order(A			,       A		)

			for i, (row_id, row_dict) in enumerate(generate_fn()		):
						UpperCAmelCase_    ,      UpperCAmelCase_      :       int          =       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   (					)    ->   Union[str, Any]:
			UpperCAmelCase_      :       str          =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
			UpperCAmelCase_      :       List[Any]          =       spark.range(1_0		).repartition(1		)
			UpperCAmelCase_      :       Optional[int]          =       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   (					)    ->   Union[str, Any]:
			UpperCAmelCase_      :       Dict          =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
			UpperCAmelCase_      :       Dict          =       spark.range(3_0		).repartition(3		)
			# Mock the generator so that shuffle reverses the partition indices.
			with patch('''numpy.random.Generator'''		) as generator_mock:
						UpperCAmelCase_      :       Any          =       lambda A		: x.reverse()
						UpperCAmelCase_      :       Dict          =       _get_expected_row_ids_and_row_dicts_for_partition_order(A			,       [2, 1, 0]		)

						UpperCAmelCase_      :       List[Any]          =       SparkExamplesIterable(A		).shuffle_data_sources(A		)
						assert shuffled_it.n_shards == 3
						for i, (row_id, row_dict) in enumerate(A		):
									UpperCAmelCase_    ,      UpperCAmelCase_      :       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 __UpperCAmelCase   (					)    ->   int:
			UpperCAmelCase_      :       Tuple          =       pyspark.sql.SparkSession.builder.master('''local[*]'''		).appName('''pyspark'''		).getOrCreate()
			UpperCAmelCase_      :       int          =       spark.range(2_0		).repartition(4		)

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


