




'''simple docstring'''

import math
import os
from copy import deepcopy

import datasets
import evaluate
import torch
import transformers
from datasets import load_dataset
from torch.utils.data import DataLoader
from transformers import AutoModelForSequenceClassification, AutoTokenizer

from accelerate import Accelerator
from accelerate.test_utils import RegressionDataset, RegressionModel
from accelerate.utils import is_tpu_available, set_seed


a  :				Optional[Any]									=  'true'


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase=82,      __UpperCAmelCase=16	)		->					str:

					'''simple docstring'''

					set_seed(42	)
					snake_case_									=				RegressionModel()
					snake_case_									=				deepcopy(__UpperCAmelCase	)
					snake_case_									=				RegressionDataset(length=__UpperCAmelCase	)
					snake_case_									=				DataLoader(__UpperCAmelCase,      batch_size=__UpperCAmelCase	)
					model.to(accelerator.device	)
					snake_case_			,snake_case_									=				accelerator.prepare(__UpperCAmelCase,      __UpperCAmelCase	)
					return model, ddp_model, dataloader


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase=False	)		->					int:

					'''simple docstring'''

					snake_case_									=				AutoTokenizer.from_pretrained('''hf-internal-testing/mrpc-bert-base-cased'''	)
					snake_case_									=				load_dataset('''glue''',      '''mrpc''',      split='''validation'''	)

					def tokenize_function(__UpperCAmelCase	):
										snake_case_									=				tokenizer(examples['''sentence1'''],      examples['''sentence2'''],      truncation=__UpperCAmelCase,      max_length=__UpperCAmelCase	)
										return outputs

					with accelerator.main_process_first():
										snake_case_									=				dataset.map(
										    __UpperCAmelCase,      batched=__UpperCAmelCase,      remove_columns=['''idx''', '''sentence1''', '''sentence2'''],      )

					snake_case_									=				tokenized_datasets.rename_column('''label''',      '''labels'''	)

					def collate_fn(__UpperCAmelCase	):
										if use_longest:
															return tokenizer.pad(__UpperCAmelCase,      padding='''longest''',      return_tensors='''pt'''	)
										return tokenizer.pad(__UpperCAmelCase,      padding='''max_length''',      max_length=128,      return_tensors='''pt'''	)

					return DataLoader(__UpperCAmelCase,      shuffle=__UpperCAmelCase,      collate_fn=__UpperCAmelCase,      batch_size=16	)


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase	)		->					int:

					'''simple docstring'''

					snake_case_									=				Accelerator(dispatch_batches=__UpperCAmelCase,      split_batches=__UpperCAmelCase	)
					snake_case_									=				get_dataloader(__UpperCAmelCase,      not dispatch_batches	)
					snake_case_									=				AutoModelForSequenceClassification.from_pretrained(
					    '''hf-internal-testing/mrpc-bert-base-cased''',      return_dict=__UpperCAmelCase	)
					snake_case_			,snake_case_									=				accelerator.prepare(__UpperCAmelCase,      __UpperCAmelCase	)
					return {"ddp": [ddp_model, ddp_dataloader, "cuda:0"], "no": [model, dataloader, accelerator.device]}, accelerator


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase	)		->					Any:

					'''simple docstring'''

					snake_case_									=				[]
					for batch in dataloader:
										snake_case_			,snake_case_									=				batch.values()
										with torch.no_grad():
															snake_case_									=				model(__UpperCAmelCase	)
															snake_case_			,snake_case_									=				accelerator.gather_for_metrics((logit, target)	)
															logits_and_targets.append((logit, target)	)
					snake_case_			,snake_case_									=				[], []
					for logit, targ in logits_and_targets:
										logits.append(__UpperCAmelCase	)
										targs.append(__UpperCAmelCase	)
					snake_case_			,snake_case_									=				torch.cat(__UpperCAmelCase	), torch.cat(__UpperCAmelCase	)
					return logits, targs


def 		__magic_name__    (	__UpperCAmelCase,      __UpperCAmelCase=82,      __UpperCAmelCase=False,      __UpperCAmelCase=False,      __UpperCAmelCase=16	)		->					Dict:

					'''simple docstring'''

					snake_case_			,snake_case_			,snake_case_									=				get_basic_setup(__UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase	)
					snake_case_			,snake_case_									=				generate_predictions(__UpperCAmelCase,      __UpperCAmelCase,      __UpperCAmelCase	)
					assert (
					    len(__UpperCAmelCase	) == num_samples
					), F"Unexpected number of inputs:\n    Expected: {num_samples}\n    Actual: {len(__UpperCAmelCase	)}"


