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']
        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 bays(self):
        # 创建LightGBM数据集，注意不要在这里指定categorical_feature，因为我们使用的是玩具数据集
        x, y = self.read_data()
        train_size = int(len(x) * 0.9)
        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)
        def lgbm_eval(num_leaves, learning_rate, feature_fraction, bagging_fraction, bagging_freq, min_child_samples):
            params = {
                'objective': 'binary',
                'metric': 'auc',
                'verbose': -1,
                'num_leaves': int(num_leaves),
                'learning_rate': learning_rate,
                'feature_fraction': feature_fraction,
                'bagging_fraction': bagging_fraction,
                'bagging_freq': int(bagging_freq),
                'min_child_samples': int(min_child_samples),
            }
            cv_result = lgb.cv(params, train_data, nfold=5, seed=42, stratified=True, metrics=['auc'])
            return max(cv_result['auc-mean'])

        param_bounds = {
            'num_leaves': (20, 40),
            'learning_rate': (1e-4, 1e-2),
            'feature_fraction': (0.5, 0.8),
            'bagging_fraction': (0.5, 0.8),
            'bagging_freq': (1, 10),
            'min_child_samples': (20, 100),
        }

        optimizer = BayesianOptimization(f=lgbm_eval, pbounds=param_bounds, random_state=42)
        optimizer.maximize(init_points=5, n_iter=25)

        print("Best Parameters:", optimizer.max['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.bays()
    # 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)