



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       _SCREAMING_SNAKE_CASE	:
     lowerCAmelCase					=			LEDConfig
     lowerCAmelCase					=			{}
     lowerCAmelCase					=			"""gelu"""





     def __init__(       self :							Dict     ,  UpperCamelCase :							Dict     ,  UpperCamelCase :							Optional[int]=1_3     ,  UpperCamelCase :							Union[str, Any]=7     ,  UpperCamelCase :							List[Any]=True     ,  UpperCamelCase :							int=False     ,  UpperCamelCase :							Optional[int]=9_9     ,  UpperCamelCase :							Union[str, Any]=3_2     ,  UpperCamelCase :							Optional[Any]=2     ,  UpperCamelCase :							int=4     ,  UpperCamelCase :							Optional[int]=3_7     ,  UpperCamelCase :							Any=0.1     ,  UpperCamelCase :							Union[str, Any]=0.1     ,  UpperCamelCase :							int=2_0     ,  UpperCamelCase :							Any=2     ,  UpperCamelCase :							Optional[Any]=1     ,  UpperCamelCase :							Any=0     ,  UpperCamelCase :							int=4     ,  )->str:
      __SCREAMING_SNAKE_CASE       :    List[str]       =	parent
      __SCREAMING_SNAKE_CASE       :    str       =	batch_size
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	seq_length
      __SCREAMING_SNAKE_CASE       :    List[str]       =	is_training
      __SCREAMING_SNAKE_CASE       :    int       =	use_labels
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	vocab_size
      __SCREAMING_SNAKE_CASE       :    int       =	hidden_size
      __SCREAMING_SNAKE_CASE       :    Any       =	num_hidden_layers
      __SCREAMING_SNAKE_CASE       :    int       =	num_attention_heads
      __SCREAMING_SNAKE_CASE       :    Tuple       =	intermediate_size

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





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

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

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






     def 	__snake_case    (       self :							Optional[Any]     ,  UpperCamelCase :							int     ,  UpperCamelCase :							Optional[Any]				)->Tuple:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	TFLEDModel(config=UpperCamelCase				).get_decoder()
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	inputs_dict["input_ids"]

      __SCREAMING_SNAKE_CASE       :    Tuple       =	input_ids[:1, :]
      __SCREAMING_SNAKE_CASE       :    str       =	inputs_dict["attention_mask"][:1, :]
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	1

      # first forward pass
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	model(UpperCamelCase     ,  attention_mask=UpperCamelCase     ,  use_cache=UpperCamelCase				)

      __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	outputs.to_tuple()

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

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

      __SCREAMING_SNAKE_CASE       :    List[Any]       =	model(UpperCamelCase     ,  attention_mask=UpperCamelCase				)[0]
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	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
      __SCREAMING_SNAKE_CASE       :    Tuple       =	int(ids_tensor((1,)     ,  output_from_past.shape[-1]				)				)
      __SCREAMING_SNAKE_CASE       :    Tuple       =	output_from_no_past[:, -3:, random_slice_idx]
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	output_from_past[:, :, random_slice_idx]

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






