
"""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


a		:					Optional[int]    	=					{}
for task, _ in pipeline_test_mapping.items():
    a		:					Union[str, Any]    	=					{'pt': None, 'tf': None}


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



def 						__magic_name__		(				UpperCamelCase    :		List[Any]    )    ->	Any:

    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 						__magic_name__		(				UpperCamelCase    :		Dict       ,		UpperCamelCase    :		Optional[int]    )    ->	Union[str, Any]:
    # Use the cached results
    if PIPELINE_TEST_MAPPING[task].get(UpperCamelCase       ,		UpperCamelCase    ) is not None:
        return PIPELINE_TEST_MAPPING[task][framework]

    a__    = pipeline_test_mapping[task]['test']
    a__    = None

    if framework == "pt":
        a__    = getattr(UpperCamelCase       ,		'model_mapping'       ,		UpperCamelCase    )
    elif framework == "tf":
        a__    = getattr(UpperCamelCase       ,		'tf_model_mapping'       ,		UpperCamelCase    )

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

    # cache the results
    a__    = mapping
    return mapping



def 						__magic_name__		(				UpperCamelCase    :		str       ,		UpperCamelCase    :		int    )    ->	List[str]:
    a__    = get_framework(UpperCamelCase    )
    if framework is None:
        return None
    a__    = get_mapping_for_task(UpperCamelCase       ,		UpperCamelCase    )
    if mapping is None:
        return None

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

    return model_class



def 						__magic_name__		(				UpperCamelCase    :		List[Any]    )    ->	int:
    a__    = [(task, get_model_for_pipeline_test(UpperCamelCase       ,		UpperCamelCase    )) for task in pipeline_test_mapping]
    a__    = sorted([(task, model) for task, model in mapping if model is not None]       ,		key=lambda UpperCamelCase    : x[0]    )

    return dict(UpperCamelCase    )



def 						__magic_name__		(				UpperCamelCase    :		str    )    ->	Tuple:
    a__    = get_framework(UpperCamelCase    )
    if framework == "pt":
        a__    = 'torch'
    a__    = '{}'

    a__    = get_pipeline_model_mapping(UpperCamelCase    )
    if len(UpperCamelCase    ) == 0:
        return ""

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

    return text



def 						__magic_name__		(				UpperCamelCase    :		Tuple    )    ->	Union[str, Any]:
    a__    = {'ModelTesterMixin', 'TFModelTesterMixin', 'FlaxModelTesterMixin'}
    if not issubclass(UpperCamelCase       ,		unittest.TestCase    ):
        return False
    return len(base_class_names.intersection([x.__name__ for x in test_class.__bases__]    )    ) > 0



def 						__magic_name__		(				UpperCamelCase    :		List[str]    )    ->	str:
    a__    = [x for x in get_test_classes(UpperCamelCase    ) if is_valid_test_class(UpperCamelCase    )]

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

    return target_test_class



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

    return end_idx



def 						__magic_name__		(				UpperCamelCase    :		Tuple       ,		UpperCamelCase    :		int=False    )    ->	Any:
    if getattr(UpperCamelCase       ,		'pipeline_model_mapping'       ,		UpperCamelCase    ) is not None:
        if not overwrite:
            return "", -1

    a__    = get_pipeline_model_mapping_string(UpperCamelCase    )
    if len(UpperCamelCase    ) == 0:
        return "", -1
    a__    = line_to_add + '\n'

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

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

    # 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
    a__    = re.compile(r'\s(is_\S+?_available\(\))\s'    )
    for line in class_lines[start_idx : end_idx + 1]:
        a__    = r.search(UpperCamelCase    )
        if backend_condition is not None:
            # replace the leading and trailing whitespaces to the space character " ".
            a__    = ' ' + backend_condition[0][1:-1] + ' '
            a__    = r.sub(UpperCamelCase       ,		UpperCamelCase    )
            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`.
        a__    = end_idx
    else:
        # `pipeline_model_mapping` is defined. The target index is set to be one **BEFORE** its start index.
        a__    = start_idx - 1
        # mark the lines of the currently existing `pipeline_model_mapping` to be removed.
        for idx in range(UpperCamelCase       ,		end_idx + 1    ):
            # These lines are going to be removed before writing to the test file.
            a__    = None  # noqa

    # Make sure the test class is a subclass of `PipelineTesterMixin`.
    a__    = [x.__name__ for x in test_class.__bases__]
    if "PipelineTesterMixin" not in parent_classes:
        # Put `PipelineTesterMixin` just before `unittest.TestCase`
        a__    = [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'    )
        a__    = ', '.join(_parent_classes    )
        for idx, line in enumerate(UpperCamelCase    ):
            # 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    ):
                    a__    = None  # noqa
                break
        # Add the new, one-line, class declaration for `test_class`
        a__    = f'class {test_class.__name__}({parent_classes}):\n'

    # Add indentation
    a__    = ' ' * indent_level + line_to_add
    # Insert `pipeline_model_mapping` to `class_lines`.
    # (The line at `target_idx` should be kept by definition!)
    a__    = class_lines[: target_idx + 1] + [line_to_add] + class_lines[target_idx + 1 :]
    # Remove the lines that are marked to be removed
    a__    = [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)
    a__    = inspect.getsourcelines(inspect.getmodule(UpperCamelCase    )    )[0]
    # Be careful with the 1-off between line numbers and array indices
    a__    = module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
    a__    = ''.join(UpperCamelCase    )

    a__    = inspect.getsourcefile(UpperCamelCase    )
    with open(UpperCamelCase       ,		'w'       ,		encoding='UTF-8'       ,		newline='\n'    ) as fp:
        fp.write(UpperCamelCase    )

    return line_to_add



def 						__magic_name__		(				UpperCamelCase    :		int       ,		UpperCamelCase    :		List[Any]=False    )    ->	Optional[Any]:
    a__    = find_test_class(UpperCamelCase    )
    if test_class:
        add_pipeline_model_mapping(UpperCamelCase       ,		overwrite=UpperCamelCase    )


if __name__ == "__main__":
    a		:					Optional[int]    	=					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`.',
    )
    a		:					Dict    	=					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.')

    a		:					int    	=					[]
    if args.test_file:
        a		:					int    	=					[args.test_file]
    else:
        a		:					Optional[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"""




