fix: map xgboost Booster labels (#120)

* fix: map xgboost Booster labels

* fix: validate class label probability shape

---------

Co-authored-by: Nanook <nanookclaw@users.noreply.github.com>
Co-authored-by: Shahid Ul Islam <65996596+Khanz9664@users.noreply.github.com>

Author: nanookclaw

tests/test_backend_xgboost.py

@@ -51,6 +51,43 @@ def test_xgboost_integration_analyze():
     assert "xgboost" in report.metadata["backend"]["resolver"]
 
 
+def test_raw_xgboost_booster_maps_ordinal_predictions_to_string_labels():
+    rng = np.random.default_rng(42)
+    class_labels = np.array(["mouse", "cat", "dog"])
+    y_ordinal = np.repeat(np.arange(3), 30)
+    X = np.column_stack(
+        [
+            y_ordinal,
+            y_ordinal == 1,
+            y_ordinal == 2,
+            rng.normal(scale=0.01, size=len(y_ordinal)),
+        ]
+    ).astype(float)
+    y_true = class_labels[y_ordinal]
+
+    dtrain = xgboost.DMatrix(X, label=y_ordinal)
+    model = xgboost.train(
+        {
+            "objective": "multi:softprob",
+            "num_class": 3,
+            "max_depth": 2,
+            "eta": 0.5,
+            "verbosity": 0,
+        },
+        dtrain,
+        num_boost_round=10,
+    )
+
+    report = analyze(model, X, y_true, class_labels=class_labels, verbose=False)
+
+    assert report.metadata["framework"] == "xgboost"
+    assert report.y_prob.shape == (90, 3)
+    assert set(np.unique(report.y_pred)).issubset(set(class_labels))
+    assert report.results["failure"]["misclassification_summary"]["__overall__"][
+        "total_errors"
+    ] == int(np.sum(y_true != report.y_pred))
+
+
 def test_xgboost_manual_override():
     from sklearn.datasets import make_classification
     from xgboost import XGBClassifier

tests/test_core_pipeline.py

@@ -0,0 +1,14 @@
+import numpy as np
+import pytest
+
+from trustlens.core.pipeline import _encode_labels_for_probability_columns
+
+
+def test_encode_labels_rejects_class_label_length_mismatch():
+    y_true = np.array(["mouse", "cat"])
+    class_labels = np.array(["mouse", "cat"])
+
+    with pytest.raises(
+        ValueError, match=r"class_labels length \(2\).*probability column shape \(3 columns\)"
+    ):
+        _encode_labels_for_probability_columns(y_true, 3, class_labels)

trustlens/api.py

@@ -82,6 +82,7 @@ def analyze(
     sensitive_features: Optional[dict[str, np.ndarray]] = None,
     modules: Optional[list[str]] = None,
     plugins: Optional[list[str]] = None,
+    class_labels: Optional[np.ndarray] = None,
     verbose: bool = True,
 ) -> TrustReport:
     """
@@ -113,6 +114,10 @@ def analyze(
       Subset of analysis modules to run.
     plugins : list[str], optional
       Names of registered plugins to activate.
+    class_labels : np.ndarray, optional
+      Semantic class labels in the order corresponding to probability columns.
+      Useful for raw backends such as ``xgboost.Booster`` that return ordinal
+      probability columns without a ``classes_`` attribute.
     verbose : bool
       Print progress updates. Default True.
 
@@ -163,7 +168,14 @@ def analyze(
         framework = "manual"
 
     resolver = get_resolver(model, framework=framework)
-    bundle = resolver(model, X, y_pred=y_pred, y_prob=y_prob)
+    resolved_class_labels = np.asarray(class_labels) if class_labels is not None else None
+    bundle = resolver(
+        model,
+        X,
+        y_pred=y_pred,
+        y_prob=y_prob,
+        class_labels=resolved_class_labels,
+    )
 
     # ------------------------------------------------------------------
     # 2. Delegate to Core Pipeline
@@ -176,6 +188,7 @@ def analyze(
         y_prob=bundle.y_prob,
         framework=bundle.framework,
         backend_metadata=bundle.metadata,
+        class_labels=bundle.class_labels,
         embeddings=embeddings,
         sensitive_features=sensitive_features,
         modules=modules,

trustlens/backends/registry.py

@@ -40,14 +40,24 @@ def manual_resolve(
     X: np.ndarray,
     y_pred: Optional[np.ndarray] = None,
     y_prob: Optional[np.ndarray] = None,
+    class_labels: Optional[np.ndarray] = None,
 ) -> PredictionBundle:
     """
     Passthrough resolver for manual overrides.
     """
     if y_pred is None:
         if y_prob is not None:
             # Derive y_pred from y_prob if missing
-            y_pred = np.argmax(np.asarray(y_prob), axis=1)
+            y_prob_arr = np.asarray(y_prob)
+            y_pred_indices = np.argmax(y_prob_arr, axis=1)
+            if (
+                class_labels is not None
+                and y_prob_arr.ndim == 2
+                and len(class_labels) == y_prob_arr.shape[1]
+            ):
+                y_pred = np.asarray(class_labels)[y_pred_indices]
+            else:
+                y_pred = y_pred_indices
         else:
             raise ValueError("Manual override requires either y_pred or y_prob.")
 
@@ -61,6 +71,7 @@ def manual_resolve(
         y_pred=np.asarray(y_pred),
         y_prob=np.asarray(y_prob) if y_prob is not None else None,
         framework="manual",
+        class_labels=np.asarray(class_labels) if class_labels is not None else None,
         metadata=metadata,
     )
 

trustlens/backends/sklearn.py

@@ -18,6 +18,7 @@ def resolve(
     X: np.ndarray,
     y_pred: Optional[np.ndarray] = None,
     y_prob: Optional[np.ndarray] = None,
+    class_labels: Optional[np.ndarray] = None,
 ) -> PredictionBundle:
     """
     Resolve predictions and probabilities from a scikit-learn model.
@@ -72,6 +73,8 @@ def resolve(
 
             # Defensive label mapping
             classes = getattr(model, "classes_", None)
+            if classes is None:
+                classes = class_labels
             if classes is not None:
                 classes_arr = np.asarray(classes)
                 if len(classes_arr) == y_prob.shape[1]:
@@ -89,9 +92,13 @@ def resolve(
             )
 
     # 5. Metadata & Labels
-    class_labels = getattr(model, "classes_", None)
-    if class_labels is not None:
-        class_labels = np.asarray(class_labels)
+    model_class_labels = getattr(model, "classes_", None)
+    if model_class_labels is not None:
+        resolved_class_labels = np.asarray(model_class_labels)
+    elif class_labels is not None:
+        resolved_class_labels = np.asarray(class_labels)
+    else:
+        resolved_class_labels = None
 
     metadata = {
         "resolver": "sklearn",
@@ -103,6 +110,6 @@ def resolve(
         y_pred=y_pred,
         y_prob=y_prob,
         framework="sklearn",
-        class_labels=class_labels,
+        class_labels=resolved_class_labels,
         metadata=metadata,
     )

trustlens/backends/types.py

@@ -42,7 +42,7 @@ def __post_init__(self) -> None:
                 "Consult docs/internal/prediction_contract.md for details."
             )
 
-        if np.any(~np.isfinite(self.y_pred)):
+        if np.issubdtype(self.y_pred.dtype, np.number) and np.any(~np.isfinite(self.y_pred)):
             raise ValueError("y_pred contains non-finite values (NaN or Inf).")
 
         if self.y_prob is not None:

trustlens/backends/xgboost.py

@@ -15,6 +15,7 @@ def resolve(
     X: np.ndarray,
     y_pred: Optional[np.ndarray] = None,
     y_prob: Optional[np.ndarray] = None,
+    class_labels: Optional[np.ndarray] = None,
 ) -> PredictionBundle:
     """
     Prediction resolver for XGBoost models.
@@ -85,6 +86,14 @@ def resolve(
         y_prob_flat = y_prob.flatten()
         y_prob = np.column_stack([1 - y_prob_flat, y_prob_flat])
 
+    model_class_labels = getattr(model, "classes_", None)
+    if model_class_labels is not None:
+        resolved_class_labels = np.asarray(model_class_labels)
+    elif class_labels is not None:
+        resolved_class_labels = np.asarray(class_labels)
+    else:
+        resolved_class_labels = None
+
     # 4. Resolve Class Predictions
     if y_pred is None:
         if hasattr(model, "predict") and not isinstance(model, xgb.Booster):
@@ -93,10 +102,9 @@ def resolve(
         else:
             # Fallback to argmax from probabilities
             y_pred_indices = np.argmax(y_prob, axis=1)
-            if hasattr(model, "classes_"):
-                classes = np.asarray(model.classes_)
-                if len(classes) == y_prob.shape[1]:
-                    y_pred = classes[y_pred_indices]
+            if resolved_class_labels is not None:
+                if len(resolved_class_labels) == y_prob.shape[1]:
+                    y_pred = resolved_class_labels[y_pred_indices]
                 else:
                     y_pred = y_pred_indices
             else:
@@ -114,5 +122,6 @@ def resolve(
         y_pred=np.asarray(y_pred),
         y_prob=y_prob,
         framework="xgboost",
+        class_labels=resolved_class_labels,
         metadata=metadata,
     )

trustlens/core/pipeline.py

@@ -35,6 +35,41 @@
 logger = logging.getLogger(__name__)
 
 
+def _as_python_label(label: Any) -> Any:
+    """Return a hashable Python scalar for NumPy scalar labels."""
+    return label.item() if hasattr(label, "item") else label
+
+
+def _encode_labels_for_probability_columns(
+    y_true: np.ndarray,
+    n_classes: int,
+    class_labels: Optional[np.ndarray],
+) -> np.ndarray:
+    """Encode semantic labels into the class-index order used by y_prob columns."""
+    if class_labels is not None:
+        class_labels_array = np.asarray(class_labels)
+        if len(class_labels_array) != n_classes:
+            raise ValueError(
+                "class_labels length "
+                f"({len(class_labels_array)}) does not match probability column shape "
+                f"({n_classes} columns)."
+            )
+
+        label_to_index = {
+            _as_python_label(label): idx for idx, label in enumerate(class_labels_array)
+        }
+        try:
+            encoded_labels: np.ndarray = np.asarray(
+                [_as_python_label(label_to_index[_as_python_label(label)]) for label in y_true],
+                dtype=int,
+            )
+            return encoded_labels
+        except KeyError as exc:
+            raise ValueError("y_true contains labels that are missing from class_labels.") from exc
+
+    return y_true.astype(int)
+
+
 def _run_analysis_pipeline(
     model: Any,
     X: np.ndarray,
@@ -47,6 +82,7 @@ def _run_analysis_pipeline(
     plugins: Optional[list[str]] = None,
     framework: Optional[str] = None,
     backend_metadata: Optional[dict[str, Any]] = None,
+    class_labels: Optional[np.ndarray] = None,
     verbose: bool = True,
 ) -> TrustReport:
     """
@@ -98,7 +134,10 @@ def _run_analysis_pipeline(
 
                 # Multiclass Brier Score: 1/N * sum(sum((p_ic - o_ic)^2))
                 # We can compute this efficiently
-                y_true_one_hot = np.eye(n_classes)[y_true.astype(int)]
+                y_true_indices = _encode_labels_for_probability_columns(
+                    y_true, n_classes, class_labels
+                )
+                y_true_one_hot = np.eye(n_classes)[y_true_indices]
                 mbrier = np.mean(np.sum((y_prob - y_true_one_hot) ** 2, axis=1))
 
                 results["calibration"] = {

trustlens/metrics/bias.py

@@ -28,6 +28,10 @@
 ]
 
 
+def _as_python_label(label):
+    return label.item() if hasattr(label, "item") else label
+
+
 def class_imbalance_report(y_true: np.ndarray) -> dict:
     """
     Summarize the class distribution in ``y_true``.
@@ -58,13 +62,15 @@ def class_imbalance_report(y_true: np.ndarray) -> dict:
     classes, counts = np.unique(y_true, return_counts=True)
     n = len(y_true)
 
-    class_counts = {int(cls): int(cnt) for cls, cnt in zip(classes, counts)}
-    class_frequencies = {int(cls): round(float(cnt / n), 4) for cls, cnt in zip(classes, counts)}
+    class_counts = {_as_python_label(cls): int(cnt) for cls, cnt in zip(classes, counts)}
+    class_frequencies = {
+        _as_python_label(cls): round(float(cnt / n), 4) for cls, cnt in zip(classes, counts)
+    }
 
     min_count = int(counts.min())
     max_count = int(counts.max())
-    minority_class = int(classes[counts.argmin()])
-    majority_class = int(classes[counts.argmax()])
+    minority_class = _as_python_label(classes[counts.argmin()])
+    majority_class = _as_python_label(classes[counts.argmax()])
     imbalance_ratio = round(max_count / min_count, 4) if min_count > 0 else float("inf")
 
     return {

trustlens/metrics/failure.py

@@ -21,6 +21,10 @@
 ]
 
 
+def _as_python_label(label):
+    return label.item() if hasattr(label, "item") else label
+
+
 def misclassification_summary(
     y_true: np.ndarray,
     y_pred: np.ndarray,
@@ -71,7 +75,7 @@ def misclassification_summary(
 
     summary: dict = {}
     for cls in classes:
-        cls_mask = y_true == int(cls)
+        cls_mask = y_true == cls
         cls_incorrect = cls_mask & incorrect_mask
 
         n_support = int(cls_mask.sum())
@@ -92,7 +96,7 @@ def misclassification_summary(
         else:
             top_mistake_indices = []
 
-        summary[int(cls)] = {
+        summary[_as_python_label(cls)] = {
             "support": n_support,
             "n_misclassified": n_misclassified,
             "error_rate": round(error_rate, 4),