

from __future__ import annotations

import unittest

from transformers import LEDConfig, is_tf_available
from transformers.testing_utils import require_tf, slow

from ...test_configuration_common import ConfigTester
from ...test_modeling_tf_common import TFModelTesterMixin, ids_tensor
from ...test_pipeline_mixin import PipelineTesterMixin


if is_tf_available():
			import tensorflow as tf

			from transformers import TFLEDForConditionalGeneration, TFLEDModel

@require_tf
class 							UpperCAmelCase_						:
							"""simple docstring"""







							UpperCAmelCase__			:		Optional[int]             =       LEDConfig
							UpperCAmelCase__			:		Union[str, Any]             =       {}
							UpperCAmelCase__			:		Dict             =       "gelu"


							def __init__(     self			,				_a			,				_a=1_3			,				_a=7			,				_a=True			,				_a=False			,				_a=9_9			,				_a=3_2			,				_a=2			,				_a=4			,				_a=3_7			,				_a=0.1			,				_a=0.1			,				_a=2_0			,				_a=2			,				_a=1			,				_a=0			,				_a=4			,				)       ->       Dict:
												_a      :       Optional[Any]   =						parent
												_a      :       str   =						batch_size
												_a      :       Dict   =						seq_length
												_a      :       Any   =						is_training
												_a      :       Any   =						use_labels
												_a      :       List[Any]   =						vocab_size
												_a      :       Tuple   =						hidden_size
												_a      :       Dict   =						num_hidden_layers
												_a      :       str   =						num_attention_heads
												_a      :       List[Any]   =						intermediate_size

												_a      :       Any   =						hidden_dropout_prob
												_a      :       Any   =						attention_probs_dropout_prob
												_a      :       Tuple   =						max_position_embeddings
												_a      :       List[str]   =						eos_token_id
												_a      :       Any   =						pad_token_id
												_a      :       int   =						bos_token_id
												_a      :       Any   =						attention_window

												# `ModelTesterMixin.test_attention_outputs` is expecting attention tensors to be of size
												# [num_attention_heads, encoder_seq_length, encoder_key_length], but TFLongformerSelfAttention
												# returns attention of shape [num_attention_heads, encoder_seq_length, self.attention_window + 1]
												# because its local attention only attends to `self.attention_window` and one before and one after
												_a      :       Dict   =						self.attention_window + 2

												# because of padding `encoder_seq_length`, is different from `seq_length`. Relevant for
												# the `test_attention_outputs` and `test_hidden_states_output` tests
												_a      :       Tuple   =						(
												    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
												)


							def 							__lowercase							(     self			)       ->       Optional[Any]:
												_a      :       Union[str, Any]   =						ids_tensor([self.batch_size, self.seq_length - 1]			,				self.vocab_size			)
												_a      :       Optional[int]   =						tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size			)			,				1			)
												_a      :       Optional[Any]   =						tf.concat([input_ids, eos_tensor]			,				axis=1			)

												_a      :       Optional[Any]   =						ids_tensor([self.batch_size, self.seq_length]			,				self.vocab_size			)

												_a      :       Dict   =						self.config_cls(
												    vocab_size=self.vocab_size			,				d_model=self.hidden_size			,				encoder_layers=self.num_hidden_layers			,				decoder_layers=self.num_hidden_layers			,				encoder_attention_heads=self.num_attention_heads			,				decoder_attention_heads=self.num_attention_heads			,				encoder_ffn_dim=self.intermediate_size			,				decoder_ffn_dim=self.intermediate_size			,				dropout=self.hidden_dropout_prob			,				attention_dropout=self.attention_probs_dropout_prob			,				max_position_embeddings=self.max_position_embeddings			,				eos_token_ids=[2]			,				bos_token_id=self.bos_token_id			,				pad_token_id=self.pad_token_id			,				decoder_start_token_id=self.pad_token_id			,				attention_window=self.attention_window			,				**self.config_updates			,				)
												_a      :       Optional[int]   =						prepare_led_inputs_dict(_a			,				_a			,				_a			)
												_a      :       int   =						tf.concat(
												    [tf.zeros_like(_a			)[:, :-1], tf.ones_like(_a			)[:, -1:]]			,				axis=-1			,				)
												_a      :       str   =						global_attention_mask
												return config, inputs_dict




							def 							__lowercase							(     self			,				_a			,				_a			)       ->       Any:
												_a      :       Dict   =						TFLEDModel(config=_a			).get_decoder()
												_a      :       Any   =						inputs_dict['''input_ids''']

												_a      :       str   =						input_ids[:1, :]
												_a      :       Optional[int]   =						inputs_dict['''attention_mask'''][:1, :]
												_a      :       Optional[Any]   =						1

												# first forward pass
												_a      :       int   =						model(_a			,				attention_mask=_a			,				use_cache=_a			)

												_a ,   _a      :       Dict   =						outputs.to_tuple()

												# create hypothetical next token and extent to next_input_ids
												_a      :       Tuple   =						ids_tensor((self.batch_size, 3)			,				config.vocab_size			)
												_a      :       int   =						tf.cast(ids_tensor((self.batch_size, 3)			,				2			)			,				tf.inta			)

												# append to next input_ids and
												_a      :       List[str]   =						tf.concat([input_ids, next_tokens]			,				axis=-1			)
												_a      :       Any   =						tf.concat([attention_mask, next_attn_mask]			,				axis=-1			)

												_a      :       int   =						model(_a			,				attention_mask=_a			)[0]
												_a      :       Optional[Any]   =						model(_a			,				attention_mask=_a			,				past_key_values=_a			)[0]

												self.parent.assertEqual(next_tokens.shape[1]			,				output_from_past.shape[1]			)

												# select random slice
												_a      :       str   =						int(ids_tensor((1,)			,				output_from_past.shape[-1]			)			)
												_a      :       List[str]   =						output_from_no_past[:, -3:, random_slice_idx]
												_a      :       Tuple   =						output_from_past[:, :, random_slice_idx]

												# test that outputs are equal for slice
												tf.debugging.assert_near(_a			,				_a			,				rtol=1e-3			)



