Spaces:
Sleeping
Sleeping
File size: 1,565 Bytes
cfc11ee d830db3 a31c8a2 cfc11ee 067ac42 cfc11ee 740d621 f267af6 cfc11ee f267af6 cfc11ee 740d621 bd9ee57 f267af6 bd9ee57 | 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 | import os
import joblib
from pathlib import Path
import imblearn
def load_models(model_name, target_column):
"""Load models from disk"""
model_root = Path("./models")
target_column_root = model_root / target_column
if target_column == "OfficerRace":
model_selected_root = target_column_root / model_name
models = [joblib.load(f) for f in model_selected_root.iterdir()]
names = [path.stem.split("_")[-1] for path in model_selected_root.iterdir()]
return names, models
if target_column == "OfficerGender":
model_path = target_column_root / f"{model_name}.pkl"
return joblib.load(model_path)
if target_column == "PenaltyCat":
model_path = target_column_root / f"{model_name}.pkl"
return joblib.load(model_path)
def get_folder_names(directory):
return [name for name in os.listdir(directory) if os.path.isdir(os.path.join(directory, name))]
def available_models(target_column):
"""Return available models for a given target column"""
if target_column == "OfficerRace":
return get_folder_names("models/OfficerRace")
elif target_column == "OfficerGender":
return [os.path.splitext(name)[0] for name in os.listdir(f"models/{target_column}") if name.endswith(".pkl")]
elif target_column == "PenaltyCat":
return [os.path.splitext(name)[0] for name in os.listdir(f"models/{target_column}") if name != ("scaler_model"
".pkl")]
|