def 	_lowerCAmelCase						(				__lowerCamelCase   :     List[Any]		,	__lowerCamelCase   :     Tuple		,	__lowerCamelCase   :     Optional[int]		,	__lowerCamelCase   :     Optional[int]=None		,	__lowerCamelCase   :     Dict=None		,	__lowerCamelCase   :     Optional[Any]=None		,	__lowerCamelCase   :     Optional[Any]=None		,	):



 """simple docstring"""





 if attention_mask is None:
  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	tf.cast(tf.math.not_equal(__lowerCamelCase		,	config.pad_token_id     )		,	tf.inta     )
 if decoder_attention_mask is None:
  __SCREAMING_SNAKE_CASE       :    Any       =	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:
  __SCREAMING_SNAKE_CASE       :    List[str]       =	tf.ones((config.encoder_layers, config.encoder_attention_heads)     )
 if decoder_head_mask is None:
  __SCREAMING_SNAKE_CASE       :    Tuple       =	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       _SCREAMING_SNAKE_CASE	(UpperCamelCase			,    UpperCamelCase			,    unittest.TestCase ):
     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 	__snake_case    (       self :							List[Any]				)->int:
      __SCREAMING_SNAKE_CASE       :    Tuple       =	TFLEDModelTester(self				)
      __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	ConfigTester(self     ,  config_class=UpperCamelCase				)





     def 	__snake_case    (       self :							Any				)->int:
      self.config_tester.run_common_tests()





     def 	__snake_case    (       self :							Any				)->Any:
      __SCREAMING_SNAKE_CASE       :    str       =	self.model_tester.prepare_config_and_inputs_for_common()
      self.model_tester.check_decoder_model_past_large_inputs(*UpperCamelCase				)





     def 	__snake_case    (       self :							List[Any]				)->Optional[int]:
      __SCREAMING_SNAKE_CASE,		__SCREAMING_SNAKE_CASE       :    str       =	self.model_tester.prepare_config_and_inputs_for_common()
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	tf.zeros_like(inputs_dict["attention_mask"]				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	2
      __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	tf.where(
          tf.range(self.model_tester.seq_length				)[None, :] < num_global_attn_indices     ,  1     ,  inputs_dict["global_attention_mask"]     ,  )

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

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

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






     def 	__snake_case    (       self :							Tuple				)->Dict:
      # TODO: Head-masking not yet implement
      pass






def 	_lowerCAmelCase						(				__lowerCamelCase   :     Tuple     ):



 """simple docstring"""





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


_lowerCamelCase						    =  1e-4






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





     def 	__snake_case    (       self :							str				)->int:
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"				).led

      # change to intended input here
      __SCREAMING_SNAKE_CASE       :    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]]				)
      __SCREAMING_SNAKE_CASE       :    Any       =	_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]]				)
      __SCREAMING_SNAKE_CASE       :    Dict       =	prepare_led_inputs_dict(model.config     ,  UpperCamelCase     ,  UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	model(**UpperCamelCase				)[0]
      __SCREAMING_SNAKE_CASE       :    Any       =	(1, 1_0_2_4, 7_6_8)
      self.assertEqual(output.shape     ,  UpperCamelCase				)
      # change to expected output here
      __SCREAMING_SNAKE_CASE       :    Tuple       =	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 	__snake_case    (       self :							Tuple				)->Optional[Any]:
      __SCREAMING_SNAKE_CASE       :    str       =	TFLEDForConditionalGeneration.from_pretrained("allenai/led-base-16384"				)

      # change to intended input here
      __SCREAMING_SNAKE_CASE       :    Optional[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]]				)
      __SCREAMING_SNAKE_CASE       :    Optional[int]       =	_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]]				)
      __SCREAMING_SNAKE_CASE       :    List[Any]       =	prepare_led_inputs_dict(model.config     ,  UpperCamelCase     ,  UpperCamelCase				)
      __SCREAMING_SNAKE_CASE       :    List[str]       =	model(**UpperCamelCase				)[0]
      __SCREAMING_SNAKE_CASE       :    int       =	(1, 1_0_2_4, model.config.vocab_size)
      self.assertEqual(output.shape     ,  UpperCamelCase				)
      # change to expected output here
      __SCREAMING_SNAKE_CASE       :    Tuple       =	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				)











# 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 json
import os

from ...utils.constants import SAGEMAKER_PARALLEL_EC2_INSTANCES, TORCH_DYNAMO_MODES
from ...utils.dataclasses import ComputeEnvironment, SageMakerDistributedType
from ...utils.imports import is_botoa_available
from .config_args import SageMakerConfig
from .config_utils import (
    DYNAMO_BACKENDS,
    _ask_field,
    _ask_options,
    _convert_dynamo_backend,
    _convert_mixed_precision,
    _convert_sagemaker_distributed_mode,
    _convert_yes_no_to_bool,
)


if is_botoa_available():
      import botoa  # noqa: F401






