nervecode 0.1.0__py3-none-any.whl

nervecode/scoring/calibrator.py

@@ -0,0 +1,396 @@
+"""Empirical percentile calibrator and state.
+
+This module provides a minimal, torch-agnostic-at-import calibrator that
+collects an empirical surprise score distribution on held-out
+in-distribution data. The calibrator stores the sorted distribution, the
+chosen threshold quantiles, and lightweight metadata sufficient to reproduce
+calibration later (e.g., index rule, sample count, dtype).
+
+Threshold comparison, percentile lookup, and persistence are implemented in
+subsequent tasks; the goal here is to make calibration runs produce a durable
+state snapshot that downstream code can consume.
+"""
+
+from __future__ import annotations
+
+from collections.abc import Iterable, Mapping, Sequence
+from dataclasses import dataclass
+from typing import Any, cast
+
+try:  # Keep import-time behavior tolerant in environments without torch
+    import torch
+except Exception:  # pragma: no cover - torch is a project dependency in tests
+    torch = cast(Any, None)
+
+__all__ = [
+    "CalibrationState",
+    "EmpiricalPercentileCalibrator",
+]
+
+
+@dataclass(frozen=True)
+class CalibrationState:
+    """Stored result of a calibration run.
+
+    Fields
+    - ``sorted_surprise``: 1D tensor of per-sample surprise scores sorted in
+      ascending order. Typically a ``torch.Tensor`` on CPU.
+    - ``threshold_quantiles``: tuple of quantiles in ``[0, 1]`` that were
+      used to choose thresholds.
+    - ``threshold_values``: 1D tensor of threshold values aligned with
+      ``threshold_quantiles`` using a deterministic index rule (see metadata).
+    - ``metadata``: free-form details needed to reproduce calibration later
+      (e.g., index rule, number of samples, dtype, aggregation hint).
+    """
+
+    sorted_surprise: Any
+    threshold_quantiles: tuple[float, ...]
+    threshold_values: Any
+    metadata: dict[str, Any]
+
+
+# Attempt to allowlist CalibrationState for PyTorch 2.6+ safe loading
+try:  # pragma: no cover - environment-dependent
+    if torch is not None:
+        from torch import serialization as _serialization
+
+        _add = getattr(_serialization, "add_safe_globals", None)
+        if callable(_add):
+            _add([CalibrationState])
+except Exception:
+    # Older torch versions or environments without the API: skip registration.
+    pass
+
+
+class EmpiricalPercentileCalibrator:
+    """Collect an empirical surprise distribution and store threshold info.
+
+    The calibrator accepts a per-sample surprise vector, sorts it, and computes
+    threshold values at configured quantiles using a deterministic index rule
+    based on ``floor(q * (n - 1))``. The full sorted distribution, the chosen
+    quantiles, and metadata sufficient for reproduction are stored in a
+    ``CalibrationState``.
+    """
+
+    def __init__(self, *, threshold_quantiles: Sequence[float] = (0.95,)) -> None:
+        self._threshold_quantiles: tuple[float, ...] = tuple(float(q) for q in threshold_quantiles)
+        self._state: CalibrationState | None = None
+
+    @property
+    def state(self) -> CalibrationState | None:
+        """Return the last computed calibration state, if available."""
+
+        return self._state
+
+    @property
+    def threshold_quantiles(self) -> tuple[float, ...]:
+        return self._threshold_quantiles
+
+    def fit(self, scores: Any, *, aggregation: str | None = None) -> CalibrationState:
+        """Calibrate from a per-sample surprise vector and store state.
+
+        Parameters
+        - scores: Per-sample surprise signal. Accepts a ``torch.Tensor`` of any
+          shape (flattened to 1D), an iterable of numbers, or an object with a
+          ``surprise`` tensor attribute.
+        - aggregation: Optional hint describing how the surprise vector was
+          produced (e.g., ``"mean"``). Stored in metadata for later reference.
+
+        Returns
+        - ``CalibrationState`` with sorted distribution, chosen quantiles,
+          threshold values, and metadata.
+        """
+
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("Calibration requires PyTorch to be installed")
+
+        vec = _as_1d_tensor(scores)
+        if vec.numel() == 0:
+            raise ValueError("Calibration requires at least one surprise score")
+
+        # Defensive copy to CPU float for stable sorting and storage.
+        vec = vec.detach().to(device="cpu", dtype=torch.float32).contiguous().view(-1)
+        sorted_vec, _ = torch.sort(vec, stable=True)
+
+        # Validate and normalize quantiles
+        qs = list(self._threshold_quantiles)
+        if not qs:
+            qs = [0.95]
+        for q in qs:
+            if not (0.0 <= q <= 1.0):  # strict validation to avoid silent misuse
+                raise ValueError(f"Quantile {q} is outside [0, 1]")
+
+        # Deterministic index rule: floor(q * (n-1)).
+        n = int(sorted_vec.numel())
+        idxs = [min(n - 1, max(0, int((n - 1) * q))) for q in qs]
+        thr = sorted_vec[torch.tensor(idxs, dtype=torch.long)]
+
+        meta: dict[str, Any] = {
+            "method": "empirical_percentile",
+            "index_rule": "floor_q_times_n_minus_1",
+            "num_scores": n,
+            "aggregation": aggregation or "unknown",
+            "dtype": str(sorted_vec.dtype).replace("torch.", ""),
+            "device": str(sorted_vec.device),
+        }
+
+        state = CalibrationState(
+            sorted_surprise=sorted_vec,
+            threshold_quantiles=tuple(qs),
+            threshold_values=thr,
+            metadata=meta,
+        )
+        self._state = state
+        return state
+
+    # --- Inference-time utilities -------------------------------------------------
+    def threshold_for(self, quantile: float | None = None) -> Any:
+        """Return the calibrated threshold value for a given quantile.
+
+        If ``quantile`` is omitted and exactly one quantile was configured at
+        construction time, that single threshold is returned. When a quantile is
+        provided that was not part of the original configuration, the value is
+        computed on the fly using the stored sorted distribution and the same
+        deterministic index rule used during fitting (``floor(q * (n-1))``).
+        """
+
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("Calibration requires PyTorch to be installed")
+        if self._state is None:
+            raise RuntimeError("Calibrator has no state; call fit(...) first")
+
+        state = self._state
+        if quantile is None:
+            if len(self._threshold_quantiles) == 1:
+                return state.threshold_values[0]
+            raise ValueError(
+                "Multiple quantiles configured; specify 'quantile' explicitly",
+            )
+
+        q = float(quantile)
+        if not (0.0 <= q <= 1.0):
+            raise ValueError(f"Quantile {q} is outside [0, 1]")
+
+        # If the quantile matches one from fit(), return the stored value.
+        try:
+            idx = state.threshold_quantiles.index(q)
+            return state.threshold_values[idx]
+        except ValueError:
+            pass  # compute ad-hoc below
+
+        # Compute ad-hoc using the stored sorted distribution and index rule.
+        sorted_vec = cast("torch.Tensor", state.sorted_surprise)
+        n = int(sorted_vec.numel())
+        idx = min(n - 1, max(0, int((n - 1) * q)))
+        return sorted_vec[idx]
+
+    def percentiles(self, scores: Any) -> Any:
+        """Map surprise values to empirical percentiles in [0, 1].
+
+        Computes an empirical CDF against the stored ID distribution and
+        returns a percentile per input. To be robust across scoring variants
+        where either tail may correspond to increased OOD-ness, this method
+        automatically selects the upper or lower tail based on the median of
+        the provided ``scores`` relative to the ID median:
+        - If ``median(scores) >= median(ID)``: use upper-tail CDF
+          ``p(x) = # {v <= x} / n`` (higher scores → higher percentiles).
+        - Else: use lower-tail complement ``p(x) = 1 - # {v <= x} / n``
+          (smaller scores → higher percentiles).
+
+        Supports scalar and batched inputs; tensor inputs preserve their shape
+        in the result.
+        """
+
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("Calibration requires PyTorch to be installed")
+        if self._state is None:
+            raise RuntimeError("Calibrator has no state; call fit(...) first")
+
+        x = _as_tensor(scores)
+        sorted_vec = cast("torch.Tensor", self._state.sorted_surprise)
+
+        # Compute counts of values <= x via searchsorted(..., right=True)
+        flat = x.reshape(-1)
+        counts = torch.searchsorted(sorted_vec, flat, right=True)
+        n = float(sorted_vec.numel())
+        cdf = counts.to(dtype=torch.float32) / n
+        # Tail selection: compare input median vs ID median
+        try:
+            id_median = sorted_vec[sorted_vec.numel() // 2]
+            x_median = torch.median(flat).to(dtype=id_median.dtype, device=id_median.device)
+            use_upper = bool(x_median >= id_median)
+        except Exception:
+            use_upper = True
+        p = cdf if use_upper else (1.0 - cdf)
+        p = torch.clamp(p, 0.0, 1.0)
+        return p.reshape(x.shape)
+
+    def is_ood(self, scores: Any, *, quantile: float | None = None) -> Any:
+        """Return a boolean mask where inputs are marked OOD.
+
+        By default, compares the empirical percentile of each input against a
+        configured threshold quantile from calibration, marking OOD when
+        ``percentile >= quantile``. If ``quantile`` is omitted:
+        - with a single configured threshold, that quantile is used;
+        - with multiple thresholds, the most conservative choice (max quantile)
+          is used to avoid ambiguity and keep a sensible default.
+        """
+
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("Calibration requires PyTorch to be installed")
+        if self._state is None:
+            raise RuntimeError("Calibrator has no state; call fit(...) first")
+
+        q: float
+        if quantile is None:
+            if len(self._threshold_quantiles) == 0:
+                q = 0.95
+            elif len(self._threshold_quantiles) == 1:
+                q = float(self._threshold_quantiles[0])
+            else:
+                # Default to the highest configured quantile for a conservative threshold
+                q = float(max(self._threshold_quantiles))
+        else:
+            q = float(quantile)
+        if not (0.0 <= q <= 1.0):
+            raise ValueError(f"Quantile {q} is outside [0, 1]")
+
+        p = self.percentiles(scores)
+        # 'p' is a tensor shaped like the input; compare against q
+        return cast("torch.Tensor", p) >= q
+
+    # --- Persistence helpers -----------------------------------------------------
+    def get_state(self) -> CalibrationState | None:
+        """Return the current CalibrationState (or None if not fitted)."""
+
+        return self._state
+
+    def set_state(self, state: CalibrationState) -> None:
+        """Assign an externally provided CalibrationState."""
+
+        if not isinstance(state, CalibrationState):  # pragma: no cover - defensive
+            raise TypeError("state must be a CalibrationState")
+        self._state = state
+
+    def state_dict(self) -> dict[str, Any]:
+        """Return a versioned dictionary suitable for checkpointing.
+
+        Contains the configured threshold quantiles and, when available, the
+        fitted CalibrationState with its tensors and metadata.
+        """
+
+        out: dict[str, Any] = {
+            "__nervecode__": "EmpiricalPercentileCalibrator",
+            "version": 1,
+            "threshold_quantiles": self._threshold_quantiles,
+        }
+        if self._state is not None:
+            out["state"] = {
+                "sorted_surprise": self._state.sorted_surprise,
+                "threshold_quantiles": self._state.threshold_quantiles,
+                "threshold_values": self._state.threshold_values,
+                "metadata": self._state.metadata,
+            }
+        return out
+
+    def load_state_dict(self, state_dict: Mapping[str, Any]) -> None:
+        """Load from a dictionary produced by state_dict()."""
+
+        if not isinstance(state_dict, Mapping):  # pragma: no cover - defensive
+            raise TypeError("state_dict must be a mapping")
+        # Threshold quantiles configuration
+        tq = state_dict.get("threshold_quantiles")
+        if tq is not None:
+            self._threshold_quantiles = tuple(float(q) for q in list(tq))
+        # Optional embedded calibration state
+        st = state_dict.get("state")
+        if isinstance(st, Mapping) and "sorted_surprise" in st and "threshold_values" in st:
+            self._state = CalibrationState(
+                sorted_surprise=st["sorted_surprise"],
+                threshold_quantiles=tuple(
+                    float(q) for q in list(st.get("threshold_quantiles", self._threshold_quantiles))
+                ),
+                threshold_values=st["threshold_values"],
+                metadata=dict(st.get("metadata", {})),
+            )
+
+    def save(self, path: str) -> None:
+        """Save state_dict() to a file using torch.save if available."""
+
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("Saving calibrator state requires PyTorch to be installed")
+        torch.save(self.state_dict(), path)
+
+    @classmethod
+    def load(cls, path: str) -> EmpiricalPercentileCalibrator:
+        """Load a calibrator from a file produced by save()."""
+
+        if torch is None:  # pragma: no cover - defensive
+            raise RuntimeError("Loading calibrator state requires PyTorch to be installed")
+        sd = torch.load(path, map_location="cpu")
+        # Attempt to recover configured quantiles; default to embedded state
+        tq = sd.get("threshold_quantiles", (0.95,)) if isinstance(sd, Mapping) else (0.95,)
+        calib = cls(threshold_quantiles=tq)
+        if isinstance(sd, Mapping):
+            calib.load_state_dict(sd)
+        return calib
+
+
+def _as_1d_tensor(scores: Any) -> Any:
+    """Return a 1D tensor view of per-sample surprise values.
+
+    Supports common inputs without importing project-internal types here.
+    - ``torch.Tensor``: flattened to 1D (keeps dtype/device for intermediate ops)
+    - object with ``surprise`` attribute: treat it as a tensor and flatten
+    - sequence/iterable of numbers: converted to a CPU float tensor
+    """
+
+    if torch is None:  # pragma: no cover - defensive
+        raise RuntimeError("Calibration requires PyTorch to be installed")
+
+    if isinstance(scores, torch.Tensor):
+        return scores.reshape(-1)
+
+    cand = getattr(scores, "surprise", None)
+    if isinstance(cand, torch.Tensor):
+        return cand.reshape(-1)
+
+    if isinstance(scores, Iterable):
+        try:
+            return torch.tensor(list(scores), dtype=torch.float32)
+        except Exception as exc:  # pragma: no cover - defensive
+            raise TypeError("Unsupported iterable for calibration scores") from exc
+
+    raise TypeError("Unsupported input type for calibration scores")
+
+
+def _as_tensor(scores: Any) -> Any:
+    """Return a tensor view of input scores preserving shape when possible.
+
+    Accepted inputs mirror ``_as_1d_tensor`` but keep the original shape for
+    tensor-like objects and scalars. Iterables are materialized into a 1D
+    tensor. The returned tensor uses ``float32`` when constructed from Python
+    values; otherwise the input dtype/device is preserved.
+    """
+
+    if torch is None:  # pragma: no cover - defensive
+        raise RuntimeError("Calibration requires PyTorch to be installed")
+
+    if isinstance(scores, torch.Tensor):
+        return scores
+
+    cand = getattr(scores, "surprise", None)
+    if isinstance(cand, torch.Tensor):
+        return cand
+
+    # Common scalar path
+    if isinstance(scores, (int, float)):
+        return torch.tensor(scores, dtype=torch.float32)
+
+    if isinstance(scores, Iterable):
+        try:
+            return torch.tensor(list(scores), dtype=torch.float32)
+        except Exception as exc:  # pragma: no cover - defensive
+            raise TypeError("Unsupported iterable for calibration scores") from exc
+
+    raise TypeError("Unsupported input type for calibration scores")

nervecode/scoring/types.py

@@ -0,0 +1,33 @@
+"""Scoring/aggregation result types.
+
+This module defines lightweight dataclasses used by the scoring/aggregation
+APIs. Keep imports torch-agnostic so the package remains importable in minimal
+environments; fields that are typically ``torch.Tensor`` are typed as ``Any``.
+"""
+
+from __future__ import annotations
+
+from dataclasses import dataclass
+from typing import Any
+
+__all__ = ["AggregatedSurprise"]
+
+
+@dataclass(frozen=True)
+class AggregatedSurprise:
+    """Result of aggregating per-layer surprise signals.
+
+    Fields
+    - ``surprise``: per-sample aggregated surprise signal (e.g., shape ``(B,)``
+      or ``(...,)``). Typically a ``torch.Tensor``.
+    - ``method``: aggregation method identifier such as ``"mean"``. Future
+      values may include ``"max"`` and ``"weighted"`` without changing the
+      public result type.
+    - ``num_layers``: number of layer signals included in the aggregation.
+    - ``details``: optional aggregation-specific details (e.g., weights).
+    """
+
+    surprise: Any
+    method: str
+    num_layers: int
+    details: dict[str, Any] | None = None

nervecode/training/__init__.py

@@ -0,0 +1,25 @@
+"""Training-time helpers and losses.
+
+This subpackage will include auxiliary losses, schedulers, and small training
+helpers used to optimize codebooks and calibrate scores.
+"""
+
+from __future__ import annotations
+
+from .diagnostics import (
+    CsvDiagnosticsLogger,
+    DiagnosticsCollector,
+    DiagnosticsRow,
+    JsonlDiagnosticsLogger,
+)
+from .loss import CodingLoss
+from .updaters import EmaCodebookUpdater
+
+__all__ = [
+    "CodingLoss",
+    "CsvDiagnosticsLogger",
+    "DiagnosticsCollector",
+    "DiagnosticsRow",
+    "EmaCodebookUpdater",
+    "JsonlDiagnosticsLogger",
+]

nervecode/training/diagnostics.py

@@ -0,0 +1,194 @@
+"""Lightweight diagnostics collection and logging backends.
+
+This module provides utilities to summarize per-layer CodingTrace objects into
+simple, machine-readable rows and to write those rows to CSV or JSONL without
+bringing extra dependencies. It complements wrapper-level convenience metrics
+by producing explicit records during training, calibration, or evaluation.
+
+Design goals
+- Minimal surface; importable without heavy dependencies.
+- Fail-open behavior: best-effort extraction; skip rows on missing tensors.
+- CSV and JSONL outputs for easy downstream processing.
+"""
+
+from __future__ import annotations
+
+import csv
+import json
+from collections.abc import Iterable, Mapping
+from dataclasses import dataclass
+from pathlib import Path
+
+try:  # optional torch for numeric reductions
+    import torch
+except Exception:  # pragma: no cover - environments without torch
+    from typing import Any
+    from typing import cast as _cast
+
+    torch = _cast(Any, None)
+
+try:
+    from nervecode.core import CodingTrace
+except Exception:  # pragma: no cover - available during normal package use
+    CodingTrace = object  # type: ignore[misc,assignment]
+
+__all__ = [
+    "CsvDiagnosticsLogger",
+    "DiagnosticsCollector",
+    "DiagnosticsRow",
+    "JsonlDiagnosticsLogger",
+]
+
+
+@dataclass(frozen=True)
+class DiagnosticsRow:
+    """Single row of layer-level diagnostics derived from a CodingTrace."""
+
+    step: int | None
+    layer: str
+    K: int | None
+    code_dim: int | None
+    utilization: float | None
+    dead_code_frac: float | None
+    mean_entropy: float | None
+    mean_best_length: float | None
+    mean_commitment: float | None
+    reduction: str | None
+
+
+class DiagnosticsCollector:
+    """Summarize per-layer traces into DiagnosticsRow entries."""
+
+    def summarize_trace(
+        self, layer_name: str, trace: CodingTrace, *, step: int | None = None
+    ) -> DiagnosticsRow:
+        # Default values
+        K: int | None = None
+        code_dim: int | None = None
+        utilization: float | None = None
+        dead_frac: float | None = None
+        mean_entropy: float | None = None
+        mean_best_len: float | None = None
+        mean_commit: float | None = None
+        reduction: str | None = None
+
+        # Resolve K and D from the trace's SoftCode and reduced vectors
+        try:
+            _k = getattr(trace.soft_code, "code_dim", 0)
+            K = int(_k) if isinstance(_k, (int,)) else None
+        except Exception:
+            K = None
+        try:
+            _d = getattr(trace, "reduced_dim", 0)
+            code_dim = int(_d) if isinstance(_d, (int,)) else None
+        except Exception:
+            code_dim = None
+
+        # Utilization and dead-code fraction
+        if torch is not None and hasattr(trace.soft_code, "best_indices"):
+            try:
+                best = trace.soft_code.best_indices
+                if not isinstance(best, torch.Tensor):
+                    best = getattr(trace, "chosen_center_indices", None)
+                if isinstance(best, torch.Tensor) and K and K > 0:
+                    used = len(torch.unique(best.reshape(-1)).tolist())
+                    utilization = float(used) / float(K)
+                    dead_frac = max(0.0, 1.0 - utilization)
+            except Exception:
+                pass
+
+        # Entropy, best-length, and commitment distance means
+        if torch is not None:
+            try:
+                ent = getattr(trace.soft_code, "entropy", None)
+                if isinstance(ent, torch.Tensor):
+                    mean_entropy = float(ent.mean().item())
+            except Exception:
+                pass
+            try:
+                bl = getattr(trace.soft_code, "best_length", None)
+                if isinstance(bl, torch.Tensor):
+                    mean_best_len = float(bl.mean().item())
+            except Exception:
+                pass
+            try:
+                cm = getattr(trace, "commitment_distances", None)
+                if isinstance(cm, torch.Tensor):
+                    mean_commit = float(cm.mean().item())
+            except Exception:
+                pass
+
+        # Reduction metadata
+        try:
+            meta = getattr(trace, "reduction_meta", {}) or {}
+            reduction = str(meta.get("method")) if isinstance(meta, Mapping) else None
+        except Exception:
+            reduction = None
+
+        return DiagnosticsRow(
+            step=step,
+            layer=layer_name,
+            K=K,
+            code_dim=code_dim,
+            utilization=utilization,
+            dead_code_frac=dead_frac,
+            mean_entropy=mean_entropy,
+            mean_best_length=mean_best_len,
+            mean_commitment=mean_commit,
+            reduction=reduction,
+        )
+
+    def collect(
+        self, traces: Mapping[str, CodingTrace], *, step: int | None = None
+    ) -> list[DiagnosticsRow]:
+        rows: list[DiagnosticsRow] = []
+        for name, trace in traces.items():
+            try:
+                rows.append(self.summarize_trace(str(name), trace, step=step))
+            except Exception:
+                continue
+        return rows
+
+
+class CsvDiagnosticsLogger:
+    """Append diagnostics rows to a CSV file with a fixed header."""
+
+    def __init__(self, path: str | Path) -> None:
+        self._path = Path(path)
+        self._header_written = False
+        self._fieldnames = [
+            "step",
+            "layer",
+            "K",
+            "code_dim",
+            "utilization",
+            "dead_code_frac",
+            "mean_entropy",
+            "mean_best_length",
+            "mean_commitment",
+            "reduction",
+        ]
+
+    def log(self, rows: Iterable[DiagnosticsRow]) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        with self._path.open("a", encoding="utf-8", newline="") as f:
+            writer = csv.DictWriter(f, fieldnames=self._fieldnames)
+            if not self._header_written and self._path.stat().st_size == 0:
+                writer.writeheader()
+                self._header_written = True
+            for r in rows:
+                writer.writerow({k: getattr(r, k) for k in self._fieldnames})
+
+
+class JsonlDiagnosticsLogger:
+    """Append diagnostics rows as JSONL objects (one per line)."""
+
+    def __init__(self, path: str | Path) -> None:
+        self._path = Path(path)
+
+    def log(self, rows: Iterable[DiagnosticsRow]) -> None:
+        self._path.parent.mkdir(parents=True, exist_ok=True)
+        with self._path.open("a", encoding="utf-8") as f:
+            for r in rows:
+                f.write(json.dumps(r.__dict__))
+                f.write("\n")