自定义函数调用train_test_split报test_size多值错误及预测逻辑咨询

一、train_test_split报错修复 #

这个报错核心原因90%以上是当前运行环境中存在同名自定义函数覆盖了sklearn官方的train_test_split接口，比如你之前写过同名自定义函数，参数列表和官方不同，就会出现参数传参错位的问题。
你可以先在调用函数前加一行代码验证：

print(train_test_split.__module__)

如果输出不是sklearn.model_selection._split，就说明确实被覆盖了，解决方案有两种：

• 重启Jupyter/代码运行内核，清空之前定义的变量和函数
• 显式指定调用sklearn的接口，避免被覆盖：

from sklearn import model_selection
X_train, X_test, y_train, y_test = model_selection.train_test_split(X, y, test_size=0.20, random_state=5)

二、其余逻辑错误修正 #

就算解决了上述报错，你的代码还有多处逻辑问题会导致无法得到预期结果：

1. 预测传参完全错误：旧版代码你给model.predict传的是target['rated_object']（就是o3、q3这类字符串，不是特征值），新版代码里你写的target['questionId']在你的数据集里根本不存在，会直接报KeyError
2. 待预测样本参与了训练，属于数据泄露：你要预测的目标样本本身混入了训练集，会导致预测结果失真
3. 没有处理缺失值：XGBoost原生不支持NaN输入，你数据集里的空值会直接触发训练报错

三、修正后可运行代码 #

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

def predict_cn(df, rated_object):
    df_copy = df.copy()
    # 单独提取待预测样本，不参与训练
    target_mask = df_copy['rated_object'] == rated_object
    target_row = df_copy[target_mask].iloc[0]
    # 提取待预测样本的特征，剔除id和标签列
    target_features = target_row.drop(['rated_object', 'rating']).values.reshape(1, -1)
    # 构造训练数据集，排除待预测样本
    train_df = df_copy[~target_mask].drop('rated_object', axis=1)
    # 填充缺失值，可根据需求更换填充逻辑
    train_df = train_df.fillna(train_df.mean(numeric_only=True))
    X = train_df.drop('rating', axis=1)
    y = train_df['rating']
    X_train, X_test, y_train, y_test = train_test_split(X, y, test_size=0.2, random_state=5)
    model = XGBClassifier()
    model.fit(X_train, y_train)
    # 给待预测样本填充缺失值后再预测
    target_features = np.nan_to_num(target_features, nan=X.mean(numeric_only=True).values)
    prediction = model.predict(target_features)
    return prediction[0]

# 用你提供的模拟数据测试
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'])
print(predict_cn(df_mock, 'q3'))

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