


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      a    :


     """simple docstring"""






     SCREAMING_SNAKE_CASE	:			Union[str, Any]							       =       LEDConfig
     SCREAMING_SNAKE_CASE	:			Optional[Any]							       =       {}
     SCREAMING_SNAKE_CASE	:			str							       =       '''gelu'''

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

          __UpperCAmelCase     :			str  							=      hidden_dropout_prob
          __UpperCAmelCase     :			Any  							=      attention_probs_dropout_prob
          __UpperCAmelCase     :			Tuple  							=      max_position_embeddings
          __UpperCAmelCase     :			Optional[int]  							=      eos_token_id
          __UpperCAmelCase     :			Any  							=      pad_token_id
          __UpperCAmelCase     :			List[str]  							=      bos_token_id
          __UpperCAmelCase     :			List[str]  							=      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     :			Optional[int]  							=      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     :			Dict  							=      (
              self.seq_length + (self.attention_window - self.seq_length % self.attention_window) % self.attention_window
          )

     def 			lowerCamelCase__			(    self					:		Optional[Any]      )						->							str:
          __UpperCAmelCase     :			Optional[Any]  							=      ids_tensor([self.batch_size, self.seq_length - 1]							,			self.vocab_size      )
          __UpperCAmelCase     :			Dict  							=      tf.expand_dims(tf.constant([self.eos_token_id] * self.batch_size      )							,			1      )
          __UpperCAmelCase     :			List[str]  							=      tf.concat([input_ids, eos_tensor]							,			axis=1      )

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

          __UpperCAmelCase     :			Any  							=      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     :			Dict  							=      prepare_led_inputs_dict(__snake_case							,			__snake_case							,			__snake_case      )
          __UpperCAmelCase     :			str  							=      tf.concat(
              [tf.zeros_like(__snake_case      )[:, :-1], tf.ones_like(__snake_case      )[:, -1:]]							,			axis=-1							,			)
          __UpperCAmelCase     :			Optional[int]  							=      global_attention_mask
          return config, inputs_dict






     def 			lowerCamelCase__			(    self					:		Optional[int]							,			snake_case					:		List[Any]							,			snake_case					:		Tuple      )						->							Dict:
          __UpperCAmelCase     :			int  							=      TFLEDModel(config=__snake_case      ).get_decoder()
          __UpperCAmelCase     :			Union[str, Any]  							=      inputs_dict['''input_ids''']

          __UpperCAmelCase     :			List[str]  							=      input_ids[:1, :]
          __UpperCAmelCase     :			Dict  							=      inputs_dict['''attention_mask'''][:1, :]
          __UpperCAmelCase     :			Union[str, Any]  							=      1

          # first forward pass
          __UpperCAmelCase     :			Any  							=      model(__snake_case							,			attention_mask=__snake_case							,			use_cache=__snake_case      )

          __UpperCAmelCase							,       __UpperCAmelCase     :			List[Any]  							=      outputs.to_tuple()

          # create hypothetical next token and extent to next_input_ids
          __UpperCAmelCase     :			Optional[Any]  							=      ids_tensor((self.batch_size, 3)							,			config.vocab_size      )
          __UpperCAmelCase     :			Dict  							=      tf.cast(ids_tensor((self.batch_size, 3)							,			2      )							,			tf.inta      )

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

          __UpperCAmelCase     :			Optional[Any]  							=      model(__snake_case							,			attention_mask=__snake_case      )[0]
          __UpperCAmelCase     :			Dict  							=      model(__snake_case							,			attention_mask=__snake_case							,			past_key_values=__snake_case      )[0]

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

          # select random slice
          __UpperCAmelCase     :			str  							=      int(ids_tensor((1,)							,			output_from_past.shape[-1]      )      )
          __UpperCAmelCase     :			Dict  							=      output_from_no_past[:, -3:, random_slice_idx]
          __UpperCAmelCase     :			Optional[Any]  							=      output_from_past[:, :, random_slice_idx]

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




