



'''simple docstring'''






import os
import tempfile
from functools import partial
from unittest import TestCase
from unittest.mock import patch

import datasets
import datasets.config

from .utils import require_beam

class 				SCREAMING_SNAKE_CASE(							datasets.BeamBasedBuilder			):







	"""simple docstring"""






	def      A	(						self			:     Dict		)			->		Tuple:
								return datasets.DatasetInfo(
								    features=datasets.Features({'''content''': datasets.Value('''string'''		)}		)		,			supervised_keys=__snake_case		,			)







	def      A	(						self			:     Union[str, Any]		,			__snake_case			:     List[str]		,			__snake_case			:     str		)			->		Optional[Any]:
								return [datasets.SplitGenerator(name=datasets.Split.TRAIN		,			gen_kwargs={'''examples''': get_test_dummy_examples()}		)]







	def      A	(						self			:     Optional[Any]		,			__snake_case			:     Optional[Any]		,			__snake_case			:     Optional[Any]		)			->		Optional[Any]:
								import apache_beam as beam

								return pipeline | "Load Examples" >> beam.Create(__snake_case		)

class 				SCREAMING_SNAKE_CASE(							datasets.BeamBasedBuilder			):







	"""simple docstring"""






	def      A	(						self			:     Tuple		)			->		Union[str, Any]:
								return datasets.DatasetInfo(
								    features=datasets.Features({'''a''': datasets.Sequence({'''b''': datasets.Value('''string'''		)}		)}		)		,			supervised_keys=__snake_case		,			)







	def      A	(						self			:     Union[str, Any]		,			__snake_case			:     Optional[int]		,			__snake_case			:     Optional[Any]		)			->		Tuple:
								return [
								    datasets.SplitGenerator(name=datasets.Split.TRAIN		,			gen_kwargs={'''examples''': get_test_nested_examples()}		)
								]







	def      A	(						self			:     Tuple		,			__snake_case			:     Tuple		,			__snake_case			:     str		)			->		str:
								import apache_beam as beam

								return pipeline | "Load Examples" >> beam.Create(__snake_case		)


def  snake_case_     (					)      ->   Optional[Any]:
							return [(i, {"content": content}) for i, content in enumerate(['''foo''', '''bar''', '''foobar''']  )]


def  snake_case_     (					)      ->   Union[str, Any]:
							return [(i, {"a": {"b": [content]}}) for i, content in enumerate(['''foo''', '''bar''', '''foobar''']  )]




