easy-ml/model/model_base.py

# -*- coding:utf-8 -*-
"""
@author: yq
@time: 2024/1/2
@desc: 模型基类
"""
import abc
from typing import Dict

import pandas as pd

from entitys import TrainConfigEntity, DataPreparedEntity, MetricFucEntity


class ModelBase(metaclass=abc.ABCMeta):

    def __init__(self, train_config: TrainConfigEntity = None, *args, **kwargs):
        if train_config is not None:
            self._train_config = train_config
        else:
            self._train_config = TrainConfigEntity(*args, **kwargs)

    @abc.abstractmethod
    def get_template_path(self, ) -> str:
        pass

    @abc.abstractmethod
    def train(self, data: DataPreparedEntity, *args, **kwargs) -> Dict[str, MetricFucEntity]:
        pass

    @abc.abstractmethod
    def predict_prob(self, x: pd.DataFrame, *args, **kwargs):
        pass

    @abc.abstractmethod
    def export_model_file(self, ):
        pass