Spaces:
Sleeping
Sleeping
| 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}") | |