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 sklearn.preprocessing import LabelEncoder
from sklearn.metrics import accuracy_score

from sklearn.model_selection import train_test_split, StratifiedKFold
from sklearn.datasets import load_breast_cancer
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 objective(self, trial):
        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:]

        dtrain = lgb.Dataset(X_train, label=Y_train)
        dvalid = lgb.Dataset(X_test, label=Y_test, reference=dtrain)

        param = {
            'objective': 'binary',  # 根据问题修改，例如'regression'或'multiclass'
            'metric': 'binary_logloss',  # 根据问题修改，例如'l2'或'multi_logloss'
            'verbosity': -1,
            'boosting_type': 'gbdt',
            'num_leaves': trial.suggest_int('num_leaves', 20, 40),
            'learning_rate': trial.suggest_float('learning_rate', 1e-4, 1e-1, log=True),
            'feature_fraction': trial.suggest_float('feature_fraction', 0.6, 0.9),
            'bagging_fraction': trial.suggest_float('bagging_fraction', 0.6, 0.9),
            'bagging_freq': trial.suggest_int('bagging_freq', 1, 10),
            'min_child_samples': trial.suggest_int('min_child_samples', 5, 100),
            'num_thread': 16
        }

        gbm = lgb.train(param, dtrain, valid_sets=[dvalid])
        preds = gbm.predict(X_test)
        pred_labels = np.rint(preds)
        accuracy = accuracy_score(Y_test, pred_labels)
        return accuracy  # 或其他优化指标

    def tune(self, n_trials=100):
        study = optuna.create_study(direction='maximize')
        study.optimize(self.objective, n_trials=n_trials)
        print('Number of finished trials:', len(study.trials))
        print('Best trial:', study.best_trial.params)

    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.tune()
    # 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)