codingutf8 2023 the huggingface inc team 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 orglicenseslicense2 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 a script to add andor update the attribute pipelinemodelmapping in model test files this script will be mostly used in the following 2 situations run within a scheduled ci job to check if model test files in the library have updated pipelinemodelmapping andor update test files and possibly open a github pull request automatically being run by a transformers member to quickly check and update some particular test files this script is not intended to be run manually by community contributors do not add item to this set unless the reason is approved infer the framework from the test class testclass if modeltestermixin in x name for x in testclass bases return pt elif tfmodeltestermixin in x name for x in testclass bases return tf elif flaxmodeltestermixin in x name for x in testclass bases return flax else return none def getmappingfortasktask framework use the cached results cache the results get the model architectures related to the test class testclass for a pipeline task framework getframeworktestclass if framework is none return none mapping getmappingfortasktask framework if mapping is none return none configclasses listmodelclass configclass for modelclass in testclass allmodelclasses if lenconfigclasses 1 raise valueerrorthere should be exactly one configuration class from testclass allmodelclasses this could be a listtuple of model classes but it s rare modelclass mapping getconfigclasses0 none if isinstancemodelclass tuple list modelclass sortedmodelclass keylambda x x name return modelclass def getpipelinemodelmappingtestclass get pipelinemodelmapping for testclass as a string to be added to the test file this will be a 1line string after this is added to a test file make style will format it beautifully a listtuple of model classes a single model class restrict to xxxmodeltestermixin and should be a subclass of unittest testcase baseclassnames modeltestermixin tfmodeltestermixin flaxmodeltestermixin if not issubclasstestclass unittest testcase return false return lenbaseclassnames intersectionx name for x in testclass bases 0 def findtestclasstestfile if a test class has defined pipelinemodelmapping let s take it take the test class with the shortest name just a heuristic outside the definition block of pipelinemodelmapping add pipelinemodelmapping to testclass if getattrtestclass pipelinemodelmapping none is not none if not overwrite return 1 linetoadd getpipelinemodelmappingstringtestclass if lenlinetoadd 0 return 1 linetoadd linetoadd n the code defined the class testclass classlines classstartlineno inspect getsourcelinestestclass inspect gives the code for an object including decorators if any we only need the exact line of the class definition for idx line in enumerateclasslines if line lstrip startswithclass classlines classlinesidx classstartlineno idx break classendlineno classstartlineno lenclasslines 1 the index in classlines that starts the definition of allmodelclasses allgenerativemodelclasses or pipelinemodelmapping this assumes they are defined in such order and we take the start index of the last block that appears in a testclass startidx none the indent level of the line at classlinesstartidx if defined indentlevel 0 to record if pipelinemodelmapping is found in testclass defline none for idx line in enumerateclasslines if line strip startswithallmodelclasses indentlevel lenline lenline lstrip startidx idx elif line strip startswithallgenerativemodelclasses indentlevel lenline lenline lstrip startidx idx elif line strip startswithpipelinemodelmapping indentlevel lenline lenline lstrip startidx idx defline line break if startidx is none return 1 find the ending index inclusive of the above found block endidx findblockendingclasslines startidx indentlevel extract isxxxavailable from existing blocks some models require specific libraries like timm and use istimmavailable instead of istorchavailable keep leading and trailing whitespaces r re compilersiss availables for line in classlinesstartidx endidx 1 backendcondition r searchline if backendcondition is not none replace the leading and trailing whitespaces to the space character target backendcondition01 1 linetoadd r subtarget linetoadd break if defline is none pipelinemodelmapping is not defined the target index is set to the ending index inclusive of allmodelclasses or allgenerativemodelclasses targetidx endidx else pipelinemodelmapping is defined the target index is set to be one before its start index targetidx startidx 1 mark the lines of the currently existing pipelinemodelmapping to be removed for idx in rangestartidx endidx 1 these lines are going to be removed before writing to the test file classlinesidx none noqa make sure the test class is a subclass of pipelinetestermixin parentclasses x name for x in testclass bases if pipelinetestermixin not in parentclasses put pipelinetestermixin just before unittest testcase parentclasses x for x in parentclasses if x testcase pipelinetestermixin if testcase in parentclasses here we assume the original string is always with unittest testcase parentclasses appendunittest testcase parentclasses joinparentclasses for idx line in enumerateclasslines find the ending of the declaration of testclass if line strip endswith mark the lines of the declaration of testclass to be removed for idx in rangeidx 1 classlinesidx none noqa break add the new oneline class declaration for testclass classlines0 fclass testclass nameparentclasses n add indentation linetoadd indentlevel linetoadd insert pipelinemodelmapping to classlines the line at targetidx should be kept by definition classlines classlines targetidx 1 linetoadd classlinestargetidx 1 remove the lines that are marked to be removed classlines x for x in classlines if x is not none move from test class to module in order to write to the test file modulelines inspect getsourcelinesinspect getmoduletestclass0 be careful with the 1off between line numbers and array indices modulelines modulelines classstartlineno 1 classlines modulelinesclassendlineno code joinmodulelines moddulefile inspect getsourcefiletestclass with openmoddulefile w encodingutf8 newlinen as fp fp writecode return linetoadd def addpipelinemodelmappingtotestfiletestfile overwritefalse flax is not concerned at this moment coding utf 8 2023 the huggingface inc team 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 a script to add and or update the attribute pipeline_model_mapping in model test files this script will be mostly used in the following 2 situations run within a scheduled ci job to check if model test files in the library have updated pipeline_model_mapping and or update test files and possibly open a github pull request automatically being run by a transformers member to quickly check and update some particular test file s this script is not intended to be run manually by community contributors do not add item to this set unless the reason is approved the pipeline test mapping is added to test_modeling_esm py infer the framework from the test class test_class get mappings defined in xxxpipelinetests for the task task use the cached results cache the results get the model architecture s related to the test class test_class for a pipeline task this could be a list tuple of model classes but it s rare get pipeline_model_mapping for test_class get pipeline_model_mapping for test_class as a string to be added to the test file this will be a 1 line string after this is added to a test file make style will format it beautifully a list tuple of model classes a single model class restrict to xxxmodeltestermixin and should be a subclass of unittest testcase find a test class in test_file to which we will add pipeline_model_mapping if a test class has defined pipeline_model_mapping let s take it take the test class with the shortest name just a heuristic outside the definition block of pipeline_model_mapping add pipeline_model_mapping to test_class the code defined the class test_class inspect gives the code for an object including decorator s if any we only need the exact line of the class definition 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 the indent level of the line at class_lines start_idx if defined to record if pipeline_model_mapping is found in test_class find the ending index inclusive of the above found block 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 replace the leading and trailing whitespaces to the space character 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 pipeline_model_mapping is defined the target index is set to be one before its start index mark the lines of the currently existing pipeline_model_mapping to be removed these lines are going to be removed before writing to the test file noqa make sure the test class is a subclass of pipelinetestermixin put pipelinetestermixin just before unittest testcase here we assume the original string is always with unittest testcase find the ending of the declaration of test_class mark the lines of the declaration of test_class to be removed noqa add the new one line class declaration for test_class add indentation insert pipeline_model_mapping to class_lines the line at target_idx should be kept by definition remove the lines that are marked to be removed move from test class to module in order to write to the test file be careful with the 1 off between line numbers and array indices add pipeline_model_mapping to test_file flax is not concerned at this moment
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


