| from __future__ import annotations |
|
|
| import json |
| from pathlib import Path |
|
|
| import joblib |
| from sklearn.datasets import load_iris |
| from sklearn.model_selection import GridSearchCV |
| from sklearn.neighbors import KNeighborsClassifier |
| from sklearn.pipeline import Pipeline |
| from sklearn.preprocessing import StandardScaler |
|
|
|
|
| ARTIFACT_PATH = Path(__file__).with_name("model.joblib") |
| METRICS_PATH = Path(__file__).with_name("metrics.json") |
|
|
|
|
| def train_and_save(artifact_path: Path = ARTIFACT_PATH) -> dict: |
| iris = load_iris(as_frame=True) |
| X = iris.data |
| y = iris.target |
|
|
| target_names = [str(name) for name in iris.target_names] |
| feature_names = [str(name) for name in iris.feature_names] |
|
|
| pipeline = Pipeline( |
| steps=[ |
| ("scaler", StandardScaler()), |
| ("knn", KNeighborsClassifier()), |
| ] |
| ) |
|
|
| param_grid = { |
| "knn__n_neighbors": list(range(1, 21)), |
| "knn__weights": ["uniform", "distance"], |
| "knn__p": [1, 2], |
| } |
|
|
| search = GridSearchCV( |
| estimator=pipeline, |
| param_grid=param_grid, |
| cv=5, |
| scoring="accuracy", |
| n_jobs=-1, |
| refit=True, |
| ) |
| search.fit(X, y) |
|
|
| best_model = search.best_estimator_ |
| joblib.dump( |
| { |
| "model": best_model, |
| "target_names": target_names, |
| "feature_names": feature_names, |
| }, |
| artifact_path, |
| ) |
|
|
| metrics = { |
| "cv_best_accuracy": float(search.best_score_), |
| "best_params": search.best_params_, |
| } |
| METRICS_PATH.write_text(json.dumps(metrics, indent=2), encoding="utf-8") |
| return metrics |
|
|
|
|
| if __name__ == "__main__": |
| metrics = train_and_save() |
| print(f"Saved model to: {ARTIFACT_PATH}") |
| print(json.dumps(metrics, indent=2)) |
|
|
