0

I'm trying to save my text classification model as a pickle file. I have a specific set of preprocessing steps that I wanted to save in my end model to apply it on unseen data for prediction. Currently I tried using a sklearn pipeline which includes preprocessing, applying count vectorizer and applying the algorithm. My question is if this is a right way to save the preprocessing steps in the end model or should I save it as a seperate file. Below is my code

from sklearn import model_selection
X_train, X_test, y_train, y_test = model_selection.train_test_split(df_train.data, df_train.label, test_size=0.2)
vect = CountVectorizer(stop_words='english', max_features=10000, , ngram_range=(1,2))
X_train_dtm = vect.fit_transform(X_train)
X_test_dtm = vect.transform(X_test)

from sklearn import metrics
from sklearn.ensemble import RandomForestClassifier
rf_classifier = RandomForestClassifier(n_estimators=100)
rf_classifier.fit(X_train_dtm, y_train)
rf_predictions = rf_classifier.predict(X_test_dtm)
print("RF Accuracy:")
metrics.accuracy_score(y_test, rf_predictions)


import pickle
from sklearn.pipeline import Pipeline
pipe = Pipeline([("prep", prep),("CV", vect), ("RF", rf_classifier)])
with open('PII_model.pickle', 'wb') as picklefile:
    pickle.dump(pipe, picklefile)

I have a method for preprocessing which I invoke it and I've included this in my sklearn pipeline

prep = prep(df_train)
dmorgan
  • 73
  • 1
  • 11

1 Answers1

4

Generally, it is the right approach. But you can improve it put everything in the pipeline from the very beginning:

vect = CountVectorizer(stop_words='english', 
                       max_features=10000, 
                       ngram_range=(1,2))
rf_classifier = RandomForestClassifier(n_estimators=100)
pipe = Pipeline([("prep", prep),("CV", vect), ("RF", rf_classifier)])    
pipe.fit(X_train_dtm, y_train)

rf_predictions = pipe.predict(X_test_dtm)
print("RF Accuracy:")
metrics.accuracy_score(y_test, rf_predictions)

with open('PII_model.pickle', 'wb') as picklefile:
    pickle.dump(pipe, picklefile)

One more benefit to putting everything in the one pipeline - you can easily use GridSearch to tune parameters of the model.

Here is you can also find official documentation on how to organize a pipeline with mixed types.

Danylo Baibak
  • 2,106
  • 1
  • 11
  • 18
  • I just have a followup question. I tried saving my pre-processing steps in the pipeline as stated above and on predicting new data I do not see my preprocessing steps being applied. Also when I try saving my model with and without preprocessing steps I do not see any difference in the pickle file size. So I was wondering if I have to first have a function to apply my preprocessing `(prep=prep(df_train))` and then load the pickle file. – dmorgan Jul 30 '20 at 10:02
  • 1
    Seems, yes, you need to apply your preprocessing function before fitting the pipeline. There is another approach - you can wrap your logic into custom transformers. Then you can include your transformer into the sklearn pipeline. For implementation, you need just extend https://scikit-learn.org/stable/modules/generated/sklearn.base.TransformerMixin.html to define a new transformer. Here is good documentation with an example (https://cloud.google.com/ai-platform/prediction/docs/custom-pipeline#create_custom_transformers) of how to implement it. – Danylo Baibak Jul 30 '20 at 10:40