141 lines
4.8 KiB
Smarty
141 lines
4.8 KiB
Smarty
# Copyright (c) Alibaba, Inc. and its affiliates.
|
||
|
||
import torch
|
||
|
||
from modelscope.models.base import TorchModel
|
||
from modelscope.preprocessors.base import Preprocessor
|
||
from modelscope.pipelines.base import Model, Pipeline
|
||
from modelscope.utils.config import Config
|
||
from modelscope.pipelines.builder import PIPELINES
|
||
from modelscope.preprocessors.builder import PREPROCESSORS
|
||
from modelscope.models.builder import MODELS
|
||
|
||
|
||
@MODELS.register_module('${task_name}', module_name='my-custom-model')
|
||
class ${model_name}(TorchModel):
|
||
|
||
def __init__(self, model_dir, *args, **kwargs):
|
||
super().__init__(model_dir, *args, **kwargs)
|
||
self.model = self.init_model(**kwargs)
|
||
|
||
def forward(self, input_tensor, **forward_params):
|
||
return self.model(input_tensor, **forward_params)
|
||
|
||
def init_model(self, **kwargs):
|
||
"""Provide default implementation based on TorchModel and user can reimplement it.
|
||
include init model and load ckpt from the model_dir, maybe include preprocessor
|
||
if nothing to do, then return lambda x: x
|
||
"""
|
||
return lambda x: x
|
||
|
||
|
||
@PREPROCESSORS.register_module('${task_name}', module_name='my-custom-preprocessor')
|
||
class ${preprocessor_name}(Preprocessor):
|
||
|
||
def __init__(self, *args, **kwargs):
|
||
super().__init__(*args, **kwargs)
|
||
self.trainsforms = self.init_preprocessor(**kwargs)
|
||
|
||
def __call__(self, results):
|
||
return self.trainsforms(results)
|
||
|
||
def init_preprocessor(self, **kwarg):
|
||
""" Provide default implementation based on preprocess_cfg and user can reimplement it.
|
||
if nothing to do, then return lambda x: x
|
||
"""
|
||
return lambda x: x
|
||
|
||
|
||
@PIPELINES.register_module('${task_name}', module_name='my-custom-pipeline')
|
||
class ${pipeline_name}(Pipeline):
|
||
""" Give simple introduction to this pipeline.
|
||
|
||
Examples:
|
||
|
||
>>> from modelscope.pipelines import pipeline
|
||
>>> input = "Hello, ModelScope!"
|
||
>>> my_pipeline = pipeline('my-task', 'my-model-id')
|
||
>>> result = my_pipeline(input)
|
||
|
||
"""
|
||
|
||
def __init__(self, model, preprocessor=None, **kwargs):
|
||
"""
|
||
use `model` and `preprocessor` to create a custom pipeline for prediction
|
||
Args:
|
||
model: model id on modelscope hub.
|
||
preprocessor: the class of method be init_preprocessor
|
||
"""
|
||
super().__init__(model=model, auto_collate=False)
|
||
assert isinstance(model, str) or isinstance(model, Model), \
|
||
'model must be a single str or Model'
|
||
if isinstance(model, str):
|
||
pipe_model = Model.from_pretrained(model)
|
||
elif isinstance(model, Model):
|
||
pipe_model = model
|
||
else:
|
||
raise NotImplementedError
|
||
pipe_model.eval()
|
||
|
||
if preprocessor is None:
|
||
preprocessor = ${preprocessor_name}()
|
||
super().__init__(model=pipe_model, preprocessor=preprocessor, **kwargs)
|
||
|
||
def _sanitize_parameters(self, **pipeline_parameters):
|
||
"""
|
||
this method should sanitize the keyword args to preprocessor params,
|
||
forward params and postprocess params on '__call__' or '_process_single' method
|
||
considered to be a normal classmethod with default implementation / output
|
||
|
||
Default Returns:
|
||
Dict[str, str]: preprocess_params = {}
|
||
Dict[str, str]: forward_params = {}
|
||
Dict[str, str]: postprocess_params = pipeline_parameters
|
||
"""
|
||
return {}, pipeline_parameters, {}
|
||
|
||
def _check_input(self, inputs):
|
||
pass
|
||
|
||
def _check_output(self, outputs):
|
||
pass
|
||
|
||
def forward(self, inputs, **forward_params):
|
||
""" Provide default implementation using self.model and user can reimplement it
|
||
"""
|
||
return super().forward(inputs, **forward_params)
|
||
|
||
def postprocess(self, inputs):
|
||
""" If current pipeline support model reuse, common postprocess
|
||
code should be write here.
|
||
|
||
Args:
|
||
inputs: input data
|
||
|
||
Return:
|
||
dict of results: a dict containing outputs of model, each
|
||
output should have the standard output name.
|
||
"""
|
||
return inputs
|
||
|
||
|
||
# Tips: usr_config_path is the temporary save configuration location， after upload modelscope hub, it is the model_id
|
||
usr_config_path = '${configuration_path}'
|
||
config = Config({
|
||
"framework": 'pytorch',
|
||
"task": '${task_name}',
|
||
"model": {'type': 'my-custom-model'},
|
||
"pipeline": {"type": "my-custom-pipeline"},
|
||
"allow_remote": True
|
||
})
|
||
config.dump('${configuration_path}' + 'configuration.json')
|
||
|
||
if __name__ == "__main__":
|
||
from modelscope.models import Model
|
||
from modelscope.pipelines import pipeline
|
||
# model = Model.from_pretrained(usr_config_path)
|
||
input = "Hello, ModelScope!"
|
||
inference = pipeline('${task_name}', model=usr_config_path)
|
||
output = inference(input)
|
||
print(output)
|