如何修改含TFIDF与GBM的sklearn Pipeline使其直接输出类别1的预测概率

你可以通过自定义类修改管道或分类器的逻辑实现需求，以下是两种可直接复用的实现方案：

方案1：自定义Pipeline类（无需修改原有分类器逻辑） #

直接继承sklearn原生Pipeline类，重写predict方法返回类别1的概率，训练逻辑和原生Pipeline完全兼容：

from sklearn.pipeline import Pipeline
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.ensemble import GradientBoostingClassifier

class CustomProbPipeline(Pipeline):
    def predict(self, X):
        # 调用原生predict_proba获取全类别概率
        all_proba = super().predict_proba(X)
        # 单样本返回浮点值，多样本返回一维概率数组
        return all_proba[0, 1] if all_proba.shape[0] == 1 else all_proba[:, 1]

# 直接用自定义Pipeline替换原有管道即可
gbm_pipeline = CustomProbPipeline([
    ('tfidf', TfidfVectorizer(use_idf=True)),
    ('gbm', GradientBoostingClassifier(random_state = 23))
])

调用gbm_pipeline.predict(输入数据)即可直接返回你需要的类别1概率值。

方案2：自定义分类器包装类 #

如果需要保留原生Pipeline的其他方法逻辑，也可以包装GradientBoostingClassifier，修改预测输出：

from sklearn.ensemble import GradientBoostingClassifier

class GBMOnlyPosProb(GradientBoostingClassifier):
    def predict(self, X):
        all_proba = super().predict_proba(X)
        return all_proba[0, 1] if all_proba.shape[0] == 1 else all_proba[:, 1]

# 替换原有管道中的GBM分类器即可
gbm_pipeline = Pipeline([
    ('tfidf', TfidfVectorizer(use_idf=True)),
    ('gbm', GBMOnlyPosProb(random_state = 23))
])

两种方案都不需要对管道输出做二次处理，训练阶段的逻辑和原生实现完全一致，不会影响模型效果。

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