





import argparse
import glob
import inspect
import os
import re
import unittest

from get_test_info import get_test_classes

from tests.test_pipeline_mixin import pipeline_test_mapping


UpperCamelCase__          ={}
for task, _ in pipeline_test_mapping.items():
       UpperCamelCase__          ={'pt': None, 'tf': None}


# DO **NOT** add item to this set (unless the reason is approved)
UpperCamelCase__          ={
    'tests/models/esm/test_modeling_esmfold.py',  # The pipeline test mapping is added to `test_modeling_esm.py`
}






def 			lowerCamelCase__	(__lowerCamelCase		):

     if "ModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
          return "pt"
     elif "TFModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
          return "tf"
     elif "FlaxModelTesterMixin" in [x.__name__ for x in test_class.__bases__]:
          return "flax"
     else:
          return None






def 			lowerCamelCase__	(__lowerCamelCase,				__lowerCamelCase		):
     # Use the cached results
     if PIPELINE_TEST_MAPPING[task].get(__lowerCamelCase,				__lowerCamelCase		) is not None:
          return PIPELINE_TEST_MAPPING[task][framework]

     _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							pipeline_test_mapping[task]["test"]
     _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							None

     if framework == "pt":
          _SCREAMING_SNAKE_CASE				:		Tuple  			=							getattr(__lowerCamelCase,				"model_mapping",				__lowerCamelCase		)
     elif framework == "tf":
          _SCREAMING_SNAKE_CASE				:		Tuple  			=							getattr(__lowerCamelCase,				"tf_model_mapping",				__lowerCamelCase		)

     if mapping is not None:
          _SCREAMING_SNAKE_CASE				:		List[Any]  			=							dict(mapping.items()		)

     # cache the results
     _SCREAMING_SNAKE_CASE				:		int  			=							mapping
     return mapping






def 			lowerCamelCase__	(__lowerCamelCase,				__lowerCamelCase		):
     _SCREAMING_SNAKE_CASE				:		int  			=							get_framework(__lowerCamelCase		)
     if framework is None:
          return None
     _SCREAMING_SNAKE_CASE				:		Any  			=							get_mapping_for_task(__lowerCamelCase,				__lowerCamelCase		)
     if mapping is None:
          return None

     _SCREAMING_SNAKE_CASE				:		Tuple  			=							list({model_class.config_class for model_class in test_class.all_model_classes}		)
     if len(__lowerCamelCase		) != 1:
          raise ValueError("There should be exactly one configuration class from `test_class.all_model_classes`."		)

     # This could be a list/tuple of model classes, but it's rare.
     _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							mapping.get(config_classes[0],				__lowerCamelCase		)
     if isinstance(__lowerCamelCase,				(tuple, list)		):
          _SCREAMING_SNAKE_CASE				:		List[str]  			=							sorted(__lowerCamelCase,				key=lambda __lowerCamelCase		: x.__name__		)

     return model_class






def 			lowerCamelCase__	(__lowerCamelCase		):
     _SCREAMING_SNAKE_CASE				:		Any  			=							[(task, get_model_for_pipeline_test(__lowerCamelCase,				__lowerCamelCase		)) for task in pipeline_test_mapping]
     _SCREAMING_SNAKE_CASE				:		List[str]  			=							sorted([(task, model) for task, model in mapping if model is not None],				key=lambda __lowerCamelCase		: x[0]		)

     return dict(__lowerCamelCase		)






def 			lowerCamelCase__	(__lowerCamelCase		):
     _SCREAMING_SNAKE_CASE				:		str  			=							get_framework(__lowerCamelCase		)
     if framework == "pt":
          _SCREAMING_SNAKE_CASE				:		Tuple  			=							"torch"
     _SCREAMING_SNAKE_CASE				:		List[str]  			=							"{}"

     _SCREAMING_SNAKE_CASE				:		int  			=							get_pipeline_model_mapping(__lowerCamelCase		)
     if len(__lowerCamelCase		) == 0:
          return ""

     _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							[]
     for task, model_classes in mapping.items():
          if isinstance(__lowerCamelCase,				(tuple, list)		):
               # A list/tuple of model classes
               _SCREAMING_SNAKE_CASE				:		Any  			=							"(" + ", ".join([x.__name__ for x in model_classes]		) + ")"
          else:
               # A single model class
               _SCREAMING_SNAKE_CASE				:		List[Any]  			=							model_classes.__name__
          texts.append(f"""\"{task}\": {value}"""		)
     _SCREAMING_SNAKE_CASE				:		Dict  			=							"{" + ", ".join(__lowerCamelCase		) + "}"
     _SCREAMING_SNAKE_CASE				:		Any  			=							f"""pipeline_model_mapping = {text} if is_{framework}_available() else {default_value}"""

     return text






