AutoAndroidController/py/venv/Lib/site-packages/modelscope/models/base/base_model.py

# Copyright (c) Alibaba, Inc. and its affiliates.
import os
import os.path as osp
from abc import ABC, abstractmethod
from typing import Any, Dict, List, Optional, Union

from modelscope.hub.snapshot_download import snapshot_download
from modelscope.metainfo import Tasks
from modelscope.models.builder import build_backbone, build_model
from modelscope.utils.automodel_utils import (can_load_by_ms,
                                              check_model_from_owner_group,
                                              try_to_load_hf_model)
from modelscope.utils.config import Config, ConfigDict
from modelscope.utils.constant import DEFAULT_MODEL_REVISION, Invoke, ModelFile
from modelscope.utils.device import verify_device
from modelscope.utils.logger import get_logger
from modelscope.utils.plugins import (register_modelhub_repo,
                                      register_plugins_repo)

logger = get_logger()

Tensor = Union['torch.Tensor', 'tf.Tensor']


class Model(ABC):
    """Base model interface.
    """

    def __init__(self, model_dir, *args, **kwargs):
        self.model_dir = model_dir
        device_name = kwargs.get('device', 'gpu')
        verify_device(device_name)
        self._device_name = device_name
        self.trust_remote_code = kwargs.get('trust_remote_code', False)

    def __call__(self, *args, **kwargs) -> Dict[str, Any]:
        return self.postprocess(self.forward(*args, **kwargs))

    def check_trust_remote_code(self,
                                info_str: Optional[str] = None,
                                model_dir: Optional[str] = None):
        """Check trust_remote_code if the model needs to import extra libs

        Args:
            info_str(str): The info showed to user if trust_remote_code is `False`.
            model_dir(`Optional[str]`): The local model directory. If is a trusted model, check remote code will pass.
        """
        info_str = info_str or (
            'This model requires `trust_remote_code` to be `True` because it needs to '
            'import extra libs or execute the code in the model repo, setting this to true '
            'means you trust the files in it.')
        if not check_model_from_owner_group(model_dir=model_dir):
            assert self.trust_remote_code, info_str

    @abstractmethod
    def forward(self, *args, **kwargs) -> Dict[str, Any]:
        """
        Run the forward pass for a model.

        Returns:
            Dict[str, Any]: output from the model forward pass
        """
        pass

    def postprocess(self, inputs: Dict[str, Any], **kwargs) -> Dict[str, Any]:
        """ Model specific postprocess and convert model output to
        standard model outputs.

        Args:
            inputs:  input data

        Return:
            dict of results:  a dict containing outputs of model, each
                output should have the standard output name.
        """
        return inputs

    @classmethod
    def _instantiate(cls, **kwargs):
        """ Define the instantiation method of a model,default method is by
            calling the constructor. Note that in the case of no loading model
            process in constructor of a task model, a load_model method is
            added, and thus this method is overloaded
        """
        return cls(**kwargs)

    @classmethod
    def from_pretrained(cls,
                        model_name_or_path: str,
                        revision: Optional[str] = DEFAULT_MODEL_REVISION,
                        cfg_dict: Config = None,
                        device: str = None,
                        trust_remote_code: Optional[bool] = False,
                        **kwargs):
        """Instantiate a model from local directory or remote model repo. Note
        that when loading from remote, the model revision can be specified.

        Args:
            model_name_or_path(str): A model dir or a model id to be loaded
            revision(str, `optional`): The revision used when the model_name_or_path is
                a model id of the remote hub. default `master`.
            cfg_dict(Config, `optional`): An optional model config. If provided, it will replace
                the config read out of the `model_name_or_path`
            device(str, `optional`): The device to load the model.
            trust_remote_code(bool, `optional`): Whether to trust and allow execution of remote code. Default is False.
            **kwargs:
                task(str, `optional`): The `Tasks` enumeration value to replace the task value
                    read out of config in the `model_name_or_path`. This is useful when the model to be loaded is not
                    equal to the model saved.
                    For example, load a `backbone` into a `text-classification` model.
                    Other kwargs will be directly fed into the `model` key, to replace the default configs.
                use_hf(bool, `optional`):
                    If set to True, it will initialize the model using AutoModel or AutoModelFor* from hf.
                    If set to False, the model is loaded using the modelscope mode.
                    If set to None, the loading mode will be automatically selected.
                ignore_file_pattern(List[str], `optional`):
                    This parameter is passed to snapshot_download
                device_map(str | Dict[str, str], `optional`):
                    This parameter is passed to AutoModel or AutoModelFor*
                torch_dtype(torch.dtype, `optional`):
                    This parameter is passed to AutoModel or AutoModelFor*
                config(PretrainedConfig, `optional`):
                    This parameter is passed to AutoModel or AutoModelFor*
        Returns:
            A model instance.

        Examples:
            >>> from modelscope.models import Model
            >>> Model.from_pretrained('damo/nlp_structbert_backbone_base_std', task='text-classification')
        """
        prefetched = kwargs.get('model_prefetched')
        if prefetched is not None:
            kwargs.pop('model_prefetched')
        invoked_by = kwargs.get(Invoke.KEY)
        if invoked_by is not None:
            kwargs.pop(Invoke.KEY)
        else:
            invoked_by = Invoke.PRETRAINED

        ignore_file_pattern = kwargs.pop('ignore_file_pattern', None)
        if osp.exists(model_name_or_path):
            local_model_dir = model_name_or_path
        else:
            if prefetched is True:
                raise RuntimeError(
                    'Expecting model is pre-fetched locally, but is not found.'
                )

            invoked_by = '%s/%s' % (Invoke.KEY, invoked_by)
            local_model_dir = snapshot_download(
                model_name_or_path,
                revision,
                user_agent=invoked_by,
                ignore_file_pattern=ignore_file_pattern)
        logger.info(f'initialize model from {local_model_dir}')

        configuration_path = osp.join(local_model_dir, ModelFile.CONFIGURATION)
        cfg = None
        if cfg_dict is not None:
            cfg = cfg_dict
        elif os.path.exists(configuration_path):
            cfg = Config.from_file(configuration_path)
        task_name = getattr(cfg, 'task', None)
        if 'task' in kwargs:
            task_name = kwargs.pop('task')
        model_cfg = getattr(cfg, 'model', ConfigDict())
        if hasattr(model_cfg, 'model_type') and not hasattr(model_cfg, 'type'):
            model_cfg.type = model_cfg.model_type
        model_type = getattr(model_cfg, 'type', None)
        if isinstance(device, str) and device.startswith('gpu'):
            device = 'cuda' + device[3:]
        use_hf = kwargs.pop('use_hf', None)
        if use_hf is None and can_load_by_ms(local_model_dir, task_name,
                                             model_type):
            use_hf = False
        model = None
        if use_hf in {True, None}:
            model = try_to_load_hf_model(local_model_dir, task_name, use_hf,
                                         **kwargs)
        if model is not None:
            device_map = kwargs.pop('device_map', None)
            if device_map is None and device is not None:
                model = model.to(device)
            return model
        # use ms
        if cfg is None:
            raise FileNotFoundError(
                f'`{ModelFile.CONFIGURATION}` file not found.')
        model_cfg.model_dir = local_model_dir

        # Security check: Only allow execution of remote code or plugins if trust_remote_code is True
        plugins = cfg.safe_get('plugins')
        if plugins and not trust_remote_code:
            raise RuntimeError(
                'Detected plugins field in the model configuration file, but '
                'trust_remote_code=True was not explicitly set.\n'
                'To prevent potential execution of malicious code, loading has been refused.\n'
                'If you trust this model repository, please pass trust_remote_code=True to from_pretrained.'
            )
        if plugins and trust_remote_code:
            logger.warning(
                'Use trust_remote_code=True. Will invoke codes or install plugins from remote model repo. '
                'Please make sure that you can trust the external codes.')
        register_modelhub_repo(local_model_dir, allow_remote=trust_remote_code)
        default_args = {}
        if trust_remote_code:
            default_args = {'trust_remote_code': trust_remote_code}
        register_plugins_repo(plugins)
        for k, v in kwargs.items():
            model_cfg[k] = v
        if device is not None:
            model_cfg.device = device
        if task_name is Tasks.backbone:
            model_cfg.init_backbone = True
            model = build_backbone(model_cfg)
        else:
            model = build_model(
                model_cfg, task_name=task_name, default_args=default_args)

        # dynamically add pipeline info to model for pipeline inference
        if hasattr(cfg, 'pipeline'):
            model.pipeline = cfg.pipeline

        if not hasattr(model, 'cfg'):
            model.cfg = cfg

        model_cfg.pop('model_dir', None)
        model.name = model_name_or_path
        model.model_dir = local_model_dir
        return model

    def save_pretrained(self,
                        target_folder: Union[str, os.PathLike],
                        save_checkpoint_names: Union[str, List[str]] = None,
                        config: Optional[dict] = None,
                        **kwargs):
        """save the pretrained model, its configuration and other related files to a directory,
            so that it can be re-loaded

        Args:
            target_folder (Union[str, os.PathLike]):
            Directory to which to save. Will be created if it doesn't exist.

            save_checkpoint_names (Union[str, List[str]]):
            The checkpoint names to be saved in the target_folder

            config (Optional[dict], optional):
            The config for the configuration.json, might not be identical with model.config
        """
        raise NotImplementedError(
            'save_pretrained method need to be implemented by the subclass.')