


'''simple docstring'''






from transformers import BertTokenizer, EncoderDecoderModel, SeqaSeqTrainer, SeqaSeqTrainingArguments
from transformers.testing_utils import TestCasePlus, require_torch, slow
from transformers.utils import is_datasets_available


if is_datasets_available():
		import datasets




class 	a__       (							_snake_case						):





							'''simple docstring'''






							@slow
							@require_torch
							def 						__SCREAMING_SNAKE_CASE							(	self     ) ->							List[str]:
														lowerCAmelCase__      =			EncoderDecoderModel.from_encoder_decoder_pretrained('''prajjwal1/bert-tiny'''      ,					'''prajjwal1/bert-tiny'''     )
														lowerCAmelCase__      =			BertTokenizer.from_pretrained('''bert-base-uncased'''     )

														lowerCAmelCase__      =			bertabert.config.encoder.vocab_size
														lowerCAmelCase__      =			tokenizer.sep_token_id
														lowerCAmelCase__      =			tokenizer.cls_token_id
														lowerCAmelCase__      =			1_28

														lowerCAmelCase__      =			datasets.load_dataset('''cnn_dailymail'''      ,					'''3.0.0'''      ,					split='''train[:1%]'''     )
														lowerCAmelCase__      =			datasets.load_dataset('''cnn_dailymail'''      ,					'''3.0.0'''      ,					split='''validation[:1%]'''     )

														lowerCAmelCase__      =			train_dataset.select(range(32     )     )
														lowerCAmelCase__      =			val_dataset.select(range(16     )     )

														lowerCAmelCase__      =			4

														def _map_to_encoder_decoder_inputs(lowerCamelCase_     ):
																					# Tokenizer will automatically set [BOS] <text> [EOS]
																					lowerCAmelCase__      =			tokenizer(batch['''article''']      ,					padding='''max_length'''      ,					truncation=lowerCAmelCase__      ,					max_length=5_12     )
																					lowerCAmelCase__      =			tokenizer(batch['''highlights''']      ,					padding='''max_length'''      ,					truncation=lowerCAmelCase__      ,					max_length=1_28     )
																					lowerCAmelCase__      =			inputs.input_ids
																					lowerCAmelCase__      =			inputs.attention_mask

																					lowerCAmelCase__      =			outputs.input_ids
																					lowerCAmelCase__      =			outputs.input_ids.copy()
																					lowerCAmelCase__      =			[
																					    [-1_00 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch['''labels''']
																					]
																					lowerCAmelCase__      =			outputs.attention_mask

																					assert all(len(lowerCAmelCase__     ) == 5_12 for x in inputs.input_ids     )
																					assert all(len(lowerCAmelCase__     ) == 1_28 for x in outputs.input_ids     )

																					return batch

														def _compute_metrics(lowerCamelCase_     ):
																					lowerCAmelCase__      =			pred.label_ids
																					lowerCAmelCase__      =			pred.predictions

																					# all unnecessary tokens are removed
																					lowerCAmelCase__      =			tokenizer.batch_decode(lowerCAmelCase__      ,					skip_special_tokens=lowerCAmelCase__     )
																					lowerCAmelCase__      =			tokenizer.batch_decode(lowerCAmelCase__      ,					skip_special_tokens=lowerCAmelCase__     )

																					lowerCAmelCase__      =			sum([int(pred_str[i] == label_str[i]     ) for i in range(len(lowerCAmelCase__     )     )]     ) / len(lowerCAmelCase__     )

																					return {"accuracy": accuracy}

														# map train dataset
														lowerCAmelCase__      =			train_dataset.map(
														    _map_to_encoder_decoder_inputs      ,					batched=lowerCAmelCase__      ,					batch_size=lowerCAmelCase__      ,					remove_columns=['''article''', '''highlights''']      ,					)
														train_dataset.set_format(
														    type='''torch'''      ,					columns=['''input_ids''', '''attention_mask''', '''decoder_input_ids''', '''decoder_attention_mask''', '''labels''']      ,					)

														# same for validation dataset
														lowerCAmelCase__      =			val_dataset.map(
														    _map_to_encoder_decoder_inputs      ,					batched=lowerCAmelCase__      ,					batch_size=lowerCAmelCase__      ,					remove_columns=['''article''', '''highlights''']      ,					)
														val_dataset.set_format(
														    type='''torch'''      ,					columns=['''input_ids''', '''attention_mask''', '''decoder_input_ids''', '''decoder_attention_mask''', '''labels''']      ,					)

														lowerCAmelCase__      =			self.get_auto_remove_tmp_dir()

														lowerCAmelCase__      =			SeqaSeqTrainingArguments(
														    output_dir=lowerCAmelCase__      ,					per_device_train_batch_size=lowerCAmelCase__      ,					per_device_eval_batch_size=lowerCAmelCase__      ,					predict_with_generate=lowerCAmelCase__      ,					evaluation_strategy='''steps'''      ,					do_train=lowerCAmelCase__      ,					do_eval=lowerCAmelCase__      ,					warmup_steps=0      ,					eval_steps=2      ,					logging_steps=2      ,					)

														# instantiate trainer
														lowerCAmelCase__      =			SeqaSeqTrainer(
														    model=lowerCAmelCase__      ,					args=lowerCAmelCase__      ,					compute_metrics=_compute_metrics      ,					train_dataset=lowerCAmelCase__      ,					eval_dataset=lowerCAmelCase__      ,					tokenizer=lowerCAmelCase__      ,					)

														# start training
														trainer.train()





'''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     (  A				,					A )		->					List[Any]:
							lowerCAmelCase__      =			[]
							for part_id in partition_order:
														lowerCAmelCase__      =			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       _snake_case     (  )		->					Tuple:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(100 ).repartition(1 )
							lowerCAmelCase__      =			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       _snake_case     (  )		->					Optional[int]:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(10 ).repartition(2 )
							lowerCAmelCase__      =			[1, 0]
							lowerCAmelCase__      =			_generate_iterable_examples(A				,					A )  # Reverse the partitions.
							lowerCAmelCase__      =			_get_expected_row_ids_and_row_dicts_for_partition_order(A				,					A )

							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       _snake_case     (  )		->					Optional[Any]:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(10 ).repartition(1 )
							lowerCAmelCase__      =			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       _snake_case     (  )		->					Union[str, Any]:
							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 A : x.reverse()
														lowerCAmelCase__      =			_get_expected_row_ids_and_row_dicts_for_partition_order(A				,					[2, 1, 0] )

														lowerCAmelCase__      =			SparkExamplesIterable(A ).shuffle_data_sources(A )
														assert shuffled_it.n_shards == 3
														for i, (row_id, row_dict) in enumerate(A ):
																					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       _snake_case     (  )		->					Dict:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(20 ).repartition(4 )

							# Partitions 0 and 2
							lowerCAmelCase__      =			SparkExamplesIterable(A ).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(A				,					[0, 2] )
							for i, (row_id, row_dict) in enumerate(A ):
														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(A ).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(A				,					[1, 3] )
							for i, (row_id, row_dict) in enumerate(A ):
														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       _snake_case     (  )		->					Dict:
							lowerCAmelCase__      =			pyspark.sql.SparkSession.builder.master('''local[*]''' ).appName('''pyspark''' ).getOrCreate()
							lowerCAmelCase__      =			spark.range(100 ).repartition(1 )
							lowerCAmelCase__      =			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() == 100