officer race predictor works

inference.py

@@ -1,19 +1,20 @@
 import numpy as np
 
+import config
 from model import load_models
 from preprocessor import preprocess
 
 
 def predict_officer_race(model_name, X):
-    models = load_models(model_name, 'OfficerRace')
+    (names, models) = load_models(model_name, 'OfficerRace')
     n_samples = 1  # single sample
     probs = np.zeros((n_samples, 4))
     for i, model in enumerate(models):
         prob = model.predict_proba(X)[:, 1]
         probs[:, i] = prob
 
-    Y = np.argmax(probs, axis=1)
-    return np.eye(4)[Y]
+    y = np.argmax(probs, axis=1)
+    return names[int(y[0])]
 
 
 def infer_officer_race(model_name, current_rank, incident_rank, previous_complaints, complaint_duration_days,
@@ -25,4 +26,5 @@ def infer_officer_race(model_name, current_rank, incident_rank, previous_complai
                              penalty_rec, penalty_cat, location_type, contact_outcome, impacted_gender, impacted_race,
                              incident_precinct)
     prediction = predict_officer_race(model_name, input_array)
+    prediction = f"The officer is predicted to be {prediction}"
     return prediction

model.py

@@ -4,9 +4,14 @@ from pathlib import Path
 
 def load_models(model_name, target_column):
     """Load models from disk"""
-    if model_name == "OfficerRace":
-        models = [joblib.load(path) for path in Path(f"models/{target_column}/{model_name}").glob("*.pkl")]
-        return models
+    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
 
 
 def get_folder_names(directory):