easy-ml/trainer/train.py

# -*- coding: utf-8 -*-
"""
@author: yq
@time: 2024/11/1
@desc: 模型训练管道
"""
from entitys import DataPreparedEntity, TrainConfigEntity
from model import f_get_model


class TrainPipeline():
    def __init__(self, train_config: TrainConfigEntity):
        self._train_config = train_config
        model_clazz = f_get_model(self._train_config.model_type)
        self.model = model_clazz(self._train_config)

    def train(self, data: DataPreparedEntity):
        metric_train = self.model.train(data)
        print(metric_train)

    def generate_report(self, data: DataPreparedEntity):
        pass


if __name__ == "__main__":
    pass