PIPELINE_TEST_MAPPING = {}
for task, _ in pipeline_test_mapping.items():
    PIPELINE_TEST_MAPPING[task] = {"pt": None, "tf": None}



TEST_FILE_TO_IGNORE = {
    "tests/models/esm/test_modeling_esmfold.py",  
}


def get_framework(test_class):
    

    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 get_mapping_for_task(task, framework):
    
    
    if PIPELINE_TEST_MAPPING[task].get(framework, None) is not None:
        return PIPELINE_TEST_MAPPING[task][framework]

    pipeline_test_class = pipeline_test_mapping[task]["test"]
    mapping = None

    if framework == "pt":
        mapping = getattr(pipeline_test_class, "model_mapping", None)
    elif framework == "tf":
        mapping = getattr(pipeline_test_class, "tf_model_mapping", None)

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

    
    PIPELINE_TEST_MAPPING[task][framework] = mapping
    return mapping


def get_model_for_pipeline_test(test_class, task):
    
    framework = get_framework(test_class)
    if framework is None:
        return None
    mapping = get_mapping_for_task(task, framework)
    if mapping is None:
        return None

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

    
    model_class = mapping.get(config_classes[0], None)
    if isinstance(model_class, (tuple, list)):
        model_class = sorted(model_class, key=lambda x: x.__name__)

    return model_class


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

    return dict(mapping)


def get_pipeline_model_mapping_string(test_class):
    
    framework = get_framework(test_class)
    if framework == "pt":
        framework = "torch"
    default_value = "{}"

    mapping = get_pipeline_model_mapping(test_class)
    if len(mapping) == 0:
        return ""

    texts = []
    for task, model_classes in mapping.items():
        if isinstance(model_classes, (tuple, list)):
            
            value = "(" + ", ".join([x.__name__ for x in model_classes]) + ")"
        else:
            
            value = model_classes.__name__
        texts.append(f'"{task}": {value}')
    text = "{" + ", ".join(texts) + "}"
    text = f"pipeline_model_mapping = {text} if is_{framework}_available() else {default_value}"

    return text


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