import os
import pickle
import unittest
from collections import UserDict, namedtuple

import torch

from accelerate.test_utils.testing import require_cuda, require_torch_min_version
from accelerate.test_utils.training import RegressionModel
from accelerate.utils import (
    convert_outputs_to_fpaa,
    extract_model_from_parallel,
    find_device,
    listify,
    patch_environment,
    recursively_apply,
    send_to_device,
)


a		:					Tuple    	=					namedtuple('ExampleNamedTuple', 'a b c')







class 						lowercase(unittest.TestCase					):



     def  lowercase__     (       self      )		->    Union[str, Any]:







         """simple docstring"""



         a__    = torch.randn(5    ,							2      )
         a__    = torch.device('cuda'      ) if torch.cuda.is_available() else torch.device('cpu'      )

         a__    = send_to_device(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta.cpu()    ,							__SCREAMING_SNAKE_CASE      )      )

         a__    = send_to_device((tensor, [tensor, tensor], 1)    ,							__SCREAMING_SNAKE_CASE      )
         self.assertIsInstance(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta[0].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertIsInstance(resulta[1]    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta[1][0].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertTrue(torch.equal(resulta[1][1].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertEqual(resulta[2]    ,							1      )

         a__    = send_to_device({'a': tensor, 'b': [tensor, tensor], 'c': 1}    ,							__SCREAMING_SNAKE_CASE      )
         self.assertIsInstance(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta['a'].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertIsInstance(resulta['b']    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta['b'][0].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertTrue(torch.equal(resulta['b'][1].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertEqual(resulta['c']    ,							1      )

         a__    = send_to_device(ExampleNamedTuple(a=__SCREAMING_SNAKE_CASE    ,							b=[tensor, tensor]    ,							c=1      )    ,							__SCREAMING_SNAKE_CASE      )
         self.assertIsInstance(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta.a.cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertIsInstance(resulta.b    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta.b[0].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertTrue(torch.equal(resulta.b[1].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertEqual(resulta.c    ,							1      )

         a__    = send_to_device(UserDict({'a': tensor, 'b': [tensor, tensor], 'c': 1}      )    ,							__SCREAMING_SNAKE_CASE      )
         self.assertIsInstance(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta['a'].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertIsInstance(resulta['b']    ,							__SCREAMING_SNAKE_CASE      )
         self.assertTrue(torch.equal(resulta['b'][0].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertTrue(torch.equal(resulta['b'][1].cpu()    ,							__SCREAMING_SNAKE_CASE      )      )
         self.assertEqual(resulta['c']    ,							1      )



     def  lowercase__     (       self      )		->    Dict:







         """simple docstring"""



         with self.assertRaises(__SCREAMING_SNAKE_CASE      ) as cm:
             a__    = recursively_apply(torch.tensor    ,							(torch.tensor(1      ), 1)    ,							error_on_other_type=__SCREAMING_SNAKE_CASE      )
         self.assertEqual(
             str(cm.exception      )    ,							'Unsupported types (<class \'int\'>) passed to `tensor`. Only nested list/tuple/dicts of objects that are valid for `is_torch_tensor` should be passed.'    ,							)



     def  lowercase__     (       self      )		->    Any:







         """simple docstring"""



         a__    = torch.tensor([1, 2, 3, 4, 5]      )
         self.assertEqual(listify(__SCREAMING_SNAKE_CASE      )    ,							[1, 2, 3, 4, 5]      )

         a__    = torch.tensor([[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]      )
         self.assertEqual(listify(__SCREAMING_SNAKE_CASE      )    ,							[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]]      )

         a__    = torch.tensor([[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]      )
         self.assertEqual(
             listify(__SCREAMING_SNAKE_CASE      )    ,							[[[1, 2, 3, 4, 5], [6, 7, 8, 9, 1_0]], [[1_1, 1_2, 1_3, 1_4, 1_5], [1_6, 1_7, 1_8, 1_9, 2_0]]]      )



     def  lowercase__     (       self      )		->    str:







         """simple docstring"""



         with patch_environment(aa=1    ,							BB=2      ):
             self.assertEqual(os.environ.get('AA'      )    ,							'1'      )
             self.assertEqual(os.environ.get('BB'      )    ,							'2'      )

         self.assertNotIn('AA'    ,							os.environ      )
         self.assertNotIn('BB'    ,							os.environ      )



     def  lowercase__     (       self      )		->    Tuple:







         """simple docstring"""



         a__    = RegressionModel()
         a__    = model.forward
         a__    = convert_outputs_to_fpaa(model.forward      )
         a__    = extract_model_from_parallel(__SCREAMING_SNAKE_CASE    ,							keep_fpaa_wrapper=__SCREAMING_SNAKE_CASE      )
         a__    = pickle.dumps(__SCREAMING_SNAKE_CASE      )



     @require_cuda
     def  lowercase__     (       self      )		->    Union[str, Any]:







         """simple docstring"""



         a__    = RegressionModel()
         a__    = model.forward
         a__    = torch.cuda.amp.autocast(dtype=torch.floataa      )(model.forward      )
         a__    = convert_outputs_to_fpaa(model.forward      )
         a__    = extract_model_from_parallel(__SCREAMING_SNAKE_CASE    ,							keep_fpaa_wrapper=__SCREAMING_SNAKE_CASE      )
         a__    = pickle.dumps(__SCREAMING_SNAKE_CASE      )



     @require_cuda
     @require_torch_min_version(version='2.0'      )
     def  lowercase__     (       self      )		->    Optional[int]:







         """simple docstring"""



         a__    = RegressionModel()
         a__    = model.forward
         a__    = torch.cuda.amp.autocast(dtype=torch.floataa      )(model.forward      )
         a__    = convert_outputs_to_fpaa(model.forward      )
         a__    = torch.compile(model.forward    ,							backend='inductor'      )
         a__    = torch.randn(4    ,							1_0      ).cuda()
         a__    = model(__SCREAMING_SNAKE_CASE      )



     def  lowercase__     (       self      )		->    str:







         """simple docstring"""



         a__    = RegressionModel()
         # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
         a__    = torch.nn.parallel.DataParallel(__SCREAMING_SNAKE_CASE      )
         a__    = extract_model_from_parallel(__SCREAMING_SNAKE_CASE      )

         self.assertEqual(__SCREAMING_SNAKE_CASE    ,							__SCREAMING_SNAKE_CASE      )



     @require_torch_min_version(version='2.0'      )
     def  lowercase__     (       self      )		->    str:







         """simple docstring"""



         a__    = RegressionModel()
         a__    = torch.compile(__SCREAMING_SNAKE_CASE      )

         # could also do a test with DistributedDataParallel, but difficult to run on CPU or single GPU
         a__    = torch.nn.parallel.DataParallel(__SCREAMING_SNAKE_CASE      )
         a__    = torch.compile(__SCREAMING_SNAKE_CASE      )
         a__    = extract_model_from_parallel(__SCREAMING_SNAKE_CASE      )

         self.assertEqual(compiled_model._orig_mod    ,							compiled_model_unwrapped._orig_mod      )




     def  lowercase__     (       self      )		->    Optional[Any]:







         """simple docstring"""



         self.assertEqual(find_device([1, 'a', torch.tensor([1, 2, 3]      )]      )    ,							torch.device('cpu'      )      )
         self.assertEqual(find_device({'a': 1, 'b': torch.tensor([1, 2, 3]      )}      )    ,							torch.device('cpu'      )      )
         self.assertIsNone(find_device([1, 'a']      )      )


