easy-ml/model/__init__.py

# -*- coding:utf-8 -*-
"""
@author: yq
@time: 2023/12/28
@desc: 模型相关
"""
from commom import GeneralException
from enums import ModelEnum, ResultCodesEnum
from .model_base import ModelBase
from .model_lr import ModelLr

__all__ = ['ModelBase', 'f_get_model']

model_map = {
    ModelEnum.LR.value: ModelLr
}


def f_get_model(model_type: str):
    if model_type not in model_map.keys():
        raise GeneralException(ResultCodesEnum.ILLEGAL_PARAMS, message=f"模型【{model_type}】没有实现")
    return model_map[model_type]


if __name__ == "__main__":
    pass