




'''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			   ={}
for task, _ in pipeline_test_mapping.items():
   A			   ={'pt': None, 'tf': None}


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





def 						snake_case_							(_a     :							int					):

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

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

    if framework == "pt":
        UpperCAmelCase      	=		getattr(_a   ,       '''model_mapping'''   ,       _a					)
    elif framework == "tf":
        UpperCAmelCase      	=		getattr(_a   ,       '''tf_model_mapping'''   ,       _a					)

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

    # cache the results
    UpperCAmelCase      	=		mapping
    return mapping





def 						snake_case_							(_a     :							Tuple   ,       _a     :							int					):
    UpperCAmelCase      	=		get_framework(_a					)
    if framework is None:
        return None
    UpperCAmelCase      	=		get_mapping_for_task(_a   ,       _a					)
    if mapping is None:
        return None

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

    return model_class





def 						snake_case_							(_a     :							Any					):
    UpperCAmelCase      	=		[(task, get_model_for_pipeline_test(_a   ,       _a					)) for task in pipeline_test_mapping]
    UpperCAmelCase      	=		sorted([(task, model) for task, model in mapping if model is not None]   ,       key=lambda _a					: x[0]					)

    return dict(_a					)





def 						snake_case_							(_a     :							Optional[int]					):
    UpperCAmelCase      	=		get_framework(_a					)
    if framework == "pt":
        UpperCAmelCase      	=		'''torch'''
    UpperCAmelCase      	=		'''{}'''

    UpperCAmelCase      	=		get_pipeline_model_mapping(_a					)
    if len(_a					) == 0:
        return ""

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

    return text





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





def 						snake_case_							(_a     :							Dict					):
    UpperCAmelCase      	=		[x for x in get_test_classes(_a					) if is_valid_test_class(_a					)]

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

    return target_test_class





def 						snake_case_							(_a     :							Any   ,       _a     :							int   ,       _a     :							int					):
    UpperCAmelCase      	=		start_idx
    for idx, line in enumerate(lines[start_idx:]					):
        UpperCAmelCase      	=		len(_a					) - len(line.lstrip()					)
        if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
            UpperCAmelCase      	=		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_							(_a     :							str   ,       _a     :							Any=False					):
    if getattr(_a   ,       '''pipeline_model_mapping'''   ,       _a					) is not None:
        if not overwrite:
            return "", -1

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

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

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

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

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

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

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

    return line_to_add





def 						snake_case_							(_a     :							Optional[int]   ,       _a     :							Union[str, Any]=False					):
    UpperCAmelCase      	=		find_test_class(_a					)
    if test_class:
        add_pipeline_model_mapping(_a   ,       overwrite=_a					)


if __name__ == "__main__":
   A			   =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			   =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			   =[]
   if args.test_file:
      A			   =[args.test_file]
   else:
      A			   =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 tempfile
import unittest
from typing import List

from transformers.models.esm.tokenization_esm import VOCAB_FILES_NAMES, EsmTokenizer
from transformers.testing_utils import require_tokenizers
from transformers.tokenization_utils import PreTrainedTokenizer
from transformers.tokenization_utils_base import PreTrainedTokenizerBase


@require_tokenizers
class       _a  (     unittest.TestCase			):
       __a    :					Optional[int]           = EsmTokenizer
       def    A				(     self						:     Tuple    ):



           '''simple docstring'''

           super().setUp()
           UpperCAmelCase      	=		tempfile.mkdtemp()
           # fmt: off
           UpperCAmelCase      	=		["<cls>", "<pad>", "<eos>", "<unk>", "L", "A", "G", "V", "S", "E", "R", "T", "I", "D", "P", "K", "Q", "N", "F", "Y", "M", "H", "W", "C", "X", "B", "U", "Z", "O", ".", "-", "<null_1>", "<mask>"]  # noqa: E501
           # fmt: on
           UpperCAmelCase      	=		os.path.join(self.tmpdirname       ,    VOCAB_FILES_NAMES['''vocab_file''']    )
           with open(self.vocab_file       ,    '''w'''       ,    encoding='''utf-8'''    ) as vocab_writer:
               vocab_writer.write(''''''.join([x + '''\n''' for x in vocab_tokens]    )    )
       def    A				(     self						:     Union[str, Any]       ,    **lowercase						:     Tuple    ):



           '''simple docstring'''

           return [self.get_tokenizer(**lowercase    )]
       def    A				(     self						:     Optional[Any]       ,    **lowercase						:     str    ):



           '''simple docstring'''

           return self.tokenizer_class.from_pretrained(self.tmpdirname       ,    **lowercase    )
       def    A				(     self						:     Optional[Any]    ):



           '''simple docstring'''

           UpperCAmelCase      	=		self.tokenizer_class(self.vocab_file    )

           UpperCAmelCase      	=		tokenizer.tokenize('''LAGVS'''    )
           self.assertListEqual(lowercase       ,    ['''L''', '''A''', '''G''', '''V''', '''S''']    )
           self.assertListEqual(tokenizer.convert_tokens_to_ids(lowercase    )       ,    [4, 5, 6, 7, 8]    )
       def    A				(     self						:     str    ):



           '''simple docstring'''

           UpperCAmelCase      	=		self.tokenizer_class(self.vocab_file    )

           UpperCAmelCase      	=		'''LAGVS'''
           self.assertListEqual(tokenizer.encode(lowercase    )       ,    [0, 4, 5, 6, 7, 8, 2]    )
       def    A				(     self						:     Optional[Any]    ):



           '''simple docstring'''

           UpperCAmelCase      	=		self.tokenizer_class(self.vocab_file    )

           UpperCAmelCase      	=		['''LAGVS''', '''WCB''']
           UpperCAmelCase      	=		tokenizer(lowercase       ,    padding=lowercase    )['''input_ids''']

           self.assertListEqual(lowercase       ,    [[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2]]    )
       def    A				(     self						:     Union[str, Any]    ):



           '''simple docstring'''

           UpperCAmelCase      	=		self.tokenizer_class(self.vocab_file    )

           UpperCAmelCase      	=		['''LAGVS''', '''WCB''']
           UpperCAmelCase      	=		tokenizer(lowercase       ,    padding=lowercase    )['''input_ids''']

           self.assertListEqual(lowercase       ,    [[0, 4, 5, 6, 7, 8, 2], [0, 22, 23, 25, 2, 1, 1]]    )







       def    A				(     self						:     Any    ):



           '''simple docstring'''

           UpperCAmelCase      	=		self.get_tokenizers(fast=lowercase    )
           for tokenizer in tokenizers:
               with self.subTest(f"{tokenizer.__class__.__name__}"    ):
                   UpperCAmelCase      	=		'''<unk>'''
                   UpperCAmelCase      	=		'''<mask>'''

                   UpperCAmelCase      	=		tokenizer.tokenize(lowercase    )
                   UpperCAmelCase      	=		tokenizer.tokenize(lowercase    )

                   self.assertEqual(len(lowercase    )       ,    1    )
                   self.assertEqual(len(lowercase    )       ,    1    )
                   self.assertEqual(token_a[0]       ,    lowercase    )
                   self.assertEqual(token_a[0]       ,    lowercase    )




