[ENH] Add and enhance average confidence distance calculation

abl/reasoning/reasoner.py

@@ -6,7 +6,7 @@ from zoopt import Dimension, Objective, Opt, Parameter, Solution
 
 from ..data.structures import ListData
 from ..reasoning import KBBase
-from ..utils.utils import confidence_dist, hamming_dist
+from ..utils.utils import hamming_dist, confidence_dist, avg_confidence_dist
 
 
 class Reasoner:
@@ -74,10 +74,10 @@ class Reasoner:
 
     def _check_valid_dist(self, dist_func):
         if isinstance(dist_func, str):
-            if dist_func not in ["hamming", "confidence"]:
+            if dist_func not in ["hamming", "confidence", "avg_confidence"]:
                 raise NotImplementedError(
                     'Valid options for predefined dist_func include "hamming" '
-                    + f'and "confidence", but got {dist_func}.'
+                    + f' "confidence" and "avg_confidence", but got {dist_func}.'
                 )
             return
         elif callable(dist_func):
@@ -167,6 +167,9 @@ class Reasoner:
         elif self.dist_func == "confidence":
             candidates_idxs = [[self.label_to_idx[x] for x in c] for c in candidates]
             return confidence_dist(data_example.pred_prob, candidates_idxs)
+        elif self.dist_func == "avg_confidence":
+            candidates_idxs = [[self.label_to_idx[x] for x in c] for c in candidates]
+            return avg_confidence_dist(data_example.pred_prob, candidates_idxs)
         else:
             candidate_idxs = [[self.label_to_idx[x] for x in c] for c in candidates]
             cost_list = self.dist_func(data_example, candidates, candidate_idxs, reasoning_results)

abl/utils/__init__.py

@@ -2,6 +2,7 @@ from .cache import Cache, abl_cache
 from .logger import ABLLogger, print_log
 from .utils import (
     confidence_dist,
+    avg_confidence_dist,
     flatten,
     hamming_dist,
     reform_list,
@@ -14,6 +15,7 @@ __all__ = [
     "ABLLogger",
     "print_log",
     "confidence_dist",
+    "avg_confidence_dist",
     "flatten",
     "hamming_dist",
     "reform_list",

abl/utils/utils.py

@@ -87,13 +87,14 @@ def hamming_dist(pred_pseudo_label: List[Any], candidates: List[List[Any]]) -> n
     return np.sum(pred_pseudo_label != candidates, axis=1)
 
 
-def confidence_dist(pred_prob: List[np.ndarray], candidates_idxs: List[List[Any]]) -> np.ndarray:
+def confidence_dist(pred_prob: np.ndarray, candidates_idxs: List[List[Any]]) -> np.ndarray:
     """
-    Compute the confidence distance between prediction probabilities and candidates.
+    Compute the confidence distance between prediction probabilities and candidates,
+    where the confidence distance is defined as 1 - the product of prediction probabilities.
 
     Parameters
     ----------
-    pred_prob : List[np.ndarray]
+    pred_prob : np.ndarray
         Prediction probability distributions, each element is an array
         representing the probability distribution of a particular prediction.
     candidates_idxs : List[List[Any]]
@@ -105,9 +106,29 @@ def confidence_dist(pred_prob: List[np.ndarray], candidates_idxs: List[List[Any]
         Confidence distances computed for each candidate.
     """
     pred_prob = np.clip(pred_prob, 1e-9, 1)
-    _, cols = np.indices((len(candidates_idxs), len(candidates_idxs[0])))
+    cols = np.arange(len(candidates_idxs[0]))[None, :]
     return 1 - np.prod(pred_prob[cols, candidates_idxs], axis=1)
 
+def avg_confidence_dist(pred_prob: np.ndarray, candidates_idxs: List[List[Any]]) -> np.ndarray:
+    """
+    Compute the average confidence distance between prediction probabilities and candidates,
+    where the confidence distance is defined as 1 - the average of prediction probabilities.
+
+    Parameters
+    ----------
+    pred_prob : np.ndarray
+        Prediction probability distributions, each element is an array
+        representing the probability distribution of a particular prediction.
+    candidates_idxs : List[List[Any]]
+        Multiple possible candidates' indices.
+
+    Returns
+    -------
+    np.ndarray
+        Confidence distances computed for each candidate.
+    """
+    cols = np.arange(len(candidates_idxs[0]))[None, :]
+    return 1 - np.average(pred_prob[cols, candidates_idxs], axis=1)
 
 def to_hashable(x: Union[List[Any], Any]) -> Union[Tuple[Any, ...], Any]:
     """

examples/zoo/main.py

@@ -8,7 +8,7 @@ from abl.bridge import SimpleBridge
 from abl.data.evaluation import ReasoningMetric, SymbolAccuracy
 from abl.learning import ABLModel
 from abl.reasoning import Reasoner
-from abl.utils import ABLLogger, confidence_dist, print_log, tab_data_to_tuple
+from abl.utils import ABLLogger, avg_confidence_dist, print_log, tab_data_to_tuple
 
 from get_dataset import load_and_preprocess_dataset, split_dataset
 from kb import ZooKB
@@ -16,7 +16,7 @@ from kb import ZooKB
 
 def consitency(data_example, candidates, candidate_idxs, reasoning_results):
     pred_prob = data_example.pred_prob
-    model_scores = confidence_dist(pred_prob, candidate_idxs)
+    model_scores = avg_confidence_dist(pred_prob, candidate_idxs)
     rule_scores = np.array(reasoning_results)
     scores = model_scores + rule_scores
     return scores

examples/zoo/zoo.ipynb

@@ -27,7 +27,7 @@
     "from abl.data.evaluation import ReasoningMetric, SymbolAccuracy\n",
     "from abl.learning import ABLModel\n",
     "from abl.reasoning import Reasoner\n",
-    "from abl.utils import ABLLogger, confidence_dist, print_log, tab_data_to_tuple\n",
+    "from abl.utils import ABLLogger, avg_confidence_dist, print_log, tab_data_to_tuple\n",
     "\n",
     "from get_dataset import load_and_preprocess_dataset, split_dataset\n",
     "from kb import ZooKB"
@@ -235,7 +235,7 @@
    "source": [
     "def consitency(data_example, candidates, candidate_idxs, reasoning_results):\n",
     "    pred_prob = data_example.pred_prob\n",
-    "    model_scores = confidence_dist(pred_prob, candidate_idxs)\n",
+    "    model_scores = avg_confidence_dist(pred_prob, candidate_idxs)\n",
     "    rule_scores = np.array(reasoning_results)\n",
     "    scores = model_scores + rule_scores\n",
     "    return scores\n",