Python使用朴素贝叶斯做文本分类时出现KeyError如何解决

model.predict()返回的结果已经是你训练时传入的category列的类别值，而非数据的行索引。你尝试用预测得到的类别字符串作为行索引，去取默认是整数索引的train.category序列的值，自然触发KeyError。

直接返回预测结果即可，无需再从训练集中取值：

# 优化后的预测函数
def predict_category(s, model=model):
    pred = model.predict([s])
    return pred[0]

如果有特殊场景需要从训练集校验取值，可使用布尔索引：

def predict_category(s, train=train, model=model):
    pred = model.predict([s])
    return train.loc[train['category'] == pred[0], 'category'].iloc[0]

额外说明 #

如果你报错的Key是首字母大写的Greeting，和你CSV里的小写greeting不一致，可以检查下CSV文件的category列是否存在大小写混杂的情况，避免后续预测匹配异常。

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