






'''simple docstring'''







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						=       LEDConfig
		_UpperCamelCase						=       {}
		_UpperCamelCase						=       'gelu'






		def __init__(   self ,_lowerCAmelCase ,_lowerCAmelCase=13 ,_lowerCAmelCase=7 ,_lowerCAmelCase=True ,_lowerCAmelCase=False ,_lowerCAmelCase=99 ,_lowerCAmelCase=32 ,_lowerCAmelCase=2 ,_lowerCAmelCase=4 ,_lowerCAmelCase=37 ,_lowerCAmelCase=0.1 ,_lowerCAmelCase=0.1 ,_lowerCAmelCase=20 ,_lowerCAmelCase=2 ,_lowerCAmelCase=1 ,_lowerCAmelCase=0 ,_lowerCAmelCase=4 ,):
								lowerCamelCase__								=	parent
								lowerCamelCase__								=	batch_size
								lowerCamelCase__								=	seq_length
								lowerCamelCase__								=	is_training
								lowerCamelCase__								=	use_labels
								lowerCamelCase__								=	vocab_size
								lowerCamelCase__								=	hidden_size
								lowerCamelCase__								=	num_hidden_layers
								lowerCamelCase__								=	num_attention_heads
								lowerCamelCase__								=	intermediate_size

								lowerCamelCase__								=	hidden_dropout_prob
								lowerCamelCase__								=	attention_probs_dropout_prob
								lowerCamelCase__								=	max_position_embeddings
								lowerCamelCase__								=	eos_token_id
								lowerCamelCase__								=	pad_token_id
								lowerCamelCase__								=	bos_token_id
								lowerCamelCase__								=	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
								lowerCamelCase__								=	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
								lowerCamelCase__								=	(
								    self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
								)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	ids_tensor([self.batch_size, self.seq_length - 1] ,self.vocab_size							)
								lowerCamelCase__								=	tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size							) ,1							)
								lowerCamelCase__								=	tf.concat([input_ids, eos_tensor] ,axis=1							)

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

								lowerCamelCase__								=	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 ,)
								lowerCamelCase__								=	prepare_led_inputs_dict(_lowerCAmelCase ,_lowerCAmelCase ,_lowerCAmelCase							)
								lowerCamelCase__								=	tf.concat(
								    [tf.zeros_like(_lowerCAmelCase							)[:, :-1], tf.ones_like(_lowerCAmelCase							)[:, -1:]] ,axis=-1 ,)
								lowerCamelCase__								=	global_attention_mask
								return config, inputs_dict






		def       UpperCamelCase_       (   self ,_lowerCAmelCase ,_lowerCAmelCase							):
								lowerCamelCase__								=	TFLEDModel(config=_lowerCAmelCase							).get_decoder()
								lowerCamelCase__								=	inputs_dict["""input_ids"""]

								lowerCamelCase__								=	input_ids[:1, :]
								lowerCamelCase__								=	inputs_dict["""attention_mask"""][:1, :]
								lowerCamelCase__								=	1

								# first forward pass
								lowerCamelCase__								=	model(_lowerCAmelCase ,attention_mask=_lowerCAmelCase ,use_cache=_lowerCAmelCase							)

								lowerCamelCase__							,      lowerCamelCase__								=	outputs.to_tuple()

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

								# append to next input_ids and
								lowerCamelCase__								=	tf.concat([input_ids, next_tokens] ,axis=-1							)
								lowerCamelCase__								=	tf.concat([attention_mask, next_attn_mask] ,axis=-1							)

								lowerCamelCase__								=	model(_lowerCAmelCase ,attention_mask=_lowerCAmelCase							)[0]
								lowerCamelCase__								=	model(_lowerCAmelCase ,attention_mask=_lowerCAmelCase ,past_key_values=_lowerCAmelCase							)[0]

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

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

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



