






"""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        lowercase__      (							snake_case__       ):


			@slow
			@require_torch
			def      UpperCAmelCase__						(				self					:  Tuple ):
										lowerCamelCase_	:							str									=EncoderDecoderModel.from_encoder_decoder_pretrained("prajjwal1/bert-tiny"       ,	"prajjwal1/bert-tiny" )
										lowerCamelCase_	:							List[Any]									=BertTokenizer.from_pretrained("bert-base-uncased" )

										lowerCamelCase_	:							Optional[int]									=bertabert.config.encoder.vocab_size
										lowerCamelCase_	:							List[str]									=tokenizer.sep_token_id
										lowerCamelCase_	:							Optional[int]									=tokenizer.cls_token_id
										lowerCamelCase_	:							int									=128

										lowerCamelCase_	:							Tuple									=datasets.load_dataset("cnn_dailymail"       ,	"3.0.0"       ,	split="train[:1%]" )
										lowerCamelCase_	:							List[Any]									=datasets.load_dataset("cnn_dailymail"       ,	"3.0.0"       ,	split="validation[:1%]" )

										lowerCamelCase_	:							str									=train_dataset.select(range(32 ) )
										lowerCamelCase_	:							Tuple									=val_dataset.select(range(16 ) )

										lowerCamelCase_	:							Tuple									=4

										def _map_to_encoder_decoder_inputs(snake_case__					:  Dict ):
																	# Tokenizer will automatically set [BOS] <text> [EOS]
																	lowerCamelCase_	:							Optional[int]									=tokenizer(batch["article"]       ,	padding="max_length"       ,	truncation=snake_case__       ,	max_length=512 )
																	lowerCamelCase_	:							List[str]									=tokenizer(batch["highlights"]       ,	padding="max_length"       ,	truncation=snake_case__       ,	max_length=128 )
																	lowerCamelCase_	:							Any									=inputs.input_ids
																	lowerCamelCase_	:							Dict									=inputs.attention_mask

																	lowerCamelCase_	:							List[Any]									=outputs.input_ids
																	lowerCamelCase_	:							str									=outputs.input_ids.copy()
																	lowerCamelCase_	:							Optional[Any]									=[
																	    [-100 if token == tokenizer.pad_token_id else token for token in labels] for labels in batch["labels"]
																	]
																	lowerCamelCase_	:							Tuple									=outputs.attention_mask

																	assert all(len(snake_case__ ) == 512 for x in inputs.input_ids )
																	assert all(len(snake_case__ ) == 128 for x in outputs.input_ids )

																	return batch

										def _compute_metrics(snake_case__					:  Optional[int] ):
																	lowerCamelCase_	:							Dict									=pred.label_ids
																	lowerCamelCase_	:							List[Any]									=pred.predictions

																	# all unnecessary tokens are removed
																	lowerCamelCase_	:							Union[str, Any]									=tokenizer.batch_decode(snake_case__       ,	skip_special_tokens=snake_case__ )
																	lowerCamelCase_	:							str									=tokenizer.batch_decode(snake_case__       ,	skip_special_tokens=snake_case__ )

																	lowerCamelCase_	:							Dict									=sum([int(pred_str[i] == label_str[i] ) for i in range(len(snake_case__ ) )] ) / len(snake_case__ )

																	return {"accuracy": accuracy}

										# map train dataset
										lowerCamelCase_	:							List[Any]									=train_dataset.map(
										    _map_to_encoder_decoder_inputs       ,	batched=snake_case__       ,	batch_size=snake_case__       ,	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_	:							Optional[Any]									=val_dataset.map(
										    _map_to_encoder_decoder_inputs       ,	batched=snake_case__       ,	batch_size=snake_case__       ,	remove_columns=["article", "highlights"]       ,	)
										val_dataset.set_format(
										    type="torch"       ,	columns=["input_ids", "attention_mask", "decoder_input_ids", "decoder_attention_mask", "labels"]       ,	)

										lowerCamelCase_	:							Union[str, Any]									=self.get_auto_remove_tmp_dir()

										lowerCamelCase_	:							List[Any]									=SeqaSeqTrainingArguments(
										    output_dir=snake_case__       ,	per_device_train_batch_size=snake_case__       ,	per_device_eval_batch_size=snake_case__       ,	predict_with_generate=snake_case__       ,	evaluation_strategy="steps"       ,	do_train=snake_case__       ,	do_eval=snake_case__       ,	warmup_steps=0       ,	eval_steps=2       ,	logging_steps=2       ,	)

										# instantiate trainer
										lowerCamelCase_	:							List[str]									=SeqaSeqTrainer(
										    model=snake_case__       ,	args=snake_case__       ,	compute_metrics=_compute_metrics       ,	train_dataset=snake_case__       ,	eval_dataset=snake_case__       ,	tokenizer=snake_case__       ,	)

										# start training
										trainer.train()











"""simple docstring"""



import pytest

from datasets.parallel import ParallelBackendConfig, parallel_backend
from datasets.utils.py_utils import map_nested

from .utils import require_dill_gt_0_3_2, require_joblibspark, require_not_windows

def 			_snake_case    (					lowerCamelCase__     :   Any  )    ->     Union[str, Any]:  # picklable for multiprocessing
							return i + 1

@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
def 			_snake_case    (					)    ->     List[Any]:
							with parallel_backend("spark"  ):
														assert ParallelBackendConfig.backend_name == "spark"

							lowerCamelCase_	:							Optional[Any]									=[1, 2, 3]
							with pytest.raises(lowerCamelCase__  ):
														with parallel_backend("unsupported backend"  ):
																					map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=2  )

							with pytest.raises(lowerCamelCase__  ):
														with parallel_backend("unsupported backend"  ):
																					map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=-1  )

@require_dill_gt_0_3_2
@require_joblibspark
@require_not_windows
@pytest.mark.parametrize("num_proc"    ,						[2, -1]  )
def 			_snake_case    (					lowerCamelCase__     :   Tuple  )    ->     Optional[Any]:
							lowerCamelCase_	:							str									=[1, 2]
							lowerCamelCase_	:							List[str]									={"a": 1, "b": 2}
							lowerCamelCase_	:							List[str]									={"a": [1, 2], "b": [3, 4]}
							lowerCamelCase_	:							Optional[int]									={"a": {"1": 1}, "b": 2}
							lowerCamelCase_	:							int									={"a": 1, "b": 2, "c": 3, "d": 4}
							lowerCamelCase_	:							Optional[int]									=[2, 3]
							lowerCamelCase_	:							List[Any]									={"a": 2, "b": 3}
							lowerCamelCase_	:							int									={"a": [2, 3], "b": [4, 5]}
							lowerCamelCase_	:							str									={"a": {"1": 2}, "b": 3}
							lowerCamelCase_	:							Dict									={"a": 2, "b": 3, "c": 4, "d": 5}

							with parallel_backend("spark"  ):
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa
														assert map_nested(lowerCamelCase__    ,						lowerCamelCase__    ,						num_proc=lowerCamelCase__  ) == expected_map_nested_sa