def 		__magic_name__    (	__UpperCAmelCase = False,      __UpperCAmelCase = False	)		->					Union[str, Any]:

					'''simple docstring'''

					snake_case_									=				evaluate.load('''glue''',      '''mrpc'''	)
					snake_case_			,snake_case_									=				get_mrpc_setup(__UpperCAmelCase,      __UpperCAmelCase	)
					# First do baseline
					snake_case_			,snake_case_			,snake_case_									=				setup['''no''']
					model.to(__UpperCAmelCase	)
					model.eval()
					for batch in dataloader:
										batch.to(__UpperCAmelCase	)
										with torch.inference_mode():
															snake_case_									=				model(**__UpperCAmelCase	)
										snake_case_									=				outputs.logits.argmax(dim=-1	)
										metric.add_batch(predictions=__UpperCAmelCase,      references=batch['''labels''']	)
					snake_case_									=				metric.compute()

					# Then do distributed
					snake_case_			,snake_case_			,snake_case_									=				setup['''ddp''']
					model.eval()
					for batch in dataloader:
										with torch.inference_mode():
															snake_case_									=				model(**__UpperCAmelCase	)
										snake_case_									=				outputs.logits.argmax(dim=-1	)
										snake_case_									=				batch['''labels''']
										snake_case_			,snake_case_									=				accelerator.gather_for_metrics((preds, references)	)
										metric.add_batch(predictions=__UpperCAmelCase,      references=__UpperCAmelCase	)
					snake_case_									=				metric.compute()

					for key in "accuracy f1".split():
										assert math.isclose(
										    baseline[key],      distributed[key]	), F"Baseline and Distributed are not the same for key {key}:\n\tBaseline: {baseline[key]}\n\tDistributed: {distributed[key]}\n"


def 		__magic_name__    (	)		->					Union[str, Any]:

					'''simple docstring'''

					snake_case_									=				Accelerator(split_batches=__UpperCAmelCase,      dispatch_batches=__UpperCAmelCase	)
					if accelerator.is_local_main_process:
										datasets.utils.logging.set_verbosity_warning()
										transformers.utils.logging.set_verbosity_warning()
					else:
										datasets.utils.logging.set_verbosity_error()
										transformers.utils.logging.set_verbosity_error()
					# These are a bit slower so they should only be ran on the GPU or TPU
					if torch.cuda.is_available() or is_tpu_available():
										if accelerator.is_local_main_process:
															print('''**Testing gather_for_metrics**'''	)
										for split_batches in [True, False]:
															for dispatch_batches in [True, False]:
																				if accelerator.is_local_main_process:
																									print(F"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`"	)
																				test_mrpc(__UpperCAmelCase,      __UpperCAmelCase	)
																				accelerator.state._reset_state()
					if accelerator.is_local_main_process:
										print('''**Test torch metrics**'''	)
					for split_batches in [True, False]:
										for dispatch_batches in [True, False]:
															snake_case_									=				Accelerator(split_batches=__UpperCAmelCase,      dispatch_batches=__UpperCAmelCase	)
															if accelerator.is_local_main_process:
																				print(F"With: `split_batches={split_batches}`, `dispatch_batches={dispatch_batches}`, length=99"	)
															test_torch_metrics(__UpperCAmelCase,      99	)
															accelerator.state._reset_state()
					if accelerator.is_local_main_process:
										print('''**Test last batch is not dropped when perfectly divisible**'''	)
					snake_case_									=				Accelerator()
					test_torch_metrics(__UpperCAmelCase,      512	)
					accelerator.state._reset_state()


def 		__magic_name__    (	__UpperCAmelCase	)		->					str:

					'''simple docstring'''

					main()


if __name__ == "__main__":
							main()





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

import pytest
from parameterized import 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.g4dn.xlarge''',
            '''results''': {'''train_runtime''': 6_5_0, '''eval_accuracy''': 0.6, '''eval_loss''': 0.9},
        },
        {
            '''framework''': '''tensorflow''',
            '''script''': '''run_tf.py''',
            '''model_name_or_path''': '''distilbert-base-cased''',
            '''instance_type''': '''ml.g4dn.xlarge''',
            '''results''': {'''train_runtime''': 6_0_0, '''eval_accuracy''': 0.3, '''eval_loss''': 0.9},
        },
    ]      )
class        snake_case_   (unittest.TestCase      ):
					def  lowerCamelCase__(      self   :Dict      )					->    Union[str, Any]:
												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=__snake_case			,)
												assert hasattr(self			,'env'      )
					def  lowerCamelCase__(      self   :List[Any]			,__snake_case   :Optional[Any]=1      )					->    int:
												# 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=F'{self.env.base_job_name}-single'			,instance_count=__snake_case			,instance_type=self.instance_type			,debugger_hook_config=__snake_case			,hyperparameters={**self.env.hyperparameters, 'model_name_or_path': self.model_name_or_path}			,metric_definitions=self.env.metric_definitions			,py_version='py36'			,)
					def  lowerCamelCase__(      self   :Dict			,__snake_case   :str      )					->    Optional[int]:
												TrainingJobAnalytics(__snake_case      ).export_csv(F'{self.env.test_path}/{job_name}_metrics.csv'      )



					def  lowerCamelCase__(      self   :Tuple      )					->    Union[str, Any]:
												# create estimator
												a__             =     self.create_estimator()

												# run training
												estimator.fit()

												# result dataframe
												a__             =     TrainingJobAnalytics(estimator.latest_training_job.name      ).dataframe()

												# extract kpis
												a__             =     list(result_metrics_df[result_metrics_df.metric_name == 'eval_accuracy']['value']      )
												a__             =     list(result_metrics_df[result_metrics_df.metric_name == 'eval_loss']['value']      )
												# get train time from SageMaker job, this includes starting, preprocessing, stopping
												a__             =     (
												    Session().describe_training_job(estimator.latest_training_job.name      ).get('TrainingTimeInSeconds'			,99_99_99      )
												)

												# 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}			,__snake_case      )




