improve custom group support

oxfordinternetinstitute · Oct 25, 2024 · 55d3c0c · 55d3c0c
1 parent 303195d
commit 55d3c0c
Showing 1 changed file with 7 additions and 1 deletion.
diff --git a/src/oxonfair/learners/fair.py b/src/oxonfair/learners/fair.py
@@ -1012,7 +1012,13 @@ def call_or_get_proba(predictor, data) -> np.ndarray:
     if isinstance(data, dict):
         data = data['data']
     if callable(predictor):
-        return np.asarray(predictor(data))
+        out = np.asarray(predictor(data))
+        if out.ndim == 1:
+            width = out.max()+1
+            new_out = np.zeros((out.shape[0], width))
+            new_out[(np.arange(out.shape[0]), out)] = 1
+            return new_out
+        return out
     return np.asarray(predictor.predict_proba(data))