class 				SCREAMING_SNAKE_CASE(							A__			):







	"""simple docstring"""






	@require_beam
	def      A	(						self			:     Tuple		)			->		Dict:
								UpperCAmelCase :       Dict 							=		len(get_test_dummy_examples()		)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
															UpperCAmelCase :       str 							=		DummyBeamDataset(cache_dir=__snake_case		,			beam_runner='''DirectRunner'''		)
															builder.download_and_prepare()
															self.assertTrue(
															    os.path.exists(
															        os.path.join(__snake_case		,			builder.name		,			'''default'''		,			'''0.0.0'''		,			F"""{builder.name}-train.arrow"""		)		)		)
															self.assertDictEqual(builder.info.features		,			datasets.Features({'''content''': datasets.Value('''string'''		)}		)		)
															UpperCAmelCase :       str 							=		builder.as_dataset()
															self.assertEqual(dset['''train'''].num_rows		,			__snake_case		)
															self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples		,			__snake_case		)
															self.assertDictEqual(dset['''train'''][0]		,			get_test_dummy_examples()[0][1]		)
															self.assertDictEqual(
															    dset['''train'''][expected_num_examples - 1]		,			get_test_dummy_examples()[expected_num_examples - 1][1]		)
															self.assertTrue(
															    os.path.exists(os.path.join(__snake_case		,			builder.name		,			'''default'''		,			'''0.0.0'''		,			'''dataset_info.json'''		)		)		)
															del dset







	@require_beam
	def      A	(						self			:     List[Any]		)			->		List[Any]:
								import apache_beam as beam

								UpperCAmelCase :       Tuple 							=		beam.io.parquetio.WriteToParquet

								UpperCAmelCase :       List[str] 							=		len(get_test_dummy_examples()		)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
															UpperCAmelCase :       Any 							=		DummyBeamDataset(cache_dir=__snake_case		,			beam_runner='''DirectRunner'''		)
															with patch('''apache_beam.io.parquetio.WriteToParquet'''		) as write_parquet_mock:
																						UpperCAmelCase :       str 							=		partial(__snake_case		,			num_shards=2		)
																						builder.download_and_prepare()
															self.assertTrue(
															    os.path.exists(
															        os.path.join(
															            __snake_case		,			builder.name		,			'''default'''		,			'''0.0.0'''		,			F"""{builder.name}-train-00000-of-00002.arrow"""		)		)		)
															self.assertTrue(
															    os.path.exists(
															        os.path.join(
															            __snake_case		,			builder.name		,			'''default'''		,			'''0.0.0'''		,			F"""{builder.name}-train-00000-of-00002.arrow"""		)		)		)
															self.assertDictEqual(builder.info.features		,			datasets.Features({'''content''': datasets.Value('''string'''		)}		)		)
															UpperCAmelCase :       Any 							=		builder.as_dataset()
															self.assertEqual(dset['''train'''].num_rows		,			__snake_case		)
															self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples		,			__snake_case		)
															# Order is not preserved when sharding, so we just check that all the elements are there
															self.assertListEqual(sorted(dset['''train''']['''content''']		)		,			sorted(['''foo''', '''bar''', '''foobar''']		)		)
															self.assertTrue(
															    os.path.exists(os.path.join(__snake_case		,			builder.name		,			'''default'''		,			'''0.0.0'''		,			'''dataset_info.json'''		)		)		)
															del dset







	@require_beam
	def      A	(						self			:     Optional[Any]		)			->		List[str]:
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
															UpperCAmelCase :       Any 							=		DummyBeamDataset(cache_dir=__snake_case		)
															self.assertRaises(datasets.builder.MissingBeamOptions		,			builder.download_and_prepare		)







	@require_beam
	def      A	(						self			:     Optional[Any]		)			->		Optional[Any]:
								UpperCAmelCase :       List[str] 							=		len(get_test_nested_examples()		)
								with tempfile.TemporaryDirectory() as tmp_cache_dir:
															UpperCAmelCase :       List[str] 							=		NestedBeamDataset(cache_dir=__snake_case		,			beam_runner='''DirectRunner'''		)
															builder.download_and_prepare()
															self.assertTrue(
															    os.path.exists(
															        os.path.join(__snake_case		,			builder.name		,			'''default'''		,			'''0.0.0'''		,			F"""{builder.name}-train.arrow"""		)		)		)
															self.assertDictEqual(
															    builder.info.features		,			datasets.Features({'''a''': datasets.Sequence({'''b''': datasets.Value('''string'''		)}		)}		)		)
															UpperCAmelCase :       List[Any] 							=		builder.as_dataset()
															self.assertEqual(dset['''train'''].num_rows		,			__snake_case		)
															self.assertEqual(dset['''train'''].info.splits['''train'''].num_examples		,			__snake_case		)
															self.assertDictEqual(dset['''train'''][0]		,			get_test_nested_examples()[0][1]		)
															self.assertDictEqual(
															    dset['''train'''][expected_num_examples - 1]		,			get_test_nested_examples()[expected_num_examples - 1][1]		)
															self.assertTrue(
															    os.path.exists(os.path.join(__snake_case		,			builder.name		,			'''default'''		,			'''0.0.0'''		,			'''dataset_info.json'''		)		)		)
															del dset











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 		__lowerCAmelCase	:
							UpperCamelCase__					     =  LEDConfig
							UpperCamelCase__					     =  {}
							UpperCamelCase__					     =  '''gelu'''
							def __init__(						self						:Optional[int]		,      __magic_name__						:Dict		,      __magic_name__						:List[str]=13		,      __magic_name__						:Union[str, Any]=7		,      __magic_name__						:str=True		,      __magic_name__						:Union[str, Any]=False		,      __magic_name__						:Union[str, Any]=99		,      __magic_name__						:List[Any]=32		,      __magic_name__						:str=2		,      __magic_name__						:List[str]=4		,      __magic_name__						:str=37		,      __magic_name__						:Any=0.1		,      __magic_name__						:Dict=0.1		,      __magic_name__						:List[str]=20		,      __magic_name__						:Union[str, Any]=2		,      __magic_name__						:List[Any]=1		,      __magic_name__						:Optional[int]=0		,      __magic_name__						:Optional[int]=4		,      ):







								'''simple docstring'''







								a										=							parent
								a										=							batch_size
								a										=							seq_length
								a										=							is_training
								a										=							use_labels
								a										=							vocab_size
								a										=							hidden_size
								a										=							num_hidden_layers
								a										=							num_attention_heads
								a										=							intermediate_size

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







								'''simple docstring'''







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

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

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





							def       lowerCamelCase__							(						self						:Union[str, Any]		,      __magic_name__						:Optional[int]		,      __magic_name__						:List[Any]				):







								'''simple docstring'''







								a										=							TFLEDModel(config=__magic_name__				).get_decoder()
								a										=							inputs_dict["""input_ids"""]

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

								# first forward pass
								a										=							model(__magic_name__		,      attention_mask=__magic_name__		,      use_cache=__magic_name__				)

								a				, a										=							outputs.to_tuple()

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

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

								a										=							model(__magic_name__		,      attention_mask=__magic_name__				)[0]
								a										=							model(__magic_name__		,      attention_mask=__magic_name__		,      past_key_values=__magic_name__				)[0]

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

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

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







