





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[Any]     ,	UpperCAmelCase    : str     ,	UpperCAmelCase    : Tuple=13     ,	UpperCAmelCase    : List[str]=7     ,	UpperCAmelCase    : Dict=True     ,	UpperCAmelCase    : Optional[Any]=False     ,	UpperCAmelCase    : Tuple=99     ,	UpperCAmelCase    : str=32     ,	UpperCAmelCase    : List[str]=2     ,	UpperCAmelCase    : List[Any]=4     ,	UpperCAmelCase    : List[str]=37     ,	UpperCAmelCase    : Union[str, Any]=0.1     ,	UpperCAmelCase    : List[str]=0.1     ,	UpperCAmelCase    : List[Any]=20     ,	UpperCAmelCase    : Dict=2     ,	UpperCAmelCase    : Optional[int]=1     ,	UpperCAmelCase    : Optional[int]=0     ,	UpperCAmelCase    : int=4     ,	)     ->      Tuple:
     lowerCamelCase__       : Any   					=						parent
     lowerCamelCase__       : Optional[Any]   					=						batch_size
     lowerCamelCase__       : int   					=						seq_length
     lowerCamelCase__       : int   					=						is_training
     lowerCamelCase__       : Optional[int]   					=						use_labels
     lowerCamelCase__       : Any   					=						vocab_size
     lowerCamelCase__       : Optional[Any]   					=						hidden_size
     lowerCamelCase__       : List[Any]   					=						num_hidden_layers
     lowerCamelCase__       : List[str]   					=						num_attention_heads
     lowerCamelCase__       : Tuple   					=						intermediate_size

     lowerCamelCase__       : int   					=						hidden_dropout_prob
     lowerCamelCase__       : Optional[int]   					=						attention_probs_dropout_prob
     lowerCamelCase__       : Union[str, Any]   					=						max_position_embeddings
     lowerCamelCase__       : Dict   					=						eos_token_id
     lowerCamelCase__       : List[Any]   					=						pad_token_id
     lowerCamelCase__       : Union[str, Any]   					=						bos_token_id
     lowerCamelCase__       : Optional[int]   					=						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__       : List[str]   					=						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__       : List[str]   					=						(
         self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
     )
 def 	A_      ( self    : Union[str, Any]      )     ->      Dict:
     lowerCamelCase__       : Any   					=						ids_tensor([self.batch_size, self.seq_length - 1]     ,	self.vocab_size      )
     lowerCamelCase__       : int   					=						tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size      )     ,	1      )
     lowerCamelCase__       : Any   					=						tf.concat([input_ids, eos_tensor]     ,	axis=1      )

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

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







 def 	A_      ( self    : Tuple     ,	UpperCAmelCase    : Tuple     ,	UpperCAmelCase    : List[Any]      )     ->      List[Any]:
     lowerCamelCase__       : Optional[Any]   					=						TFLEDModel(config=UpperCAmelCase      ).get_decoder()
     lowerCamelCase__       : Dict   					=						inputs_dict['input_ids']

     lowerCamelCase__       : Optional[int]   					=						input_ids[:1, :]
     lowerCamelCase__       : str   					=						inputs_dict['attention_mask'][:1, :]
     lowerCamelCase__       : Tuple   					=						1

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

     lowerCamelCase__					,				lowerCamelCase__       : List[str]   					=						outputs.to_tuple()

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

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

     lowerCamelCase__       : Any   					=						model(UpperCAmelCase     ,	attention_mask=UpperCAmelCase      )[0]
     lowerCamelCase__       : Tuple   					=						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
     lowerCamelCase__       : Dict   					=						int(ids_tensor((1,)     ,	output_from_past.shape[-1]      )      )
     lowerCamelCase__       : Union[str, Any]   					=						output_from_no_past[:, -3:, random_slice_idx]
     lowerCamelCase__       : List[Any]   					=						output_from_past[:, :, random_slice_idx]

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


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase		,      _UpperCAmelCase=None		,      _UpperCAmelCase=None		,      _UpperCAmelCase=None		,      _UpperCAmelCase=None		,      )     ->	List[str]:
    if attention_mask is None:
        lowerCamelCase__       : Dict   					=						tf.cast(tf.math.not_equal(_UpperCAmelCase		,      config.pad_token_id					)		,      tf.inta					)
    if decoder_attention_mask is None:
        lowerCamelCase__       : Optional[int]   					=						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__       : Optional[Any]   					=						tf.ones((config.encoder_layers, config.encoder_attention_heads)					)
    if decoder_head_mask is None:
        lowerCamelCase__       : Dict   					=						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				(   __UpperCamelCase,    __UpperCamelCase,    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 	A_      ( self    : int      )     ->      Tuple:
     lowerCamelCase__       : Any   					=						TFLEDModelTester(self      )
     lowerCamelCase__       : Tuple   					=						ConfigTester(self     ,	config_class=UpperCAmelCase      )
 def 	A_      ( self    : Tuple      )     ->      Optional[Any]:
     self.config_tester.run_common_tests()
 def 	A_      ( self    : Tuple      )     ->      List[Any]:
     lowerCamelCase__       : int   					=						self.model_tester.prepare_config_and_inputs_for_common()
     self.model_tester.check_decoder_model_past_large_inputs(*UpperCAmelCase      )
 def 	A_      ( self    : Union[str, Any]      )     ->      int:
     lowerCamelCase__					,				lowerCamelCase__       : str   					=						self.model_tester.prepare_config_and_inputs_for_common()
     lowerCamelCase__       : Optional[int]   					=						tf.zeros_like(inputs_dict['attention_mask']      )
     lowerCamelCase__       : Optional[int]   					=						2
     lowerCamelCase__       : Optional[int]   					=						tf.where(
         tf.range(self.model_tester.seq_length      )[None, :] < num_global_attn_indices     ,	1     ,	inputs_dict['global_attention_mask']     ,	)

     lowerCamelCase__       : Optional[int]   					=						True
     lowerCamelCase__       : List[str]   					=						self.model_tester.seq_length
     lowerCamelCase__       : Optional[Any]   					=						self.model_tester.encoder_seq_length

     def check_decoder_attentions_output(UpperCAmelCase    : int      ):
         lowerCamelCase__       : Optional[Any]   					=						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]      ):
         lowerCamelCase__       : Tuple   					=						[t.numpy() for t in outputs.encoder_attentions]
         lowerCamelCase__       : str   					=						[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:
         lowerCamelCase__       : List[str]   					=						True
         lowerCamelCase__       : Any   					=						False
         lowerCamelCase__       : List[Any]   					=						False
         lowerCamelCase__       : Union[str, Any]   					=						model_class(UpperCAmelCase      )
         lowerCamelCase__       : int   					=						model(self._prepare_for_class(UpperCAmelCase     ,	UpperCAmelCase      )      )
         lowerCamelCase__       : Optional[Any]   					=						len(UpperCAmelCase      )
         self.assertEqual(config.output_hidden_states     ,	UpperCAmelCase      )
         check_encoder_attentions_output(UpperCAmelCase      )

         if self.is_encoder_decoder:
             lowerCamelCase__       : Any   					=						model_class(UpperCAmelCase      )
             lowerCamelCase__       : Optional[Any]   					=						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"]
         lowerCamelCase__       : int   					=						True
         lowerCamelCase__       : List[str]   					=						model_class(UpperCAmelCase      )
         lowerCamelCase__       : Optional[int]   					=						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
         lowerCamelCase__       : Optional[Any]   					=						True
         lowerCamelCase__       : Tuple   					=						True
         lowerCamelCase__       : Optional[int]   					=						model_class(UpperCAmelCase      )
         lowerCamelCase__       : str   					=						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 	A_      ( self    : str      )     ->      List[str]:
     pass







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


def 			SCREAMING_SNAKE_CASE (			_UpperCAmelCase					)     ->	List[str]:
    return tf.constant(_UpperCAmelCase		,      dtype=tf.intaa					)


_UpperCAmelCase							:			Dict        =							1E-4


@slow
@require_tf
class   lowerCAmelCase				(   unittest.TestCase     ):
 def 	A_      ( self    : str      )     ->      Union[str, Any]:
     lowerCamelCase__       : Any   					=						TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'      ).led

     # change to intended input here
     lowerCamelCase__       : Tuple   					=						_long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]      )
     lowerCamelCase__       : Union[str, Any]   					=						_long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]      )
     lowerCamelCase__       : Optional[int]   					=						prepare_led_inputs_dict(model.config     ,	UpperCAmelCase     ,	UpperCAmelCase      )
     lowerCamelCase__       : Dict   					=						model(**UpperCAmelCase      )[0]
     lowerCamelCase__       : List[Any]   					=						(1, 1024, 768)
     self.assertEqual(output.shape     ,	UpperCAmelCase      )
     # change to expected output here
     lowerCamelCase__       : Optional[int]   					=						tf.convert_to_tensor(
         [[2.3_0_5_0, 2.8_2_7_9, 0.6_5_3_1], [-1.8_4_5_7, -0.1_4_5_5, -3.5_6_6_1], [-1.0_1_8_6, 0.4_5_8_6, -2.2_0_4_3]]     ,	)
     tf.debugging.assert_near(output[:, :3, :3]     ,	UpperCAmelCase     ,	atol=1e-3      )







 def 	A_      ( self    : Tuple      )     ->      Dict:
     lowerCamelCase__       : Union[str, Any]   					=						TFLEDForConditionalGeneration.from_pretrained('allenai/led-base-16384'      )

     # change to intended input here
     lowerCamelCase__       : int   					=						_long_tensor([512 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]      )
     lowerCamelCase__       : Any   					=						_long_tensor([128 * [0, 31414, 232, 328, 740, 1140, 12695, 69]]      )
     lowerCamelCase__       : int   					=						prepare_led_inputs_dict(model.config     ,	UpperCAmelCase     ,	UpperCAmelCase      )
     lowerCamelCase__       : Dict   					=						model(**UpperCAmelCase      )[0]
     lowerCamelCase__       : List[Any]   					=						(1, 1024, model.config.vocab_size)
     self.assertEqual(output.shape     ,	UpperCAmelCase      )
     # change to expected output here
     lowerCamelCase__       : int   					=						tf.convert_to_tensor(
         [[3_3.6_5_0_7, 6.4_5_7_2, 1_6.8_0_8_9], [5.8_7_3_9, -2.4_2_3_8, 1_1.2_9_0_2], [-3.2_1_3_9, -4.3_1_4_9, 4.2_7_8_3]]     ,	)
     tf.debugging.assert_near(output[:, :3, :3]     ,	UpperCAmelCase     ,	atol=1e-3     ,	rtol=1e-3      )