def   __UpperCAmelCase					(      __a :     int						,__a :     str						,__a :     Optional[Any]						,__a :     Optional[int]=None						,__a :     Optional[int]=None						,__a :     Dict=None						,__a :     List[Any]=None						,)     ->	str:

					"""simple docstring"""

					if attention_mask is None:
										_a      :       Optional[int]   =						tf.cast(tf.math.not_equal(__a						,config.pad_token_id      )						,tf.inta      )
					if decoder_attention_mask is None:
										_a      :       List[Any]   =						tf.concat(
										    [
										        tf.ones(decoder_input_ids[:, :1].shape						,dtype=tf.inta      ),
										        tf.cast(tf.math.not_equal(decoder_input_ids[:, 1:]						,config.pad_token_id      )						,tf.inta      ),
										    ]						,axis=-1						,)
					if head_mask is None:
										_a      :       List[Any]   =						tf.ones((config.encoder_layers, config.encoder_attention_heads)      )
					if decoder_head_mask is None:
										_a      :       Optional[Any]   =						tf.ones((config.decoder_layers, config.decoder_attention_heads)      )
					return {
					    "input_ids": input_ids,
					    "attention_mask": attention_mask,
					    "decoder_input_ids": decoder_input_ids,
					    "decoder_attention_mask": decoder_attention_mask,
					    "head_mask": head_mask,
					    "decoder_head_mask": decoder_head_mask,
					}

