rov-offline/ad_xgboost_predict.py

import pandas as pd
import xgboost as xgb
from xgboost.sklearn import XGBClassifier
from utils import RedisHelper
redis_helper = RedisHelper()


# 1. 模型加载
model = XGBClassifier()
booster = xgb.Booster()
booster.load_model('./data/ad_xgb.model')
model._Booster = booster
# 2. 预测：ad_status = 0, 不出广告
df_0 = pd.read_csv('./data/predict_data/predict_data_0.csv')
columns_0 = df_0.columns.values.tolist()
columns_0.remove('videoid')
y_pred_proba_0 = model.predict_proba(df_0[columns_0[2:]])
df_0['y_0'] = [x[1] for x in y_pred_proba_0]
pre_df_0 = df_0[['apptype', 'mid', 'videoid', 'y_0']].copy()

# 3. 预测：ad_status = 1, 不出广告
df_1 = pd.read_csv('./data/predict_data/predict_data_1.csv')
columns_1 = df_1.columns.values.tolist()
columns_1.remove('videoid')
y_pred_proba_1 = model.predict_proba(df_1[columns_1[2:]])
df_1['y_1'] = [x[1] for x in y_pred_proba_1]
pre_df_1 = df_1[['apptype', 'mid', 'videoid', 'y_1']].copy()

# 4. merge 结果
res_df = pd.merge(pre_df_0, pre_df_1, how='left', on=['apptype', 'mid', 'videoid'])
res_df['res_predict'] = res_df['y_0'] - res_df['y_1']
print(res_df.head())

# 5. to csv
res_df.to_csv('./data/predict_data/predict_res.csv', index=False)
print("to csv finished!")

# 6. to redis
for ind, row in res_df.iterrows():
    app_type = row['apptype']
    mid = row['mid']
    video_id = row['videoid']
    pre_res = row['res_predict']
    key = f"ad:xgb:predict:{app_type}:{mid}:{video_id}"
    redis_helper.set_data_to_redis(key_name=key, value=pre_res, expire_time=48*3600)
print("to redis finished!")