如何在垃圾邮件分类任务中仅对spam类数据做训练测试拆分

你不需要用if-elif语句区分数据，直接用Pandas的布尔索引筛选出所有spam类数据，再对这个子集做训练测试拆分即可，操作更简洁高效。

实现步骤 #

1. 从完整数据集中筛选仅包含spam类的子集
2. 对该spam子集使用train_test_split完成拆分

修改后的代码 #

替换原代码中train_test_split相关段落为以下内容：

# 筛选所有spam类数据
spam_data = data[data['class'] == 'spam']

# 仅对spam数据执行训练测试拆分
X_train, X_test, y_train, y_test = train_test_split(
    spam_data, 
    spam_data['class'],  # y应为标签列，而非整个数据集
    random_state=104,
    test_size=0.20,
    shuffle=True
)

print('X_train: ')
print(X_train.head())
print('')
print('X_test: ')
print(X_test.head())
print('')
print('y_train: ')
print(y_train.head())
print('')
print('y_test: ')
print(y_test.head())

额外优化提示 #

你当前的dataFrameFromDirectory函数只存储了邮件分类标签，未将邮件内容（message）存入DataFrame，这会导致后续无法用邮件文本训练分类模型。若需用邮件内容作为特征，需修改该函数：

def dataFrameFromDirectory(path, classification):
    rows = []
    index = []
    for filename, message in readFiles(path):
        # 新增message字段存储邮件内容
        rows.append({'class': classification, 'message': message})
        index.append(filename)

    return DataFrame(rows, index=index)

修改后，拆分出的X_train['message']可作为模型输入特征，y_train作为分类标签。

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