|
|
|
import copy |
|
|
|
from transformers import PretrainedConfig |
|
|
|
from lmdeploy.turbomind.deploy.target_model.base import TurbomindModelConfig |
|
from lmdeploy.version import __version__ as lm_version |
|
|
|
|
|
class LMDeployConfig(PretrainedConfig): |
|
"""Lmdeploy config.""" |
|
|
|
def __init__(self, turbomind: dict = None, **kwargs): |
|
default_tm_cfg = copy.deepcopy( |
|
TurbomindModelConfig.from_dict({}, allow_none=True).__dict__) |
|
if turbomind is not None: |
|
default_tm_cfg.update(turbomind) |
|
self.turbomind = default_tm_cfg |
|
self.lmdeploy_version = lm_version |
|
super().__init__(**kwargs) |
|
|
|
@classmethod |
|
def from_pretrained(cls, pretrained_model_name_or_path, **kwargs): |
|
return_unused_kwargs = kwargs.pop('return_unused_kwargs', False) |
|
config, kwargs = super().from_pretrained(pretrained_model_name_or_path, |
|
return_unused_kwargs=True, |
|
**kwargs) |
|
for k, v in kwargs.items(): |
|
if k in config.turbomind.keys(): |
|
config.turbomind[k] = v |
|
if 'tp' in kwargs: |
|
config.turbomind['tensor_para_size'] = kwargs['tp'] |
|
if return_unused_kwargs: |
|
return config, kwargs |
|
else: |
|
return config |
|
|
|
|
