# 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    a__       (						__UpperCamelCase ):
      SCREAMING_SNAKE_CASE_      =     botoa.client("iam" )

      SCREAMING_SNAKE_CASE_      =     {
          "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=__UpperCamelCase   ,					AssumeRolePolicyDocument=json.dumps(__UpperCamelCase   ,					indent=2 ) )
            SCREAMING_SNAKE_CASE_      =     {
                "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=__UpperCamelCase   ,					PolicyName=F'''{role_name}_policy_permission'''   ,					PolicyDocument=json.dumps(__UpperCamelCase   ,					indent=2 )   ,					)
      except iam_client.exceptions.EntityAlreadyExistsException:
            print(F'''role {role_name} already exists. Using existing one''' )


def    a__       (						__UpperCamelCase ):
      SCREAMING_SNAKE_CASE_      =     botoa.client("iam" )
      return iam_client.get_role(RoleName=__UpperCamelCase )["Role"]["Arn"]





def    a__       (						):
      SCREAMING_SNAKE_CASE_      =     _ask_options(
          "How do you want to authorize?"   ,					["AWS Profile", "Credentials (AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY) "]   ,					__UpperCamelCase   ,					)
      SCREAMING_SNAKE_CASE_      =     None
      if credentials_configuration == 0:
            SCREAMING_SNAKE_CASE_      =     _ask_field("Enter your AWS Profile name: [default] "   ,					default="default" )
            SCREAMING_SNAKE_CASE_      =     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_      =     _ask_field("AWS Access Key ID: " )
            SCREAMING_SNAKE_CASE_      =     aws_access_key_id

            SCREAMING_SNAKE_CASE_      =     _ask_field("AWS Secret Access Key: " )
            SCREAMING_SNAKE_CASE_      =     aws_secret_access_key

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

      SCREAMING_SNAKE_CASE_      =     _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"]   ,					__UpperCamelCase   ,					)
      if role_management == 0:
            SCREAMING_SNAKE_CASE_      =     _ask_field("Enter your IAM role name: " )
      else:
            SCREAMING_SNAKE_CASE_      =     "accelerate_sagemaker_execution_role"
            print(F'''Accelerate will create an iam role "{iam_role_name}" using the provided credentials''' )
            _create_iam_role_for_sagemaker(__UpperCamelCase )

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

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

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

      SCREAMING_SNAKE_CASE_      =     _ask_options(
          "What is the distributed mode?"   ,					["No distributed training", "Data parallelism"]   ,					_convert_sagemaker_distributed_mode   ,					)
      SCREAMING_SNAKE_CASE_      =     {}
      SCREAMING_SNAKE_CASE_      =     _ask_field(
          "Do you wish to optimize your script with torch dynamo?[yes/NO]:"   ,					_convert_yes_no_to_bool   ,					default=__UpperCamelCase   ,					error_message="Please enter yes or no."   ,					)
      if use_dynamo:
            SCREAMING_SNAKE_CASE_      =     "dynamo_"
            SCREAMING_SNAKE_CASE_      =     _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_      =     _ask_field(
                "Do you want to customize the defaults sent to torch.compile? [yes/NO]: "   ,					_convert_yes_no_to_bool   ,					default=__UpperCamelCase   ,					error_message="Please enter yes or no."   ,					)

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

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

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



import argparse
import os

from transformers.utils import direct_transformers_import


# All paths are set with the intent you should run this script from the root of the repo with the command
# python utils/check_task_guides.py
A   :     Tuple       =				"src/transformers"
A   :     Optional[Any]       =				"docs/source/en/tasks"


def    a__       (						__UpperCamelCase   ,					__UpperCamelCase   ,					__UpperCamelCase ):
      with open(__UpperCamelCase   ,					"r"   ,					encoding="utf-8"   ,					newline="\n" ) as f:
            SCREAMING_SNAKE_CASE_      =     f.readlines()
      # Find the start prompt.
      SCREAMING_SNAKE_CASE_      =     0
      while not lines[start_index].startswith(__UpperCamelCase ):
            start_index += 1
      start_index += 1

      SCREAMING_SNAKE_CASE_      =     start_index
      while not lines[end_index].startswith(__UpperCamelCase ):
            end_index += 1
      end_index -= 1

      while len(lines[start_index] ) <= 1:
            start_index += 1
      while len(lines[end_index] ) <= 1:
            end_index -= 1
      end_index += 1
      return "".join(lines[start_index:end_index] ), start_index, end_index, lines


# This is to make sure the transformers module imported is the one in the repo.
A   :     List[str]       =				direct_transformers_import(TRANSFORMERS_PATH)

A   :     List[Any]       =				{
    "asr.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_CTC_MAPPING_NAMES,
    "audio_classification.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_AUDIO_CLASSIFICATION_MAPPING_NAMES,
    "language_modeling.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_CAUSAL_LM_MAPPING_NAMES,
    "image_classification.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_IMAGE_CLASSIFICATION_MAPPING_NAMES,
    "masked_language_modeling.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_MASKED_LM_MAPPING_NAMES,
    "multiple_choice.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_MULTIPLE_CHOICE_MAPPING_NAMES,
    "object_detection.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_OBJECT_DETECTION_MAPPING_NAMES,
    "question_answering.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_QUESTION_ANSWERING_MAPPING_NAMES,
    "semantic_segmentation.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_SEMANTIC_SEGMENTATION_MAPPING_NAMES,
    "sequence_classification.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_SEQUENCE_CLASSIFICATION_MAPPING_NAMES,
    "summarization.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
    "token_classification.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_TOKEN_CLASSIFICATION_MAPPING_NAMES,
    "translation.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_SEQ_TO_SEQ_CAUSAL_LM_MAPPING_NAMES,
    "video_classification.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_VIDEO_CLASSIFICATION_MAPPING_NAMES,
    "document_question_answering.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_DOCUMENT_QUESTION_ANSWERING_MAPPING_NAMES,
    "monocular_depth_estimation.md": transformers_module.models.auto.modeling_auto.MODEL_FOR_DEPTH_ESTIMATION_MAPPING_NAMES,
}

# This list contains model types used in some task guides that are not in `CONFIG_MAPPING_NAMES` (therefore not in any
# `MODEL_MAPPING_NAMES` or any `MODEL_FOR_XXX_MAPPING_NAMES`).
A   :     Any       =				{
    "summarization.md": ("nllb",),
    "translation.md": ("nllb",),
}


def    a__       (						__UpperCamelCase ):
      SCREAMING_SNAKE_CASE_      =     TASK_GUIDE_TO_MODELS[task_guide]
      SCREAMING_SNAKE_CASE_      =     SPECIAL_TASK_GUIDE_TO_MODEL_TYPES.get(__UpperCamelCase   ,					set() )
      SCREAMING_SNAKE_CASE_      =     {
          code: name
          for code, name in transformers_module.MODEL_NAMES_MAPPING.items()
          if (code in model_maping_names or code in special_model_types)
      }
      return ", ".join([F'''[{name}](../model_doc/{code})''' for code, name in model_names.items()] ) + "\n"


def    a__       (						__UpperCamelCase   ,					__UpperCamelCase=False ):
      SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_   ,       SCREAMING_SNAKE_CASE_      =     _find_text_in_file(
          filename=os.path.join(__UpperCamelCase   ,					__UpperCamelCase )   ,					start_prompt="<!--This tip is automatically generated by `make fix-copies`, do not fill manually!-->"   ,					end_prompt="<!--End of the generated tip-->"   ,					)

      SCREAMING_SNAKE_CASE_      =     get_model_list_for_task(__UpperCamelCase )

      if current_list != new_list:
            if overwrite:
                  with open(os.path.join(__UpperCamelCase   ,					__UpperCamelCase )   ,					"w"   ,					encoding="utf-8"   ,					newline="\n" ) as f:
                        f.writelines(lines[:start_index] + [new_list] + lines[end_index:] )
            else:
                  raise ValueError(
                      F'''The list of models that can be used in the {task_guide} guide needs an update. Run `make fix-copies`'''
                      " to fix this." )


if __name__ == "__main__":
      A   :     Tuple       =				argparse.ArgumentParser()
      parser.add_argument("--fix_and_overwrite", action="store_true", help="Whether to fix inconsistencies.")
      A   :     Dict       =				parser.parse_args()

      for task_guide in TASK_GUIDE_TO_MODELS.keys():
            check_model_list_for_task(task_guide, args.fix_and_overwrite)