def 		__A		(	__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase							,			__lowerCamelCase=None							,			__lowerCamelCase=None							,			__lowerCamelCase=None							,			__lowerCamelCase=None							,			)      ->       List[str]:
	if attention_mask is None:
		a										=							tf.cast(tf.math.not_equal(__lowerCamelCase							,			config.pad_token_id    )							,			tf.inta    )
	if decoder_attention_mask is None:
		a										=							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										=							tf.ones((config.encoder_layers, config.encoder_attention_heads)    )
	if decoder_head_mask is None:
		a										=							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 		__lowerCAmelCase	(	__magic_name__      ,       __magic_name__      ,       unittest.TestCase ):
							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       lowerCamelCase__							(						self						:Tuple				):







								'''simple docstring'''







								a										=							TFLEDModelTester(self				)
								a										=							ConfigTester(self		,      config_class=__magic_name__				)
							def       lowerCamelCase__							(						self						:int				):







								'''simple docstring'''







								self.config_tester.run_common_tests()
							def       lowerCamelCase__							(						self						:str				):







								'''simple docstring'''







								a										=							self.model_tester.prepare_config_and_inputs_for_common()
								self.model_tester.check_decoder_model_past_large_inputs(*__magic_name__				)
							def       lowerCamelCase__							(						self						:Optional[int]				):







								'''simple docstring'''







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

								a										=							True
								a										=							self.model_tester.seq_length
								a										=							self.model_tester.encoder_seq_length

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

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

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

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

									self.assertEqual(out_len + (2 if self.is_encoder_decoder else 1)		,      len(__magic_name__				)				)
									self.assertEqual(model.config.output_hidden_states		,      __magic_name__				)
									check_encoder_attentions_output(__magic_name__				)
							@unittest.skip("""LED keeps using potentially symbolic tensors in conditionals and breaks tracing."""				)
							def       lowerCamelCase__							(						self						:Union[str, Any]				):







								'''simple docstring'''







								pass





							def       lowerCamelCase__							(						self						:int				):







								'''simple docstring'''







								pass







def 		__A		(	__lowerCamelCase    )      ->       int:
	return tf.constant(__lowerCamelCase							,			dtype=tf.intaa    )


__UpperCamelCase					:       int    					=      1E-4



@slow
@require_tf
class 		__lowerCAmelCase	(	unittest.TestCase ):
							def       lowerCamelCase__							(						self						:Union[str, Any]				):







								'''simple docstring'''







								a										=							TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""				).led

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





							def       lowerCamelCase__							(						self						:str				):







								'''simple docstring'''







								a										=							TFLEDForConditionalGeneration.from_pretrained("""allenai/led-base-16384"""				)

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

