alg/main_userupload.py

"""
针对用户类型数据单独训练模型
"""
import os
import sys

from sklearn.preprocessing import LabelEncoder

sys.path.append(os.getcwd())

import pandas as pd
import lightgbm as lgb

from scipy.stats import randint as sp_randint
from scipy.stats import uniform as sp_uniform
from sklearn.model_selection import RandomizedSearchCV, train_test_split
from sklearn.metrics import roc_auc_score, accuracy_score


class LightGBM(object):
    """
    LightGBM model for classification
    """

    def __init__(self, flag, dt):
        self.label_encoder = LabelEncoder()
        self.my_c = [
            "channel",
            "user_fans",
            "user_view_30",
            "user_share_30",
            "user_return_30",
            "user_rov",
            "user_str",
            "user_return_videos_30",
            "user_return_videos_3",
            "user_return_3",
            "user_view_3",
            "user_share_3",
            "address",
            "tag1",
            "tag2",
            "tag3",
            "rov_30",
            "rov_3",
            "ros_30",
            "ros_3"
        ]
        self.str_columns = ["channel", "address", "tag1", "tag2", "tag3"]
        self.float_columns = [
            "user_fans",
            "user_view_30",
            "user_share_30",
            "user_return_30",
            "user_rov",
            "user_str",
            "user_return_videos_30",
            "user_return_videos_3",
            "user_return_3",
            "user_share_3",
            "user_view_3",
            "rov_30",
            "rov_3",
            "ros_30",
            "ros_3"
        ]
        self.split_c = 0.7
        self.yc = 0.8
        self.model = "models/lightgbm_0402_user.bin"
        self.flag = flag
        self.dt = dt

    def read_data(self, path, yc=None):
        """
        Read data from local
        :return:
        """
        df = pd.read_json(path)
        df = df.dropna(subset=['label'])
        labels = df['label']
        temp = sorted(labels)
        if yc is None:
            yc = temp[int(len(temp) * 0.8)]
        print("阈值", yc)
        labels = [0 if i < yc else 1 for i in labels]
        features = df.drop("label", axis=1)
        for key in self.float_columns:
            features[key] = pd.to_numeric(features[key], errors="coerce")
        for key in self.str_columns:
            features[key] = self.label_encoder.fit_transform(features[key])
        return features, labels

    def best_params(self):
        path = "data/train_data/user_train_20240402.json"
        X, y = self.read_data(path)
        X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=42)

        lgbm = lgb.LGBMClassifier(objective='binary')

        # 设置搜索的参数范围
        param_dist = {
            'num_leaves': sp_randint(20, 40),
            'learning_rate': sp_uniform(0.001, 0.1),
            'feature_fraction': sp_uniform(0.5, 0.9),
            'bagging_fraction': sp_uniform(0.5, 0.9),
            'bagging_freq': sp_randint(1, 10),
            'min_child_samples': sp_randint(5, 100),
        }

        # 定义 RandomizedSearchCV
        rsearch = RandomizedSearchCV(
            estimator=lgbm,
            param_distributions=param_dist,
            n_iter=100,
            cv=3,
            scoring='roc_auc',
            random_state=42, verbose=2
        )

        # 开始搜索
        rsearch.fit(X_train, y_train)

        # 打印最佳参数和对应的AUC得分
        print("Best parameters found: ", rsearch.best_params_)
        print("Best AUC found: ", rsearch.best_score_)

        # 使用最佳参数在测试集上的表现
        best_model = rsearch.best_estimator_
        y_pred = best_model.predict_proba(X_test)[:, 1]
        auc = roc_auc_score(y_test, y_pred)
        print("AUC on test set: ", auc)

    def train_model(self):
        """
        Load dataset
        :return:
        """
        path = "data/train_data/user_train_20240402.json"
        x, y = self.read_data(path)
        train_size = int(len(x) * self.split_c)
        X_train, X_test = x[:train_size], x[train_size:]
        Y_train, Y_test = y[:train_size], y[train_size:]
        train_data = lgb.Dataset(
            X_train,
            label=Y_train,
            categorical_feature=self.str_columns,
        )
        test_data = lgb.Dataset(X_test, label=Y_test, reference=train_data)
        params = {
            'bagging_fraction': 0.9778411249854227,
            'bagging_freq': 6,
            'feature_fraction': 0.7178497443603504,
            'learning_rate': 0.06821355474058786,
            'min_child_samples': 79,
            'num_leaves': 36,
            'num_threads': 16
        }

        # 训练模型
        num_round = 100
        print("开始训练......")
        bst = lgb.train(params, train_data, num_round, valid_sets=[test_data])
        bst.save_model(self.model)
        print("模型训练完成✅")

    def evaluate_model(self):
        """
        评估模型性能
        :return:
        """
        fw = open("result/summary_{}.txt".format(dt), "a+", encoding="utf-8")
        path = 'data/predict_data/user_predict_{}.json'.format(dt)
        x, y = self.read_data(path, yc=23)
        bst = lgb.Booster(model_file=self.model)
        y_pred = bst.predict(x, num_iteration=bst.best_iteration)
        temp = sorted(list(y_pred))
        yuzhi = temp[int(len(temp) * 0.7) - 1]
        y_pred_binary = [0 if i <= yuzhi else 1 for i in list(y_pred)]
        # 转换为二进制输出
        score_list = []
        for index, item in enumerate(list(y_pred)):
            real_label = y[index]
            score = item
            prid_label = y_pred_binary[index]
            print(real_label, "\t", prid_label, "\t", score)
            fw.write("{}\t{}\t{}\n".format(real_label, prid_label, score))
            score_list.append(score)
        print("预测样本总量: {}".format(len(score_list)))
        data_series = pd.Series(score_list)
        print("统计 score 信息")
        print(data_series.describe())
        # 评估模型
        accuracy = accuracy_score(y, y_pred_binary)
        print(f"Accuracy: {accuracy}")
        fw.close()

    def feature_importance(self):
        """
        Get the importance of each feature
        :return:
        """
        lgb_model = lgb.Booster(model_file=self.model)
        importance = lgb_model.feature_importance(importance_type='split')
        feature_name = lgb_model.feature_name()
        feature_importance = sorted(zip(feature_name, importance), key=lambda x: x[1], reverse=True)

        # 打印特征重要性
        for name, imp in feature_importance:
            print(name, imp)


if __name__ == "__main__":
    i = int(input("输入 1 训练， 输入 2 预测:\n"))
    if i == 1:
        f = "train"
        dt = "whole"
        L = LightGBM(flag=f, dt=dt)
        L.train_model()
    elif i == 2:
        f = "predict"
        dt = int(input("输入日期， 20240316-21:\n"))
        L = LightGBM(flag=f, dt=dt)
        L.evaluate_model()
        L.feature_importance()
    elif i == 3:
        L = LightGBM("train", "whole")
        L.best_params()