| import joblib | |
| from sklearn.datasets import fetch_openml | |
| from sklearn.preprocessing import StandardScaler, OneHotEncoder | |
| from sklearn.compose import make_column_transformer | |
| from sklearn.pipeline import make_pipeline | |
| from sklearn.model_selection import train_test_split | |
| from sklearn.tree import DecisionTreeRegressor | |
| dataset = fetch_openml(data_id=43355, as_frame=True, parser='auto') | |
| diamond_prices = dataset.data | |
| target = ['price'] | |
| numeric_features = ['carat'] | |
| categorical_features = ['shape', 'cut', 'color', 'clarity', 'report', 'type'] | |
| X = diamond_prices.drop(columns=target) | |
| y = diamond_prices[target] | |
| Xtrain, Xtest, ytrain, ytest = train_test_split( | |
| X, y, | |
| test_size=0.2, | |
| random_state=42 | |
| ) | |
| preprocessor = make_column_transformer( | |
| (StandardScaler(), numeric_features), | |
| (OneHotEncoder(handle_unknown='ignore'), categorical_features) | |
| ) | |
| model_pipeline = make_pipeline(preprocessor, DecisionTreeRegressor()) | |
| model_pipeline.fit(Xtrain, ytrain) | |
| joblib.dump(model_pipeline, 'model-v1.joblib') |