'''simple docstring'''




import argparse
import os
import re


_A								=   'src/transformers'

# Pattern that looks at the indentation in a line.
_A								=   re.compile(R'^(\s*)\S')
# Pattern that matches `"key":" and puts `key` in group 0.
_A								=   re.compile(R'^\s*"([^"]+)":')
# Pattern that matches `_import_structure["key"]` and puts `key` in group 0.
_A								=   re.compile(R'^\s*_import_structure\["([^"]+)"\]')
# Pattern that matches `"key",` and puts `key` in group 0.
_A								=   re.compile(R'^\s*"([^"]+)",\s*$')
# Pattern that matches any `[stuff]` and puts `stuff` in group 0.
_A								=   re.compile(R'\[([^\]]+)\]')

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_							):
     lowercase_					:       str     				=					_re_indent.search(SCREAMING_SNAKE_CASE_							)
     return "" if search is None else search.groups()[0]

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_=""  ,       SCREAMING_SNAKE_CASE_=None  ,       SCREAMING_SNAKE_CASE_=None							):
     lowercase_					:       List[str]     				=					0
     lowercase_					:       Tuple     				=					code.split('\n'							)
     if start_prompt is not None:
          while not lines[index].startswith(SCREAMING_SNAKE_CASE_							):
               index += 1
          lowercase_					:       Dict     				=					['\n'.join(lines[:index]							)]
     else:
          lowercase_					:       Optional[int]     				=					[]

     # We split into blocks until we get to the `end_prompt` (or the end of the block).
     lowercase_					:       List[str]     				=					[lines[index]]
     index += 1
     while index < len(SCREAMING_SNAKE_CASE_							) and (end_prompt is None or not lines[index].startswith(SCREAMING_SNAKE_CASE_							)):
          if len(lines[index]							) > 0 and get_indent(lines[index]							) == indent_level:
               if len(SCREAMING_SNAKE_CASE_							) > 0 and get_indent(current_block[-1]							).startswith(indent_level + ' '							):
                    current_block.append(lines[index]							)
                    blocks.append('\n'.join(SCREAMING_SNAKE_CASE_							)							)
                    if index < len(SCREAMING_SNAKE_CASE_							) - 1:
                         lowercase_					:       List[str]     				=					[lines[index + 1]]
                         index += 1
                    else:
                         lowercase_					:       Tuple     				=					[]
               else:
                    blocks.append('\n'.join(SCREAMING_SNAKE_CASE_							)							)
                    lowercase_					:       Dict     				=					[lines[index]]
          else:
               current_block.append(lines[index]							)
          index += 1

     # Adds current block if it's nonempty.
     if len(SCREAMING_SNAKE_CASE_							) > 0:
          blocks.append('\n'.join(SCREAMING_SNAKE_CASE_							)							)

     # Add final block after end_prompt if provided.
     if end_prompt is not None and index < len(SCREAMING_SNAKE_CASE_							):
          blocks.append('\n'.join(lines[index:]							)							)

     return blocks

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_							):

     def _inner(SCREAMING_SNAKE_CASE_							):
          return key(SCREAMING_SNAKE_CASE_							).lower().replace('_'  ,       ''							)

     return _inner

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_=None							):

     # If no key is provided, we use a noop.
     def noop(SCREAMING_SNAKE_CASE_							):
          return x

     if key is None:
          lowercase_					:       Optional[int]     				=					noop
     # Constants are all uppercase, they go first.
     lowercase_					:       Optional[int]     				=					[obj for obj in objects if key(SCREAMING_SNAKE_CASE_							).isupper()]
     # Classes are not all uppercase but start with a capital, they go second.
     lowercase_					:       Optional[Any]     				=					[obj for obj in objects if key(SCREAMING_SNAKE_CASE_							)[0].isupper() and not key(SCREAMING_SNAKE_CASE_							).isupper()]
     # Functions begin with a lowercase, they go last.
     lowercase_					:       Optional[int]     				=					[obj for obj in objects if not key(SCREAMING_SNAKE_CASE_							)[0].isupper()]

     lowercase_					:       str     				=					ignore_underscore(SCREAMING_SNAKE_CASE_							)
     return sorted(SCREAMING_SNAKE_CASE_  ,       key=SCREAMING_SNAKE_CASE_							) + sorted(SCREAMING_SNAKE_CASE_  ,       key=SCREAMING_SNAKE_CASE_							) + sorted(SCREAMING_SNAKE_CASE_  ,       key=SCREAMING_SNAKE_CASE_							)

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_							):

     # This inner function sort imports between [ ].
     def _replace(SCREAMING_SNAKE_CASE_							):
          lowercase_					:       int     				=					match.groups()[0]
          if "," not in imports:
               return f'''[{imports}]'''
          lowercase_					:       Optional[Any]     				=					[part.strip().replace('"'  ,       ''							) for part in imports.split(','							)]
          # We will have a final empty element if the line finished with a comma.
          if len(keys[-1]							) == 0:
               lowercase_					:       str     				=					keys[:-1]
          return "[" + ", ".join([f'''"{k}"''' for k in sort_objects(SCREAMING_SNAKE_CASE_							)]							) + "]"

     lowercase_					:       Tuple     				=					import_statement.split('\n'							)
     if len(SCREAMING_SNAKE_CASE_							) > 3:
          # Here we have to sort internal imports that are on several lines (one per name):
          # key: [
          #     "object1",
          #     "object2",
          #     ...
          # ]

          # We may have to ignore one or two lines on each side.
          lowercase_					:       Optional[Any]     				=					2 if lines[1].strip() == '[' else 1
          lowercase_					:       Dict     				=					[(i, _re_strip_line.search(SCREAMING_SNAKE_CASE_							).groups()[0]) for i, line in enumerate(lines[idx:-idx]							)]
          lowercase_					:       Optional[int]     				=					sort_objects(SCREAMING_SNAKE_CASE_  ,       key=lambda SCREAMING_SNAKE_CASE_							: x[1]							)
          lowercase_					:       List[Any]     				=					[lines[x[0] + idx] for x in sorted_indices]
          return "\n".join(lines[:idx] + sorted_lines + lines[-idx:]							)
     elif len(SCREAMING_SNAKE_CASE_							) == 3:
          # Here we have to sort internal imports that are on one separate line:
          # key: [
          #     "object1", "object2", ...
          # ]
          if _re_bracket_content.search(lines[1]							) is not None:
               lowercase_					:       Optional[Any]     				=					_re_bracket_content.sub(_replace  ,       lines[1]							)
          else:
               lowercase_					:       Optional[int]     				=					[part.strip().replace('"'  ,       ''							) for part in lines[1].split(','							)]
               # We will have a final empty element if the line finished with a comma.
               if len(keys[-1]							) == 0:
                    lowercase_					:       List[str]     				=					keys[:-1]
               lowercase_					:       Optional[int]     				=					get_indent(lines[1]							) + ', '.join([f'''"{k}"''' for k in sort_objects(SCREAMING_SNAKE_CASE_							)]							)
          return "\n".join(SCREAMING_SNAKE_CASE_							)
     else:
          # Finally we have to deal with imports fitting on one line
          lowercase_					:       Union[str, Any]     				=					_re_bracket_content.sub(_replace  ,       SCREAMING_SNAKE_CASE_							)
          return import_statement

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_=True							):
     with open(SCREAMING_SNAKE_CASE_  ,       encoding='utf-8'							) as f:
          lowercase_					:       Union[str, Any]     				=					f.read()

     if "_import_structure" not in code:
          return

     # Blocks of indent level 0
     lowercase_					:       List[str]     				=					split_code_in_indented_blocks(
         SCREAMING_SNAKE_CASE_  ,       start_prompt='_import_structure = {'  ,       end_prompt='if TYPE_CHECKING:'							)

     # We ignore block 0 (everything untils start_prompt) and the last block (everything after end_prompt).
     for block_idx in range(1  ,       len(SCREAMING_SNAKE_CASE_							) - 1							):
          # Check if the block contains some `_import_structure`s thingy to sort.
          lowercase_					:       List[Any]     				=					main_blocks[block_idx]
          lowercase_					:       List[Any]     				=					block.split('\n'							)

          # Get to the start of the imports.
          lowercase_					:       Union[str, Any]     				=					0
          while line_idx < len(SCREAMING_SNAKE_CASE_							) and "_import_structure" not in block_lines[line_idx]:
               # Skip dummy import blocks
               if "import dummy" in block_lines[line_idx]:
                    lowercase_					:       str     				=					len(SCREAMING_SNAKE_CASE_							)
               else:
                    line_idx += 1
          if line_idx >= len(SCREAMING_SNAKE_CASE_							):
               continue

          # Ignore beginning and last line: they don't contain anything.
          lowercase_					:       int     				=					'\n'.join(block_lines[line_idx:-1]							)
          lowercase_					:       Tuple     				=					get_indent(block_lines[1]							)
          # Slit the internal block into blocks of indent level 1.
          lowercase_					:       List[str]     				=					split_code_in_indented_blocks(SCREAMING_SNAKE_CASE_  ,       indent_level=SCREAMING_SNAKE_CASE_							)
          # We have two categories of import key: list or _import_structure[key].append/extend
          lowercase_					:       Optional[int]     				=					_re_direct_key if '_import_structure = {' in block_lines[0] else _re_indirect_key
          # Grab the keys, but there is a trap: some lines are empty or just comments.
          lowercase_					:       List[str]     				=					[(pattern.search(SCREAMING_SNAKE_CASE_							).groups()[0] if pattern.search(SCREAMING_SNAKE_CASE_							) is not None else None) for b in internal_blocks]
          # We only sort the lines with a key.
          lowercase_					:       Optional[int]     				=					[(i, key) for i, key in enumerate(SCREAMING_SNAKE_CASE_							) if key is not None]
          lowercase_					:       List[Any]     				=					[x[0] for x in sorted(SCREAMING_SNAKE_CASE_  ,       key=lambda SCREAMING_SNAKE_CASE_							: x[1]							)]

          # We reorder the blocks by leaving empty lines/comments as they were and reorder the rest.
          lowercase_					:       List[str]     				=					0
          lowercase_					:       Dict     				=					[]
          for i in range(len(SCREAMING_SNAKE_CASE_							)							):
               if keys[i] is None:
                    reorderded_blocks.append(internal_blocks[i]							)
               else:
                    lowercase_					:       Optional[Any]     				=					sort_objects_in_import(internal_blocks[sorted_indices[count]]							)
                    reorderded_blocks.append(SCREAMING_SNAKE_CASE_							)
                    count += 1

        # And we put our main block back together with its first and last line.
          lowercase_					:       Dict     				=					'\n'.join(block_lines[:line_idx] + reorderded_blocks + [block_lines[-1]]							)

     if code != "\n".join(SCREAMING_SNAKE_CASE_							):
          if check_only:
               return True
          else:
               print(f'''Overwriting {file}.'''							)
               with open(SCREAMING_SNAKE_CASE_  ,       'w'  ,       encoding='utf-8'							) as f:
                    f.write('\n'.join(SCREAMING_SNAKE_CASE_							)							)

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_=True							):
     lowercase_					:       List[str]     				=					[]
     for root, _, files in os.walk(SCREAMING_SNAKE_CASE_							):
          if "__init__.py" in files:
               lowercase_					:       Dict     				=					sort_imports(os.path.join(SCREAMING_SNAKE_CASE_  ,       '__init__.py'							)  ,       check_only=SCREAMING_SNAKE_CASE_							)
               if result:
                    lowercase_					:       Optional[int]     				=					[os.path.join(SCREAMING_SNAKE_CASE_  ,       '__init__.py'							)]
     if len(SCREAMING_SNAKE_CASE_							) > 0:
          raise ValueError(f'''Would overwrite {len(SCREAMING_SNAKE_CASE_							)} files, run `make style`.'''							)


