validmind 2.13.7__tar.gz → 2.13.9__tar.gz

{validmind-2.13.7 → validmind-2.13.9}/PKG-INFO

@@ -1,6 +1,6 @@
 Metadata-Version: 2.4
 Name: validmind
-Version: 2.13.7
+Version: 2.13.9
 Summary: ValidMind Library
 Author-email: Andres Rodriguez <andres@validmind.ai>, Juan Martinez <juan@validmind.ai>, Anil Sorathiya <anil@validmind.ai>, Luis Pallares <luis@validmind.ai>, John Walz <john@validmind.ai>
 License:                          DUAL LICENSE NOTICE

{validmind-2.13.7 → validmind-2.13.9}/pyproject.toml

@@ -1,6 +1,6 @@
 [project]
 name = "validmind"
-version = "2.13.7"
+version = "2.13.9"
 description = "ValidMind Library"
 readme = "README.pypi.md"
 requires-python = ">=3.9,<3.15"

validmind-2.13.9/validmind/__version__.py

@@ -0,0 +1 @@
+__version__ = "2.13.9"

{validmind-2.13.7 → validmind-2.13.9}/validmind/api_client.py

@@ -11,6 +11,7 @@ import asyncio
 import atexit
 import json
 import os
+import threading
 from io import BytesIO
 from typing import Any, Dict, List, Optional, Tuple, Union
 from urllib.parse import urlencode, urljoin
@@ -42,6 +43,10 @@ _monitoring = False
 _auth_mode = "api_key"
 _access_token: Optional[str] = None
 _oidc_login_context: Optional[Dict[str, str]] = None
+# Expiry (ISO-8601) of the in-memory OIDC access token, for cheap request-path
+# staleness checks; guarded together with the token by _oidc_refresh_lock.
+_oidc_expires_at: Optional[str] = None
+_oidc_refresh_lock = threading.Lock()
 
 __api_session: Optional[aiohttp.ClientSession] = None
 
