File size: 1,776 Bytes
dc1dc20 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 | 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))
|