




'''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 						snake_case_       (UpperCamelCase			:    Union[str, Any]					,		UpperCamelCase			:    Union[str, Any] ):





							'''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 						snake_case_       ():





							'''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 						snake_case_       ():





							'''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 						snake_case_       ():





							'''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 						snake_case_       ():





							'''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 						snake_case_       ():





							'''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 						snake_case_       ():





							'''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 json
from typing import List, Optional, Tuple

from tokenizers import normalizers

from ...tokenization_utils_fast import PreTrainedTokenizerFast
from ...utils import logging
from .tokenization_squeezebert import SqueezeBertTokenizer


_snake_case : Tuple  				=						logging.get_logger(__name__)

_snake_case : Optional[int]  				=						{'vocab_file': 'vocab.txt', 'tokenizer_file': 'tokenizer.json'}

_snake_case : List[Any]  				=						{
    'vocab_file': {
        'squeezebert/squeezebert-uncased': (
            'https://huggingface.co/squeezebert/squeezebert-uncased/resolve/main/vocab.txt'
        ),
        'squeezebert/squeezebert-mnli': 'https://huggingface.co/squeezebert/squeezebert-mnli/resolve/main/vocab.txt',
        'squeezebert/squeezebert-mnli-headless': (
            'https://huggingface.co/squeezebert/squeezebert-mnli-headless/resolve/main/vocab.txt'
        ),
    },
    'tokenizer_file': {
        'squeezebert/squeezebert-uncased': (
            'https://huggingface.co/squeezebert/squeezebert-uncased/resolve/main/tokenizer.json'
        ),
        'squeezebert/squeezebert-mnli': (
            'https://huggingface.co/squeezebert/squeezebert-mnli/resolve/main/tokenizer.json'
        ),
        'squeezebert/squeezebert-mnli-headless': (
            'https://huggingface.co/squeezebert/squeezebert-mnli-headless/resolve/main/tokenizer.json'
        ),
    },
}

_snake_case : Union[str, Any]  				=						{
    'squeezebert/squeezebert-uncased': 512,
    'squeezebert/squeezebert-mnli': 512,
    'squeezebert/squeezebert-mnli-headless': 512,
}


_snake_case : Tuple  				=						{
    'squeezebert/squeezebert-uncased': {'do_lower_case': True},
    'squeezebert/squeezebert-mnli': {'do_lower_case': True},
    'squeezebert/squeezebert-mnli-headless': {'do_lower_case': True},
}

class      A       (						_a       ):
							lowercase_							     =							VOCAB_FILES_NAMES
							lowercase_							     =							PRETRAINED_VOCAB_FILES_MAP
							lowercase_							     =							PRETRAINED_INIT_CONFIGURATION
							lowercase_							     =							PRETRAINED_POSITIONAL_EMBEDDINGS_SIZES
							lowercase_							     =							SqueezeBertTokenizer







							def __init__(    self			:    str				,     lowerCAmelCase_			:    str=None				,     lowerCAmelCase_			:    List[str]=None				,     lowerCAmelCase_			:    str=True				,     lowerCAmelCase_			:    List[str]="[UNK]"				,     lowerCAmelCase_			:    Union[str, Any]="[SEP]"				,     lowerCAmelCase_			:    Optional[Any]="[PAD]"				,     lowerCAmelCase_			:    Any="[CLS]"				,     lowerCAmelCase_			:    List[str]="[MASK]"				,     lowerCAmelCase_			:    int=True				,     lowerCAmelCase_			:    List[Any]=None				,     **lowerCAmelCase_			:    Optional[int]				,     )     ->				int:


														"""simple docstring"""


														super().__init__(
														    lowerCAmelCase_				,     tokenizer_file=lowerCAmelCase_				,     do_lower_case=lowerCAmelCase_				,     unk_token=lowerCAmelCase_				,     sep_token=lowerCAmelCase_				,     pad_token=lowerCAmelCase_				,     cls_token=lowerCAmelCase_				,     mask_token=lowerCAmelCase_				,     tokenize_chinese_chars=lowerCAmelCase_				,     strip_accents=lowerCAmelCase_				,     **lowerCAmelCase_				,     )

														_a      	=							json.loads(self.backend_tokenizer.normalizer.__getstate__()   )
														if (
														    normalizer_state.get('''lowercase'''				,     lowerCAmelCase_   ) != do_lower_case
														    or normalizer_state.get('''strip_accents'''				,     lowerCAmelCase_   ) != strip_accents
														    or normalizer_state.get('''handle_chinese_chars'''				,     lowerCAmelCase_   ) != tokenize_chinese_chars
														):
																					_a      	=							getattr(lowerCAmelCase_				,     normalizer_state.pop('''type'''   )   )
																					_a      	=							do_lower_case
																					_a      	=							strip_accents
																					_a      	=							tokenize_chinese_chars
																					_a      	=							normalizer_class(**lowerCAmelCase_   )

														_a      	=							do_lower_case







							def   __lowerCAmelCase     (    self			:    Union[str, Any]				,     lowerCAmelCase_			:    int				,     lowerCAmelCase_			:    Optional[Any]=None   )     ->				List[str]:


														"""simple docstring"""


														_a      	=							[self.cls_token_id] + token_ids_a + [self.sep_token_id]

														if token_ids_a:
																					output += token_ids_a + [self.sep_token_id]

														return output







							def   __lowerCAmelCase     (    self			:    Any				,     lowerCAmelCase_			:    List[int]				,     lowerCAmelCase_			:    Optional[List[int]] = None   )     ->				List[int]:


														"""simple docstring"""


														_a      	=							[self.sep_token_id]
														_a      	=							[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   ) * [0] + len(token_ids_a + sep   ) * [1]







							def   __lowerCAmelCase     (    self			:    Optional[Any]				,     lowerCAmelCase_			:    str				,     lowerCAmelCase_			:    Optional[str] = None   )     ->				Tuple[str]:


														"""simple docstring"""


														_a      	=							self._tokenizer.model.save(lowerCAmelCase_				,     name=lowerCAmelCase_   )
														return tuple(lowerCAmelCase_   )