def find_test_class(test_file):
    
    test_classes = [x for x in get_test_classes(test_file) if is_valid_test_class(x)]

    target_test_class = None
    for test_class in test_classes:
        
        if getattr(test_class, "pipeline_model_mapping", None) is not None:
            target_test_class = test_class
            break
    
    if target_test_class is None and len(test_classes) > 0:
        target_test_class = sorted(test_classes, key=lambda x: (len(x.__name__), x.__name__))[0]

    return target_test_class


def find_block_ending(lines, start_idx, indent_level):
    end_idx = start_idx
    for idx, line in enumerate(lines[start_idx:]):
        indent = len(line) - len(line.lstrip())
        if idx == 0 or indent > indent_level or (indent == indent_level and line.strip() == ")"):
            end_idx = start_idx + idx
        elif idx > 0 and indent <= indent_level:
            
            break

    return end_idx


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

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

    
    class_lines, class_start_line_no = inspect.getsourcelines(test_class)
    
    
    for idx, line in enumerate(class_lines):
        if line.lstrip().startswith("class "):
            class_lines = class_lines[idx:]
            class_start_line_no += idx
            break
    class_end_line_no = class_start_line_no + len(class_lines) - 1

    
    
    
    start_idx = None
    
    indent_level = 0
    
    def_line = None
    for idx, line in enumerate(class_lines):
        if line.strip().startswith("all_model_classes = "):
            indent_level = len(line) - len(line.lstrip())
            start_idx = idx
        elif line.strip().startswith("all_generative_model_classes = "):
            indent_level = len(line) - len(line.lstrip())
            start_idx = idx
        elif line.strip().startswith("pipeline_model_mapping = "):
            indent_level = len(line) - len(line.lstrip())
            start_idx = idx
            def_line = line
            break

    if start_idx is None:
        return "", -1
    
    end_idx = find_block_ending(class_lines, start_idx, indent_level)

    
    
    
    r = re.compile(r"\s(is_\S+?_available\(\))\s")
    for line in class_lines[start_idx : end_idx + 1]:
        backend_condition = r.search(line)
        if backend_condition is not None:
            
            target = " " + backend_condition[0][1:-1] + " "
            line_to_add = r.sub(target, line_to_add)
            break

    if def_line is None:
        
        
        target_idx = end_idx
    else:
        
        target_idx = start_idx - 1
        
        for idx in range(start_idx, end_idx + 1):
            
            class_lines[idx] = None  

    
    parent_classes = [x.__name__ for x in test_class.__bases__]
    if "PipelineTesterMixin" not in parent_classes:
        
        _parent_classes = [x for x in parent_classes if x != "TestCase"] + ["PipelineTesterMixin"]
        if "TestCase" in parent_classes:
            
            _parent_classes.append("unittest.TestCase")
        parent_classes = ", ".join(_parent_classes)
        for idx, line in enumerate(class_lines):
            
            if line.strip().endswith("):"):
                
                for _idx in range(idx + 1):
                    class_lines[_idx] = None  
                break
        
        class_lines[0] = f"class {test_class.__name__}({parent_classes}):\n"

    
    line_to_add = " " * indent_level + line_to_add
    
    
    class_lines = class_lines[: target_idx + 1] + [line_to_add] + class_lines[target_idx + 1 :]
    
    class_lines = [x for x in class_lines if x is not None]

    
    module_lines = inspect.getsourcelines(inspect.getmodule(test_class))[0]
    
    module_lines = module_lines[: class_start_line_no - 1] + class_lines + module_lines[class_end_line_no:]
    code = "".join(module_lines)

    moddule_file = inspect.getsourcefile(test_class)
    with open(moddule_file, "w", encoding="UTF-8", newline="\n") as fp:
        fp.write(code)

    return line_to_add


def add_pipeline_model_mapping_to_test_file(test_file, overwrite=False):
    
    test_class = find_test_class(test_file)
    if test_class:
        add_pipeline_model_mapping(test_class, overwrite=overwrite)


if __name__ == "__main__":
    parser = 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`.",
    )
    args = 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.")

    test_files = []
    if args.test_file:
        test_files = [args.test_file]
    else:
        pattern = os.path.join("tests", "models", "**", "test_modeling_*.py")
        for test_file in glob.glob(pattern):
            
            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)