| from sklearn.base import BaseEstimator, TransformerMixin | |
| import pandas as pd | |
| class ProductTypeCategorizer(BaseEstimator, TransformerMixin): | |
| def __init__(self): | |
| self.perishables = [ | |
| "Dairy", | |
| "Meat", | |
| "Fruits and Vegetables", | |
| "Breakfast", | |
| "Breads", | |
| "Seafood", | |
| ] | |
| def fit(self, X, y=None): | |
| return self | |
| def transform(self, X): | |
| X = X.copy() | |
| X = pd.Series(X.squeeze()) # ensure 1D | |
| return pd.DataFrame( | |
| X.apply(lambda x: "Perishables" if x in self.perishables else "Non Perishables") | |
| ) | |