easy-ml/model/model_lr.py

# -*- coding: utf-8 -*-
"""
@author: yq
@time: 2024/11/1
@desc: 
"""
import pandas as pd
from sklearn.linear_model import LogisticRegression

from entitys import DataFeatureEntity, MetricTrainEntity
from .model_base import ModelBase


class ModelLr(ModelBase):
    def __init__(self, ):
        self.lr = LogisticRegression(penalty='l1', C=0.9, solver='saga', n_jobs=-1)

    def train(self, data: DataFeatureEntity, *args, **kwargs) -> MetricTrainEntity:
        self.lr.fit(data.get_Xdata(), data.get_Ydata())
        return MetricTrainEntity(0.7, 0.4)

    def predict_prob(self, x: pd.DataFrame, *args, **kwargs):
        return self.lr.predict_proba(x)[:, 1]

    def predict(self, x: pd.DataFrame, *args, **kwargs):
        pass

    def export_model_file(self):
        pass


if __name__ == "__main__":
    pass