40 lines
1.3 KiB
Python
40 lines
1.3 KiB
Python
# These codes are copied from modelscope revision c58451baead80d83281f063d12fb377fad415257
|
|
# Copyright (c) Alibaba, Inc. and its affiliates.
|
|
from transformers import BloomConfig
|
|
from transformers import BloomModel as BloomModelTransform
|
|
|
|
from modelscope.metainfo import Models
|
|
from modelscope.models import TorchModel
|
|
from modelscope.models.builder import BACKBONES
|
|
from modelscope.utils.constant import Tasks
|
|
|
|
|
|
class MsModelMixin:
|
|
|
|
@classmethod
|
|
def _instantiate(cls, **kwargs):
|
|
"""Instantiate the model.
|
|
Args:
|
|
kwargs: Input args.
|
|
model_dir: The model dir used to load the checkpoint and the label information.
|
|
Returns:
|
|
The loaded model, which is initialized by transformers.PreTrainedModel.from_pretrained
|
|
"""
|
|
|
|
model_dir = kwargs.pop('model_dir', None)
|
|
kwargs.pop('device', None)
|
|
if model_dir is None:
|
|
config = BloomConfig(**kwargs)
|
|
model = cls(config)
|
|
else:
|
|
model = super(MsModelMixin, cls).from_pretrained(
|
|
pretrained_model_name_or_path=model_dir, **kwargs)
|
|
model.model_dir = model_dir
|
|
return model
|
|
|
|
|
|
@BACKBONES.register_module(group_key=Tasks.backbone, module_name=Models.bloom)
|
|
class BloomModel(MsModelMixin, BloomModelTransform, TorchModel):
|
|
|
|
pass
|