Unable to pickle sklearn pipeline in Databricks

Viewed 1883

I was trying to pickle sklearn pipeline in databricks. but it ended up in error. The same code worked fine in local jupyter notebook

Tried with joblib also. But the same error occurs

#Build the baseline for Random forest
from sklearn.base import BaseEstimator, TransformerMixin
from sklearn import preprocessing
from sklearn.feature_extraction.text import CountVectorizer
from sklearn.feature_extraction.text import TfidfTransformer
from sklearn.feature_extraction.text import TfidfVectorizer
from sklearn.pipeline import Pipeline
from sklearn.pipeline import FeatureUnion
from sklearn.ensemble import RandomForestClassifier

#Build Baseline
class Select_Column(BaseEstimator, TransformerMixin):
    def __init__(self, column_name):  # no *args or **kargs
        self.column_name = column_name
    def fit(self, X, y=None):
        return self  # nothing else to do

    def transform(self, X, ):
        return X[self.column_name].values.reshape(-1, 1)


class Text_features(BaseEstimator, TransformerMixin):
    def __init__(self, column_name):  # no *args or **kargs
        self.column_name = column_name

    def fit(self, X, y=None):
        return self  # nothing else to do

    def transform(self, X, ):
        return X[self.column_name].values.astype('U')

class Binarized_features(BaseEstimator, TransformerMixin):
    def __init__(self, column_name):  # no *args or **kargs
                self.column_name = column_name

    def fit(self, X, y=None):
            return self  # nothing else to do

    def transform(self, X, ):
            return preprocessing.Binarizer(threshold=0.0).transform(X[self.column_name].values.reshape(-1, 1))



full_pipeline = FeatureUnion(transformer_list=[
            ("AB", Select_Column(["ABC"])),
            ("AC", Select_Column(["ACD"])),
            ("SCc", Select_Column(["SCc"])),
            ("Text_Features1", Pipeline([
                ('selector', Text_features(column_name='Text')),
                ('tfidf_vectorizer', TfidfVectorizer(min_df=1, 
                                     norm='l2', 
                                     max_features=100000, 
                                     smooth_idf=True, 
                                     use_idf=True, 
                                     ngram_range=(1,2)))  # TFID
            ])) 
        ])

vectorizer = full_pipeline.fit(data)
import pickle
pickle.dump(vectorizer_fit, open("vector.pickle", 'wb'))

The above code cause error in databricks, PicklingError: Can't pickle : attribute lookup Select_Column on main failed.

But the same code worked fine in local jupyter notebook

Databricks python Version 3.5.2

0 Answers
Related