Function _patch_pretrained_class

modelscope/utils/hf_util/patcher.py:158–477

Patch all class to download from modelscope Args: wrap: Wrap the class or monkey patch the original class Returns: The classes after patched

(all_imported_modules, wrap=False)

def _patch_pretrained_class(all_imported_modules, wrap=False):
    """Patch all class to download from modelscope

    Args:
        wrap: Wrap the class or monkey patch the original class

    Returns:
        The classes after patched
    """

    def get_model_dir(pretrained_model_name_or_path,
                      ignore_file_pattern=None,
                      allow_file_pattern=None,
                      **kwargs):
        from modelscope import snapshot_download
        subfolder = kwargs.pop('subfolder', None)
        file_filter = None
        if subfolder:
            file_filter = f'{subfolder}/*'
        if not os.path.exists(pretrained_model_name_or_path):
            revision = kwargs.pop('revision', None)
            if revision is None or revision == 'main':
                revision = 'master'
            if file_filter is not None:
                allow_file_pattern = file_filter
            local_files_only = kwargs.pop('local_files_only', False)
            model_dir = snapshot_download(
                pretrained_model_name_or_path,
                revision=revision,
                local_files_only=local_files_only,
                ignore_file_pattern=ignore_file_pattern,
                allow_file_pattern=allow_file_pattern)
            if subfolder:
                model_dir = os.path.join(model_dir, subfolder)
        else:
            model_dir = pretrained_model_name_or_path
        return model_dir

    def patch_pretrained_model_name_or_path(cls, pretrained_model_name_or_path,
                                            *model_args, **kwargs):
        """Patch all from_pretrained"""
        model_dir = get_model_dir(pretrained_model_name_or_path,
                                  kwargs.pop('ignore_file_pattern', None),
                                  kwargs.pop('allow_file_pattern', None),
                                  **kwargs)
        return cls._from_pretrained_origin.__func__(cls, model_dir,
                                                    *model_args, **kwargs)

    def patch_get_config_dict(cls, pretrained_model_name_or_path, *model_args,
                              **kwargs):
        """Patch all get_config_dict"""
        model_dir = get_model_dir(pretrained_model_name_or_path,
                                  kwargs.pop('ignore_file_pattern', None),
                                  kwargs.pop('allow_file_pattern', None),
                                  **kwargs)
        return cls._get_config_dict_origin.__func__(cls, model_dir,
                                                    *model_args, **kwargs)