def     A__      (					__lowerCAmelCase		:      str ,     __lowerCAmelCase		:      Union[str, Any] ,     __lowerCAmelCase		:      List[Any] ,     __lowerCAmelCase		:      Optional[int]=None ,     __lowerCAmelCase		:      Dict=None ,     __lowerCAmelCase		:      Optional[int]=None ,     __lowerCAmelCase		:      Optional[Any]=None ,     ):
						if attention_mask is None:
												lowerCamelCase__								=	tf.cast(tf.math.not_equal(__lowerCAmelCase ,     config.pad_token_id    ) ,     tf.inta    )
						if decoder_attention_mask is None:
												lowerCamelCase__								=	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:
												lowerCamelCase__								=	tf.ones((config.encoder_layers, config.encoder_attention_heads)    )
						if decoder_head_mask is None:
												lowerCamelCase__								=	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__     (a ,a ,unittest.TestCase      ):




		'''simple docstring'''





		_UpperCamelCase						=       (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
		_UpperCamelCase						=       (TFLEDForConditionalGeneration,) if is_tf_available() else ()
		_UpperCamelCase						=       (
		    {
		        'conversational': TFLEDForConditionalGeneration,
		        'feature-extraction': TFLEDModel,
		        'summarization': TFLEDForConditionalGeneration,
		        'text2text-generation': TFLEDForConditionalGeneration,
		        'translation': TFLEDForConditionalGeneration,
		    }
		    if is_tf_available()
		    else {}
		)
		_UpperCamelCase						=       True
		_UpperCamelCase						=       False
		_UpperCamelCase						=       False
		_UpperCamelCase						=       False






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	TFLEDModelTester(self							)
								lowerCamelCase__								=	ConfigTester(self ,config_class=_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								self.config_tester.run_common_tests()






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.check_decoder_model_past_large_inputs(*_lowerCAmelCase							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__							,      lowerCamelCase__								=	self.model_tester.prepare_config_and_inputs_for_common()
								lowerCamelCase__								=	tf.zeros_like(inputs_dict["""attention_mask"""]							)
								lowerCamelCase__								=	2
								lowerCamelCase__								=	tf.where(
								    tf.range(self.model_tester.seq_length							)[None, :] < num_global_attn_indices ,1 ,inputs_dict["""global_attention_mask"""] ,)

								lowerCamelCase__								=	True
								lowerCamelCase__								=	self.model_tester.seq_length
								lowerCamelCase__								=	self.model_tester.encoder_seq_length

								def check_decoder_attentions_output(_lowerCAmelCase							):
														lowerCamelCase__								=	outputs.decoder_attentions
														self.assertEqual(len(_lowerCAmelCase							) ,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(_lowerCAmelCase							):
														lowerCamelCase__								=	[t.numpy() for t in outputs.encoder_attentions]
														lowerCamelCase__								=	[t.numpy() for t in outputs.encoder_global_attentions]
														self.assertEqual(len(_lowerCAmelCase							) ,self.model_tester.num_hidden_layers							)
														self.assertEqual(len(_lowerCAmelCase							) ,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:
														lowerCamelCase__								=	True
														lowerCamelCase__								=	False
														lowerCamelCase__								=	False
														lowerCamelCase__								=	model_class(_lowerCAmelCase							)
														lowerCamelCase__								=	model(self._prepare_for_class(_lowerCAmelCase ,_lowerCAmelCase							)							)
														lowerCamelCase__								=	len(_lowerCAmelCase							)
														self.assertEqual(config.output_hidden_states ,_lowerCAmelCase							)
														check_encoder_attentions_output(_lowerCAmelCase							)

														if self.is_encoder_decoder:
																				lowerCamelCase__								=	model_class(_lowerCAmelCase							)
																				lowerCamelCase__								=	model(self._prepare_for_class(_lowerCAmelCase ,_lowerCAmelCase							)							)
																				self.assertEqual(config.output_hidden_states ,_lowerCAmelCase							)
																				check_decoder_attentions_output(_lowerCAmelCase							)

														# Check that output attentions can also be changed via the config
														del inputs_dict["output_attentions"]
														lowerCamelCase__								=	True
														lowerCamelCase__								=	model_class(_lowerCAmelCase							)
														lowerCamelCase__								=	model(self._prepare_for_class(_lowerCAmelCase ,_lowerCAmelCase							)							)
														self.assertEqual(config.output_hidden_states ,_lowerCAmelCase							)
														check_encoder_attentions_output(_lowerCAmelCase							)

														# Check attention is always last and order is fine
														lowerCamelCase__								=	True
														lowerCamelCase__								=	True
														lowerCamelCase__								=	model_class(_lowerCAmelCase							)
														lowerCamelCase__								=	model(self._prepare_for_class(_lowerCAmelCase ,_lowerCAmelCase							)							)

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






		@unittest.skip("""LED keeps using potentially symbolic tensors in conditionals and breaks tracing."""							)
		def       UpperCamelCase_       (   self							):
								pass






		def       UpperCamelCase_       (   self							):
								# TODO: Head-masking not yet implement
								pass



def     A__      (					__lowerCAmelCase		:      Optional[int]    ):
						return tf.constant(__lowerCAmelCase ,     dtype=tf.intaa    )


UpperCamelCase     :   List[Any]	      =		1E-4





@slow
@require_tf
class 		UpperCamelCase__     (unittest.TestCase      ):




		'''simple docstring'''





		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""							).led

								# change to intended input here
								lowerCamelCase__								=	_long_tensor([5_12 * [0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69]]							)
								lowerCamelCase__								=	_long_tensor([1_28 * [0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69]]							)
								lowerCamelCase__								=	prepare_led_inputs_dict(model.config ,_lowerCAmelCase ,_lowerCAmelCase							)
								lowerCamelCase__								=	model(**_lowerCAmelCase							)[0]
								lowerCamelCase__								=	(1, 10_24, 7_68)
								self.assertEqual(output.shape ,_lowerCAmelCase							)
								# change to expected output here
								lowerCamelCase__								=	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] ,_lowerCAmelCase ,atol=1E-3							)






		def       UpperCamelCase_       (   self							):
								lowerCamelCase__								=	TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""							)

								# change to intended input here
								lowerCamelCase__								=	_long_tensor([5_12 * [0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69]]							)
								lowerCamelCase__								=	_long_tensor([1_28 * [0, 3_14_14, 2_32, 3_28, 7_40, 11_40, 1_26_95, 69]]							)
								lowerCamelCase__								=	prepare_led_inputs_dict(model.config ,_lowerCAmelCase ,_lowerCAmelCase							)
								lowerCamelCase__								=	model(**_lowerCAmelCase							)[0]
								lowerCamelCase__								=	(1, 10_24, model.config.vocab_size)
								self.assertEqual(output.shape ,_lowerCAmelCase							)
								# change to expected output here
								lowerCamelCase__								=	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] ,_lowerCAmelCase ,atol=1E-3 ,rtol=1E-3							)


