

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

import argparse
import os

from accelerate.utils import ComputeEnvironment

from .cluster import get_cluster_input
from .config_args import cache_dir, default_config_file, default_yaml_config_file, load_config_from_file  # noqa: F401
from .config_utils import _ask_field, _ask_options, _convert_compute_environment  # noqa: F401
from .sagemaker import get_sagemaker_input


lowerCAmelCase__											=	'''Launches a series of prompts to create and save a `default_config.yaml` configuration file for your training system. Should always be ran first on your machine'''


def 		_A     (    ):






 """simple docstring"""



 __lowercase							  = _ask_options(
     '''In which compute environment are you running?'''     ,				['''This machine''', '''AWS (Amazon SageMaker)''']     ,				_convert_compute_environment     ,				)
 if compute_environment == ComputeEnvironment.AMAZON_SAGEMAKER:
  __lowercase							  = get_sagemaker_input()
 else:
  __lowercase							  = get_cluster_input()
 return config


def 		_A     (    A__=None      ):






 """simple docstring"""



 if subparsers is not None:
  __lowercase							  = subparsers.add_parser('''config'''     ,				description=A__      )
 else:
  __lowercase							  = argparse.ArgumentParser('''Accelerate config command'''     ,				description=A__      )

 parser.add_argument(
     '''--config_file'''     ,				default=A__     ,				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\'.'''
     )     ,				)

 if subparsers is not None:
  parser.set_defaults(func=A__      )
 return parser


def 		_A     (    A__      ):






 """simple docstring"""



 __lowercase							  = get_user_input()
 if args.config_file is not None:
  __lowercase							  = args.config_file
 else:
  if not os.path.isdir(A__      ):
   os.makedirs(A__      )
  __lowercase							  = default_yaml_config_file

 if config_file.endswith('''.json'''      ):
  config.to_json_file(A__      )
 else:
  config.to_yaml_file(A__      )
 print(F"accelerate configuration saved at {config_file}"      )


def 		_A     (    ):






 """simple docstring"""



 __lowercase							  = config_command_parser()
 __lowercase							  = parser.parse_args()
 config_command(A__      )