if __name__ == "__main__":
 _A								=   argparse.ArgumentParser()
 parser.add_argument('--check_only', action='store_true', help='Whether to only check or fix style.')
 _A								=   parser.parse_args()

 sort_imports_in_all_inits(check_only=args.check_only)





'''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 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 _UpperCamelCase   (       SCREAMING_SNAKE_CASE_							):
     lowercase_					:       str     				=					botoa.client('iam'							)

     lowercase_					:       Optional[int]     				=					{
         '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=SCREAMING_SNAKE_CASE_  ,       AssumeRolePolicyDocument=json.dumps(SCREAMING_SNAKE_CASE_  ,       indent=2							)							)
          lowercase_					:       Optional[Any]     				=					{
              '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=SCREAMING_SNAKE_CASE_  ,       PolicyName=f'''{role_name}_policy_permission'''  ,       PolicyDocument=json.dumps(SCREAMING_SNAKE_CASE_  ,       indent=2							)  ,       )
     except iam_client.exceptions.EntityAlreadyExistsException:
          print(f'''role {role_name} already exists. Using existing one'''							)

def _UpperCamelCase   (       SCREAMING_SNAKE_CASE_							):
     lowercase_					:       List[Any]     				=					botoa.client('iam'							)
     return iam_client.get_role(RoleName=SCREAMING_SNAKE_CASE_							)["Role"]["Arn"]






def _UpperCamelCase   (       ):
     lowercase_					:       Dict     				=					_ask_options(
         'How do you want to authorize?'  ,       ['AWS Profile', 'Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) ']  ,       SCREAMING_SNAKE_CASE_  ,       )
     lowercase_					:       List[Any]     				=					None
     if credentials_configuration == 0:
          lowercase_					:       List[Any]     				=					_ask_field('Enter your AWS Profile name: [default] '  ,       default='default'							)
          lowercase_					:       Optional[int]     				=					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`'							)
          lowercase_					:       Optional[Any]     				=					_ask_field('AWS Access Key ID: '							)
          lowercase_					:       Optional[int]     				=					aws_access_key_id

          lowercase_					:       Optional[Any]     				=					_ask_field('AWS Secret Access Key: '							)
          lowercase_					:       Optional[int]     				=					aws_secret_access_key

     lowercase_					:       Any     				=					_ask_field('Enter your AWS Region: [us-east-1]'  ,       default='us-east-1'							)
     lowercase_					:       Tuple     				=					aws_region

     lowercase_					:       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']  ,       SCREAMING_SNAKE_CASE_  ,       )
     if role_management == 0:
          lowercase_					:       int     				=					_ask_field('Enter your IAM role name: '							)
     else:
          lowercase_					:       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(SCREAMING_SNAKE_CASE_							)

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

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

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

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

          if use_custom_options:
               lowercase_					:       Union[str, Any]     				=					_ask_options(
                   'Which mode do you want to use?'  ,       SCREAMING_SNAKE_CASE_  ,       lambda SCREAMING_SNAKE_CASE_							: TORCH_DYNAMO_MODES[int(SCREAMING_SNAKE_CASE_							)]  ,       default='default'  ,       )
               lowercase_					:       Union[str, Any]     				=					_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=SCREAMING_SNAKE_CASE_  ,       error_message='Please enter yes or no.'  ,       )
               lowercase_					:       Tuple     				=					_ask_field(
                   'Do you want to enable dynamic shape tracing? [yes/NO]: '  ,       _convert_yes_no_to_bool  ,       default=SCREAMING_SNAKE_CASE_  ,       error_message='Please enter yes or no.'  ,       )
     lowercase_					:       List[str]     				=					'Which EC2 instance type you want to use for your training?'
     if distributed_type != SageMakerDistributedType.NO:
          lowercase_					:       Optional[Any]     				=					_ask_options(
              SCREAMING_SNAKE_CASE_  ,       SCREAMING_SNAKE_CASE_  ,       lambda SCREAMING_SNAKE_CASE_							: SAGEMAKER_PARALLEL_EC2_INSTANCES[int(SCREAMING_SNAKE_CASE_							)]							)
     else:
          eca_instance_query += "? [ml.p3.2xlarge]:"
          lowercase_					:       Tuple     				=					_ask_field(SCREAMING_SNAKE_CASE_  ,       lambda SCREAMING_SNAKE_CASE_							: str(SCREAMING_SNAKE_CASE_							).lower()  ,       default='ml.p3.2xlarge'							)

     lowercase_					:       str     				=					1
     if distributed_type in (SageMakerDistributedType.DATA_PARALLEL, SageMakerDistributedType.MODEL_PARALLEL):
          lowercase_					:       Union[str, Any]     				=					_ask_field(
              'How many machines do you want use? [1]: '  ,       SCREAMING_SNAKE_CASE_  ,       default=1  ,       )

     lowercase_					:       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=SCREAMING_SNAKE_CASE_  ,       compute_environment=ComputeEnvironment.AMAZON_SAGEMAKER  ,       distributed_type=SCREAMING_SNAKE_CASE_  ,       use_cpu=SCREAMING_SNAKE_CASE_  ,       dynamo_config=SCREAMING_SNAKE_CASE_  ,       eca_instance_type=SCREAMING_SNAKE_CASE_  ,       profile=SCREAMING_SNAKE_CASE_  ,       region=SCREAMING_SNAKE_CASE_  ,       iam_role_name=SCREAMING_SNAKE_CASE_  ,       mixed_precision=SCREAMING_SNAKE_CASE_  ,       num_machines=SCREAMING_SNAKE_CASE_  ,       sagemaker_inputs_file=SCREAMING_SNAKE_CASE_  ,       sagemaker_metrics_file=SCREAMING_SNAKE_CASE_  ,       )




