How do I incorporate SelectKBest in an SKlearn pipeline

Question

I'm trying to build a text classifier with sklearn. The idea is to:

1. Vectorize training corpus using TfidfVectorizer
2. Select the top 20,000 features that result (or using all features if the resultant number is below 20k) using SelectKBest
3. Feed these features into a Logistic Regression Classifier

I've set it up successfully as follows:

from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.linear_model import LogisticRegression

vectorizer = TfidfVectorizer()
x_train = vectorizer.fit_transform(df_train["input"])
selector = SelectKBest(f_classif, k=min(20000, x_train.shape[1]))
selector.fit(x_train, df_train["label"].values)
x_train = selector.transform(x_train)
classifier = LogisticRegression()
classifier.fit(x_train, df_train["label"])

I would now like to wrap all this up into a pipeline, and share the pipeline so it can be used by others for their own text data. Yet, I can't figure how to get SelectKBest to achieve the same behavior as it did above, ie accept min(20000, n_features from vectorizer output) as k. If I were to simply leave it as k=20000 as below, the pipeline doesn't work (throws an error) when fitting new corpora with less than 20k vectorized features.

pipe = Pipeline([
            ("vect",TfidfVectorizer()),
            ("selector",SelectKBest(f_classif, k=20000)),
            ("clf",LogisticRegression())])

Answer 1

As @vivek kumar pointed out you need to override the _check_params method of SelectKBest and add your logic to it as shown below:

class MySelectKBest(SelectKBest):
    def _check_params(self, X, y):
        if (self.k >= X.shape[1]):
            warnings.warn("Less than %d number of features found, so setting k as %d" % (self.k, X.shape[1]),
                      UserWarning)
            self.k = X.shape[1]
        if not (self.k == "all" or 0 <= self.k):
            raise ValueError("k should be >=0, <= n_features = %d; got %r. "
                             "Use k='all' to return all features."
                             % (X.shape[1], self.k)) 

I have also set a warning in case if the number of features found is less than the threshold set. Now let's see a working example of the same:

from sklearn.datasets import fetch_20newsgroups
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.feature_selection import SelectKBest, f_classif
from sklearn.linear_model import LogisticRegression
from sklearn.pipeline import Pipeline
import warnings

categories = ['alt.atheism', 'comp.graphics',
              'comp.sys.ibm.pc.hardware', 'comp.sys.mac.hardware',
              'comp.windows.x', 'misc.forsale', 'rec.autos']
newsgroups = fetch_20newsgroups(categories=categories)
y_true = newsgroups.target

# newsgroups result in 47K odd features after performing TFIDF vectorizer

# Case 1: When K < No. of features - the regular case
pipe = Pipeline([
            ("vect",TfidfVectorizer()),
            ("selector",MySelectKBest(f_classif, k=30000)),
            ("clf",LogisticRegression())])

pipe.fit(newsgroups.data, y_true)
pipe.score(newsgroups.data, y_true)
#0.968

#Case 2: When K > No. of cases - the one with an issue

pipe = Pipeline([
            ("vect",TfidfVectorizer()),
            ("selector",MySelectKBest(f_classif, k=50000)),
            ("clf",LogisticRegression())])

pipe.fit(newsgroups.data, y_true)
UserWarning: Less than 50000 number of features found, so setting k as 47407

pipe.score(newsgroups.data, y_true)
#0.9792

Hope this helps!