



# Copyright 2022 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
import subprocess

from packaging.version import Version, parse

from accelerate.commands.config.config_args import default_config_file, load_config_from_file


__lowerCamelCase     :							Dict  							=				"Run commands across TPU VMs for initial setup before running `accelerate launch`."

def    lowerCamelCase_(lowerCamelCase_=None    ) ->       Optional[int]:
						if subparsers is not None:
												UpperCAmelCase         =     subparsers.add_parser("tpu-config"    ,   description=_description    )
						else:
												UpperCAmelCase         =     argparse.ArgumentParser("Accelerate tpu-config command"    ,   description=_description    )
						# Core arguments
						UpperCAmelCase         =     parser.add_argument_group(
						    "Config Arguments"    ,   "Arguments that can be configured through `accelerate config`."    )
						config_args.add_argument(
						    "--config_file"    ,   type=lowerCamelCase_    ,   default=lowerCamelCase_    ,   help="Path to the config file to use for accelerate."    ,   )
						config_args.add_argument(
						    "--tpu_name"    ,   default=lowerCamelCase_    ,   help="The name of the TPU to use. If not specified, will use the TPU specified in the config file."    ,   )
						config_args.add_argument(
						    "--tpu_zone"    ,   default=lowerCamelCase_    ,   help="The zone of the TPU to use. If not specified, will use the zone specified in the config file."    ,   )
						UpperCAmelCase         =     parser.add_argument_group("TPU Arguments"    ,   "Arguments for options ran inside the TPU."    )
						pod_args.add_argument(
						    "--use_alpha"    ,   action="store_true"    ,   help="Whether to use `gcloud alpha` when running the TPU training script instead of `gcloud`."    ,   )
						pod_args.add_argument(
						    "--command_file"    ,   default=lowerCamelCase_    ,   help="The path to the file containing the commands to run on the pod on startup."    ,   )
						pod_args.add_argument(
						    "--command"    ,   action="append"    ,   nargs="+"    ,   help="A command to run on the pod. Can be passed multiple times."    ,   )
						pod_args.add_argument(
						    "--install_accelerate"    ,   action="store_true"    ,   help="Whether to install accelerate on the pod. Defaults to False."    ,   )
						pod_args.add_argument(
						    "--accelerate_version"    ,   default="latest"    ,   help="The version of accelerate to install on the pod. If not specified, will use the latest pypi version. Specify 'dev' to install from GitHub."    ,   )
						pod_args.add_argument(
						    "--debug"    ,   action="store_true"    ,   help="If set, will print the command that would be run instead of running it."    )

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

def    lowerCamelCase_(lowerCamelCase_    ) ->       Optional[int]:
						UpperCAmelCase         =     None

						# Get the default from the config file if it exists.
						if args.config_file is not None or os.path.isfile(lowerCamelCase_    ):
												UpperCAmelCase         =     load_config_from_file(args.config_file    )
												if not args.command_file and defaults.command_file is not None and not args.command:
																		UpperCAmelCase         =     defaults.command_file
												if not args.command and defaults.commands is not None:
																		UpperCAmelCase         =     defaults.commands
												if not args.tpu_name:
																		UpperCAmelCase         =     defaults.tpu_name
												if not args.tpu_zone:
																		UpperCAmelCase         =     defaults.tpu_zone
						if args.accelerate_version == "dev":
												UpperCAmelCase         =     "git+https://github.com/huggingface/accelerate.git"
						elif args.accelerate_version == "latest":
												UpperCAmelCase         =     "accelerate -U"
						elif isinstance(parse(args.accelerate_version    )    ,   lowerCamelCase_    ):
												UpperCAmelCase         =     F'accelerate=={args.accelerate_version}'

						if not args.command_file and not args.command:
												raise ValueError("You must specify either a command file or a command to run on the pod."    )

						if args.command_file:
												with open(args.command_file    ,   "r"    ) as f:
																		UpperCAmelCase         =     [f.read().splitlines()]

    # To turn list of lists into list of strings
						if isinstance(args.command[0]    ,   lowerCamelCase_    ):
												UpperCAmelCase         =     [line for cmd in args.command for line in cmd]
						# Default to the shared folder and install accelerate
						UpperCAmelCase         =     ["cd /usr/share"]
						if args.install_accelerate:
												new_cmd += [F'pip install {args.accelerate_version}']
						new_cmd += args.command
						UpperCAmelCase         =     "; ".join(lowerCamelCase_    )

						# Then send it to gcloud
						# Eventually try to use google-api-core to do this instead of subprocess
						UpperCAmelCase         =     ["gcloud"]
						if args.use_alpha:
												cmd += ["alpha"]
						cmd += [
						    "compute",
						    "tpus",
						    "tpu-vm",
						    "ssh",
						    args.tpu_name,
						    "--zone",
						    args.tpu_zone,
						    "--command",
						    args.command,
						    "--worker",
						    "all",
						]
						if args.debug:
												print(F'Running {" ".join(lowerCamelCase_    )}'    )
												return
						subprocess.run(lowerCamelCase_    )
						print("Successfully setup pod."    )





def    lowerCamelCase_() ->       Union[str, Any]:
						UpperCAmelCase         =     tpu_command_parser()
						UpperCAmelCase         =     parser.parse_args()

						tpu_command_launcher(lowerCamelCase_    )








from typing import TYPE_CHECKING

from ...utils import (
    OptionalDependencyNotAvailable,
    _LazyModule,
    is_torch_available,
)


__lowerCamelCase     :							Any  							=				{
    "configuration_mega": ["MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP", "MegaConfig", "MegaOnnxConfig"],
}

try:
		if not is_torch_available():
				raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
		pass
else:
		__lowerCamelCase     :							Optional[int]  							=				[
		    "MEGA_PRETRAINED_MODEL_ARCHIVE_LIST",
		    "MegaForCausalLM",
		    "MegaForMaskedLM",
		    "MegaForMultipleChoice",
		    "MegaForQuestionAnswering",
		    "MegaForSequenceClassification",
		    "MegaForTokenClassification",
		    "MegaModel",
		    "MegaPreTrainedModel",
		]

if TYPE_CHECKING:
		from .configuration_mega import MEGA_PRETRAINED_CONFIG_ARCHIVE_MAP, MegaConfig, MegaOnnxConfig

		try:
				if not is_torch_available():
						raise OptionalDependencyNotAvailable()
		except OptionalDependencyNotAvailable:
				pass
		else:
				from .modeling_mega import (
				    MEGA_PRETRAINED_MODEL_ARCHIVE_LIST,
				    MegaForCausalLM,
				    MegaForMaskedLM,
				    MegaForMultipleChoice,
				    MegaForQuestionAnswering,
				    MegaForSequenceClassification,
				    MegaForTokenClassification,
				    MegaModel,
				    MegaPreTrainedModel,
				)

else:
		import sys

		__lowerCamelCase     :							int  							=				_LazyModule(__name__, globals()["__file__"], _import_structure, module_spec=__spec__)



