fix: add test for single-sample categorical prediction and fix bug in…

… get_categorical_conf()

lightwood/analysis/nc/util.py

@@ -156,13 +156,15 @@ def get_numeric_conf_range(
 
 def get_categorical_conf(raw_confs: np.ndarray):
     """
-    Gets ICP confidence estimation for categorical targets from raw p-values per class.
+    ICP confidence estimation for categorical targets from raw p-values:
+        1.0 minus 2nd highest p-value yields confidence for predicted label.
     :param all_confs: p-value for each class per data point
-    :return: confidence for each data instance
+    :return: confidence for each data point
     """
-    # one minus 2nd best p-value yields confidence for predicted label
-    second_best = np.sort(raw_confs, axis=1)[:, -2]
-    confs = np.clip(np.subtract(1, second_best), 0.0001, 0.9999)
+    if len(raw_confs.shape) == 1:
+        raw_confs = np.expand_dims(raw_confs, axis=0)
+    second_p = np.sort(raw_confs, axis=1)[:, -2]
+    confs = np.clip(np.subtract(1, second_p), 0.0001, 0.9999)
     return confs
 
 

tests/integration/basic/test_categorical.py

@@ -54,6 +54,9 @@ def test_1_categorical(self):
         predictions = predictor.predict(df[:10], args={'all_mixers': True})
         assert '__mdb_mixer_Neural' in predictions.columns
 
+        # predict single sample
+        predictor.predict(df.iloc[[0]])
+
     def test_2_binary_no_analysis(self):
         df = pd.read_csv('tests/data/ionosphere.csv')[:100]
         mask = np.random.rand(len(df)) < 0.8