def 			lowerCamelCase__	(__lowerCamelCase		):
     _SCREAMING_SNAKE_CASE				:		Any  			=							{"ModelTesterMixin", "TFModelTesterMixin", "FlaxModelTesterMixin"}
     if not issubclass(__lowerCamelCase,				unittest.TestCase		):
          return False
     return len(base_class_names.intersection([x.__name__ for x in test_class.__bases__]		)		) > 0






def 			lowerCamelCase__	(__lowerCamelCase		):
     _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							[x for x in get_test_classes(__lowerCamelCase		) if is_valid_test_class(__lowerCamelCase		)]

     _SCREAMING_SNAKE_CASE				:		str  			=							None
     for test_class in test_classes:
          # If a test class has defined `pipeline_model_mapping`, let's take it
          if getattr(__lowerCamelCase,				"pipeline_model_mapping",				__lowerCamelCase		) is not None:
               _SCREAMING_SNAKE_CASE				:		int  			=							test_class
               break
    # Take the test class with the shortest name (just a heuristic)
     if target_test_class is None and len(__lowerCamelCase		) > 0:
          _SCREAMING_SNAKE_CASE				:		Dict  			=							sorted(__lowerCamelCase,				key=lambda __lowerCamelCase		: (len(x.__name__		), x.__name__)		)[0]

     return target_test_class






def 			lowerCamelCase__	(__lowerCamelCase,				__lowerCamelCase,				__lowerCamelCase		):
     _SCREAMING_SNAKE_CASE				:		int  			=							start_idx
     for idx, line in enumerate(lines[start_idx:]		):
          _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							len(__lowerCamelCase		) - len(line.lstrip()		)
          if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
               _SCREAMING_SNAKE_CASE				:		int  			=							start_idx + idx
          elif idx > 0 and indent <= indent_level:
               # Outside the definition block of `pipeline_model_mapping`
               break

     return end_idx