def      _a     (						_lowercase :     int      , _lowercase :     Optional[Any]      , _lowercase :     str      , _lowercase :     str=None      , _lowercase :     List[str]=None      , _lowercase :     List[Any]=None      , _lowercase :     str=None      , ):





     '''simple docstring'''






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


     """simple docstring"""






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

     def 			lowerCamelCase__			(    self					:		str      )						->							Union[str, Any]:
          __UpperCAmelCase     :			int  							=      TFLEDModelTester(self      )
          __UpperCAmelCase     :			List[Any]  							=      ConfigTester(self							,			config_class=__snake_case      )

     def 			lowerCamelCase__			(    self					:		int      )						->							List[str]:
          self.config_tester.run_common_tests()

     def 			lowerCamelCase__			(    self					:		Optional[Any]      )						->							Any:
          __UpperCAmelCase     :			Any  							=      self.model_tester.prepare_config_and_inputs_for_common()
          self.model_tester.check_decoder_model_past_large_inputs(*__snake_case      )

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

          __UpperCAmelCase     :			Union[str, Any]  							=      True
          __UpperCAmelCase     :			List[str]  							=      self.model_tester.seq_length
          __UpperCAmelCase     :			Tuple  							=      self.model_tester.encoder_seq_length

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

               if self.is_encoder_decoder:
                    __UpperCAmelCase     :			Optional[Any]  							=      model_class(__snake_case      )
                    __UpperCAmelCase     :			Optional[int]  							=      model(self._prepare_for_class(__snake_case							,			__snake_case      )      )
                    self.assertEqual(config.output_hidden_states							,			__snake_case      )
                    check_decoder_attentions_output(__snake_case      )

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

               # Check attention is always last and order is fine
               __UpperCAmelCase     :			Optional[int]  							=      True
               __UpperCAmelCase     :			Any  							=      True
               __UpperCAmelCase     :			Tuple  							=      model_class(__snake_case      )
               __UpperCAmelCase     :			Any  							=      model(self._prepare_for_class(__snake_case							,			__snake_case      )      )

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

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






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




def      _a     (						_lowercase :     Dict  ):





     '''simple docstring'''






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


