



import unittest

import numpy as np

from transformers import BertConfig, is_flax_available
from transformers.testing_utils import require_flax, slow

from ...test_modeling_flax_common import FlaxModelTesterMixin, floats_tensor, ids_tensor, random_attention_mask


if is_flax_available():
							from transformers.models.bert.modeling_flax_bert import (
							    FlaxBertForMaskedLM,
							    FlaxBertForMultipleChoice,
							    FlaxBertForNextSentencePrediction,
							    FlaxBertForPreTraining,
							    FlaxBertForQuestionAnswering,
							    FlaxBertForSequenceClassification,
							    FlaxBertForTokenClassification,
							    FlaxBertModel,
							)







class       __lowerCAmelCase   (      unittest.TestCase						):



			def __init__(						self       :  Tuple			,		A       :  Union[str, Any]			,		A       :  Union[str, Any]=13			,		A       :  List[Any]=7			,		A       :  int=True			,		A       :  List[str]=True			,		A       :  Dict=True			,		A       :  Union[str, Any]=True			,		A       :  Tuple=99			,		A       :  Tuple=32			,		A       :  Tuple=5			,		A       :  int=4			,		A       :  List[str]=37			,		A       :  Tuple="gelu"			,		A       :  int=0.1			,		A       :  Tuple=0.1			,		A       :  Any=5_12			,		A       :  int=16			,		A       :  Dict=2			,		A       :  Tuple=0.0_2			,		A       :  Dict=4			,		)       ->   Any:



						"""simple docstring"""






						_UpperCAmelCase			=       parent
						_UpperCAmelCase			=       batch_size
						_UpperCAmelCase			=       seq_length
						_UpperCAmelCase			=       is_training
						_UpperCAmelCase			=       use_attention_mask
						_UpperCAmelCase			=       use_token_type_ids
						_UpperCAmelCase			=       use_labels
						_UpperCAmelCase			=       vocab_size
						_UpperCAmelCase			=       hidden_size
						_UpperCAmelCase			=       num_hidden_layers
						_UpperCAmelCase			=       num_attention_heads
						_UpperCAmelCase			=       intermediate_size
						_UpperCAmelCase			=       hidden_act
						_UpperCAmelCase			=       hidden_dropout_prob
						_UpperCAmelCase			=       attention_probs_dropout_prob
						_UpperCAmelCase			=       max_position_embeddings
						_UpperCAmelCase			=       type_vocab_size
						_UpperCAmelCase			=       type_sequence_label_size
						_UpperCAmelCase			=       initializer_range
						_UpperCAmelCase			=       num_choices



			def    _lowerCamelCase	(						self       :  List[Any])       ->   Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase			=       ids_tensor([self.batch_size, self.seq_length]			,		self.vocab_size)

						_UpperCAmelCase			=       None
						if self.use_attention_mask:
									_UpperCAmelCase			=       random_attention_mask([self.batch_size, self.seq_length])

						_UpperCAmelCase			=       None
						if self.use_token_type_ids:
									_UpperCAmelCase			=       ids_tensor([self.batch_size, self.seq_length]			,		self.type_vocab_size)

						_UpperCAmelCase			=       BertConfig(
						    vocab_size=self.vocab_size			,		hidden_size=self.hidden_size			,		num_hidden_layers=self.num_hidden_layers			,		num_attention_heads=self.num_attention_heads			,		intermediate_size=self.intermediate_size			,		hidden_act=self.hidden_act			,		hidden_dropout_prob=self.hidden_dropout_prob			,		attention_probs_dropout_prob=self.attention_probs_dropout_prob			,		max_position_embeddings=self.max_position_embeddings			,		type_vocab_size=self.type_vocab_size			,		is_decoder=A			,		initializer_range=self.initializer_range			,		)

						return config, input_ids, token_type_ids, attention_mask



			def    _lowerCamelCase	(						self       :  int)       ->   str:



						"""simple docstring"""






						_UpperCAmelCase			=       self.prepare_config_and_inputs()
						_UpperCAmelCase	,					_UpperCAmelCase	,					_UpperCAmelCase	,					_UpperCAmelCase			=       config_and_inputs
						_UpperCAmelCase			=       {'input_ids': input_ids, 'token_type_ids': token_type_ids, 'attention_mask': attention_mask}
						return config, inputs_dict



			def    _lowerCamelCase	(						self       :  Tuple)       ->   List[str]:



						"""simple docstring"""






						_UpperCAmelCase			=       self.prepare_config_and_inputs()
						_UpperCAmelCase	,					_UpperCAmelCase	,					_UpperCAmelCase	,					_UpperCAmelCase			=       config_and_inputs

						_UpperCAmelCase			=       True
						_UpperCAmelCase			=       floats_tensor([self.batch_size, self.seq_length, self.hidden_size])
						_UpperCAmelCase			=       ids_tensor([self.batch_size, self.seq_length]			,		vocab_size=2)

						return (
						    config,
						    input_ids,
						    attention_mask,
						    encoder_hidden_states,
						    encoder_attention_mask,
						)







@require_flax
class       __lowerCAmelCase   (      A    ,						unittest.TestCase						):
			UpperCamelCase        =					True

			UpperCamelCase        =					(
			    (
			        FlaxBertModel,
			        FlaxBertForPreTraining,
			        FlaxBertForMaskedLM,
			        FlaxBertForMultipleChoice,
			        FlaxBertForQuestionAnswering,
			        FlaxBertForNextSentencePrediction,
			        FlaxBertForSequenceClassification,
			        FlaxBertForTokenClassification,
			        FlaxBertForQuestionAnswering,
			    )
			    if is_flax_available()
			    else ()
			)



			def    _lowerCamelCase	(						self       :  str)       ->   Dict:



						"""simple docstring"""






						_UpperCAmelCase			=       FlaxBertModelTester(self)



			@slow
			def    _lowerCamelCase	(						self       :  Optional[Any])       ->   Optional[int]:



						"""simple docstring"""






						_UpperCAmelCase			=       FlaxBertModel.from_pretrained('bert-base-cased')
						_UpperCAmelCase			=       model(np.ones((1, 1)))
						self.assertIsNotNone(A)





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       __lowerCAmelCase   (      unittest.TestCase						):



			def    _lowerCamelCase	(						self       :  List[Any])       ->   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=A			,		)
						assert hasattr(self			,		'env')



			def    _lowerCamelCase	(						self       :  Any			,		A       :  Tuple=1)       ->   List[str]:



						"""simple docstring"""






						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=A			,		instance_type=self.instance_type			,		debugger_hook_config=A			,		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			,		A       :  int)       ->   str:



						"""simple docstring"""






						TrainingJobAnalytics(A).export_csv(F"{self.env.test_path}/{job_name}_metrics.csv")



			def    _lowerCamelCase	(						self       :  Union[str, Any])       ->   Dict:



						"""simple docstring"""






						_UpperCAmelCase			=       self.create_estimator()

						# run training
						estimator.fit()

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

						# extract kpis
						_UpperCAmelCase			=       list(result_metrics_df[result_metrics_df.metric_name == 'eval_accuracy']['value'])
						_UpperCAmelCase			=       list(result_metrics_df[result_metrics_df.metric_name == 'eval_loss']['value'])
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						_UpperCAmelCase			=       (
						    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}			,		A)