'''simple docstring'''




from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_tf_available,
    is_tokenizers_available,
    is_torch_available,
)


__UpperCAmelCase					={
    "configuration_mobilebert": [
        "MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP",
        "MobileBertConfig",
        "MobileBertOnnxConfig",
    ],
    "tokenization_mobilebert": ["MobileBertTokenizer"],
}

try:
  if not is_tokenizers_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase					=["MobileBertTokenizerFast"]

try:
  if not is_torch_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase					=[
      "MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "MobileBertForMaskedLM",
      "MobileBertForMultipleChoice",
      "MobileBertForNextSentencePrediction",
      "MobileBertForPreTraining",
      "MobileBertForQuestionAnswering",
      "MobileBertForSequenceClassification",
      "MobileBertForTokenClassification",
      "MobileBertLayer",
      "MobileBertModel",
      "MobileBertPreTrainedModel",
      "load_tf_weights_in_mobilebert",
  ]

try:
  if not is_tf_available():
    raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
  pass
else:
  __UpperCAmelCase					=[
      "TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST",
      "TFMobileBertForMaskedLM",
      "TFMobileBertForMultipleChoice",
      "TFMobileBertForNextSentencePrediction",
      "TFMobileBertForPreTraining",
      "TFMobileBertForQuestionAnswering",
      "TFMobileBertForSequenceClassification",
      "TFMobileBertForTokenClassification",
      "TFMobileBertMainLayer",
      "TFMobileBertModel",
      "TFMobileBertPreTrainedModel",
  ]


if TYPE_CHECKING:
  from .configuration_mobilebert import (
      MOBILEBERT_PRETRAINED_CONFIG_ARCHIVE_MAP,
      MobileBertConfig,
      MobileBertOnnxConfig,
  )
  from .tokenization_mobilebert import MobileBertTokenizer

  try:
    if not is_tokenizers_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .tokenization_mobilebert_fast import MobileBertTokenizerFast

  try:
    if not is_torch_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_mobilebert import (
        MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        MobileBertForMaskedLM,
        MobileBertForMultipleChoice,
        MobileBertForNextSentencePrediction,
        MobileBertForPreTraining,
        MobileBertForQuestionAnswering,
        MobileBertForSequenceClassification,
        MobileBertForTokenClassification,
        MobileBertLayer,
        MobileBertModel,
        MobileBertPreTrainedModel,
        load_tf_weights_in_mobilebert,
    )

  try:
    if not is_tf_available():
      raise OptionalDependencyNotAvailable()
  except OptionalDependencyNotAvailable:
    pass
  else:
    from .modeling_tf_mobilebert import (
        TF_MOBILEBERT_PRETRAINED_MODEL_ARCHIVE_LIST,
        TFMobileBertForMaskedLM,
        TFMobileBertForMultipleChoice,
        TFMobileBertForNextSentencePrediction,
        TFMobileBertForPreTraining,
        TFMobileBertForQuestionAnswering,
        TFMobileBertForSequenceClassification,
        TFMobileBertForTokenClassification,
        TFMobileBertMainLayer,
        TFMobileBertModel,
        TFMobileBertPreTrainedModel,
    )

else:
  import sys

  __UpperCAmelCase					=_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)
