


'''simple docstring'''


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_	:     List[str]							     =	{}
for task, _ in pipeline_test_mapping.items():
				UpperCAmelCase_	:     List[Any]							     =	{'pt': None, 'tf': None}


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







def     snake_case_     (						SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





				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     snake_case_     (						SCREAMING_SNAKE_CASE__							,   SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





				if PIPELINE_TEST_MAPPING[task].get(lowercase_							,   lowercase_  ) is not None:
								return PIPELINE_TEST_MAPPING[task][framework]

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

				if framework == "pt":
								_SCREAMING_SNAKE_CASE				: str  	=   getattr(lowercase_							,   """model_mapping"""							,   lowercase_  )
				elif framework == "tf":
								_SCREAMING_SNAKE_CASE				: Any  	=   getattr(lowercase_							,   """tf_model_mapping"""							,   lowercase_  )

				if mapping is not None:
								_SCREAMING_SNAKE_CASE				: Tuple  	=   dict(mapping.items()  )

				# cache the results
				_SCREAMING_SNAKE_CASE				: Tuple  	=   mapping
				return mapping







def     snake_case_     (						SCREAMING_SNAKE_CASE__							,   SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE				: Dict  	=   get_framework(lowercase_  )
				if framework is None:
								return None
				_SCREAMING_SNAKE_CASE				: Any  	=   get_mapping_for_task(lowercase_							,   lowercase_  )
				if mapping is None:
								return None

				_SCREAMING_SNAKE_CASE				: List[str]  	=   list({model_class.config_class for model_class in test_class.all_model_classes}  )
				if len(lowercase_  ) != 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				: str  	=   mapping.get(config_classes[0]							,   lowercase_  )
				if isinstance(lowercase_							,   (tuple, list)  ):
								_SCREAMING_SNAKE_CASE				: Dict  	=   sorted(lowercase_							,   key=lambda SCREAMING_SNAKE_CASE__  : x.__name__  )

				return model_class







def     snake_case_     (						SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





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

				return dict(lowercase_  )







def     snake_case_     (						SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE				: Tuple  	=   get_framework(lowercase_  )
				if framework == "pt":
								_SCREAMING_SNAKE_CASE				: Any  	=   """torch"""
				_SCREAMING_SNAKE_CASE				: Tuple  	=   """{}"""

				_SCREAMING_SNAKE_CASE				: Tuple  	=   get_pipeline_model_mapping(lowercase_  )
				if len(lowercase_  ) == 0:
								return ""

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

				return text







def     snake_case_     (						SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





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







def     snake_case_     (						SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE				: Optional[int]  	=   [x for x in get_test_classes(lowercase_  ) if is_valid_test_class(lowercase_  )]

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

				return target_test_class







def     snake_case_     (						SCREAMING_SNAKE_CASE__							,   SCREAMING_SNAKE_CASE__							,   SCREAMING_SNAKE_CASE__  ):







				"""simple docstring"""





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

				return end_idx







def     snake_case_     (						SCREAMING_SNAKE_CASE__							,   SCREAMING_SNAKE_CASE__=False  ):







				"""simple docstring"""





				if getattr(lowercase_							,   """pipeline_model_mapping"""							,   lowercase_  ) is not None:
								if not overwrite:
												return "", -1

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

				# The code defined the class `test_class`
				_SCREAMING_SNAKE_CASE	, _SCREAMING_SNAKE_CASE				: int  	=   inspect.getsourcelines(lowercase_  )
				# `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(lowercase_  ):
								if line.lstrip().startswith("""class """  ):
												_SCREAMING_SNAKE_CASE				: Optional[int]  	=   class_lines[idx:]
												class_start_line_no += idx
												break
				_SCREAMING_SNAKE_CASE				: str  	=   class_start_line_no + len(lowercase_  ) - 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				: List[str]  	=   None
				# The indent level of the line at `class_lines[start_idx]` (if defined)
				_SCREAMING_SNAKE_CASE				: str  	=   0
				# To record if `pipeline_model_mapping` is found in `test_class`.
				_SCREAMING_SNAKE_CASE				: Tuple  	=   None
				for idx, line in enumerate(lowercase_  ):
								if line.strip().startswith("""all_model_classes = """  ):
												_SCREAMING_SNAKE_CASE				: Tuple  	=   len(lowercase_  ) - len(line.lstrip()  )
												_SCREAMING_SNAKE_CASE				: List[str]  	=   idx
								elif line.strip().startswith("""all_generative_model_classes = """  ):
												_SCREAMING_SNAKE_CASE				: List[str]  	=   len(lowercase_  ) - len(line.lstrip()  )
												_SCREAMING_SNAKE_CASE				: int  	=   idx
								elif line.strip().startswith("""pipeline_model_mapping = """  ):
												_SCREAMING_SNAKE_CASE				: Tuple  	=   len(lowercase_  ) - len(line.lstrip()  )
												_SCREAMING_SNAKE_CASE				: str  	=   idx
												_SCREAMING_SNAKE_CASE				: List[Any]  	=   line
												break

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

				# 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				: int  	=   re.compile(R"""\s(is_\S+?_available\(\))\s"""  )
				for line in class_lines[start_idx : end_idx + 1]:
								_SCREAMING_SNAKE_CASE				: Union[str, Any]  	=   r.search(lowercase_  )
								if backend_condition is not None:
												# replace the leading and trailing whitespaces to the space character " ".
												_SCREAMING_SNAKE_CASE				: int  	=   """ """ + backend_condition[0][1:-1] + """ """
												_SCREAMING_SNAKE_CASE				: Dict  	=   r.sub(lowercase_							,   lowercase_  )
												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				: Optional[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(lowercase_							,   end_idx + 1  ):
												# These lines are going to be removed before writing to the test file.
												_SCREAMING_SNAKE_CASE				: List[str]  	=   None  # noqa

    # Make sure the test class is a subclass of `PipelineTesterMixin`.
				_SCREAMING_SNAKE_CASE				: 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				: str  	=   """, """.join(_parent_classes  )
								for idx, line in enumerate(lowercase_  ):
												# 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				: Union[str, Any]  	=   None  # noqa
																break
        # Add the new, one-line, class declaration for `test_class`
								_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   f"""class {test_class.__name__}({parent_classes}):\n"""

				# Add indentation
				_SCREAMING_SNAKE_CASE				: List[Any]  	=   """ """ * 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				: int  	=   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				: str  	=   [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				: str  	=   inspect.getsourcelines(inspect.getmodule(lowercase_  )  )[0]
				# Be careful with the 1-off between line numbers and array indices
				_SCREAMING_SNAKE_CASE				: Tuple  	=   module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
				_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   """""".join(lowercase_  )

				_SCREAMING_SNAKE_CASE				: str  	=   inspect.getsourcefile(lowercase_  )
				with open(lowercase_							,   """w"""							,   encoding="""UTF-8"""							,   newline="""\n"""  ) as fp:
								fp.write(lowercase_  )

				return line_to_add







def     snake_case_     (						SCREAMING_SNAKE_CASE__							,   SCREAMING_SNAKE_CASE__=False  ):







				"""simple docstring"""





				_SCREAMING_SNAKE_CASE				: Optional[Any]  	=   find_test_class(lowercase_  )
				if test_class:
								add_pipeline_model_mapping(lowercase_							,   overwrite=lowercase_  )


if __name__ == "__main__":
				UpperCAmelCase_	:     List[str]							     =	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_	:     Union[str, Any]							     =	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_	:     Optional[Any]							     =	[]
				if args.test_file:
								UpperCAmelCase_	:     List[Any]							     =	[args.test_file]
				else:
								UpperCAmelCase_	:     Union[str, Any]							     =	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)








'''simple docstring'''


from typing import TYPE_CHECKING

from ...utils import OptionalDependencyNotAvailable, _LazyModule, is_torch_available


UpperCAmelCase_	:     Any							     =	{
    'configuration_musicgen': [
        'MUSICGEN_PRETRAINED_CONFIG_ARCHIVE_MAP',
        'MusicgenConfig',
        'MusicgenDecoderConfig',
    ],
    'processing_musicgen': ['MusicgenProcessor'],
}

try:
				if not is_torch_available():
								raise OptionalDependencyNotAvailable()
except OptionalDependencyNotAvailable:
				pass
else:
				UpperCAmelCase_	:     Any							     =	[
				    'MUSICGEN_PRETRAINED_MODEL_ARCHIVE_LIST',
				    'MusicgenForConditionalGeneration',
				    'MusicgenForCausalLM',
				    'MusicgenModel',
				    'MusicgenPreTrainedModel',
				]

if TYPE_CHECKING:
				from .configuration_musicgen import (
				    MUSICGEN_PRETRAINED_CONFIG_ARCHIVE_MAP,
				    MusicgenConfig,
				    MusicgenDecoderConfig,
				)
				from .processing_musicgen import MusicgenProcessor

				try:
								if not is_torch_available():
												raise OptionalDependencyNotAvailable()
				except OptionalDependencyNotAvailable:
								pass
				else:
								from .modeling_musicgen import (
								    MUSICGEN_PRETRAINED_MODEL_ARCHIVE_LIST,
								    MusicgenForCausalLM,
								    MusicgenForConditionalGeneration,
								    MusicgenModel,
								    MusicgenPreTrainedModel,
								)

else:
				import sys

				UpperCAmelCase_	:     Tuple							     =	_LazyModule(__name__, globals()['__file__'], _import_structure, module_spec=__spec__)




