Function create_tiny_models

utils/create_dummy_models.py:1683–1798

(
    output_path,
    all,
    model_types,
    models_to_skip,
    no_check,
    upload,
    organization,
    token,
    num_workers=1,
)

def create_tiny_models(
    output_path,
    all,
    model_types,
    models_to_skip,
    no_check,
    upload,
    organization,
    token,
    num_workers=1,
):
    clone_path = os.path.abspath(os.path.dirname(os.path.dirname(__file__)))
    if os.getcwd() != clone_path:
        raise ValueError(f"This script should be run from the root of the clone of `transformers` {clone_path}")

    report_path = os.path.join(output_path, "reports")
    os.makedirs(report_path, exist_ok=True)

    _pytorch_arch_mappings = [x for x in dir(transformers_module) if x.startswith("MODEL_") and x.endswith("_MAPPING")]

    pytorch_arch_mappings = [getattr(transformers_module, x) for x in _pytorch_arch_mappings]

    config_classes = CONFIG_MAPPING.values()
    if not all:
        config_classes = [CONFIG_MAPPING[model_type] for model_type in model_types]

    # TODO: we should add information to the reports instead of skip them
    config_classes = [x for x in config_classes if x.__name__ not in no_model_tester_at_all]
    config_classes = [x for x in config_classes if x.__name__ not in configs_requiring_too_exotic_dependency]
    config_classes = [x for x in config_classes if x.__name__ not in deprecated_models]
    config_classes = [x for x in config_classes if x.__name__ not in config_without_meaningful_model_class]

    # A map from config classes to tuples of processors (tokenizer, feature extractor, processor) classes
    processor_type_map = {c: get_processor_types_from_config_class(c) for c in config_classes}

    to_create = {}
    for c in config_classes:
        processors = processor_type_map[c]
        models = get_architectures_from_config_class(c, pytorch_arch_mappings, models_to_skip)
        if len(models) > 0:
            to_create[c] = {"processor": processors, "pytorch": models}

    results = {}
    if num_workers <= 1:
        for c, models_to_create in list(to_create.items()):
            print(f"Create models for {c.__name__} ...")
            result = build(c, models_to_create, output_dir=os.path.join(output_path, c.model_type))
            results[c.__name__] = result
            print("=" * 40)
    else:
        all_build_args = []
        for c, models_to_create in list(to_create.items()):
            all_build_args.append((c, models_to_create, os.path.join(output_path, c.model_type)))
        with multiprocessing.Pool(processes=num_workers) as pool:
            results = pool.starmap(build, all_build_args)
            results = {build_args[0].__name__: result for build_args, result in zip(all_build_args, results)}

    print(results)