如何将XGBoost分类模型封装为接收dataframe与rated_object返回预测评分的函数

问题根因和修复方案 #

该需求完全可以实现，下面是具体的问题排查和修复步骤：

1. 解决train_test_split报错 #

这个报错的核心原因是命名冲突：你本地环境之前大概率自定义过同名的train_test_split函数，覆盖了从sklearn导入的官方函数，导致参数匹配异常。
解决方法很简单，导入时给官方函数取别名避开冲突即可：

from sklearn.model_selection import train_test_split as tts

后续调用时用tts()代替原来的train_test_split()即可。

2. 修复函数逻辑错误 #

你编写的函数还有几个明显的逻辑问题：

• 预测时错误调用了不存在的target['questionId']字段，应该传入目标对象的特征列数据
• 原代码中predict方法传入了不支持的verbose参数，会触发额外报错
• 没有将目标对象的特征转为符合模型输入要求的二维数组格式

3. 最终可运行的完整代码 #

import pandas as pd
import xgboost as xgb
from sklearn.model_selection import train_test_split as tts
from xgboost import XGBClassifier
from numpy import nan

# 模拟数据
data_mock = [['q1', 10.93, 20, 1, 0], ['q2', nan, 12, 0, 1], ['q3', 14.34, 30, 0, 1], ['q4', 12.93, 20, 0, 1], ['q5', nan, 62, 1, 0], ['q6', 14.34, 60, 0, 0], ['q7', 16.93, 28, 1, 1], ['q8', nan, 12, 1, 1], ['q9', 10.34, 50, 0, 0], ['q10', 10.93, 20, 0, 0], ['q11', nan, 57, 1, 1], ['q12', 89.34, 30, 0, 0]]
df_mock = pd.DataFrame(data_mock, columns = ['rated_object', 'feature_1', 'feature_2', 'feature_n', 'rating'])

def predict_cn(df, rated_object):
    # 深拷贝避免污染原数据
    df_copy = df.copy()
    # 提取目标对象行
    target_row = df_copy[df_copy['rated_object'] == rated_object].iloc[0]
    # 清理非特征列
    df_copy.drop('rated_object', axis=1, inplace=True)
    # 拆分特征和标签
    X = df_copy.drop('rating', axis=1)
    y = df_copy['rating']
    # 拆分训练测试集
    X_train, X_test, y_train, y_test = tts(X, y, test_size=0.20, random_state=5)
    # 训练模型，添加参数避免警告
    model = XGBClassifier(use_label_encoder=False, eval_metric='logloss')
    model.fit(X_train, y_train)
    # 提取目标特征转为二维数组格式
    target_features = target_row[X.columns].values.reshape(1, -1)
    # 预测返回标量结果
    prediction = model.predict(target_features)[0]
    return prediction

# 测试调用
print(predict_cn(df_mock, 'q3'))

优化建议 #

如果函数需要被多次调用，建议不要每次调用都重新训练模型，可以将训练好的模型缓存起来，仅在数据更新时重新训练，大幅提升运行效率。

内容的提问来源于stack exchange，提问作者johnnydoe