GSLC-MLOps / src /train.py
github-actions
Deploy from GitHub Actions: 4025db1ea73f71f9c2a3953aa74dd93cdedb2084
d7f4720
Raw
History Blame Contribute Delete
1.55 kB
import json
from joblib import dump
from sklearn.datasets import load_breast_cancer
from sklearn.pipeline import Pipeline
from sklearn.preprocessing import MinMaxScaler
from sklearn.ensemble import RandomForestClassifier
from sklearn.model_selection import train_test_split, GridSearchCV
from sklearn.metrics import recall_score
from src.utils import MODEL_PATH, META_PATH
def load_data():
data = load_breast_cancer()
return data.data, data.target, list(data.feature_names), list(data.target_names)
def train(smoke=False):
X, y, feature_names, target_names = load_data()
Xtr, Xte, ytr, yte = train_test_split(
X, y, test_size=0.2, random_state=42, stratify=y
)
pipe = Pipeline([
("scaler", MinMaxScaler()),
("rf", RandomForestClassifier(random_state=42))
])
if smoke:
params = {"rf__n_estimators": [50]}
cv = 2
else:
params = {"rf__n_estimators": [100, 200], "rf__max_depth": [None, 8, 12]}
cv = 5
grid = GridSearchCV(pipe, params, scoring="recall", cv=cv, n_jobs=-1)
grid.fit(Xtr, ytr)
best = grid.best_estimator_
yhat = best.predict(Xte)
rec = recall_score(yte, yhat)
dump(best, MODEL_PATH)
META_PATH.write_text(json.dumps({
"best_params": grid.best_params_,
"cv": cv,
"recall_test": rec,
"feature_names": feature_names,
"target_names": target_names
}, indent=2))
return rec
if __name__ == "__main__":
r = train(smoke=False)
print(f"Recall (test): {r:.4f}")