@require_tf
class 							UpperCAmelCase_						(     __lowercase							,			__lowercase							,			unittest.TestCase  ):
							"""simple docstring"""







							UpperCAmelCase__			:		Union[str, Any]             =       (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
							UpperCAmelCase__			:		Optional[int]             =       (TFLEDForConditionalGeneration,) if is_tf_available() else ()
							UpperCAmelCase__			:		List[Any]             =       (
							    {
							        "conversational": TFLEDForConditionalGeneration,
							        "feature-extraction": TFLEDModel,
							        "summarization": TFLEDForConditionalGeneration,
							        "text2text-generation": TFLEDForConditionalGeneration,
							        "translation": TFLEDForConditionalGeneration,
							    }
							    if is_tf_available()
							    else {}
							)
							UpperCAmelCase__			:		str             =       True
							UpperCAmelCase__			:		Dict             =       False
							UpperCAmelCase__			:		Optional[Any]             =       False
							UpperCAmelCase__			:		Dict             =       False


							def 							__lowercase							(     self			)       ->       Optional[Any]:
												_a      :       Dict   =						TFLEDModelTester(self			)
												_a      :       Any   =						ConfigTester(self			,				config_class=_a			)


							def 							__lowercase							(     self			)       ->       Union[str, Any]:
												self.config_tester.run_common_tests()


							def 							__lowercase							(     self			)       ->       int:
												_a      :       Optional[Any]   =						self.model_tester.prepare_config_and_inputs_for_common()
												self.model_tester.check_decoder_model_past_large_inputs(*_a			)


							def 							__lowercase							(     self			)       ->       str:
												_a ,   _a      :       List[Any]   =						self.model_tester.prepare_config_and_inputs_for_common()
												_a      :       int   =						tf.zeros_like(inputs_dict['''attention_mask''']			)
												_a      :       Tuple   =						2
												_a      :       Optional[Any]   =						tf.where(
												    tf.range(self.model_tester.seq_length			)[None, :] < num_global_attn_indices			,				1			,				inputs_dict['''global_attention_mask''']			,				)

												_a      :       Tuple   =						True
												_a      :       Union[str, Any]   =						self.model_tester.seq_length
												_a      :       Dict   =						self.model_tester.encoder_seq_length

												def check_decoder_attentions_output(_a			):
																	_a      :       List[str]   =						outputs.decoder_attentions
																	self.assertEqual(len(_a			)			,				self.model_tester.num_hidden_layers			)
																	self.assertListEqual(
																	    list(decoder_attentions[0].shape[-3:]			)			,				[self.model_tester.num_attention_heads, seq_length, seq_length]			,				)

												def check_encoder_attentions_output(_a			):
																	_a      :       Optional[int]   =						[t.numpy() for t in outputs.encoder_attentions]
																	_a      :       Optional[Any]   =						[t.numpy() for t in outputs.encoder_global_attentions]
																	self.assertEqual(len(_a			)			,				self.model_tester.num_hidden_layers			)
																	self.assertEqual(len(_a			)			,				self.model_tester.num_hidden_layers			)
																	self.assertListEqual(
																	    list(attentions[0].shape[-3:]			)			,				[self.model_tester.num_attention_heads, seq_length, seq_length]			,				)
																	self.assertListEqual(
																	    list(global_attentions[0].shape[-3:]			)			,				[self.model_tester.num_attention_heads, encoder_seq_length, num_global_attn_indices]			,				)

												for model_class in self.all_model_classes:
																	_a      :       List[Any]   =						True
																	_a      :       str   =						False
																	_a      :       Any   =						False
																	_a      :       List[Any]   =						model_class(_a			)
																	_a      :       Union[str, Any]   =						model(self._prepare_for_class(_a			,				_a			)			)
																	_a      :       Any   =						len(_a			)
																	self.assertEqual(config.output_hidden_states			,				_a			)
																	check_encoder_attentions_output(_a			)

																	if self.is_encoder_decoder:
																						_a      :       str   =						model_class(_a			)
																						_a      :       Dict   =						model(self._prepare_for_class(_a			,				_a			)			)
																						self.assertEqual(config.output_hidden_states			,				_a			)
																						check_decoder_attentions_output(_a			)

																	# Check that output attentions can also be changed via the config
																	del inputs_dict["output_attentions"]
																	_a      :       Union[str, Any]   =						True
																	_a      :       Tuple   =						model_class(_a			)
																	_a      :       Optional[Any]   =						model(self._prepare_for_class(_a			,				_a			)			)
																	self.assertEqual(config.output_hidden_states			,				_a			)
																	check_encoder_attentions_output(_a			)

																	# Check attention is always last and order is fine
																	_a      :       Tuple   =						True
																	_a      :       Dict   =						True
																	_a      :       Union[str, Any]   =						model_class(_a			)
																	_a      :       Optional[Any]   =						model(self._prepare_for_class(_a			,				_a			)			)

																	self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)			,				len(_a			)			)
																	self.assertEqual(model.config.output_hidden_states			,				_a			)
																	check_encoder_attentions_output(_a			)


							@unittest.skip('''LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'''			)
							def 							__lowercase							(     self			)       ->       Union[str, Any]:
												pass




							def 							__lowercase							(     self			)       ->       List[str]:
												# TODO: Head-masking not yet implement
												pass



