





'''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"""







		lowerCAmelCase_					=				LEDConfig
		lowerCAmelCase_					=				{}
		lowerCAmelCase_					=				"""gelu"""


		def __init__(				self    :					Dict							,	UpperCamelCase__    :					Tuple							,	UpperCamelCase__    :					Optional[int]=13							,	UpperCamelCase__    :					List[str]=7							,	UpperCamelCase__    :					Dict=True							,	UpperCamelCase__    :					str=False							,	UpperCamelCase__    :					Optional[Any]=99							,	UpperCamelCase__    :					Tuple=32							,	UpperCamelCase__    :					Optional[int]=2							,	UpperCamelCase__    :					Tuple=4							,	UpperCamelCase__    :					List[str]=37							,	UpperCamelCase__    :					Optional[Any]=0.1							,	UpperCamelCase__    :					List[str]=0.1							,	UpperCamelCase__    :					List[str]=20							,	UpperCamelCase__    :					str=2							,	UpperCamelCase__    :					List[str]=1							,	UpperCamelCase__    :					Optional[Any]=0							,	UpperCamelCase__    :					List[str]=4							,	)     ->						Any:
				_UpperCamelCase       			=parent
				_UpperCamelCase       			=batch_size
				_UpperCamelCase       			=seq_length
				_UpperCamelCase       			=is_training
				_UpperCamelCase       			=use_labels
				_UpperCamelCase       			=vocab_size
				_UpperCamelCase       			=hidden_size
				_UpperCamelCase       			=num_hidden_layers
				_UpperCamelCase       			=num_attention_heads
				_UpperCamelCase       			=intermediate_size

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


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

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

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




		def   UpperCamelCase__      (				self    :					str							,	UpperCamelCase__    :					List[Any]							,	UpperCamelCase__    :					List[Any]  )     ->						Any:
				_UpperCamelCase       			=TFLEDModel(config=UpperCamelCase__  ).get_decoder()
				_UpperCamelCase       			=inputs_dict['''input_ids''']

				_UpperCamelCase       			=input_ids[:1, :]
				_UpperCamelCase       			=inputs_dict['''attention_mask'''][:1, :]
				_UpperCamelCase       			=1

				# first forward pass
				_UpperCamelCase       			=model(UpperCamelCase__							,	attention_mask=UpperCamelCase__							,	use_cache=UpperCamelCase__  )

				_UpperCamelCase			,     _UpperCamelCase       			=outputs.to_tuple()

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

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

				_UpperCamelCase       			=model(UpperCamelCase__							,	attention_mask=UpperCamelCase__  )[0]
				_UpperCamelCase       			=model(UpperCamelCase__							,	attention_mask=UpperCamelCase__							,	past_key_values=UpperCamelCase__  )[0]

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

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

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




def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=None  ,      __SCREAMING_SNAKE_CASE=None  ,      ):


		"""simple docstring"""

		if attention_mask is None:
				_UpperCamelCase       			=tf.cast(tf.math.not_equal(__SCREAMING_SNAKE_CASE  ,      config.pad_token_id   )  ,      tf.inta   )
		if decoder_attention_mask is None:
				_UpperCamelCase       			=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:
				_UpperCamelCase       			=tf.ones((config.encoder_layers, config.encoder_attention_heads)   )
		if decoder_head_mask is None:
				_UpperCamelCase       			=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"""







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


		def   UpperCamelCase__      (				self    :					List[str]  )     ->						Dict:
				_UpperCamelCase       			=TFLEDModelTester(self  )
				_UpperCamelCase       			=ConfigTester(self							,	config_class=UpperCamelCase__  )


		def   UpperCamelCase__      (				self    :					str  )     ->						List[Any]:
				self.config_tester.run_common_tests()


		def   UpperCamelCase__      (				self    :					List[Any]  )     ->						List[Any]:
				_UpperCamelCase       			=self.model_tester.prepare_config_and_inputs_for_common()
				self.model_tester.check_decoder_model_past_large_inputs(*UpperCamelCase__  )


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

				_UpperCamelCase       			=True
				_UpperCamelCase       			=self.model_tester.seq_length
				_UpperCamelCase       			=self.model_tester.encoder_seq_length

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

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

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

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

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


		@unittest.skip('''LED keeps using potentially symbolic tensors in conditionals and breaks tracing.'''  )
		def   UpperCamelCase__      (				self    :					int  )     ->						Optional[Any]:
				pass




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




def 					_a   (__SCREAMING_SNAKE_CASE   ):


		"""simple docstring"""

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


__lowerCamelCase						: Union[str, Any]		=   1E-4







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


		"""simple docstring"""







		def   UpperCamelCase__      (				self    :					Dict  )     ->						Optional[int]:
				_UpperCamelCase       			=TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''  ).led

				# change to intended input here
				_UpperCamelCase       			=_long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]  )
				_UpperCamelCase       			=_long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]  )
				_UpperCamelCase       			=prepare_led_inputs_dict(model.config							,	UpperCamelCase__							,	UpperCamelCase__  )
				_UpperCamelCase       			=model(**UpperCamelCase__  )[0]
				_UpperCamelCase       			=(1, 1024, 768)
				self.assertEqual(output.shape							,	UpperCamelCase__  )
				# change to expected output here
				_UpperCamelCase       			=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]							,	UpperCamelCase__							,	atol=1E-3  )




		def   UpperCamelCase__      (				self    :					List[Any]  )     ->						Union[str, Any]:
				_UpperCamelCase       			=TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''  )

				# change to intended input here
				_UpperCamelCase       			=_long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]  )
				_UpperCamelCase       			=_long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]  )
				_UpperCamelCase       			=prepare_led_inputs_dict(model.config							,	UpperCamelCase__							,	UpperCamelCase__  )
				_UpperCamelCase       			=model(**UpperCamelCase__  )[0]
				_UpperCamelCase       			=(1, 1024, model.config.vocab_size)
				self.assertEqual(output.shape							,	UpperCamelCase__  )
				# change to expected output here
				_UpperCamelCase       			=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]							,	UpperCamelCase__							,	atol=1E-3							,	rtol=1E-3  )













'''simple docstring'''






import contextlib
from multiprocessing import Pool, RLock

from tqdm.auto import tqdm

from ..utils import experimental, logging


__lowerCamelCase						: Union[str, Any]		=   logging.get_logger(__name__)







class 		UpperCAmelCase   :


  """simple docstring"""







  lowerCAmelCase_					=				None




@experimental
def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  if ParallelBackendConfig.backend_name is None:
    return _map_with_multiprocessing_pool(
        __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )

  return _map_with_joblib(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=num_proc if num_proc <= len(__SCREAMING_SNAKE_CASE   ) else len(__SCREAMING_SNAKE_CASE   )
  _UpperCamelCase       			=[]  # We organize the splits ourselve (contiguous splits)
  for index in range(__SCREAMING_SNAKE_CASE   ):
    _UpperCamelCase       			=len(__SCREAMING_SNAKE_CASE   ) // num_proc
    _UpperCamelCase       			=len(__SCREAMING_SNAKE_CASE   ) % num_proc
    _UpperCamelCase       			=div * index + min(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
    _UpperCamelCase       			=start + div + (1 if index < mod else 0)
    split_kwds.append((function, iterable[start:end], types, index, disable_tqdm, desc)   )

  if len(__SCREAMING_SNAKE_CASE   ) != sum(len(i[1]   ) for i in split_kwds   ):
    raise ValueError(
        f'''Error dividing inputs iterable among processes. '''
        f'''Total number of objects {len(__SCREAMING_SNAKE_CASE   )}, '''
        f'''length: {sum(len(i[1]   ) for i in split_kwds   )}'''   )

  logger.info(
      f'''Spawning {num_proc} processes for {len(__SCREAMING_SNAKE_CASE   )} objects in slices of {[len(i[1]   ) for i in split_kwds]}'''   )
  _UpperCamelCase			,     _UpperCamelCase       			=None, None
  if not disable_tqdm:
    _UpperCamelCase			,     _UpperCamelCase       			=(RLock(),), tqdm.set_lock
  with Pool(__SCREAMING_SNAKE_CASE  ,      initargs=__SCREAMING_SNAKE_CASE  ,      initializer=__SCREAMING_SNAKE_CASE   ) as pool:
    _UpperCamelCase       			=pool.map(__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   )
  logger.info(f'''Finished {num_proc} processes'''   )
  _UpperCamelCase       			=[obj for proc_res in mapped for obj in proc_res]
  logger.info(f'''Unpacked {len(__SCREAMING_SNAKE_CASE   )} objects'''   )

  return mapped


def 					_a   (__SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE  ,      __SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  import joblib

  with joblib.parallel_backend(ParallelBackendConfig.backend_name  ,      n_jobs=__SCREAMING_SNAKE_CASE   ):
    return joblib.Parallel()(
        joblib.delayed(__SCREAMING_SNAKE_CASE   )((function, obj, types, None, True, None)   ) for obj in iterable   )


@experimental
@contextlib.contextmanager
def 					_a   (__SCREAMING_SNAKE_CASE   ):


  """simple docstring"""

  _UpperCamelCase       			=backend_name

  if backend_name == "spark":
    from joblibspark import register_spark

    register_spark()

    # TODO: call create_cache_and_write_probe if "download" in steps
    # TODO: raise NotImplementedError when Dataset.map etc is called

  try:
    yield
  finally:
    _UpperCamelCase       			=None






