from typing import Any, Dict, List, Optional, Tuple, Union

import torch
from torch import nn
from torch.utils.data import DistributedSampler, RandomSampler

from transformers import PreTrainedModel, Trainer, logging
from transformers.integrations import is_fairscale_available
from transformers.models.fsmt.configuration_fsmt import FSMTConfig
from transformers.optimization import (
    Adafactor,
    AdamW,
    get_constant_schedule,
    get_constant_schedule_with_warmup,
    get_cosine_schedule_with_warmup,
    get_cosine_with_hard_restarts_schedule_with_warmup,
    get_linear_schedule_with_warmup,
    get_polynomial_decay_schedule_with_warmup,
)
from transformers.trainer_pt_utils import get_tpu_sampler
from transformers.training_args import ParallelMode
from transformers.utils import is_torch_tpu_available


if is_fairscale_available():
					from fairscale.optim import OSS


__UpperCamelCase	: Dict    		=  logging.get_logger(__name__)

__UpperCamelCase	: Union[str, Any]    		=  {
    'linear': get_linear_schedule_with_warmup,
    'cosine': get_cosine_schedule_with_warmup,
    'cosine_w_restarts': get_cosine_with_hard_restarts_schedule_with_warmup,
    'polynomial': get_polynomial_decay_schedule_with_warmup,
    'constant': get_constant_schedule,
    'constant_w_warmup': get_constant_schedule_with_warmup,
}





