alg/main_spider.py

"""
针对爬虫类型数据单独训练模型
"""
import os
import sys
import json
import optuna

from sklearn.linear_model import LogisticRegression

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

import numpy as np
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
from bayes_opt import BayesianOptimization


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

    def __init__(self, flag, dt):
        self.label_encoder = LabelEncoder()
        self.my_c = [
            "channel",
            "out_user_id",
            "mode",
            "out_play_cnt",
            "out_like_cnt",
            "out_share_cnt",
            "lop",
            "duration",
            "tag1",
            "tag2",
            "tag3"
        ]
        self.str_columns = ["channel", "mode", "out_user_id", "tag1", "tag2", "tag3"]
        self.float_columns = [
            "out_play_cnt",
            "out_like_cnt",
            "out_share_cnt",
            "lop",
            "duration"
        ]
        self.split_c = 0.7
        self.yc = 0.8
        self.model = "models/lightgbm_0401_spider.bin"
        self.flag = flag
        self.dt = dt

    def read_data(self):
        """
        Read data from local
        :return:
        """
        path = "data/train_data/spider_data_240401.json"
        df = pd.read_json(path)
        df = df.dropna(subset=['label'])
        labels = df['label']
        temp = sorted(labels)
        yc = temp[int(len(temp) * self.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):
        X, y = self.read_data()
        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:
        """
        x, y = self.read_data()
        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=["channel", "mode", "out_user_id", "tag1", "tag2", "tag3"],
        )
        test_data = lgb.Dataset(X_test, label=Y_test, reference=train_data)
        params = {
            'num_leaves': 31,
            'learning_rate': 0.00020616904432655601,
            'feature_fraction': 0.6508847259863764,
            'bagging_fraction': 0.7536774652478249,
            'bagging_freq': 6,
            'min_child_samples': 99,
            '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("summary_tag_03{}_spider.txt".format(self.dt), "a+", encoding="utf-8")
        # 测试数据
        with open("data/produce_data/x_data_total_return_predict_{}_spider.json".format(self.dt)) as f1:
            x_list = json.loads(f1.read())

        # 测试 label
        with open("data/produce_data/y_data_total_return_predict_{}_spider.json".format(self.dt)) as f2:
            Y_test = json.loads(f2.read())

        Y_test = [0 if i <= 19 else 1 for i in Y_test]
        X_test = pd.DataFrame(x_list, columns=self.my_c)
        for key in self.str_columns:
            X_test[key] = self.label_encoder.fit_transform(X_test[key])
        for key in self.float_columns:
            X_test[key] = pd.to_numeric(X_test[key], errors="coerce")
        bst = lgb.Booster(model_file=self.model)
        y_pred = bst.predict(X_test, 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_test[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_test, 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("输入日期， 16-21:\n"))
    #     L = LightGBM(flag=f, dt=dt)
    #     L.evaluate_model()
    #     L.feature_importance()
    L = LightGBM("train", "whole")
    L.best_params()
    # study = optuna.create_study(direction='maximize')
    # study.optimize(L.bays_params, n_trials=100)
    # print('Number of finished trials:', len(study.trials))
    # print('Best trial:', study.best_trial.params)