def 	_lowerCAmelCase						(				__lowerCamelCase   :     Optional[int]     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	botoa.client("iam"     )

 __SCREAMING_SNAKE_CASE       :    List[Any]       =	{
     "Version": "2012-10-17",
     "Statement": [
         {"Effect": "Allow", "Principal": {"Service": "sagemaker.amazonaws.com"}, "Action": "sts:AssumeRole"}
     ],
 }
 try:
  # create the role, associated with the chosen trust policy
  iam_client.create_role(
      RoleName=__lowerCamelCase		,	AssumeRolePolicyDocument=json.dumps(__lowerCamelCase		,	indent=2     )     )
  __SCREAMING_SNAKE_CASE       :    str       =	{
      "Version": "2012-10-17",
      "Statement": [
          {
              "Effect": "Allow",
              "Action": [
                  "sagemaker:*",
                  "ecr:GetDownloadUrlForLayer",
                  "ecr:BatchGetImage",
                  "ecr:BatchCheckLayerAvailability",
                  "ecr:GetAuthorizationToken",
                  "cloudwatch:PutMetricData",
                  "cloudwatch:GetMetricData",
                  "cloudwatch:GetMetricStatistics",
                  "cloudwatch:ListMetrics",
                  "logs:CreateLogGroup",
                  "logs:CreateLogStream",
                  "logs:DescribeLogStreams",
                  "logs:PutLogEvents",
                  "logs:GetLogEvents",
                  "s3:CreateBucket",
                  "s3:ListBucket",
                  "s3:GetBucketLocation",
                  "s3:GetObject",
                  "s3:PutObject",
              ],
              "Resource": "*",
          }
      ],
  }
  # attach policy to role
  iam_client.put_role_policy(
      RoleName=__lowerCamelCase		,	PolicyName=F"""{role_name}_policy_permission"""		,	PolicyDocument=json.dumps(__lowerCamelCase		,	indent=2     )		,	)
 except iam_client.exceptions.EntityAlreadyExistsException:
  print(F"""role {role_name} already exists. Using existing one"""     )






def 	_lowerCAmelCase						(				__lowerCamelCase   :     Dict     ):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    List[str]       =	botoa.client("iam"     )
 return iam_client.get_role(RoleName=__lowerCamelCase     )["Role"]["Arn"]






def 	_lowerCAmelCase						(				):



 """simple docstring"""





 __SCREAMING_SNAKE_CASE       :    Any       =	_ask_options(
     "How do you want to authorize?"		,	["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "]		,	__lowerCamelCase		,	)
 __SCREAMING_SNAKE_CASE       :    str       =	None
 if credentials_configuration == 0:
  __SCREAMING_SNAKE_CASE       :    List[str]       =	_ask_field("Enter your AWS Profile name: [default] "		,	default="default"     )
  __SCREAMING_SNAKE_CASE       :    Any       =	aws_profile
 else:
  print(
      "Note you will need to provide AWS_ACCESS_KEY_ID and AWS_SECRET_ACCESS_KEY when you launch you training script with,"
      "`accelerate launch --aws_access_key_id XXX --aws_secret_access_key YYY`"     )
  __SCREAMING_SNAKE_CASE       :    Dict       =	_ask_field("AWS Access Key ID: "     )
  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	aws_access_key_id

  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	_ask_field("AWS Secret Access Key: "     )
  __SCREAMING_SNAKE_CASE       :    Optional[int]       =	aws_secret_access_key

 __SCREAMING_SNAKE_CASE       :    List[Any]       =	_ask_field("Enter your AWS Region: [us-east-1]"		,	default="us-east-1"     )
 __SCREAMING_SNAKE_CASE       :    Any       =	aws_region

 __SCREAMING_SNAKE_CASE       :    int       =	_ask_options(
     "Do you already have an IAM Role for executing Amazon SageMaker Training Jobs?"		,	["Provide IAM Role name", "Create new IAM role using credentials"]		,	__lowerCamelCase		,	)
 if role_management == 0:
  __SCREAMING_SNAKE_CASE       :    List[str]       =	_ask_field("Enter your IAM role name: "     )
 else:
  __SCREAMING_SNAKE_CASE       :    Any       =	"accelerate_sagemaker_execution_role"
  print(F"""Accelerate will create an iam role \"{iam_role_name}\" using the provided credentials"""     )
  _create_iam_role_for_sagemaker(__lowerCamelCase     )

 __SCREAMING_SNAKE_CASE       :    Any       =	_ask_field(
     "Do you want to use custom Docker image? [yes/NO]: "		,	_convert_yes_no_to_bool		,	default=__lowerCamelCase		,	error_message="Please enter yes or no."		,	)
 __SCREAMING_SNAKE_CASE       :    Tuple       =	None
 if is_custom_docker_image:
  __SCREAMING_SNAKE_CASE       :    List[Any]       =	_ask_field("Enter your Docker image: "		,	lambda __lowerCamelCase     : str(__lowerCamelCase     ).lower()     )

 __SCREAMING_SNAKE_CASE       :    Dict       =	_ask_field(
     "Do you want to provide SageMaker input channels with data locations? [yes/NO]: "		,	_convert_yes_no_to_bool		,	default=__lowerCamelCase		,	error_message="Please enter yes or no."		,	)
 __SCREAMING_SNAKE_CASE       :    List[Any]       =	None
 if is_sagemaker_inputs_enabled:
  __SCREAMING_SNAKE_CASE       :    List[str]       =	_ask_field(
      "Enter the path to the SageMaker inputs TSV file with columns (channel_name, data_location): "		,	lambda __lowerCamelCase     : str(__lowerCamelCase     ).lower()		,	)

 __SCREAMING_SNAKE_CASE       :    Dict       =	_ask_field(
     "Do you want to enable SageMaker metrics? [yes/NO]: "		,	_convert_yes_no_to_bool		,	default=__lowerCamelCase		,	error_message="Please enter yes or no."		,	)
 __SCREAMING_SNAKE_CASE       :    int       =	None
 if is_sagemaker_metrics_enabled:
  __SCREAMING_SNAKE_CASE       :    Dict       =	_ask_field(
      "Enter the path to the SageMaker metrics TSV file with columns (metric_name, metric_regex): "		,	lambda __lowerCamelCase     : str(__lowerCamelCase     ).lower()		,	)

 __SCREAMING_SNAKE_CASE       :    Any       =	_ask_options(
     "What is the distributed mode?"		,	["No distributed training", "Data parallelism"]		,	_convert_sagemaker_distributed_mode		,	)
 __SCREAMING_SNAKE_CASE       :    Tuple       =	{}
 __SCREAMING_SNAKE_CASE       :    str       =	_ask_field(
     "Do you wish to optimize your script with torch dynamo?[yes/NO]:"		,	_convert_yes_no_to_bool		,	default=__lowerCamelCase		,	error_message="Please enter yes or no."		,	)
 if use_dynamo:
  __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	"dynamo_"
  __SCREAMING_SNAKE_CASE       :    Optional[Any]       =	_ask_options(
      "Which dynamo backend would you like to use?"		,	[x.lower() for x in DYNAMO_BACKENDS]		,	_convert_dynamo_backend		,	default=2		,	)
  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	_ask_field(
      "Do you want to customize the defaults sent to torch.compile? [yes/NO]: "		,	_convert_yes_no_to_bool		,	default=__lowerCamelCase		,	error_message="Please enter yes or no."		,	)

  if use_custom_options:
   __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	_ask_options(
       "Which mode do you want to use?"		,	__lowerCamelCase		,	lambda __lowerCamelCase     : TORCH_DYNAMO_MODES[int(__lowerCamelCase     )]		,	default="default"		,	)
   __SCREAMING_SNAKE_CASE       :    str       =	_ask_field(
       "Do you want the fullgraph mode or it is ok to break model into several subgraphs? [yes/NO]: "		,	_convert_yes_no_to_bool		,	default=__lowerCamelCase		,	error_message="Please enter yes or no."		,	)
   __SCREAMING_SNAKE_CASE       :    Any       =	_ask_field(
       "Do you want to enable dynamic shape tracing? [yes/NO]: "		,	_convert_yes_no_to_bool		,	default=__lowerCamelCase		,	error_message="Please enter yes or no."		,	)
 __SCREAMING_SNAKE_CASE       :    Optional[int]       =	"Which EC2 instance type you want to use for your training?"
 if distributed_type != SageMakerDistributedType.NO:
  __SCREAMING_SNAKE_CASE       :    List[Any]       =	_ask_options(
      __lowerCamelCase		,	__lowerCamelCase		,	lambda __lowerCamelCase     : SAGEMAKER_PARALLEL_EC2_INSTANCES[int(__lowerCamelCase     )]     )
 else:
  eca_instance_query += "? [ml.p3.2xlarge]:"
  __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	_ask_field(__lowerCamelCase		,	lambda __lowerCamelCase     : str(__lowerCamelCase     ).lower()		,	default="ml.p3.2xlarge"     )

 __SCREAMING_SNAKE_CASE       :    List[Any]       =	1
 if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
  __SCREAMING_SNAKE_CASE       :    Any       =	_ask_field(
      "How many machines do you want use? [1]: "		,	__lowerCamelCase		,	default=1		,	)

 __SCREAMING_SNAKE_CASE       :    Union[str, Any]       =	_ask_options(
     "Do you wish to use FP16 or BF16 (mixed precision)?"		,	["no", "fp16", "bf16", "fp8"]		,	_convert_mixed_precision		,	)

 if use_dynamo and mixed_precision == "no":
  print(
      "Torch dynamo used without mixed precision requires TF32 to be efficient. Accelerate will enable it by default when launching your scripts."     )

 return SageMakerConfig(
     image_uri=__lowerCamelCase		,	compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER		,	distributed_type=__lowerCamelCase		,	use_cpu=__lowerCamelCase		,	dynamo_config=__lowerCamelCase		,	eca_instance_type=__lowerCamelCase		,	profile=__lowerCamelCase		,	region=__lowerCamelCase		,	iam_role_name=__lowerCamelCase		,	mixed_precision=__lowerCamelCase		,	num_machines=__lowerCamelCase		,	sagemaker_inputs_file=__lowerCamelCase		,	sagemaker_metrics_file=__lowerCamelCase		,	)