class       lowercase__		(     UpperCamelCase_):


						def __init__( self			:       Union[str, Any]					,	UpperCamelCase__			:       Union[str, Any]=None					,	UpperCamelCase__			:       Optional[int]=None					,	*UpperCamelCase__			:       Dict					,	**UpperCamelCase__			:       str		):





												'''simple docstring'''





												super().__init__(*UpperCamelCase__					,	**UpperCamelCase__		)

												if config is None:
																		assert isinstance(self.model					,	UpperCamelCase__		), (
																		    "If no `config` is passed the model to be trained has to be of type `PreTrainedModel`, but is"
																		    f""" {self.model.__class__}"""
																		)
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							self.model.config
												else:
																		SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							config

												SCREAMING_SNAKE_CASE							:					List[Any]           =							data_args
												SCREAMING_SNAKE_CASE							:					str           =							self.config.tgt_vocab_size if isinstance(self.config					,	UpperCamelCase__		) else self.config.vocab_size

												if self.args.label_smoothing != 0 or (self.data_args is not None and self.data_args.ignore_pad_token_for_loss):
																		assert self.config.pad_token_id is not None, (
																		    "Make sure that `config.pad_token_id` is correcly defined when ignoring `pad_token` for loss"
																		    " calculation or doing label smoothing."
																		)

												if self.config.pad_token_id is None and self.config.eos_token_id is not None:
																		logger.warning(
																		    f"""The `config.pad_token_id` is `None`. Using `config.eos_token_id` = {self.config.eos_token_id} for"""
																		    ''' padding..'''		)

												if self.args.label_smoothing == 0:
																		SCREAMING_SNAKE_CASE							:					int           =							torch.nn.CrossEntropyLoss(ignore_index=self.config.pad_token_id		)
												else:
																		# dynamically import label_smoothed_nll_loss
																		from utils import label_smoothed_nll_loss

																		SCREAMING_SNAKE_CASE							:					List[str]           =							label_smoothed_nll_loss


						def  __A ( self			:       Any					,	UpperCamelCase__			:       int		):





												'''simple docstring'''





												if self.optimizer is None:
																		SCREAMING_SNAKE_CASE							:					str           =							['''bias''', '''LayerNorm.weight''']
																		SCREAMING_SNAKE_CASE							:					Optional[int]           =							[
																		    {
																		        '''params''': [p for n, p in self.model.named_parameters() if not any(nd in n for nd in no_decay		)],
																		        '''weight_decay''': self.args.weight_decay,
																		    },
																		    {
																		        '''params''': [p for n, p in self.model.named_parameters() if any(nd in n for nd in no_decay		)],
																		        '''weight_decay''': 0.0,
																		    },
																		]
																		SCREAMING_SNAKE_CASE							:					int           =							Adafactor if self.args.adafactor else AdamW
																		if self.args.adafactor:
																								SCREAMING_SNAKE_CASE							:					str           =							Adafactor
																								SCREAMING_SNAKE_CASE							:					Optional[Any]           =							{'''scale_parameter''': False, '''relative_step''': False}
																		else:
																								SCREAMING_SNAKE_CASE							:					Optional[int]           =							AdamW
																								SCREAMING_SNAKE_CASE							:					Any           =							{
																								    '''betas''': (self.args.adam_betaa, self.args.adam_betaa),
																								    '''eps''': self.args.adam_epsilon,
																								}
																		SCREAMING_SNAKE_CASE							:					Dict           =							self.args.learning_rate
																		if self.sharded_ddp:
																								SCREAMING_SNAKE_CASE							:					Any           =							OSS(
																								    params=UpperCamelCase__					,	optim=UpperCamelCase__					,	**UpperCamelCase__					,	)
																		else:
																								SCREAMING_SNAKE_CASE							:					Any           =							optimizer_cls(UpperCamelCase__					,	**UpperCamelCase__		)

												if self.lr_scheduler is None:
																		SCREAMING_SNAKE_CASE							:					List[Any]           =							self._get_lr_scheduler(UpperCamelCase__		)
												else:  # ignoring --lr_scheduler
																		logger.warning('''scheduler is passed to `Seq2SeqTrainer`, `--lr_scheduler` arg is ignored.'''		)


						def  __A ( self			:       Union[str, Any]					,	UpperCamelCase__			:       List[Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							arg_to_scheduler[self.args.lr_scheduler]
												if self.args.lr_scheduler == "constant":
																		SCREAMING_SNAKE_CASE							:					List[str]           =							schedule_func(self.optimizer		)
												elif self.args.lr_scheduler == "constant_w_warmup":
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							schedule_func(self.optimizer					,	num_warmup_steps=self.args.warmup_steps		)
												else:
																		SCREAMING_SNAKE_CASE							:					str           =							schedule_func(
																		    self.optimizer					,	num_warmup_steps=self.args.warmup_steps					,	num_training_steps=UpperCamelCase__		)
												return scheduler


						def  __A ( self			:       Any		):





												'''simple docstring'''





												if isinstance(self.train_dataset					,	torch.utils.data.IterableDataset		):
																		return None
												elif is_torch_tpu_available():
																		return get_tpu_sampler(self.train_dataset		)
												else:
																		if self.args.sortish_sampler:
																								self.train_dataset.make_sortish_sampler(
																								    self.args.per_device_train_batch_size					,	distributed=(self.args.parallel_mode == ParallelMode.DISTRIBUTED)					,	)

																		return (
																		    RandomSampler(self.train_dataset		)
																		    if self.args.local_rank == -1
																		    else DistributedSampler(self.train_dataset		)
																		)


						def  __A ( self			:       Union[str, Any]					,	UpperCamelCase__			:       Tuple					,	UpperCamelCase__			:       Union[str, Any]					,	UpperCamelCase__			:       List[str]		):





												'''simple docstring'''





												if self.args.label_smoothing == 0:
																		if self.data_args is not None and self.data_args.ignore_pad_token_for_loss:
																								# force training to ignore pad token
																								SCREAMING_SNAKE_CASE							:					List[str]           =							model(**UpperCamelCase__					,	use_cache=UpperCamelCase__		)[0]
																								SCREAMING_SNAKE_CASE							:					List[str]           =							self.loss_fn(logits.view(-1					,	logits.shape[-1]		)					,	labels.view(-1		)		)
																		else:
																								# compute usual loss via models
																								SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							model(**UpperCamelCase__					,	labels=UpperCamelCase__					,	use_cache=UpperCamelCase__		)[:2]
												else:
																		# compute label smoothed loss
																		SCREAMING_SNAKE_CASE							:					List[str]           =							model(**UpperCamelCase__					,	use_cache=UpperCamelCase__		)[0]
																		SCREAMING_SNAKE_CASE							:					int           =							torch.nn.functional.log_softmax(UpperCamelCase__					,	dim=-1		)
																		SCREAMING_SNAKE_CASE							:					int           =							self.loss_fn(UpperCamelCase__					,	UpperCamelCase__					,	self.args.label_smoothing					,	ignore_index=self.config.pad_token_id		)
												return loss, logits


						def  __A ( self			:       str					,	UpperCamelCase__			:       Dict					,	UpperCamelCase__			:       Union[str, Any]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							inputs.pop('''labels'''		)
												SCREAMING_SNAKE_CASE							:					List[Any]           =							self._compute_loss(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)
												return loss


						def  __A ( self			:       str					,	UpperCamelCase__			:       nn.Module					,	UpperCamelCase__			:       Dict[str, Union[torch.Tensor, Any]]					,	UpperCamelCase__			:       bool					,	UpperCamelCase__			:       Optional[List[str]] = None					,	):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							self._prepare_inputs(UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Dict           =							{
												    '''max_length''': self.data_args.val_max_target_length
												    if self.data_args is not None
												    else self.config.max_length,
												    '''num_beams''': self.data_args.eval_beams if self.data_args is not None else self.config.num_beams,
												}

												if self.args.predict_with_generate and not self.args.prediction_loss_only:
																		SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self.model.generate(
																		    inputs['''input_ids''']					,	attention_mask=inputs['''attention_mask''']					,	**UpperCamelCase__					,	)
																		# in case the batch is shorter than max length, the output should be padded
																		if generated_tokens.shape[-1] < gen_kwargs["max_length"]:
																								SCREAMING_SNAKE_CASE							:					Optional[Any]           =							self._pad_tensors_to_max_len(UpperCamelCase__					,	gen_kwargs['''max_length''']		)

												SCREAMING_SNAKE_CASE							:					Tuple           =							inputs.pop('''labels'''		)
												with torch.no_grad():
																		# compute loss on predict data
																		SCREAMING_SNAKE_CASE							:					int           =							self._compute_loss(UpperCamelCase__					,	UpperCamelCase__					,	UpperCamelCase__		)

												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							loss.mean().detach()
												if self.args.prediction_loss_only:
																		return (loss, None, None)

												SCREAMING_SNAKE_CASE							:					int           =							generated_tokens if self.args.predict_with_generate else logits

												if labels.shape[-1] < gen_kwargs["max_length"]:
																		SCREAMING_SNAKE_CASE							:					Dict           =							self._pad_tensors_to_max_len(UpperCamelCase__					,	gen_kwargs['''max_length''']		)

												return (loss, logits, labels)



						def  __A ( self			:       List[str]					,	UpperCamelCase__			:       Optional[int]					,	UpperCamelCase__			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Union[str, Any]           =							self.config.pad_token_id if self.config.pad_token_id is not None else self.config.eos_token_id

												if pad_token_id is None:
																		raise ValueError(
																		    '''Make sure that either `config.pad_token_id` or `config.eos_token_id` is defined if tensor has to be'''
																		    f""" padded to `max_length`={max_length}"""		)

												SCREAMING_SNAKE_CASE							:					List[Any]           =							pad_token_id * torch.ones(
												    (tensor.shape[0], max_length)					,	dtype=tensor.dtype					,	device=tensor.device		)
												SCREAMING_SNAKE_CASE							:					Tuple           =							tensor
												return padded_tensor






import json
import os
import subprocess
import unittest
from ast import literal_eval

import pytest
from parameterized import parameterized, parameterized_class

from . import is_sagemaker_available


if is_sagemaker_available():
					from sagemaker import Session, TrainingJobAnalytics
					from sagemaker.huggingface import HuggingFace





@pytest.mark.skipif(
    literal_eval(os.getenv("""TEST_SAGEMAKER"""			,							"""False""")) is not True			,							reason="""Skipping test because should only be run when releasing minor transformers version"""			,							)
@pytest.mark.usefixtures("""sm_env""")
@parameterized_class(
    [
        {
            """framework""": """pytorch""",
            """script""": """run_glue.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.p3.16xlarge""",
            """results""": {"""train_runtime""": 650, """eval_accuracy""": 0.7, """eval_loss""": 0.6},
        },
        {
            """framework""": """pytorch""",
            """script""": """run_ddp.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.p3.16xlarge""",
            """results""": {"""train_runtime""": 600, """eval_accuracy""": 0.7, """eval_loss""": 0.6},
        },
        {
            """framework""": """tensorflow""",
            """script""": """run_tf_dist.py""",
            """model_name_or_path""": """distilbert-base-cased""",
            """instance_type""": """ml.p3.16xlarge""",
            """results""": {"""train_runtime""": 600, """eval_accuracy""": 0.6, """eval_loss""": 0.7},
        },
    ])
class       lowercase__		(     unittest.TestCase):


						def  __A ( self			:       Any		):





												'''simple docstring'''





												if self.framework == "pytorch":
																		subprocess.run(
																		    f"""cp ./examples/pytorch/text-classification/run_glue.py {self.env.test_path}/run_glue.py""".split()					,	encoding='''utf-8'''					,	check=UpperCamelCase__					,	)
												assert hasattr(self					,	'''env'''		)


						def  __A ( self			:       str					,	UpperCamelCase__			:       Dict		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					Any           =							f"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
												# distributed data settings
												SCREAMING_SNAKE_CASE							:					Any           =							{'''smdistributed''': {'''dataparallel''': {'''enabled''': True}}} if self.script != '''run_ddp.py''' else None

												# creates estimator
												return HuggingFace(
												    entry_point=self.script					,	source_dir=self.env.test_path					,	role=self.env.role					,	image_uri=self.env.image_uri					,	base_job_name=UpperCamelCase__					,	instance_count=UpperCamelCase__					,	instance_type=self.instance_type					,	debugger_hook_config=UpperCamelCase__					,	hyperparameters={**self.env.distributed_hyperparameters, '''model_name_or_path''': self.model_name_or_path}					,	metric_definitions=self.env.metric_definitions					,	distribution=UpperCamelCase__					,	py_version='''py36'''					,	)


						def  __A ( self			:       Optional[Any]					,	UpperCamelCase__			:       List[str]		):





												'''simple docstring'''





												TrainingJobAnalytics(UpperCamelCase__		).export_csv(f"""{self.env.test_path}/{job_name}_metrics.csv"""		)



						@parameterized.expand([(2,)]		)
						def  __A ( self			:       Tuple					,	UpperCamelCase__			:       Optional[int]		):





												'''simple docstring'''





												SCREAMING_SNAKE_CASE							:					List[str]           =							self.create_estimator(UpperCamelCase__		)

												# run training
												estimator.fit()

												# result dataframe
												SCREAMING_SNAKE_CASE							:					Optional[Any]           =							TrainingJobAnalytics(estimator.latest_training_job.name		).dataframe()

												# extract kpis
												SCREAMING_SNAKE_CASE							:					List[str]           =							list(result_metrics_df[result_metrics_df.metric_name == '''eval_accuracy''']['''value''']		)
												SCREAMING_SNAKE_CASE							:					Dict           =							list(result_metrics_df[result_metrics_df.metric_name == '''eval_loss''']['''value''']		)
												# get train time from SageMaker job, this includes starting, preprocessing, stopping
												SCREAMING_SNAKE_CASE							:					List[Any]           =							(
												    Session().describe_training_job(estimator.latest_training_job.name		).get('''TrainingTimeInSeconds'''					,	99_9999		)
												)

												# assert kpis
												assert train_runtime <= self.results["train_runtime"]
												assert all(t >= self.results['''eval_accuracy'''] for t in eval_accuracy		)
												assert all(t <= self.results['''eval_loss'''] for t in eval_loss		)

												# dump tests result into json file to share in PR
												with open(f"""{estimator.latest_training_job.name}.json"""					,	'''w'''		) as outfile:
																		json.dump({'''train_time''': train_runtime, '''eval_accuracy''': eval_accuracy, '''eval_loss''': eval_loss}					,	UpperCamelCase__		)