@@ -147,12 +152,21 @@ def _get_url(
 async def _get(
     endpoint: str, params: Optional[Dict[str, str]] = None
 ) -> Dict[str, Any]:
+    _ensure_fresh_oidc_token()
     url = _get_url(endpoint, params)
     session = _get_session()
 
     async with session.get(url) as r:
+        # A 401 can still slip through if the token was revoked ahead of its
+        # stated expiry; force a refresh and retry once (safe — GET has no body).
+        if r.status == 401 and _ensure_fresh_oidc_token(force=True):
+            session = _get_session()
+            async with session.get(url) as retry:
+                if retry.status != 200:
+                    _raise_for_api_error(retry.status, await retry.text())
+                return await retry.json()
         if r.status != 200:
-            raise_api_error(await r.text())
+            _raise_for_api_error(r.status, await r.text())
 
         return await r.json()
 
@@ -163,6 +177,7 @@ async def _post(
     data: Optional[Union[dict, FormData]] = None,
     files: Optional[Dict[str, Tuple[str, BytesIO, str]]] = None,
 ) -> Dict[str, Any]:
+    _ensure_fresh_oidc_token()
     url = _get_url(endpoint, params)
     session = _get_session()
 
@@ -186,20 +201,29 @@ async def _post(
         _data = data
 
     async with session.post(url, data=_data) as r:
+        # Not retried on 401: the request body / upload stream may already be
+        # consumed. Proactive refresh above covers the expiry case; a 401 here
+        # means a genuinely rejected token, surfaced as a clear auth error.
         if r.status != 200:
-            raise_api_error(await r.text())
+            _raise_for_api_error(r.status, await r.text())
 
         return await r.json()
 
 
 def _ping() -> Dict[str, Any]:
     """Validates that we can connect to the ValidMind API (does not use the async session)."""
+    _ensure_fresh_oidc_token()
     r = requests.get(
         url=_get_url("ping"),
         headers=_get_api_headers(),
     )
+    if r.status_code == 401 and _ensure_fresh_oidc_token(force=True):
+        r = requests.get(
+            url=_get_url("ping"),
+            headers=_get_api_headers(),
+        )
     if r.status_code != 200:
-        raise_api_error(r.text)
+        _raise_for_api_error(r.status_code, r.text)
 
     client_info = r.json()
 
@@ -224,6 +248,36 @@ def _ping() -> Dict[str, Any]:
         )
 
 
+def _refresh_oidc_from_cache(
+    issuer: str,
+    client_id: str,
+    refresh_token: str,
+    scope: Optional[str],
+    audience: Optional[str] = None,
+) -> Optional[Dict[str, Any]]:
+    """Exchange a cached refresh token for new tokens, persisting the result.
+
+    Shared by the ``init()`` path (``_obtain_oidc_tokens``) and the request path
+    (``_ensure_fresh_oidc_token``) so their refresh behaviour can't drift. On
+    success the refreshed entry is cached and returned. On a failed refresh (e.g.
+    revoked token / ``invalid_grant``) the stale entry is deleted and None is
+    returned, so a dead refresh token isn't re-attempted; the caller decides how
+    to recover (``init()`` falls back to device flow, the request path stops).
+    """
+    from .credentials_store import delete_cached_entry, upsert_cached_entry
+    from .oidc_device import try_refresh_cached_tokens
+
+    try:
+        new_tokens = try_refresh_cached_tokens(
+            issuer, client_id, refresh_token, scope, audience=audience
+        )
+    except ValidMindAuthError:
+        delete_cached_entry(issuer, client_id, audience=audience)
+        return None
+    upsert_cached_entry(issuer, client_id, new_tokens, audience=audience)
+    return new_tokens
+
+
 def _obtain_oidc_tokens(
     issuer: str,
     client_id: str,
@@ -232,14 +286,13 @@ def _obtain_oidc_tokens(
 ) -> Dict[str, Any]:
     """Return a credentials entry dict with access_token, expires_at, refresh_token, etc."""
     from .credentials_store import (
-        delete_cached_entry,
         get_cached_entry,
         is_expired,
         normalize_client_id,
         normalize_issuer,
         upsert_cached_entry,
     )
-    from .oidc_device import run_device_flow, try_refresh_cached_tokens
+    from .oidc_device import run_device_flow
 
     norm_issuer = normalize_issuer(issuer)
     norm_client_id = normalize_client_id(client_id)
@@ -247,20 +300,11 @@ def _obtain_oidc_tokens(
     if cached and not is_expired(cached):
         return cached
     if cached and cached.get("refresh_token"):
-        try:
-            new_tokens = try_refresh_cached_tokens(
-                norm_issuer,
-                norm_client_id,
-                cached["refresh_token"],
-                scope,
-                audience=audience,
-            )
-            upsert_cached_entry(
-                norm_issuer, norm_client_id, new_tokens, audience=audience
-            )
+        new_tokens = _refresh_oidc_from_cache(
+            norm_issuer, norm_client_id, cached["refresh_token"], scope, audience
+        )
+        if new_tokens is not None:
             return new_tokens
-        except ValidMindAuthError:
-            delete_cached_entry(norm_issuer, norm_client_id, audience=audience)
     tokens = run_device_flow(norm_issuer, norm_client_id, scope, audience=audience)
     upsert_cached_entry(norm_issuer, norm_client_id, tokens, audience=audience)
     return tokens
@@ -276,6 +320,100 @@ def _select_oidc_bearer_token(entry: Dict[str, Any]) -> str:
     return entry["access_token"]
 
 
+def _set_oidc_access_token(entry: Dict[str, Any]) -> None:
+    """Adopt the bearer token from a credentials entry as the active token.
+
+    Records its expiry for cheap request-path staleness checks and drops the
+    pooled aiohttp session so the next request rebuilds it with the new token.
+    """
+    global _access_token, _oidc_expires_at
+    _access_token = _select_oidc_bearer_token(entry)
+    _oidc_expires_at = entry.get("expires_at")
+    _invalidate_async_session()
+
+
+def _clear_oidc_access_token() -> None:
+    """Drop the in-memory OIDC token and pooled session after an unrecoverable refresh.
+
+    Pairs with deleting the cached entry: with no usable token in memory, the next
+    request fails fast at header build with the clear "run vm.init()" error instead
+    of sending a doomed request just to get a 401 back.
+    """
+    global _access_token, _oidc_expires_at
+    _access_token = None
+    _oidc_expires_at = None
+    _invalidate_async_session()
+
+
+def _oidc_token_is_stale() -> bool:
+    """Cheap in-memory expiry check (no file I/O) using the 120s skew in is_expired."""
+    from .credentials_store import is_expired
+
+    return is_expired({"expires_at": _oidc_expires_at})
+
+
+def _ensure_fresh_oidc_token(force: bool = False) -> bool:
+    """Refresh the OIDC access token on the request path when it has expired.
+
+    The token captured at ``init()`` is otherwise reused unchanged for the life of
+    the process, so a session that outlives the provider's access-token lifetime
+    starts failing on every call until ``init()`` is re-run. This uses the cached
+    refresh token (requires ``offline_access``, the default scope) to obtain a new
+    access token in place. Returns True when a usable token is in place, False when
+    no refresh was possible (non-OIDC mode, no cached refresh token, or a failed
+    refresh — in which case the cached entry is dropped, matching init()).
+    """
+    if _auth_mode != "oidc" or _oidc_login_context is None:
+        return False
+    # Hot path: skip file I/O and the lock while the current token is still valid.
+    if not force and not _oidc_token_is_stale():
+        return True
+
+    from .credentials_store import get_cached_entry, is_expired
+
+    ctx = _oidc_login_context
+    issuer = ctx["issuer"]
+    client_id = ctx["client_id"]
+    scope = ctx.get("scope")
+    audience = ctx.get("audience") or None
+
+    with _oidc_refresh_lock:
+        entry = get_cached_entry(issuer, client_id, audience=audience)
+        if entry is None:
+            return False
+        # Another caller may have refreshed the token while we waited on the lock.
+        if not force and not is_expired(entry):
+            _set_oidc_access_token(entry)
+            return True
+        refresh_token = entry.get("refresh_token")
+        if not refresh_token:
+            return False
+        # Shared with init(): on failure the cached entry is dropped so a
+        # bad/revoked refresh token isn't re-attempted on every request. Recovery
+        # is re-running vm.init() (surfaced by the clear 401 auth error); unlike
+        # init() the request path can't fall back to interactive device flow.
+        new_tokens = _refresh_oidc_from_cache(
+            issuer, client_id, refresh_token, scope, audience
+        )
+        if new_tokens is None:
+            # Hard failure: the cache entry was dropped; clear the in-memory token
+            # too so the next request fails fast with the clear re-auth message.
+            _clear_oidc_access_token()
+            return False
+        _set_oidc_access_token(new_tokens)
+        return True
+
+
+def _raise_for_api_error(status: int, text: str) -> None:
+    """Raise a clear, actionable auth error for an OIDC 401, else defer to raise_api_error."""
+    if status == 401 and _auth_mode == "oidc":
+        raise ValidMindAuthError(
+            "ValidMind rejected the OAuth access token (HTTP 401); it may have "
+            "expired or been revoked. Run vm.init() again to re-authenticate."
+        )
+    raise_api_error(text)
+
+
 def init(
     api_key: Optional[str] = None,
     api_secret: Optional[str] = None,
@@ -317,7 +455,7 @@ def init(
             Can be set via env ``VM_OIDC_ISSUER``.
         client_id (str, optional): OAuth public client id for device flow. Can be
             set via env ``VM_OIDC_CLIENT_ID``.
-        scope (str, optional): OAuth scopes (default ``openid profile email``).
+        scope (str, optional): OAuth scopes (default ``openid profile email offline_access``).
             Can be set via env ``VM_OIDC_SCOPE``.
         audience (str, optional): Resource / API identifier for the access token
             (e.g. Auth0 API Identifier). Use the same value the ValidMind backend
@@ -330,7 +468,7 @@ def init(
         ValidMindAuthError: If OIDC configuration conflicts or login fails.
     """
     global _api_key, _api_secret, _api_host, _model_cuid, _monitoring, _document
-    global _auth_mode, _access_token, _oidc_login_context
+    global _auth_mode, _access_token, _oidc_login_context, _oidc_expires_at
 
     if api_key == "...":
         # special case to detect when running a notebook placeholder (...)
@@ -392,7 +530,7 @@ def init(
         _api_key = None
         _api_secret = None
         _api_host = resolved_host
-        scope_val = oidc_scope or "openid profile email"
+        scope_val = oidc_scope or "openid profile email offline_access"
         from .credentials_store import normalize_audience
 
         oidc_audience_val = normalize_audience(
@@ -402,20 +540,21 @@ def init(
         entry = _obtain_oidc_tokens(
             oidc_issuer, oidc_client_id, scope_val, audience=oidc_audience_opt
         )
-        _access_token = _select_oidc_bearer_token(entry)
         _oidc_login_context = {
             "issuer": entry["issuer"],
             "client_id": entry["client_id"],
             "scope": scope_val,
             "audience": entry.get("audience") or oidc_audience_val,
         }
-        _invalidate_async_session()
+        # Sets _access_token / _oidc_expires_at and invalidates the async session.
+        _set_oidc_access_token(entry)
     else:
         if env_key is None or env_secret is None:
             raise MissingAPICredentialsError()
         _auth_mode = "api_key"
         _access_token = None
         _oidc_login_context = None
+        _oidc_expires_at = None
         _api_key = env_key
         _api_secret = env_secret
         _api_host = resolved_host
@@ -656,6 +795,7 @@ def _validate_log_text_context(
 
 def generate_qualitative_text(text_generation_data: Dict[str, Any]) -> Dict[str, Any]:
     """Generate qualitative text using the ValidMind AI API."""
+    _ensure_fresh_oidc_token()
     r = requests.post(
         url=_get_url("ai/generate/qualitative_text_generation"),
         headers=_get_api_headers(),
@@ -663,7 +803,7 @@ def generate_qualitative_text(text_generation_data: Dict[str, Any]) -> Dict[str,
     )
 
     if r.status_code != 200:
-        raise_api_error(r.text)
+        _raise_for_api_error(r.status_code, r.text)
 
     return r.json()
 
@@ -901,6 +1041,7 @@ def log_metric(
 
 
 def generate_test_result_description(test_result_data: Dict[str, Any]) -> str:
+    _ensure_fresh_oidc_token()
     r = requests.post(
         url=_get_url("ai/generate/test_result_description"),
         headers=_get_api_headers(),
@@ -908,6 +1049,6 @@ def generate_test_result_description(test_result_data: Dict[str, Any]) -> str:
     )
 
     if r.status_code != 200:
-        raise_api_error(r.text)
+        _raise_for_api_error(r.status_code, r.text)
 
     return r.json()

{validmind-2.13.7 → validmind-2.13.9}/validmind/tests/model_validation/sklearn/CalibrationCurve.py

@@ -2,6 +2,7 @@
 # Refer to the LICENSE file in the root of this repository for details.
 # SPDX-License-Identifier: AGPL-3.0 AND ValidMind Commercial
 
+import inspect
 from typing import Tuple
 
 import numpy as np
@@ -80,8 +81,19 @@ def CalibrationCurve(
             "Calibration Curve is only supported for binary classification models"
         )
 
+    # A two-class target encoded outside {0, 1} (e.g. {0, 4} or string labels)
+    # passes the guard above but hits sklearn's "pos_label is not specified" error.
+    # Pass the largest label as the positive class, matching the "largest label is
+    # positive" convention in _diagnosis_metrics.resolve_averaging; for {0, 1} this
+    # is byte-identical to sklearn's default. pos_label was added to
+    # calibration_curve in scikit-learn 1.1, so guard by signature inspection to
+    # keep pre-1.1 installs working (pyproject sets no scikit-learn floor).
+    kwargs = {}
+    if "pos_label" in inspect.signature(calibration_curve).parameters:
+        kwargs["pos_label"] = np.unique(dataset.y).max()
+
     prob_true, prob_pred = calibration_curve(
-        dataset.y, dataset.y_prob(model), n_bins=n_bins
+        dataset.y, dataset.y_prob(model), n_bins=n_bins, **kwargs
     )
 
     # Create DataFrame for raw data

{validmind-2.13.7 → validmind-2.13.9}/validmind/tests/model_validation/sklearn/PopulationStabilityIndex.py

@@ -14,6 +14,8 @@ from validmind.errors import SkipTestError
 from validmind.logging import get_logger
 from validmind.vm_models import VMDataset, VMModel
 
+from ._multiclass_proba import multiclass_proba, proba_matrix
+
 logger = get_logger(__name__)
 
 _PSI_PALETTE = ["#DE257E", "#1F77B4", "#2CA02C", "#FF7F0E", "#9467BD", "#8C564B"]
@@ -104,54 +106,37 @@ def _psi_table_rows(psi_results):
     return rows_with_total, table_rows
 
 
-def _multiclass_psi(datasets, model, classes, num_bins, mode):
+def _multiclass_psi(datasets, model, num_bins, mode):
     """One-vs-rest PSI for a multiclass model.
 
     PSI needs a 1-D score distribution to compare across the two datasets. The
-    stored single probability column cannot represent every class, so we ask the
-    underlying estimator for the full per-class probability matrix (mirroring the
-    ROC/PR curve tests). Models without a usable ``predict_proba`` (metadata-only
-    / precomputed single-column predictions) are skipped rather than crashed.
+    stored single probability column cannot represent every class, so the shared
+    helper reaches the underlying estimator for the full per-class probability
+    matrix (mirroring the ROC/PR curve tests), aligning columns to the training
+    class order. The initial dataset drives the class alignment and the set of
+    classes actually present; the new dataset's matrix is validated against the
+    same class list. Models that cannot supply a matching matrix are skipped.
     """
-    raw_model = getattr(model, "model", None)
-    proba_fn = getattr(raw_model, "predict_proba", None)
-    if not callable(proba_fn):
-        raise SkipTestError(
-            "Multiclass Population Stability Index requires per-class "
-            "probabilities from the underlying model's predict_proba, which is "
-            "not available for this model (e.g. metadata-only / precomputed "
-            "predictions). Skipping."
-        )
-    try:
-        prob_initial = np.asarray(proba_fn(datasets[0].x_df()))
-        prob_new = np.asarray(proba_fn(datasets[1].x_df()))
-    except Exception as e:
-        raise SkipTestError(
-            "Multiclass Population Stability Index could not compute per-class "
-            f"probabilities ({type(e).__name__}). Skipping."
-        ) from e
-
-    n_classes = len(classes)
-    for prob in (prob_initial, prob_new):
-        if prob.ndim != 2 or prob.shape[1] != n_classes:
-            raise SkipTestError(
-                "Multiclass Population Stability Index requires a per-class "
-                f"probability matrix with one column per class (got shape "
-                f"{getattr(prob, 'shape', None)} for {n_classes} classes). Skipping."
-            )
-
-    # predict_proba columns are ordered by sorted class label == np.unique.
+    aligned = multiclass_proba(model, datasets[0], "Population Stability Index")
+    prob_initial = aligned.y_prob
+    prob_new = proba_matrix(
+        model, datasets[1], aligned.class_list, "Population Stability Index"
+    )
+
+    classes_present = aligned.classes_present
+    n_present = len(classes_present)
+
     fig = make_subplots(
-        rows=n_classes,
+        rows=n_present,
         cols=1,
-        specs=[[{"secondary_y": True}] for _ in range(n_classes)],
-        subplot_titles=[f"Class {cls}" for cls in classes],
+        specs=[[{"secondary_y": True}] for _ in range(n_present)],
+        subplot_titles=[f"Class {cls}" for cls in classes_present],
         vertical_spacing=0.08,
     )
 
     tables = {}
     raw_per_class = {}
-    for i, cls in enumerate(classes):
+    for plot_i, (i, cls) in enumerate(zip(aligned.present_indices, classes_present)):
         psi_results = calculate_psi(
             prob_initial[:, i].copy(),
             prob_new[:, i].copy(),
@@ -159,17 +144,17 @@ def _multiclass_psi(datasets, model, classes, num_bins, mode):
             mode=mode,
         )
         x = list(range(len(psi_results)))
-        color = _PSI_PALETTE[i % len(_PSI_PALETTE)]
+        color = _PSI_PALETTE[plot_i % len(_PSI_PALETTE)]
         fig.add_trace(
             go.Bar(
                 x=x,
                 y=[d["percent_initial"] for d in psi_results],
                 name="Initial",
                 marker=dict(color="#DE257E"),
-                showlegend=i == 0,
+                showlegend=plot_i == 0,
                 legendgroup="initial",
             ),
-            row=i + 1,
+            row=plot_i + 1,
             col=1,
             secondary_y=False,
         )
@@ -179,10 +164,10 @@ def _multiclass_psi(datasets, model, classes, num_bins, mode):
                 y=[d["percent_new"] for d in psi_results],
                 name="New",
                 marker=dict(color="#E8B1F8"),
-                showlegend=i == 0,
+                showlegend=plot_i == 0,
                 legendgroup="new",
             ),
-            row=i + 1,
+            row=plot_i + 1,
             col=1,
             secondary_y=False,
         )
@@ -192,10 +177,10 @@ def _multiclass_psi(datasets, model, classes, num_bins, mode):
                 y=[d["psi"] for d in psi_results],
                 name="PSI",
                 line=dict(color=color),
-                showlegend=i == 0,
+                showlegend=plot_i == 0,
                 legendgroup="psi",
             ),
-            row=i + 1,
+            row=plot_i + 1,
             col=1,
             secondary_y=True,
         )
@@ -211,7 +196,7 @@ def _multiclass_psi(datasets, model, classes, num_bins, mode):
     fig.update_layout(
         title="Population Stability Index (PSI) — one-vs-rest per class",
         barmode="group",
-        height=300 * n_classes,
+        height=300 * n_present,
     )
 
     return (
@@ -291,7 +276,7 @@ def PopulationStabilityIndex(
 
     classes = np.unique(datasets[0].y)
     if len(classes) > 2:
-        return _multiclass_psi(datasets, model, classes, num_bins, mode)
+        return _multiclass_psi(datasets, model, num_bins, mode)
 
     psi_results = calculate_psi(
         datasets[0].y_prob(model).copy(),

{validmind-2.13.7 → validmind-2.13.9}/validmind/tests/model_validation/sklearn/PrecisionRecallCurve.py

@@ -7,13 +7,14 @@ from typing import Tuple
 import numpy as np
 import plotly.graph_objects as go
 from sklearn.metrics import average_precision_score, precision_recall_curve
-from sklearn.preprocessing import label_binarize
 
 from validmind import RawData, tags, tasks
 from validmind.errors import SkipTestError
 from validmind.models import FoundationModel
 from validmind.vm_models import VMDataset, VMModel
 
+from ._multiclass_proba import multiclass_proba
+
 
 @tags(
     "sklearn",
@@ -75,7 +76,7 @@ def PrecisionRecallCurve(
     classes = np.unique(y_true)
 
     if len(classes) > 2:
-        return _multiclass_pr_curve(model, dataset, classes)
+        return _multiclass_pr_curve(model, dataset)
 
     precision, recall, _ = precision_recall_curve(y_true, dataset.y_prob(model))
 
@@ -105,53 +106,27 @@ def PrecisionRecallCurve(
 
 
 def _multiclass_pr_curve(
-    model: VMModel, dataset: VMDataset, classes: np.ndarray
+    model: VMModel, dataset: VMDataset
 ) -> Tuple[go.Figure, RawData]:
     """One-vs-rest precision-recall curves for a multiclass model.
 
     Needs the full per-class probability matrix, which the stored single
-    probability column cannot provide, so we ask the model for it directly.
-    Models without a usable ``predict_proba`` (Foundation/metadata-only,
-    precomputed single-column probabilities) are skipped rather than crashed.
+    probability column cannot provide; the shared helper reaches the underlying
+    estimator, aligns the probability columns to the training class order and
+    skips models that cannot supply a matching matrix.
     """
-    # The VMModel wrapper's predict_proba is binary-only (it returns just the
-    # positive-class column), so reach the underlying estimator for the full
-    # per-class probability matrix.
-    raw_model = getattr(model, "model", None)
-    proba_fn = getattr(raw_model, "predict_proba", None)
-    if not callable(proba_fn):
-        raise SkipTestError(
-            "Multiclass Precision-Recall Curve requires per-class probabilities "
-            "from the underlying model's predict_proba, which is not available "
-            "for this model (e.g. Foundation / metadata-only / precomputed "
-            "predictions). Skipping."
-        )
-    try:
-        y_prob = np.asarray(proba_fn(dataset.x_df()))
-    except Exception as e:
-        raise SkipTestError(
-            "Multiclass Precision-Recall Curve could not compute per-class "
-            f"probabilities ({type(e).__name__}). Skipping."
-        ) from e
-
-    n_classes = len(classes)
-    if y_prob.ndim != 2 or y_prob.shape[1] != n_classes:
-        raise SkipTestError(
-            "Multiclass Precision-Recall Curve requires a per-class probability "
-            f"matrix with one column per class (got shape "
-            f"{getattr(y_prob, 'shape', None)} for {n_classes} classes). Skipping."
-        )
-
-    # One-hot the true labels in the same class order predict_proba columns use
-    # (sklearn orders predict_proba columns by sorted class label == np.unique).
-    y_bin = label_binarize(dataset.y.flatten(), classes=classes)
+    aligned = multiclass_proba(model, dataset, "Precision-Recall Curve")
+    y_bin = aligned.y_bin
+    y_prob = aligned.y_prob
 
     traces = []
     raw_precision = {}
     raw_recall = {}
     raw_ap = {}
     palette = ["#DE257E", "#1F77B4", "#2CA02C", "#FF7F0E", "#9467BD", "#8C564B"]
-    for i, cls in enumerate(classes):
+    for plot_i, (i, cls) in enumerate(
+        zip(aligned.present_indices, aligned.classes_present)
+    ):
         precision, recall, _ = precision_recall_curve(y_bin[:, i], y_prob[:, i])
         ap = average_precision_score(y_bin[:, i], y_prob[:, i])
         key = str(cls)
@@ -164,11 +139,14 @@ def _multiclass_pr_curve(
                 y=precision,
                 mode="lines",
                 name=f"Class {key} (AP = {ap:.2f})",
-                line=dict(color=palette[i % len(palette)]),
+                line=dict(color=palette[plot_i % len(palette)]),
             )
         )
 
-    # Micro-average across all one-vs-rest decisions.
+    # Micro-average across the one-vs-rest decisions of the present classes.
+    present = aligned.present_indices
+    y_bin = y_bin[:, present]
+    y_prob = y_prob[:, present]
     micro_precision, micro_recall, _ = precision_recall_curve(
         y_bin.ravel(), y_prob.ravel()
     )