if __name__ == "__main__":
 main()



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__  :
   _SCREAMING_SNAKE_CASE      :							Union[str, Any]	       =      LEDConfig
   _SCREAMING_SNAKE_CASE      :							Optional[int]	       =      {}
   _SCREAMING_SNAKE_CASE      :							int	       =      "gelu"
   def __init__(self :					List[str]					,				snake_case_ :					Union[str, Any]					,				snake_case_ :					Union[str, Any]=1_3					,				snake_case_ :					Optional[Any]=7					,				snake_case_ :					Any=True					,				snake_case_ :					List[Any]=False					,				snake_case_ :					str=9_9					,				snake_case_ :					Any=3_2					,				snake_case_ :					Dict=2					,				snake_case_ :					List[Any]=4					,				snake_case_ :					Optional[int]=3_7					,				snake_case_ :					Dict=0.1					,				snake_case_ :					int=0.1					,				snake_case_ :					Optional[Any]=2_0					,				snake_case_ :					Optional[Any]=2					,				snake_case_ :					Optional[int]=1					,				snake_case_ :					Optional[int]=0					,				snake_case_ :					str=4					,				):
    __a :  List[Any]						=  parent
    __a :  Union[str, Any]						=  batch_size
    __a :  List[str]						=  seq_length
    __a :  Any						=  is_training
    __a :  Tuple						=  use_labels
    __a :  List[Any]						=  vocab_size
    __a :  Optional[Any]						=  hidden_size
    __a :  int						=  num_hidden_layers
    __a :  Optional[int]						=  num_attention_heads
    __a :  int						=  intermediate_size

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

    __a :  List[str]						=  ids_tensor([self.batch_size, self.seq_length]					,				self.vocab_size   )

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







   def       lowerCAmelCase			(self :					List[Any]					,				snake_case_ :					Dict					,				snake_case_ :					int   ):
    __a :  List[str]						=  TFLEDModel(config=snake_case_   ).get_decoder()
    __a :  Dict						=  inputs_dict['''input_ids''']

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

    # first forward pass
    __a :  Optional[Any]						=  model(snake_case_					,				attention_mask=snake_case_					,				use_cache=snake_case_   )

    __a				,       __a :  Tuple						=  outputs.to_tuple()

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

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

    __a :  Optional[int]						=  model(snake_case_					,				attention_mask=snake_case_   )[0]
    __a :  int						=  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
    __a :  Any						=  int(ids_tensor((1,)					,				output_from_past.shape[-1]   )   )
    __a :  List[Any]						=  output_from_no_past[:, -3:, random_slice_idx]
    __a :  Dict						=  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    __UpperCamelCase     (				lowerCAmelCase__							:				Union[str, Any]  ,   lowerCAmelCase__							:				List[Any]  ,   lowerCAmelCase__							:				int  ,   lowerCAmelCase__							:				List[Any]=None  ,   lowerCAmelCase__							:				int=None  ,   lowerCAmelCase__							:				Union[str, Any]=None  ,   lowerCAmelCase__							:				List[str]=None  ,   ):
 if attention_mask is None:
  __a :  Any						=  tf.cast(tf.math.not_equal(lowerCAmelCase__  ,   config.pad_token_id  )  ,   tf.inta  )
 if decoder_attention_mask is None:
  __a :  Dict						=  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 :  Dict						=  tf.ones((config.encoder_layers, config.encoder_attention_heads)  )
 if decoder_head_mask is None:
  __a :  str						=  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    ):
   _SCREAMING_SNAKE_CASE      :							Dict	       =      (TFLEDForConditionalGeneration, TFLEDModel) if is_tf_available() else ()
   _SCREAMING_SNAKE_CASE      :							Tuple	       =      (TFLEDForConditionalGeneration,) if is_tf_available() else ()
   _SCREAMING_SNAKE_CASE      :							Optional[Any]	       =      (
       {
           "conversational": TFLEDForConditionalGeneration,
           "feature-extraction": TFLEDModel,
           "summarization": TFLEDForConditionalGeneration,
           "text2text-generation": TFLEDForConditionalGeneration,
           "translation": TFLEDForConditionalGeneration,
       }
       if is_tf_available()
       else {}
   )
   _SCREAMING_SNAKE_CASE      :							Union[str, Any]	       =      True
   _SCREAMING_SNAKE_CASE      :							Optional[int]	       =      False
   _SCREAMING_SNAKE_CASE      :							Optional[Any]	       =      False
   _SCREAMING_SNAKE_CASE      :							Optional[Any]	       =      False
   def       lowerCAmelCase			(self :					Optional[int]   ):
    __a :  List[str]						=  TFLEDModelTester(self   )
    __a :  Optional[int]						=  ConfigTester(self					,				config_class=snake_case_   )
   def       lowerCAmelCase			(self :					Any   ):
    self.config_tester.run_common_tests()
   def       lowerCAmelCase			(self :					Optional[Any]   ):
    __a :  Union[str, Any]						=  self.model_tester.prepare_config_and_inputs_for_common()
    self.model_tester.check_decoder_model_past_large_inputs(*snake_case_   )
   def       lowerCAmelCase			(self :					Any   ):
    __a				,       __a :  Tuple						=  self.model_tester.prepare_config_and_inputs_for_common()
    __a :  Any						=  tf.zeros_like(inputs_dict['''attention_mask''']   )
    __a :  Tuple						=  2
    __a :  Dict						=  tf.where(
        tf.range(self.model_tester.seq_length   )[None, :] < num_global_attn_indices					,				1					,				inputs_dict['''global_attention_mask''']					,				)

    __a :  List[str]						=  True
    __a :  Tuple						=  self.model_tester.seq_length
    __a :  Any						=  self.model_tester.encoder_seq_length

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

     if self.is_encoder_decoder:
      __a :  List[str]						=  model_class(snake_case_   )
      __a :  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"]
     __a :  List[Any]						=  True
     __a :  Dict						=  model_class(snake_case_   )
     __a :  Tuple						=  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
     __a :  List[str]						=  True
     __a :  Any						=  True
     __a :  Tuple						=  model_class(snake_case_   )
     __a :  int						=  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 :					List[str]   ):
    pass







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





def    __UpperCamelCase     (				lowerCAmelCase__							:				Optional[Any]  ):
 return tf.constant(lowerCAmelCase__  ,   dtype=tf.intaa  )


lowercase__       =1e-4



@slow
@require_tf
class 		UpperCamelCase__  (       unittest.TestCase    ):
   def       lowerCAmelCase			(self :					Any   ):
    __a :  Dict						=  TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''   ).led

    # change to intended input here
    __a :  Union[str, Any]						=  _long_tensor([5_1_2 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
    __a :  Dict						=  _long_tensor([1_2_8 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
    __a :  List[str]						=  prepare_led_inputs_dict(model.config					,				snake_case_					,				snake_case_   )
    __a :  List[str]						=  model(**snake_case_   )[0]
    __a :  Any						=  (1, 1_0_2_4, 7_6_8)
    self.assertEqual(output.shape					,				snake_case_   )
    # change to expected output here
    __a :  Dict						=  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]					,				snake_case_					,				atol=1E-3   )







   def       lowerCAmelCase			(self :					int   ):
    __a :  Optional[Any]						=  TFLEDForConditionalGeneration.from_pretrained('''allenai/led-base-16384'''   )

    # change to intended input here
    __a :  int						=  _long_tensor([5_1_2 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
    __a :  Tuple						=  _long_tensor([1_2_8 * [0, 3_1_4_1_4, 2_3_2, 3_2_8, 7_4_0, 1_1_4_0, 1_2_6_9_5, 6_9]]   )
    __a :  Dict						=  prepare_led_inputs_dict(model.config					,				snake_case_					,				snake_case_   )
    __a :  List[str]						=  model(**snake_case_   )[0]
    __a :  List[Any]						=  (1, 1_0_2_4, model.config.vocab_size)
    self.assertEqual(output.shape					,				snake_case_   )
    # change to expected output here
    __a :  str						=  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]					,				snake_case_					,				atol=1E-3					,				rtol=1E-3   )