def   __UpperCAmelCase					(      __a :     Tuple      )     ->	Any:

					"""simple docstring"""

					return tf.constant(__a						,dtype=tf.intaa      )


a__												=						1E-4







@slow
@require_tf
class 							UpperCAmelCase_						(     unittest.TestCase  ):
							"""simple docstring"""







							def 							__lowercase							(     self			)       ->       List[Any]:
												_a      :       Dict   =						TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''			).led

												# change to intended input here
												_a      :       str   =						_long_tensor([5_1_2 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]			)
												_a      :       Tuple   =						_long_tensor([1_2_8 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]			)
												_a      :       Any   =						prepare_led_inputs_dict(model.config			,				_a			,				_a			)
												_a      :       Optional[int]   =						model(**_a			)[0]
												_a      :       str   =						(1, 1_0_2_4, 7_6_8)
												self.assertEqual(output.shape			,				_a			)
												# change to expected output here
												_a      :       str   =						tf.convert_to_tensor(
												    [[2.3050, 2.8279, 0.6531], [-1.8457, -0.1455, -3.5661], [-1.0186, 0.4586, -2.2043]]			,				)
												tf.debugging.assert_near(output[:, :3, :3]			,				_a			,				atol=1e-3			)




							def 							__lowercase							(     self			)       ->       str:
												_a      :       Any   =						TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''			)

												# change to intended input here
												_a      :       List[str]   =						_long_tensor([5_1_2 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]			)
												_a      :       Dict   =						_long_tensor([1_2_8 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]			)
												_a      :       Union[str, Any]   =						prepare_led_inputs_dict(model.config			,				_a			,				_a			)
												_a      :       Union[str, Any]   =						model(**_a			)[0]
												_a      :       Optional[int]   =						(1, 1_0_2_4, model.config.vocab_size)
												self.assertEqual(output.shape			,				_a			)
												# change to expected output here
												_a      :       List[str]   =						tf.convert_to_tensor(
												    [[33.6507, 6.4572, 16.8089], [5.8739, -2.4238, 11.2902], [-3.2139, -4.3149, 4.2783]]			,				)
												tf.debugging.assert_near(output[:, :3, :3]			,				_a			,				atol=1e-3			,				rtol=1e-3			)






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''': 6_50, '''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''': 6_00, '''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''': 6_00, '''eval_accuracy''': 0.6, '''eval_loss''': 0.7},
        },
    ]							)
class        _a		(unittest.TestCase							):


					'''simple docstring'''






					def     __A     (					self     ):
						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     __A     (					self       ,			A__     ):
						A__				:       int	  =						F"""{self.env.base_job_name}-{instance_count}-{'ddp' if 'ddp' in self.script else 'smd'}"""
						# distributed data settings
						A__				:       str	  =						{"""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=A__       ,			instance_count=A__       ,			instance_type=self.instance_type       ,			debugger_hook_config=A__       ,			hyperparameters={**self.env.distributed_hyperparameters, """model_name_or_path""": self.model_name_or_path}       ,			metric_definitions=self.env.metric_definitions       ,			distribution=A__       ,			py_version="""py36"""       ,			)






					def     __A     (					self       ,			A__     ):
						TrainingJobAnalytics(A__     ).export_csv(F"""{self.env.test_path}/{job_name}_metrics.csv"""     )






					@parameterized.expand([(2,)]     )
					def     __A     (					self       ,			A__     ):
						# create estimator
						A__				:       str	  =						self.create_estimator(A__     )

						# run training
						estimator.fit()

						# result dataframe
						A__				:       str	  =						TrainingJobAnalytics(estimator.latest_training_job.name     ).dataframe()

						# extract kpis
						A__				:       List[Any]	  =						list(result_metrics_df[result_metrics_df.metric_name == """eval_accuracy"""]["""value"""]     )
						A__				:       Optional[Any]	  =						list(result_metrics_df[result_metrics_df.metric_name == """eval_loss"""]["""value"""]     )
						# get train time from SageMaker job, this includes starting, preprocessing, stopping
						A__				:       str	  =						(
						    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}       ,			A__     )