import itertools
from dataclasses import dataclass
from typing import List, Optional

import pyarrow as pa
import pyarrow.parquet as pq

import datasets
from datasets.table import table_cast


_UpperCAmelCase							:			Optional[Any]        =							datasets.utils.logging.get_logger(__name__)

@dataclass
class   lowerCAmelCase				(   datasets.BuilderConfig     ):
 UpperCAmelCase__						=						1_00_00
 UpperCAmelCase__						=						None
 UpperCAmelCase__						=						None


class   lowerCAmelCase				(   datasets.ArrowBasedBuilder     ):
 UpperCAmelCase__						=						ParquetConfig
 def 	A_      ( self    : int      )     ->      Optional[int]:
     return datasets.DatasetInfo(features=self.config.features      )
 def 	A_      ( self    : int     ,	UpperCAmelCase    : str      )     ->      Any:
     if not self.config.data_files:
         raise ValueError(F"""At least one data file must be specified, but got data_files={self.config.data_files}"""      )
     lowerCamelCase__       : List[Any]   					=						dl_manager.download_and_extract(self.config.data_files      )
     if isinstance(UpperCAmelCase     ,	(str, list, tuple)      ):
         lowerCamelCase__       : int   					=						data_files
         if isinstance(UpperCAmelCase     ,	UpperCAmelCase      ):
             lowerCamelCase__       : int   					=						[files]
         # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
         lowerCamelCase__       : str   					=						[dl_manager.iter_files(UpperCAmelCase      ) for file in files]
         return [datasets.SplitGenerator(name=datasets.Split.TRAIN     ,	gen_kwargs={'files': files}      )]
     lowerCamelCase__       : Optional[Any]   					=						[]
     for split_name, files in data_files.items():
         if isinstance(UpperCAmelCase     ,	UpperCAmelCase      ):
             lowerCamelCase__       : Optional[Any]   					=						[files]
         # Use `dl_manager.iter_files` to skip hidden files in an extracted archive
         lowerCamelCase__       : Dict   					=						[dl_manager.iter_files(UpperCAmelCase      ) for file in files]
         # Infer features is they are stoed in the arrow schema
         if self.info.features is None:
             for file in itertools.chain.from_iterable(UpperCAmelCase      ):
                 with open(UpperCAmelCase     ,	'rb'      ) as f:
                     lowerCamelCase__       : Optional[Any]   					=						datasets.Features.from_arrow_schema(pq.read_schema(UpperCAmelCase      )      )
                 break
         splits.append(datasets.SplitGenerator(name=UpperCAmelCase     ,	gen_kwargs={'files': files}      )      )
     return splits
 def 	A_      ( self    : Optional[int]     ,	UpperCAmelCase    : pa.Table      )     ->      pa.Table:
     if self.info.features is not None:
         # more expensive cast to support nested features with keys in a different order
         # allows str <-> int/float or str to Audio for example
         lowerCamelCase__       : Any   					=						table_cast(UpperCAmelCase     ,	self.info.features.arrow_schema      )
     return pa_table







 def 	A_      ( self    : Any     ,	UpperCAmelCase    : Tuple      )     ->      List[str]:
     lowerCamelCase__       : Optional[Any]   					=						self.info.features.arrow_schema if self.info.features is not None else None
     if self.info.features is not None and self.config.columns is not None:
         if sorted(field.name for field in schema      ) != sorted(self.config.columns      ):
             raise ValueError(
                 F"""Tried to load parquet data with columns '{self.config.columns}' with mismatching features '{self.info.features}'"""      )
     for file_idx, file in enumerate(itertools.chain.from_iterable(UpperCAmelCase      )      ):
         with open(UpperCAmelCase     ,	'rb'      ) as f:
             lowerCamelCase__       : List[Any]   					=						pq.ParquetFile(UpperCAmelCase      )
             try:
                 for batch_idx, record_batch in enumerate(
                     parquet_file.iter_batches(batch_size=self.config.batch_size     ,	columns=self.config.columns      )      ):
                     lowerCamelCase__       : Tuple   					=						pa.Table.from_batches([record_batch]      )
                     # Uncomment for debugging (will print the Arrow table size and elements)
                     # logger.warning(f"pa_table: {pa_table} num rows: {pa_table.num_rows}")
                     # logger.warning('\n'.join(str(pa_table.slice(i, 1).to_pydict()) for i in range(pa_table.num_rows)))
                     yield F"""{file_idx}_{batch_idx}""", self._cast_table(UpperCAmelCase      )
             except ValueError as e:
                 logger.error(F"""Failed to read file '{file}' with error {type(UpperCAmelCase      )}: {e}"""      )
                 raise