def 			lowerCamelCase__	(__lowerCamelCase,				__lowerCamelCase=False		):
     if getattr(__lowerCamelCase,				"pipeline_model_mapping",				__lowerCamelCase		) is not None:
          if not overwrite:
               return "", -1

     _SCREAMING_SNAKE_CASE				:		Tuple  			=							get_pipeline_model_mapping_string(__lowerCamelCase		)
     if len(__lowerCamelCase		) == 0:
          return "", -1
     _SCREAMING_SNAKE_CASE				:		Tuple  			=							line_to_add + "\n"

     # The code defined the class `test_class`
     _SCREAMING_SNAKE_CASE			,	_SCREAMING_SNAKE_CASE				:		Dict  			=							inspect.getsourcelines(__lowerCamelCase		)
     # `inspect` gives the code for an object, including decorator(s) if any.
     # We (only) need the exact line of the class definition.
     for idx, line in enumerate(__lowerCamelCase		):
          if line.lstrip().startswith("class "		):
               _SCREAMING_SNAKE_CASE				:		int  			=							class_lines[idx:]
               class_start_line_no += idx
               break
     _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							class_start_line_no + len(__lowerCamelCase		) - 1

     # The index in `class_lines` that starts the definition of `all_model_classes`, `all_generative_model_classes` or
     # `pipeline_model_mapping`. This assumes they are defined in such order, and we take the start index of the last
     # block that appears in a `test_class`.
     _SCREAMING_SNAKE_CASE				:		str  			=							None
     # The indent level of the line at `class_lines[start_idx]` (if defined)
     _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							0
     # To record if `pipeline_model_mapping` is found in `test_class`.
     _SCREAMING_SNAKE_CASE				:		int  			=							None
     for idx, line in enumerate(__lowerCamelCase		):
          if line.strip().startswith("all_model_classes = "		):
               _SCREAMING_SNAKE_CASE				:		str  			=							len(__lowerCamelCase		) - len(line.lstrip()		)
               _SCREAMING_SNAKE_CASE				:		Any  			=							idx
          elif line.strip().startswith("all_generative_model_classes = "		):
               _SCREAMING_SNAKE_CASE				:		Any  			=							len(__lowerCamelCase		) - len(line.lstrip()		)
               _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							idx
          elif line.strip().startswith("pipeline_model_mapping = "		):
               _SCREAMING_SNAKE_CASE				:		List[Any]  			=							len(__lowerCamelCase		) - len(line.lstrip()		)
               _SCREAMING_SNAKE_CASE				:		Any  			=							idx
               _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							line
               break

     if start_idx is None:
          return "", -1
     # Find the ending index (inclusive) of the above found block.
     _SCREAMING_SNAKE_CASE				:		int  			=							find_block_ending(__lowerCamelCase,				__lowerCamelCase,				__lowerCamelCase		)

     # Extract `is_xxx_available()` from existing blocks: some models require specific libraries like `timm` and use
     # `is_timm_available()` instead of `is_torch_available()`.
     # Keep leading and trailing whitespaces
     _SCREAMING_SNAKE_CASE				:		Tuple  			=							re.compile(R"\s(is_\S+?_available\(\))\s"		)
     for line in class_lines[start_idx : end_idx + 1]:
          _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							r.search(__lowerCamelCase		)
          if backend_condition is not None:
               # replace the leading and trailing whitespaces to the space character " ".
               _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							" " + backend_condition[0][1:-1] + " "
               _SCREAMING_SNAKE_CASE				:		List[str]  			=							r.sub(__lowerCamelCase,				__lowerCamelCase		)
               break

     if def_line is None:
          # `pipeline_model_mapping` is not defined. The target index is set to the ending index (inclusive) of
          # `all_model_classes` or `all_generative_model_classes`.
          _SCREAMING_SNAKE_CASE				:		Any  			=							end_idx
     else:
          # `pipeline_model_mapping` is defined. The target index is set to be one **BEFORE** its start index.
          _SCREAMING_SNAKE_CASE				:		Tuple  			=							start_idx - 1
          # mark the lines of the currently existing `pipeline_model_mapping` to be removed.
          for idx in range(__lowerCamelCase,				end_idx + 1		):
               # These lines are going to be removed before writing to the test file.
               _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							None  # noqa

    # Make sure the test class is a subclass of `PipelineTesterMixin`.
     _SCREAMING_SNAKE_CASE				:		List[Any]  			=							[x.__name__ for x in test_class.__bases__]
     if "PipelineTesterMixin" not in parent_classes:
          # Put `PipelineTesterMixin` just before `unittest.TestCase`
          _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							[x for x in parent_classes if x != "TestCase"] + ["PipelineTesterMixin"]
          if "TestCase" in parent_classes:
               # Here we **assume** the original string is always with `unittest.TestCase`.
               _parent_classes.append("unittest.TestCase"		)
          _SCREAMING_SNAKE_CASE				:		List[str]  			=							", ".join(_parent_classes		)
          for idx, line in enumerate(__lowerCamelCase		):
               # Find the ending of the declaration of `test_class`
               if line.strip().endswith("):"		):
                    # mark the lines of the declaration of `test_class` to be removed
                    for _idx in range(idx + 1		):
                         _SCREAMING_SNAKE_CASE				:		Tuple  			=							None  # noqa
                    break
        # Add the new, one-line, class declaration for `test_class`
          _SCREAMING_SNAKE_CASE				:		Any  			=							f"""class {test_class.__name__}({parent_classes}):\n"""

     # Add indentation
     _SCREAMING_SNAKE_CASE				:		Tuple  			=							" " * indent_level + line_to_add
     # Insert `pipeline_model_mapping` to `class_lines`.
     # (The line at `target_idx` should be kept by definition!)
     _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							class_lines[: target_idx + 1] + [line_to_add] + class_lines[target_idx + 1 :]
     # Remove the lines that are marked to be removed
     _SCREAMING_SNAKE_CASE				:		List[Any]  			=							[x for x in class_lines if x is not None]

     # Move from test class to module (in order to write to the test file)
     _SCREAMING_SNAKE_CASE				:		Union[str, Any]  			=							inspect.getsourcelines(inspect.getmodule(__lowerCamelCase		)		)[0]
     # Be careful with the 1-off between line numbers and array indices
     _SCREAMING_SNAKE_CASE				:		Optional[int]  			=							module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
     _SCREAMING_SNAKE_CASE				:		Dict  			=							"".join(__lowerCamelCase		)

     _SCREAMING_SNAKE_CASE				:		Optional[Any]  			=							inspect.getsourcefile(__lowerCamelCase		)
     with open(__lowerCamelCase,				"w",				encoding="UTF-8",				newline="\n"		) as fp:
          fp.write(__lowerCamelCase		)

     return line_to_add