__UpperCAmelCase      :int			    =       1e-4



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


     """simple docstring"""






     def 			lowerCamelCase__			(    self					:		Tuple      )						->							Optional[Any]:
          __UpperCAmelCase     :			str  							=      TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''      ).led

          # change to intended input here
          __UpperCAmelCase     :			int  							=      _long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
          __UpperCAmelCase     :			Dict  							=      _long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
          __UpperCAmelCase     :			Any  							=      prepare_led_inputs_dict(model.config							,			__snake_case							,			__snake_case      )
          __UpperCAmelCase     :			List[Any]  							=      model(**__snake_case      )[0]
          __UpperCAmelCase     :			Optional[Any]  							=      (1, 1024, 768)
          self.assertEqual(output.shape							,			__snake_case      )
          # change to expected output here
          __UpperCAmelCase     :			Dict  							=      tf.convert_to_tensor(
              [[2.3_050, 2.8_279, 0.6_531], [-1.8_457, -0.1_455, -3.5_661], [-1.0_186, 0.4_586, -2.2_043]]							,			)
          tf.debugging.assert_near(output[:, :3, :3]							,			__snake_case							,			atol=1E-3      )






     def 			lowerCamelCase__			(    self					:		int      )						->							List[str]:
          __UpperCAmelCase     :			Any  							=      TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''      )

          # change to intended input here
          __UpperCAmelCase     :			Optional[Any]  							=      _long_tensor([512 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
          __UpperCAmelCase     :			int  							=      _long_tensor([128 * [0, 3_1414, 232, 328, 740, 1140, 1_2695, 69]]      )
          __UpperCAmelCase     :			List[str]  							=      prepare_led_inputs_dict(model.config							,			__snake_case							,			__snake_case      )
          __UpperCAmelCase     :			Tuple  							=      model(**__snake_case      )[0]
          __UpperCAmelCase     :			List[str]  							=      (1, 1024, model.config.vocab_size)
          self.assertEqual(output.shape							,			__snake_case      )
          # change to expected output here
          __UpperCAmelCase     :			List[str]  							=      tf.convert_to_tensor(
              [[33.6_507, 6.4_572, 16.8_089], [5.8_739, -2.4_238, 11.2_902], [-3.2_139, -4.3_149, 4.2_783]]							,			)
          tf.debugging.assert_near(output[:, :3, :3]							,			__snake_case							,			atol=1E-3							,			rtol=1E-3      )



'''simple docstring'''



# Copyright 2021 The HuggingFace Team. All rights reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

from pathlib import Path

import torch

from ...utils import is_npu_available, is_xpu_available
from .config_args import ClusterConfig, default_json_config_file
from .config_utils import SubcommandHelpFormatter


__UpperCAmelCase      :Tuple			    =       "Create a default config file for Accelerate with only a few flags set."




def      _a     (						_lowercase :     List[Any]="no"      , _lowercase :     str = default_json_config_file      , _lowercase :     bool = False  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Dict  							=      Path(_lowercase  )
     path.parent.mkdir(parents=_lowercase      , exist_ok=_lowercase  )
     if path.exists():
          print(
              F'Configuration already exists at {save_location}, will not override. Run `accelerate config` manually or pass a different `save_location`.'  )
          return False
     __UpperCAmelCase     :			List[str]  							=      mixed_precision.lower()
     if mixed_precision not in ["no", "fp16", "bf16", "fp8"]:
          raise ValueError(
              F'`mixed_precision` should be one of \'no\', \'fp16\', \'bf16\', or \'fp8\'. Received {mixed_precision}'  )
     __UpperCAmelCase     :			int  							=      {
         '''compute_environment''': '''LOCAL_MACHINE''',
         '''mixed_precision''': mixed_precision,
     }
     if torch.cuda.is_available():
          __UpperCAmelCase     :			Optional[Any]  							=      torch.cuda.device_count()
          __UpperCAmelCase     :			List[str]  							=      num_gpus
          __UpperCAmelCase     :			int  							=      False
          if num_gpus > 1:
               __UpperCAmelCase     :			Any  							=      '''MULTI_GPU'''
          else:
               __UpperCAmelCase     :			int  							=      '''NO'''
     elif is_xpu_available() and use_xpu:
          __UpperCAmelCase     :			List[Any]  							=      torch.xpu.device_count()
          __UpperCAmelCase     :			List[Any]  							=      num_xpus
          __UpperCAmelCase     :			Optional[int]  							=      False
          if num_xpus > 1:
               __UpperCAmelCase     :			Any  							=      '''MULTI_XPU'''
          else:
               __UpperCAmelCase     :			Optional[Any]  							=      '''NO'''
     elif is_npu_available():
          __UpperCAmelCase     :			Dict  							=      torch.npu.device_count()
          __UpperCAmelCase     :			Any  							=      num_npus
          __UpperCAmelCase     :			Any  							=      False
          if num_npus > 1:
               __UpperCAmelCase     :			Dict  							=      '''MULTI_NPU'''
          else:
               __UpperCAmelCase     :			Optional[int]  							=      '''NO'''
     else:
          __UpperCAmelCase     :			List[str]  							=      0
          __UpperCAmelCase     :			Dict  							=      True
          __UpperCAmelCase     :			Dict  							=      1
          __UpperCAmelCase     :			Tuple  							=      '''NO'''
     __UpperCAmelCase     :			List[Any]  							=      ClusterConfig(**_lowercase  )
     config.to_json_file(_lowercase  )
     return path




def      _a     (						_lowercase :     Union[str, Any]      , _lowercase :     str  ):





     '''simple docstring'''






     __UpperCAmelCase     :			Optional[int]  							=      parser.add_parser('''default'''      , parents=_lowercase      , help=_lowercase      , formatter_class=_lowercase  )
     parser.add_argument(
         '''--config_file'''      , default=_lowercase      , help=(
             '''The path to use to store the config file. Will default to a file named default_config.yaml in the cache '''
             '''location, which is the content of the environment `HF_HOME` suffixed with \'accelerate\', or if you don\'t have '''
             '''such an environment variable, your cache directory (\'~/.cache\' or the content of `XDG_CACHE_HOME`) suffixed '''
             '''with \'huggingface\'.'''
         )      , dest='''save_location'''      , )

     parser.add_argument(
         '''--mixed_precision'''      , choices=['''no''', '''fp16''', '''bf16''']      , type=_lowercase      , help='''Whether or not to use mixed precision training. '''
         '''Choose between FP16 and BF16 (bfloat16) training. '''
         '''BF16 training is only supported on Nvidia Ampere GPUs and PyTorch 1.10 or later.'''      , default='''no'''      , )
     parser.set_defaults(func=_lowercase  )
     return parser




def      _a     (						_lowercase :     List[Any]  ):





     '''simple docstring'''






     __UpperCAmelCase     :			str  							=      write_basic_config(args.mixed_precision      , args.save_location  )
     if config_file:
          print(F'accelerate configuration saved at {config_file}'  )