def 			lowerCamelCase__	(__lowerCamelCase,				__lowerCamelCase=False		):
     _SCREAMING_SNAKE_CASE				:		str  			=							find_test_class(__lowerCamelCase		)
     if test_class:
          add_pipeline_model_mapping(__lowerCamelCase,				overwrite=__lowerCamelCase		)


if __name__ == "__main__":
       UpperCamelCase__          =argparse.ArgumentParser()
       parser.add_argument(
           '--test_file', type=str, help='A path to the test file, starting with the repository\'s `tests` directory.'
       )
       parser.add_argument(
           '--all',
           action='store_true',
           help='If to check and modify all test files.',
       )
       parser.add_argument(
           '--overwrite',
           action='store_true',
           help='If to overwrite a test class if it has already defined `pipeline_model_mapping`.',
       )
       UpperCamelCase__          =parser.parse_args()

       if not args.all and not args.test_file:
              raise ValueError('Please specify either `test_file` or pass `--all` to check/modify all test files.')
       elif args.all and args.test_file:
              raise ValueError('Only one of `--test_file` and `--all` could be specified.')

       UpperCamelCase__          =[]
       if args.test_file:
              UpperCamelCase__          =[args.test_file]
       else:
              UpperCamelCase__          =os.path.join('tests', 'models', '**', 'test_modeling_*.py')
              for test_file in glob.glob(pattern):
                     # `Flax` is not concerned at this moment
                     if not test_file.startswith('test_modeling_flax_'):
                            test_files.append(test_file)

       for test_file in test_files:
              if test_file in TEST_FILE_TO_IGNORE:
                     print(f"[SKIPPED] {test_file} is skipped as it is in `TEST_FILE_TO_IGNORE` in the file {__file__}.")
                     continue
              add_pipeline_model_mapping_to_test_file(test_file, overwrite=args.overwrite)






import warnings
from dataclasses import dataclass, field
from typing import Optional, Tuple

from .training_args import TrainingArguments
from .utils import cached_property, is_tf_available, logging, requires_backends


UpperCamelCase__          =logging.get_logger(__name__)

if is_tf_available():
       import tensorflow as tf





@dataclass
class 	lowerCAmelCase__(       __lowercase			):






       '''simple docstring'''


       __snake_case       	=			'tf'
       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Name of TPU'}     ,   )

       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Zone of TPU'}     ,   )

       __snake_case       	=			field(
           default=__lowercase     ,   metadata={'help': 'Name of Cloud TPU-enabled project'}     ,   )

       __snake_case       	=			field(
           default=1.0     ,   metadata={'help': 'Power for the Polynomial decay LR scheduler.'}     ,   )

       __snake_case       	=			field(default=__lowercase     ,   metadata={'help': 'Whether to activate the XLA compilation or not'}			)




       @cached_property
       def 			UpperCamelCase_	(					self							)	->   Tuple["tf.distribute.Strategy", int]:
            requires_backends(self			,  ["tf"]							)
            logger.info("Tensorflow: setting up strategy"							)

            _SCREAMING_SNAKE_CASE				:		List[Any]  			=							tf.config.list_physical_devices("GPU"							)

            # Set to float16 at first
            if self.fpaa:
                 tf.keras.mixed_precision.set_global_policy("mixed_float16"							)

            if self.no_cuda:
                 _SCREAMING_SNAKE_CASE				:		List[Any]  			=							tf.distribute.OneDeviceStrategy(device="/cpu:0"							)
            else:
                 try:
                      if self.tpu_name:
                           _SCREAMING_SNAKE_CASE				:		Any  			=							tf.distribute.cluster_resolver.TPUClusterResolver(
                               self.tpu_name			,  zone=self.tpu_zone			,  project=self.gcp_project							)
                      else:
                           _SCREAMING_SNAKE_CASE				:		str  			=							tf.distribute.cluster_resolver.TPUClusterResolver()
                 except ValueError:
                      if self.tpu_name:
                           raise RuntimeError(F"""Couldn't connect to TPU {self.tpu_name}!"""							)
                      else:
                           _SCREAMING_SNAKE_CASE				:		Any  			=							None

                 if tpu:
                      # Set to bfloat16 in case of TPU
                      if self.fpaa:
                           tf.keras.mixed_precision.set_global_policy("mixed_bfloat16"							)

                      tf.config.experimental_connect_to_cluster(__lowerCamelCase							)
                      tf.tpu.experimental.initialize_tpu_system(__lowerCamelCase							)

                      _SCREAMING_SNAKE_CASE				:		List[str]  			=							tf.distribute.TPUStrategy(__lowerCamelCase							)

                 elif len(__lowerCamelCase							) == 0:
                      _SCREAMING_SNAKE_CASE				:		Dict  			=							tf.distribute.OneDeviceStrategy(device="/cpu:0"							)
                 elif len(__lowerCamelCase							) == 1:
                      _SCREAMING_SNAKE_CASE				:		Dict  			=							tf.distribute.OneDeviceStrategy(device="/gpu:0"							)
                 elif len(__lowerCamelCase							) > 1:
                      # If you only want to use a specific subset of GPUs use `CUDA_VISIBLE_DEVICES=0`
                      _SCREAMING_SNAKE_CASE				:		Any  			=							tf.distribute.MirroredStrategy()
                 else:
                      raise ValueError("Cannot find the proper strategy, please check your environment properties."							)

            return strategy




       @property
       def 			UpperCamelCase_	(					self							)	->   "tf.distribute.Strategy":
            requires_backends(self			,  ["tf"]							)
            return self._setup_strategy




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            requires_backends(self			,  ["tf"]							)
            return self._setup_strategy.num_replicas_in_sync




       @property
       def 			UpperCamelCase_	(					self							)	->   List[Any]:
            return False  # TF Logging is handled by Keras not the Trainer




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            if self.per_gpu_train_batch_size:
                 logger.warning(
                     "Using deprecated `--per_gpu_train_batch_size` argument which will be removed in a future "
                     "version. Using `--per_device_train_batch_size` is preferred."							)
            _SCREAMING_SNAKE_CASE				:		Tuple  			=							self.per_gpu_train_batch_size or self.per_device_train_batch_size
            return per_device_batch_size * self.n_replicas




       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            if self.per_gpu_eval_batch_size:
                 logger.warning(
                     "Using deprecated `--per_gpu_eval_batch_size` argument which will be removed in a future "
                     "version. Using `--per_device_eval_batch_size` is preferred."							)
            _SCREAMING_SNAKE_CASE				:		int  			=							self.per_gpu_eval_batch_size or self.per_device_eval_batch_size
            return per_device_batch_size * self.n_replicas







       @property
       def 			UpperCamelCase_	(					self							)	->   int:
            requires_backends(self			,  ["tf"]							)
            warnings.warn(
                "The n_gpu argument is deprecated and will be removed in a future version, use n_replicas instead."			,  __lowerCamelCase			,  )
            return self._setup_strategy.num_replicas_in_sync