Probes

steerability.algorithms.core.internals.probes

Probes: calibrated affine readouts over model internals.

Probe is model-free feature math with canonical polarity (a score at or above zero means present). fit_probe and calibrate_bias fit and calibrate probes from contrastive pairs; ProbeSet scores many probes in one read-only forward and returns a ProbeReadings; ProbeSetFit defers fitting to the model a pipeline provides at steer time.

CalibrationSpec = Literal['max_f1', 'midpoint'] | tuple[Literal['target_fpr'], float] module-attribute

Probe dataclass

Calibrated affine readout of a concept from pooled hidden states.

A probe scores pooled per-layer features as sum_l features[l] @ weights[l] + bias, and its decision is score >= 0 (ties open). Polarity is canonical, i.e., positives score at or above zero, and the operating threshold is folded into bias at calibration time, so no comparator or threshold field exists. Weights are stored in raw activation coordinates (any standardization used during fitting is folded into the weights), so scoring needs no activation statistics.

The class is model-free feature math. It runs no forwards and registers no hooks; batched reads over a live model go through ProbeSet.

Attributes:

Name	Type	Description
model_type	str	HuggingFace model_type of the model the probe was fitted on.
location	str	Residual-stream boundary the weights were fitted at ("layer_input" or "layer_output").
pooling	Literal['mean', 'last']	Token aggregation the probe expects ("mean" or "last", mask-aware).
layer_ids	list[int]	Layers whose features the decision function consumes.
weights	dict[int, Tensor]	Per-layer weight vectors, [H] float32, raw-activation coordinates.
bias	float	Calibrated offset; the decision is raw_score + bias >= 0.
meta	dict	Provenance record (fit method, sample counts, calibration record, layer sweep, fingerprints, package version, polarity marker). May be empty for hand-constructed probes, which disarms all fingerprint checks.

Source code in steerability/algorithms/core/internals/probes/probe.py

@dataclass
class Probe:
    """Calibrated affine readout of a concept from pooled hidden states.

    A probe scores pooled per-layer features as `sum_l features[l] @ weights[l] + bias`, and its
    decision is `score >= 0` (ties open). Polarity is canonical, i.e., positives score at or
    above zero, and the operating threshold is folded into `bias` at calibration time, so no
    comparator or threshold field exists. Weights are stored in raw activation coordinates (any
    standardization used during fitting is folded into the weights), so scoring needs no
    activation statistics.

    The class is model-free feature math. It runs no forwards and registers no hooks; batched
    reads over a live model go through `ProbeSet`.

    Attributes:
        model_type: HuggingFace `model_type` of the model the probe was fitted on.
        location: Residual-stream boundary the weights were fitted at (`"layer_input"` or
            `"layer_output"`).
        pooling: Token aggregation the probe expects (`"mean"` or `"last"`, mask-aware).
        layer_ids: Layers whose features the decision function consumes.
        weights: Per-layer weight vectors, `[H]` float32, raw-activation coordinates.
        bias: Calibrated offset; the decision is `raw_score + bias >= 0`.
        meta: Provenance record (fit method, sample counts, calibration record, layer sweep,
            fingerprints, package version, polarity marker). May be empty for hand-constructed
            probes, which disarms all fingerprint checks.
    """

    model_type: str
    location: str
    pooling: Literal["mean", "last"]
    layer_ids: list[int]
    weights: dict[int, torch.Tensor]
    bias: float
    meta: dict = field(default_factory=dict)

    def __post_init__(self):
        if self.pooling not in ("mean", "last"):
            raise ValueError(f"pooling must be 'mean' or 'last', got {self.pooling!r}.")
        if not self.layer_ids:
            raise ValueError("layer_ids must name at least one layer.")
        self.layer_ids = [int(lid) for lid in self.layer_ids]
        missing = [lid for lid in self.layer_ids if lid not in self.weights]
        if missing:
            raise ValueError(f"weights missing for layer(s) {missing}.")
        self.weights = {
            int(lid): torch.as_tensor(w).detach().reshape(-1).to(torch.float32)
            for lid, w in self.weights.items()
        }
        self.bias = float(self.bias)

    def decision_function(self, features: Mapping[int, torch.Tensor]) -> torch.Tensor:
        """Signed decision scores for pooled features.

        Args:
            features: Mapping from layer id to pooled features of shape `[B, H]` (a `[H]` tensor
                is treated as one row).

        Returns:
            Float32 tensor of shape `[B]`: the sum over `layer_ids` of `features[l] @ weights[l]`,
            plus `bias`.

        Raises:
            KeyError: If a layer in `layer_ids` is absent from `features`; the message names the
                missing layer and the layers supplied.
        """
        scores: torch.Tensor | None = None
        for lid in self.layer_ids:
            if lid not in features:
                raise KeyError(
                    f"Probe requires features for layer {lid}, but only layers "
                    f"{sorted(features)} were supplied; capture or pool every layer in "
                    f"probe.layer_ids ({self.layer_ids})."
                )
            feats = torch.as_tensor(features[lid]).to(torch.float32)
            if feats.ndim == 1:
                feats = feats.unsqueeze(0)
            contribution = feats @ self.weights[lid].to(feats.device)
            scores = contribution if scores is None else scores + contribution
        return (scores + self.bias).cpu()

    def predict(self, features: Mapping[int, torch.Tensor]) -> torch.Tensor:
        """Boolean decisions for pooled features: `decision_function(features) >= 0` (ties open).

        Args:
            features: Mapping from layer id to pooled features of shape `[B, H]`.

        Returns:
            Bool tensor of shape `[B]`.
        """
        return self.decision_function(features) >= 0

    def score_hidden(
        self,
        hidden: Mapping[int, torch.Tensor],
        prompt_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Pool per-layer hidden states and apply the decision function.

        Each layer's `[B, T, H]` tensor is aggregated to `[B, H]` per `self.pooling` using the
        mask, so `"last"` selects the last real (non-pad) token per row.

        Args:
            hidden: Mapping from layer id to hidden states of shape `[B, T, H]`.
            prompt_mask: Attention mask of shape `[B, T]` (1 at real tokens), or None to treat
                every position as real.

        Returns:
            Float32 tensor of signed decision scores, shape `[B]`.

        Raises:
            KeyError: If a layer in `layer_ids` is absent from `hidden`.
        """
        features: dict[int, torch.Tensor] = {}
        for lid in self.layer_ids:
            if lid not in hidden:
                raise KeyError(
                    f"Probe requires hidden states for layer {lid}, but only layers "
                    f"{sorted(hidden)} were supplied; capture every layer in probe.layer_ids "
                    f"({self.layer_ids})."
                )
            features[lid] = aggregate_condition_hidden(
                hidden[lid], self.pooling, attention_mask=prompt_mask
            )
        return self.decision_function(features)

    def save(self, dir_path: str | Path) -> None:
        """Save the probe to a directory (safetensors for weights, JSON sidecar for metadata).

        Args:
            dir_path: Directory to write to; created if missing.
        """
        dir_path = Path(dir_path)
        dir_path.mkdir(parents=True, exist_ok=True)

        tensors = {
            f"weights.{lid}": self.weights[lid].to(torch.float32).contiguous()
            for lid in sorted(self.weights)
        }
        save_file(tensors, str(dir_path / _TENSOR_FILE))

        metadata = {
            "format_version": PROBE_FORMAT_VERSION,
            "polarity": POLARITY_MARKER,
            "model_type": self.model_type,
            "location": self.location,
            "pooling": self.pooling,
            "layer_ids": list(self.layer_ids),
            "weight_layer_ids": sorted(self.weights),
            "bias": self.bias,
            "meta": self.meta,
        }
        with open(dir_path / _META_FILE, "w") as f:
            json.dump(metadata, f, indent=2, sort_keys=True)
        logger.debug("Saved Probe to %s", dir_path)

    @classmethod
    def load(cls, dir_path: str | Path) -> "Probe":
        """Load a probe saved by `save`.

        Args:
            dir_path: Directory containing the safetensors file and JSON sidecar.

        Returns:
            The loaded `Probe`.

        Raises:
            ValueError: If the artifact-format version is unsupported or the recorded polarity
                marker is not canonical.
        """
        dir_path = Path(dir_path)
        with open(dir_path / _META_FILE) as f:
            metadata = json.load(f)

        version = metadata.get("format_version")
        if version != PROBE_FORMAT_VERSION:
            raise ValueError(
                f"Unsupported Probe format version {version!r}; expected {PROBE_FORMAT_VERSION}."
            )
        polarity = metadata.get("polarity")
        if polarity != POLARITY_MARKER:
            raise ValueError(
                f"Probe artifact records polarity {polarity!r}; expected {POLARITY_MARKER!r} "
                "(positives score at or above zero)."
            )

        tensors = load_file(str(dir_path / _TENSOR_FILE))
        weights = {int(lid): tensors[f"weights.{lid}"] for lid in metadata["weight_layer_ids"]}

        logger.debug("Loaded Probe from %s with layers %s", dir_path, metadata["layer_ids"])
        return cls(
            model_type=metadata["model_type"],
            location=metadata["location"],
            pooling=metadata["pooling"],
            layer_ids=[int(lid) for lid in metadata["layer_ids"]],
            weights=weights,
            bias=float(metadata["bias"]),
            meta=metadata["meta"],
        )

    def as_gate(self, *, allow_model_mismatch: bool = False) -> Gate:
        """A steering gate reproducing this probe's decision, for gated interventions.

        The gate reads the probe's layers at the probe's pooling, scores each pooled state
        against the probe's weights, and opens where the summed contributions plus the
        calibrated bias are at or above zero (ties open). The decision is evaluated on the
        prompt and holds for the whole generation.

        Args:
            allow_model_mismatch: When True, the intervention's model-identity check for this
                probe is disarmed.

        Returns:
            A `Gate` for an `Intervention`'s gate slot (e.g. `ActivationAdapter`'s `gate=`).
        """
        # function-local import to avoid a core/internals -> category package import at module load
        from steerability.algorithms.state_control.common.gating import gate_from_probe

        return gate_from_probe(self, allow_model_mismatch=allow_model_mismatch)

bias instance-attribute

layer_ids instance-attribute

location instance-attribute

meta = field(default_factory=dict) class-attribute instance-attribute

model_type instance-attribute

pooling instance-attribute

weights instance-attribute

as_gate(*, allow_model_mismatch=False)

A steering gate reproducing this probe's decision, for gated interventions.

The gate reads the probe's layers at the probe's pooling, scores each pooled state against the probe's weights, and opens where the summed contributions plus the calibrated bias are at or above zero (ties open). The decision is evaluated on the prompt and holds for the whole generation.

Parameters:

Name	Type	Description	Default
allow_model_mismatch	bool	When True, the intervention's model-identity check for this probe is disarmed.	False

Returns:

Type	Description
Gate	A Gate for an Intervention's gate slot (e.g. ActivationAdapter's gate=).

Source code in steerability/algorithms/core/internals/probes/probe.py

def as_gate(self, *, allow_model_mismatch: bool = False) -> Gate:
    """A steering gate reproducing this probe's decision, for gated interventions.

    The gate reads the probe's layers at the probe's pooling, scores each pooled state
    against the probe's weights, and opens where the summed contributions plus the
    calibrated bias are at or above zero (ties open). The decision is evaluated on the
    prompt and holds for the whole generation.

    Args:
        allow_model_mismatch: When True, the intervention's model-identity check for this
            probe is disarmed.

    Returns:
        A `Gate` for an `Intervention`'s gate slot (e.g. `ActivationAdapter`'s `gate=`).
    """
    # function-local import to avoid a core/internals -> category package import at module load
    from steerability.algorithms.state_control.common.gating import gate_from_probe

    return gate_from_probe(self, allow_model_mismatch=allow_model_mismatch)

decision_function(features)

Signed decision scores for pooled features.

Parameters:

Name	Type	Description	Default
features	Mapping[int, Tensor]	Mapping from layer id to pooled features of shape [B, H] (a [H] tensor is treated as one row).	required

Returns:

Type	Description
Tensor	Float32 tensor of shape [B]: the sum over layer_ids of features[l] @ weights[l],
Tensor	plus bias.

Raises:

Type	Description
KeyError	If a layer in layer_ids is absent from features; the message names the missing layer and the layers supplied.

Source code in steerability/algorithms/core/internals/probes/probe.py

def decision_function(self, features: Mapping[int, torch.Tensor]) -> torch.Tensor:
    """Signed decision scores for pooled features.

    Args:
        features: Mapping from layer id to pooled features of shape `[B, H]` (a `[H]` tensor
            is treated as one row).

    Returns:
        Float32 tensor of shape `[B]`: the sum over `layer_ids` of `features[l] @ weights[l]`,
        plus `bias`.

    Raises:
        KeyError: If a layer in `layer_ids` is absent from `features`; the message names the
            missing layer and the layers supplied.
    """
    scores: torch.Tensor | None = None
    for lid in self.layer_ids:
        if lid not in features:
            raise KeyError(
                f"Probe requires features for layer {lid}, but only layers "
                f"{sorted(features)} were supplied; capture or pool every layer in "
                f"probe.layer_ids ({self.layer_ids})."
            )
        feats = torch.as_tensor(features[lid]).to(torch.float32)
        if feats.ndim == 1:
            feats = feats.unsqueeze(0)
        contribution = feats @ self.weights[lid].to(feats.device)
        scores = contribution if scores is None else scores + contribution
    return (scores + self.bias).cpu()

load(dir_path) classmethod

Load a probe saved by save.

Parameters:

Name	Type	Description	Default
dir_path	str | Path	Directory containing the safetensors file and JSON sidecar.	required

Returns:

Type	Description
'Probe'	The loaded Probe.

Raises:

Type	Description
ValueError	If the artifact-format version is unsupported or the recorded polarity marker is not canonical.

Source code in steerability/algorithms/core/internals/probes/probe.py

@classmethod
def load(cls, dir_path: str | Path) -> "Probe":
    """Load a probe saved by `save`.

    Args:
        dir_path: Directory containing the safetensors file and JSON sidecar.

    Returns:
        The loaded `Probe`.

    Raises:
        ValueError: If the artifact-format version is unsupported or the recorded polarity
            marker is not canonical.
    """
    dir_path = Path(dir_path)
    with open(dir_path / _META_FILE) as f:
        metadata = json.load(f)

    version = metadata.get("format_version")
    if version != PROBE_FORMAT_VERSION:
        raise ValueError(
            f"Unsupported Probe format version {version!r}; expected {PROBE_FORMAT_VERSION}."
        )
    polarity = metadata.get("polarity")
    if polarity != POLARITY_MARKER:
        raise ValueError(
            f"Probe artifact records polarity {polarity!r}; expected {POLARITY_MARKER!r} "
            "(positives score at or above zero)."
        )

    tensors = load_file(str(dir_path / _TENSOR_FILE))
    weights = {int(lid): tensors[f"weights.{lid}"] for lid in metadata["weight_layer_ids"]}

    logger.debug("Loaded Probe from %s with layers %s", dir_path, metadata["layer_ids"])
    return cls(
        model_type=metadata["model_type"],
        location=metadata["location"],
        pooling=metadata["pooling"],
        layer_ids=[int(lid) for lid in metadata["layer_ids"]],
        weights=weights,
        bias=float(metadata["bias"]),
        meta=metadata["meta"],
    )

predict(features)

Boolean decisions for pooled features: decision_function(features) >= 0 (ties open).

Parameters:

Name	Type	Description	Default
features	Mapping[int, Tensor]	Mapping from layer id to pooled features of shape [B, H].	required

Returns:

Type	Description
Tensor	Bool tensor of shape [B].

Source code in steerability/algorithms/core/internals/probes/probe.py

def predict(self, features: Mapping[int, torch.Tensor]) -> torch.Tensor:
    """Boolean decisions for pooled features: `decision_function(features) >= 0` (ties open).

    Args:
        features: Mapping from layer id to pooled features of shape `[B, H]`.

    Returns:
        Bool tensor of shape `[B]`.
    """
    return self.decision_function(features) >= 0

save(dir_path)

Save the probe to a directory (safetensors for weights, JSON sidecar for metadata).

Parameters:

Name	Type	Description	Default
dir_path	str | Path	Directory to write to; created if missing.	required

Source code in steerability/algorithms/core/internals/probes/probe.py

def save(self, dir_path: str | Path) -> None:
    """Save the probe to a directory (safetensors for weights, JSON sidecar for metadata).

    Args:
        dir_path: Directory to write to; created if missing.
    """
    dir_path = Path(dir_path)
    dir_path.mkdir(parents=True, exist_ok=True)

    tensors = {
        f"weights.{lid}": self.weights[lid].to(torch.float32).contiguous()
        for lid in sorted(self.weights)
    }
    save_file(tensors, str(dir_path / _TENSOR_FILE))

    metadata = {
        "format_version": PROBE_FORMAT_VERSION,
        "polarity": POLARITY_MARKER,
        "model_type": self.model_type,
        "location": self.location,
        "pooling": self.pooling,
        "layer_ids": list(self.layer_ids),
        "weight_layer_ids": sorted(self.weights),
        "bias": self.bias,
        "meta": self.meta,
    }
    with open(dir_path / _META_FILE, "w") as f:
        json.dump(metadata, f, indent=2, sort_keys=True)
    logger.debug("Saved Probe to %s", dir_path)

score_hidden(hidden, prompt_mask=None)

Pool per-layer hidden states and apply the decision function.

Each layer's [B, T, H] tensor is aggregated to [B, H] per self.pooling using the mask, so "last" selects the last real (non-pad) token per row.

Parameters:

Name	Type	Description	Default
hidden	Mapping[int, Tensor]	Mapping from layer id to hidden states of shape [B, T, H].	required
prompt_mask	Tensor | None	Attention mask of shape [B, T] (1 at real tokens), or None to treat every position as real.	None

Returns:

Type	Description
Tensor	Float32 tensor of signed decision scores, shape [B].

Raises:

Type	Description
KeyError	If a layer in layer_ids is absent from hidden.

Source code in steerability/algorithms/core/internals/probes/probe.py

def score_hidden(
    self,
    hidden: Mapping[int, torch.Tensor],
    prompt_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Pool per-layer hidden states and apply the decision function.

    Each layer's `[B, T, H]` tensor is aggregated to `[B, H]` per `self.pooling` using the
    mask, so `"last"` selects the last real (non-pad) token per row.

    Args:
        hidden: Mapping from layer id to hidden states of shape `[B, T, H]`.
        prompt_mask: Attention mask of shape `[B, T]` (1 at real tokens), or None to treat
            every position as real.

    Returns:
        Float32 tensor of signed decision scores, shape `[B]`.

    Raises:
        KeyError: If a layer in `layer_ids` is absent from `hidden`.
    """
    features: dict[int, torch.Tensor] = {}
    for lid in self.layer_ids:
        if lid not in hidden:
            raise KeyError(
                f"Probe requires hidden states for layer {lid}, but only layers "
                f"{sorted(hidden)} were supplied; capture every layer in probe.layer_ids "
                f"({self.layer_ids})."
            )
        features[lid] = aggregate_condition_hidden(
            hidden[lid], self.pooling, attention_mask=prompt_mask
        )
    return self.decision_function(features)

ProbeEvaluation dataclass

Held-out scores of a probe on labeled data.

Attributes:

Name	Type	Description
positive_scores	Tensor	Signed decision scores of the positive class, shape [N_pos], ordered by descending rendered text length rather than input order.
negative_scores	Tensor	Signed decision scores of the negative class, shape [N_neg], in the same order.
accuracy	float	Fraction correct at the calibrated point (score >= 0 is positive).
f1	float	F1 of the same decision.

Source code in steerability/algorithms/core/internals/probes/fitting.py

@dataclass(frozen=True, slots=True)
class ProbeEvaluation:
    """Held-out scores of a probe on labeled data.

    Attributes:
        positive_scores: Signed decision scores of the positive class, shape `[N_pos]`, ordered by
            descending rendered text length rather than input order.
        negative_scores: Signed decision scores of the negative class, shape `[N_neg]`, in the same
            order.
        accuracy: Fraction correct at the calibrated point (`score >= 0` is positive).
        f1: F1 of the same decision.
    """

    positive_scores: torch.Tensor
    negative_scores: torch.Tensor
    accuracy: float
    f1: float

accuracy instance-attribute

f1 instance-attribute

negative_scores instance-attribute

positive_scores instance-attribute

ProbeFitSpec dataclass

Configuration for fitting a Probe from contrastive pairs.

Attributes:

Name	Type	Description
method	Literal['lda', 'mean_diff', 'logreg', 'fisher']	Direction estimation method. "lda" (default) computes the difference in class means on standardized features (diagonal LDA) and requires stats; "logreg" is L2 logistic regression on standardized features and also requires stats; "mean_diff" computes the raw difference in class means and never consults stats; "fisher" is a full-covariance discriminant computed from the class features alone (the pooled within-class covariance, pseudo-inverted with singular values truncated at 1e-6, applied to the class-mean difference) and never consults stats. The truncated pseudo-inverse handles rank deficiency when sample counts are below the hidden size. "fisher" weights are unit-normalized; the scale is part of the method's contract, since downstream consumers score raw margins.
pooling	Literal['mean', 'last']	Token aggregation for feature extraction ("mean" or "last", mask-aware).
location	str	Residual-stream boundary features are captured at.
prompt_format	PromptFormat	How pairs are rendered into model-ready text before tokenization.
candidate_layers	Sequence[int] | None	Explicit layer ids to sweep. None sweeps all decoder layers, optionally windowed by layer_range.
layer_range	tuple[float, float] | None	Fractional-depth window (lo, hi) with 0 <= lo < hi <= 1, applied to the full layer count when candidate_layers is None.
calibration	CalibrationSpec	Operating-point rule folded into the bias (see calibrate_bias).
C	float	Inverse regularization strength for "logreg".
seed	int	Random seed forwarded to "logreg".

Source code in steerability/algorithms/core/internals/probes/fitting.py

@dataclass(frozen=True)
class ProbeFitSpec:
    """Configuration for fitting a `Probe` from contrastive pairs.

    Attributes:
        method: Direction estimation method. `"lda"` (default) computes the difference in class
            means on standardized features (diagonal LDA) and requires `stats`; `"logreg"` is L2
            logistic regression on standardized features and also requires `stats`;
            `"mean_diff"` computes the raw difference in class means and never consults `stats`;
            `"fisher"` is a full-covariance discriminant computed from the class features alone
            (the pooled within-class covariance, pseudo-inverted with singular values truncated
            at `1e-6`, applied to the class-mean difference) and never consults `stats`. The
            truncated pseudo-inverse handles rank deficiency when sample counts are below the
            hidden size. `"fisher"` weights are unit-normalized; the scale is part of the
            method's contract, since downstream consumers score raw margins.
        pooling: Token aggregation for feature extraction (`"mean"` or `"last"`, mask-aware).
        location: Residual-stream boundary features are captured at.
        prompt_format: How pairs are rendered into model-ready text before tokenization.
        candidate_layers: Explicit layer ids to sweep. None sweeps all decoder layers, optionally
            windowed by `layer_range`.
        layer_range: Fractional-depth window `(lo, hi)` with `0 <= lo < hi <= 1`, applied to the
            full layer count when `candidate_layers` is None.
        calibration: Operating-point rule folded into the bias (see `calibrate_bias`).
        C: Inverse regularization strength for `"logreg"`.
        seed: Random seed forwarded to `"logreg"`.
    """

    method: Literal["lda", "mean_diff", "logreg", "fisher"] = "lda"
    pooling: Literal["mean", "last"] = "last"
    location: str = "layer_input"
    prompt_format: PromptFormat = "chat_prompt"
    candidate_layers: Sequence[int] | None = None
    layer_range: tuple[float, float] | None = None
    calibration: CalibrationSpec = "max_f1"
    C: float = 1.0
    seed: int = 0

    def __post_init__(self):
        if self.method not in ("lda", "mean_diff", "logreg", "fisher"):
            raise ValueError(
                f"method must be 'lda', 'mean_diff', 'logreg', or 'fisher', got {self.method!r}."
            )
        if self.pooling not in ("mean", "last"):
            raise ValueError(f"pooling must be 'mean' or 'last', got {self.pooling!r}.")
        if self.location not in ("layer_input", "layer_output"):
            raise ValueError(
                f"location must be 'layer_input' or 'layer_output', got {self.location!r}."
            )
        if self.prompt_format not in ("raw", "chat_completion", "chat_prompt"):
            raise ValueError(
                f"prompt_format must be one of raw/chat_completion/chat_prompt, "
                f"got {self.prompt_format!r}."
            )
        if self.candidate_layers is not None and len(self.candidate_layers) == 0:
            raise ValueError("candidate_layers must name at least one layer when supplied.")
        if self.layer_range is not None:
            lo, hi = self.layer_range
            if not 0.0 <= lo < hi <= 1.0:
                raise ValueError(
                    f"layer_range must satisfy 0 <= lo < hi <= 1, got ({lo}, {hi})."
                )
        if self.C <= 0:
            raise ValueError(f"C must be positive, got {self.C}.")
        _validate_calibration(self.calibration)

C = 1.0 class-attribute instance-attribute

calibration = 'max_f1' class-attribute instance-attribute

candidate_layers = None class-attribute instance-attribute

layer_range = None class-attribute instance-attribute

location = 'layer_input' class-attribute instance-attribute

method = 'lda' class-attribute instance-attribute

pooling = 'last' class-attribute instance-attribute

prompt_format = 'chat_prompt' class-attribute instance-attribute

seed = 0 class-attribute instance-attribute

ProbeReadings dataclass

One ProbeSet.read() call's result.

Attributes:

Name	Type	Description
scores	dict[str, Tensor]	Mapping from probe name to signed decision scores, [B] per probe.
decisions	dict[str, Tensor]	Mapping from probe name to boolean decisions, [B] per probe (scores >= 0, ties open).

Source code in steerability/algorithms/core/internals/probes/probe_set.py

@dataclass
class ProbeReadings:
    """One `ProbeSet.read()` call's result.

    Attributes:
        scores: Mapping from probe name to signed decision scores, `[B]` per probe.
        decisions: Mapping from probe name to boolean decisions, `[B]` per probe
            (`scores >= 0`, ties open).
    """

    scores: dict[str, torch.Tensor]
    decisions: dict[str, torch.Tensor]

decisions instance-attribute

scores instance-attribute

ProbeSet

Named probes scored on a batch of prompts in one read-only forward.

All probes must share one model_type and one capture location, since a single forward supplies every probe's features. read() registers forward pre-hooks on the union of the probes' layers (the layer_input boundary), runs one forward, and returns per-probe signed scores and decisions. The call is self-contained, i.e., it registers no standing hooks and never edits hidden states.

Parameters:

Name	Type	Description	Default
probes	Mapping[str, Probe]	Mapping from probe name to a fitted Probe.	required

Attributes:

Name	Type	Description
probes	dict[str, Probe]	The probes, keyed by name.
names	tuple[str, ...]	The probe names, in mapping order.
layer_ids	list[int]	Sorted union of the probes' layers.
latest	ProbeReadings | None	The most recent ProbeReadings, overwritten by each read() call; None before the first. Diagnostics only.

Raises:

Type	Description
ValueError	If probes is empty, the probes disagree on model_type or location, or two probes record different model fingerprints (probes with no recorded fingerprint mix freely).

Source code in steerability/algorithms/core/internals/probes/probe_set.py

class ProbeSet:
    """Named probes scored on a batch of prompts in one read-only forward.

    All probes must share one `model_type` and one capture `location`, since a single forward
    supplies every probe's features. `read()` registers forward pre-hooks on the union of the
    probes' layers (the `layer_input` boundary), runs one forward, and returns per-probe signed
    scores and decisions. The call is self-contained, i.e., it registers no standing hooks and
    never edits hidden states.

    Args:
        probes: Mapping from probe name to a fitted `Probe`.

    Attributes:
        probes: The probes, keyed by name.
        names: The probe names, in mapping order.
        layer_ids: Sorted union of the probes' layers.
        latest: The most recent `ProbeReadings`, overwritten by each `read()` call; None
            before the first. Diagnostics only.

    Raises:
        ValueError: If `probes` is empty, the probes disagree on `model_type` or `location`, or
            two probes record different model fingerprints (probes with no recorded fingerprint
            mix freely).
    """

    def __init__(self, probes: Mapping[str, Probe]):
        if not probes:
            raise ValueError("ProbeSet requires at least one probe.")
        self.probes: dict[str, Probe] = dict(probes)

        model_types = {probe.model_type for probe in self.probes.values()}
        if len(model_types) > 1:
            raise ValueError(
                f"probes disagree on model_type ({sorted(model_types)}); one read() capture "
                "cannot serve probes fitted on different model families."
            )
        locations = {probe.location for probe in self.probes.values()}
        if len(locations) > 1:
            raise ValueError(
                f"probes disagree on location ({sorted(locations)}); one read() capture "
                "cannot serve probes fitted at different boundaries."
            )
        fingerprints = {
            probe.meta["model_fingerprint"]
            for probe in self.probes.values()
            if probe.meta.get("model_fingerprint") is not None
        }
        if len(fingerprints) > 1:
            raise ValueError(
                f"probes record different model fingerprints ({sorted(fingerprints)}); "
                "mixed-model probes cannot share one read() capture."
            )

        self.names: tuple[str, ...] = tuple(self.probes)
        self.layer_ids: list[int] = sorted({lid for probe in self.probes.values() for lid in probe.layer_ids})
        self.latest: ProbeReadings | None = None

    @property
    def model_type(self) -> str:
        """The shared `model_type` of the probes."""
        return next(iter(self.probes.values())).model_type

    @property
    def location(self) -> str:
        """The shared capture boundary of the probes."""
        return next(iter(self.probes.values())).location

    @classmethod
    def fit(
        cls,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase,
        *,
        data: Mapping[str, ContrastivePairs],
        spec: ProbeFitSpec | Mapping[str, ProbeFitSpec] | None = None,
        stats: ActivationStats | None = None,
        calibration_data: Mapping[str, ContrastivePairs] | None = None,
        session=None,
    ) -> "ProbeSet":
        """Fit one probe per name and return the set.

        Args:
            model: Model the probes are fitted on.
            tokenizer: Tokenizer for rendering and encoding the pairs.
            data: Mapping from probe name to the contrastive pairs it is fitted on.
            spec: A shared `ProbeFitSpec` (None uses the defaults), or a mapping from probe name
                to a per-name spec.
            stats: Ambient activation statistics for methods that require them.
            calibration_data: Optional mapping from probe name to calibration pairs.

        Returns:
            The fitted `ProbeSet`.

        Raises:
            ValueError: If `data` is empty, a spec mapping misses a name, or `fit_probe` raises
                for any probe.
        """
        if not data:
            raise ValueError("data must map at least one probe name to contrastive pairs.")
        if isinstance(spec, Mapping):
            missing = [name for name in data if name not in spec]
            if missing:
                raise ValueError(f"spec mapping is missing entries for probe(s) {missing}.")

        probes: dict[str, Probe] = {}
        for name, pairs in data.items():
            probe_spec = spec[name] if isinstance(spec, Mapping) else (spec or ProbeFitSpec())
            probes[name] = fit_probe(
                model,
                tokenizer,
                data=pairs,
                spec=probe_spec,
                stats=stats,
                calibration_data=(calibration_data or {}).get(name),
                session=session,
            )
        return cls(probes)

    def read(
        self,
        model: PreTrainedModel | None,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        session=None,
    ) -> ProbeReadings:
        """Score a batch of prompts against every probe in one read-only forward.

        Registers one forward pre-hook per layer in the union (the `layer_input` boundary, the
        only boundary `read()` serves), runs a single forward under `torch.no_grad()` with
        `use_cache=False` wrapped in `auxiliary_pass(aligned=True)`, removes the hooks, then
        pools each probe's layers per its `pooling` with the attention mask and applies its
        decision function. The auxiliary marking means co-resident state controls' condition
        scorers, gates, and position counters ignore the pass, while `"all"`-scoped behavior
        transforms still apply, so scores measure the stream as deployed. A set fitted at
        `"layer_output"` raises, since a pre-hook capture would score features one block away
        from where those probes were fitted; score such probes offline via
        `layerwise_tokenwise_hidden` and `Probe.score_hidden`.

        Args:
            model: The model to read (must match the probes' `model_type`).
            input_ids: Prompt token ids of shape `[B, T]` (a 1-D tensor is treated as one row).
            attention_mask: Prompt attention mask matching `input_ids` (1 at real tokens). When
                None, every position is treated as real.

        Returns:
            A `ProbeReadings` with per-probe signed scores and decisions, also stashed on `latest`.

        Raises:
            ValueError: If the model's `model_type` does not match the probes', the set's
                `location` is not `"layer_input"`, a probe layer is out of range, or the
                architecture is unrecognized.
        """
        if model is not None:
            live_model_type = getattr(model.config, "model_type", "unknown")
            if live_model_type != self.model_type:
                raise ValueError(
                    f"ProbeSet was fitted on model_type {self.model_type!r} but read() received "
                    f"{live_model_type!r}."
                )
        elif session is None:
            raise ValueError("ProbeSet.read() requires a live model or a capture-capable session.")
        if self.location != "layer_input":
            raise ValueError(
                f"ProbeSet.read() serves the 'layer_input' boundary, but this set was fitted at "
                f"{self.location!r}; refit with location='layer_input', or score offline via "
                "layerwise_tokenwise_hidden and Probe.score_hidden."
            )

        device = next(model.parameters()).device if model is not None else torch.device("cpu")
        ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.as_tensor(input_ids, dtype=torch.long)
        if ids.ndim == 1:
            ids = ids.unsqueeze(0)
        ids = ids.to(device)

        if attention_mask is not None:
            mask = attention_mask if isinstance(attention_mask, torch.Tensor) else torch.as_tensor(attention_mask)
            if mask.ndim == 1:
                mask = mask.unsqueeze(0)
            mask = mask.to(device)
        else:
            mask = torch.ones_like(ids)

        if model is None:
            return self._read_via_session(session, ids, mask)

        layer_names = resolve_model_layout(model).layer_names
        for lid in self.layer_ids:
            if not 0 <= lid < len(layer_names):
                raise ValueError(f"probe layer {lid} out of range [0, {len(layer_names)}).")

        captured: dict[int, torch.Tensor] = {}

        def _pre_capture(layer_id: int):
            def _pre_hook(module, input_args, input_kwargs):
                hidden = input_args[0] if input_args else input_kwargs.get("hidden_states")
                if hidden is not None:
                    captured[layer_id] = hidden.detach().to("cpu", torch.float32)
                return input_args, input_kwargs

            return _pre_hook

        handles = []
        try:
            for lid in self.layer_ids:
                module = model.get_submodule(layer_names[lid])
                handles.append(module.register_forward_pre_hook(_pre_capture(lid), with_kwargs=True))
            with torch.no_grad(), auxiliary_pass(aligned=True):
                # explicit `cache_position` lets co-resident state hooks position this aligned
                # pass since transformers v5 threads the kwarg into decoder layers only when the
                # caller passes it
                model(
                    input_ids=ids,
                    attention_mask=mask,
                    use_cache=False,
                    cache_position=torch.arange(ids.size(1), device=ids.device),
                )
        finally:
            for handle in handles:
                handle.remove()

        cpu_mask = mask.to("cpu")
        scores: dict[str, torch.Tensor] = {}
        decisions: dict[str, torch.Tensor] = {}
        for name, probe in self.probes.items():
            features = {
                lid: aggregate_condition_hidden(captured[lid], probe.pooling, attention_mask=cpu_mask)
                for lid in probe.layer_ids
            }
            probe_scores = probe.decision_function(features)
            scores[name] = probe_scores
            decisions[name] = probe_scores >= 0

        readings = ProbeReadings(scores=scores, decisions=decisions)
        self.latest = readings
        return readings

    def _read_via_session(self, session, ids: torch.Tensor, mask: torch.Tensor) -> ProbeReadings:
        """Score through a capture-capable session's `capture` at the layer-input boundary."""
        from steerability.algorithms.core.execution.payloads import PreparedPrompt

        prompts = [
            PreparedPrompt.from_token_ids(ids[index:index + 1], mask[index:index + 1])
            for index in range(ids.size(0))
        ]
        result = session.capture(
            prompts, layers=list(self.layer_ids), mode="all_tokens", location="layer_input"
        )
        cpu_mask = result.attention_mask
        scores: dict[str, torch.Tensor] = {}
        decisions: dict[str, torch.Tensor] = {}
        for name, probe in self.probes.items():
            features = {
                lid: aggregate_condition_hidden(
                    result.hidden[lid].to(torch.float32), probe.pooling, attention_mask=cpu_mask
                )
                for lid in probe.layer_ids
            }
            probe_scores = probe.decision_function(features)
            scores[name] = probe_scores
            decisions[name] = probe_scores >= 0
        readings = ProbeReadings(scores=scores, decisions=decisions)
        self.latest = readings
        return readings

    def summary(self) -> dict[str, dict]:
        """Per-probe diagnostic table.

        Returns:
            Mapping from probe name to a dict with keys `"layer_ids"`, `"method"`, `"bias"`,
            `"calibration"`, `"f1"`, `"model_fingerprint"`, and `"stats_fingerprint"` (all but
            `"layer_ids"` and `"bias"` are None when the probe's `meta` does not record them).
        """
        return {
            name: {
                "layer_ids": list(probe.layer_ids),
                "method": probe.meta.get("method"),
                "bias": probe.bias,
                "calibration": probe.meta.get("calibration"),
                "f1": (probe.meta.get("calibration") or {}).get("f1"),
                "model_fingerprint": probe.meta.get("model_fingerprint"),
                "stats_fingerprint": probe.meta.get("stats_fingerprint"),
            }
            for name, probe in self.probes.items()
        }

latest = None instance-attribute

layer_ids = sorted({lid for probe in self.probes.values() for lid in probe.layer_ids}) instance-attribute

location property

The shared capture boundary of the probes.

model_type property

The shared model_type of the probes.

names = tuple(self.probes) instance-attribute

probes = dict(probes) instance-attribute

fit(model, tokenizer, *, data, spec=None, stats=None, calibration_data=None, session=None) classmethod

Fit one probe per name and return the set.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model the probes are fitted on.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the pairs.	required
data	Mapping[str, ContrastivePairs]	Mapping from probe name to the contrastive pairs it is fitted on.	required
spec	ProbeFitSpec | Mapping[str, ProbeFitSpec] | None	A shared ProbeFitSpec (None uses the defaults), or a mapping from probe name to a per-name spec.	None
stats	ActivationStats | None	Ambient activation statistics for methods that require them.	None
calibration_data	Mapping[str, ContrastivePairs] | None	Optional mapping from probe name to calibration pairs.	None

Returns:

Type	Description
ProbeSet	The fitted ProbeSet.

Raises:

Type	Description
ValueError	If data is empty, a spec mapping misses a name, or fit_probe raises for any probe.

Source code in steerability/algorithms/core/internals/probes/probe_set.py

@classmethod
def fit(
    cls,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: Mapping[str, ContrastivePairs],
    spec: ProbeFitSpec | Mapping[str, ProbeFitSpec] | None = None,
    stats: ActivationStats | None = None,
    calibration_data: Mapping[str, ContrastivePairs] | None = None,
    session=None,
) -> "ProbeSet":
    """Fit one probe per name and return the set.

    Args:
        model: Model the probes are fitted on.
        tokenizer: Tokenizer for rendering and encoding the pairs.
        data: Mapping from probe name to the contrastive pairs it is fitted on.
        spec: A shared `ProbeFitSpec` (None uses the defaults), or a mapping from probe name
            to a per-name spec.
        stats: Ambient activation statistics for methods that require them.
        calibration_data: Optional mapping from probe name to calibration pairs.

    Returns:
        The fitted `ProbeSet`.

    Raises:
        ValueError: If `data` is empty, a spec mapping misses a name, or `fit_probe` raises
            for any probe.
    """
    if not data:
        raise ValueError("data must map at least one probe name to contrastive pairs.")
    if isinstance(spec, Mapping):
        missing = [name for name in data if name not in spec]
        if missing:
            raise ValueError(f"spec mapping is missing entries for probe(s) {missing}.")

    probes: dict[str, Probe] = {}
    for name, pairs in data.items():
        probe_spec = spec[name] if isinstance(spec, Mapping) else (spec or ProbeFitSpec())
        probes[name] = fit_probe(
            model,
            tokenizer,
            data=pairs,
            spec=probe_spec,
            stats=stats,
            calibration_data=(calibration_data or {}).get(name),
            session=session,
        )
    return cls(probes)

read(model, input_ids, attention_mask=None, session=None)

Score a batch of prompts against every probe in one read-only forward.

Registers one forward pre-hook per layer in the union (the layer_input boundary, the only boundary read() serves), runs a single forward under torch.no_grad() with use_cache=False wrapped in auxiliary_pass(aligned=True), removes the hooks, then pools each probe's layers per its pooling with the attention mask and applies its decision function. The auxiliary marking means co-resident state controls' condition scorers, gates, and position counters ignore the pass, while "all"-scoped behavior transforms still apply, so scores measure the stream as deployed. A set fitted at "layer_output" raises, since a pre-hook capture would score features one block away from where those probes were fitted; score such probes offline via layerwise_tokenwise_hidden and Probe.score_hidden.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	The model to read (must match the probes' model_type).	required
input_ids	Tensor	Prompt token ids of shape [B, T] (a 1-D tensor is treated as one row).	required
attention_mask	Tensor | None	Prompt attention mask matching input_ids (1 at real tokens). When None, every position is treated as real.	None

Returns:

Type	Description
ProbeReadings	A ProbeReadings with per-probe signed scores and decisions, also stashed on latest.

Raises:

Type	Description
ValueError	If the model's model_type does not match the probes', the set's location is not "layer_input", a probe layer is out of range, or the architecture is unrecognized.

Source code in steerability/algorithms/core/internals/probes/probe_set.py

def read(
    self,
    model: PreTrainedModel | None,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    session=None,
) -> ProbeReadings:
    """Score a batch of prompts against every probe in one read-only forward.

    Registers one forward pre-hook per layer in the union (the `layer_input` boundary, the
    only boundary `read()` serves), runs a single forward under `torch.no_grad()` with
    `use_cache=False` wrapped in `auxiliary_pass(aligned=True)`, removes the hooks, then
    pools each probe's layers per its `pooling` with the attention mask and applies its
    decision function. The auxiliary marking means co-resident state controls' condition
    scorers, gates, and position counters ignore the pass, while `"all"`-scoped behavior
    transforms still apply, so scores measure the stream as deployed. A set fitted at
    `"layer_output"` raises, since a pre-hook capture would score features one block away
    from where those probes were fitted; score such probes offline via
    `layerwise_tokenwise_hidden` and `Probe.score_hidden`.

    Args:
        model: The model to read (must match the probes' `model_type`).
        input_ids: Prompt token ids of shape `[B, T]` (a 1-D tensor is treated as one row).
        attention_mask: Prompt attention mask matching `input_ids` (1 at real tokens). When
            None, every position is treated as real.

    Returns:
        A `ProbeReadings` with per-probe signed scores and decisions, also stashed on `latest`.

    Raises:
        ValueError: If the model's `model_type` does not match the probes', the set's
            `location` is not `"layer_input"`, a probe layer is out of range, or the
            architecture is unrecognized.
    """
    if model is not None:
        live_model_type = getattr(model.config, "model_type", "unknown")
        if live_model_type != self.model_type:
            raise ValueError(
                f"ProbeSet was fitted on model_type {self.model_type!r} but read() received "
                f"{live_model_type!r}."
            )
    elif session is None:
        raise ValueError("ProbeSet.read() requires a live model or a capture-capable session.")
    if self.location != "layer_input":
        raise ValueError(
            f"ProbeSet.read() serves the 'layer_input' boundary, but this set was fitted at "
            f"{self.location!r}; refit with location='layer_input', or score offline via "
            "layerwise_tokenwise_hidden and Probe.score_hidden."
        )

    device = next(model.parameters()).device if model is not None else torch.device("cpu")
    ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.as_tensor(input_ids, dtype=torch.long)
    if ids.ndim == 1:
        ids = ids.unsqueeze(0)
    ids = ids.to(device)

    if attention_mask is not None:
        mask = attention_mask if isinstance(attention_mask, torch.Tensor) else torch.as_tensor(attention_mask)
        if mask.ndim == 1:
            mask = mask.unsqueeze(0)
        mask = mask.to(device)
    else:
        mask = torch.ones_like(ids)

    if model is None:
        return self._read_via_session(session, ids, mask)

    layer_names = resolve_model_layout(model).layer_names
    for lid in self.layer_ids:
        if not 0 <= lid < len(layer_names):
            raise ValueError(f"probe layer {lid} out of range [0, {len(layer_names)}).")

    captured: dict[int, torch.Tensor] = {}

    def _pre_capture(layer_id: int):
        def _pre_hook(module, input_args, input_kwargs):
            hidden = input_args[0] if input_args else input_kwargs.get("hidden_states")
            if hidden is not None:
                captured[layer_id] = hidden.detach().to("cpu", torch.float32)
            return input_args, input_kwargs

        return _pre_hook

    handles = []
    try:
        for lid in self.layer_ids:
            module = model.get_submodule(layer_names[lid])
            handles.append(module.register_forward_pre_hook(_pre_capture(lid), with_kwargs=True))
        with torch.no_grad(), auxiliary_pass(aligned=True):
            # explicit `cache_position` lets co-resident state hooks position this aligned
            # pass since transformers v5 threads the kwarg into decoder layers only when the
            # caller passes it
            model(
                input_ids=ids,
                attention_mask=mask,
                use_cache=False,
                cache_position=torch.arange(ids.size(1), device=ids.device),
            )
    finally:
        for handle in handles:
            handle.remove()

    cpu_mask = mask.to("cpu")
    scores: dict[str, torch.Tensor] = {}
    decisions: dict[str, torch.Tensor] = {}
    for name, probe in self.probes.items():
        features = {
            lid: aggregate_condition_hidden(captured[lid], probe.pooling, attention_mask=cpu_mask)
            for lid in probe.layer_ids
        }
        probe_scores = probe.decision_function(features)
        scores[name] = probe_scores
        decisions[name] = probe_scores >= 0

    readings = ProbeReadings(scores=scores, decisions=decisions)
    self.latest = readings
    return readings

summary()

Per-probe diagnostic table.

Returns:

Type	Description
dict[str, dict]	Mapping from probe name to a dict with keys "layer_ids", "method", "bias",
dict[str, dict]	"calibration", "f1", "model_fingerprint", and "stats_fingerprint" (all but
dict[str, dict]	"layer_ids" and "bias" are None when the probe's meta does not record them).

Source code in steerability/algorithms/core/internals/probes/probe_set.py

def summary(self) -> dict[str, dict]:
    """Per-probe diagnostic table.

    Returns:
        Mapping from probe name to a dict with keys `"layer_ids"`, `"method"`, `"bias"`,
        `"calibration"`, `"f1"`, `"model_fingerprint"`, and `"stats_fingerprint"` (all but
        `"layer_ids"` and `"bias"` are None when the probe's `meta` does not record them).
    """
    return {
        name: {
            "layer_ids": list(probe.layer_ids),
            "method": probe.meta.get("method"),
            "bias": probe.bias,
            "calibration": probe.meta.get("calibration"),
            "f1": (probe.meta.get("calibration") or {}).get("f1"),
            "model_fingerprint": probe.meta.get("model_fingerprint"),
            "stats_fingerprint": probe.meta.get("stats_fingerprint"),
        }
        for name, probe in self.probes.items()
    }

ProbeSetFit dataclass

Deferred fitting recipe; holds every fitting input except the model.

fit(model, tokenizer) estimates stats first when it is a StatsSpec, fits one probe per name, and returns a ProbeSet. RoutedDecoding.steer() calls fit on the model the pipeline provides, so pipelines whose structural controls produce the final weights fit on those weights, and ProbeFitSpec fields stay grid-searchable through ControlSpec.

Attributes:

Name	Type	Description
data	Mapping[str, ContrastivePairs]	Mapping from probe name to the contrastive pairs it is fitted on.
spec	ProbeFitSpec | Mapping[str, ProbeFitSpec]	A shared ProbeFitSpec, or a mapping from probe name to a per-name spec covering every name in data.
stats	ActivationStats | StatsSpec | None	Ambient activation statistics, or a StatsSpec estimated on the model at fit time. Required when any effective spec's method is "lda" or "logreg".
calibration_data	Mapping[str, ContrastivePairs] | None	Optional mapping from probe name to calibration pairs; names absent from the mapping calibrate on their fit pairs.

Source code in steerability/algorithms/core/internals/probes/probe_set.py

@dataclass
class ProbeSetFit:
    """Deferred fitting recipe; holds every fitting input except the model.

    `fit(model, tokenizer)` estimates `stats` first when it is a `StatsSpec`, fits one probe
    per name, and returns a `ProbeSet`. `RoutedDecoding.steer()` calls `fit` on the model the
    pipeline provides, so pipelines whose structural controls produce the final weights fit on
    those weights, and `ProbeFitSpec` fields stay grid-searchable through `ControlSpec`.

    Attributes:
        data: Mapping from probe name to the contrastive pairs it is fitted on.
        spec: A shared `ProbeFitSpec`, or a mapping from probe name to a per-name spec covering
            every name in `data`.
        stats: Ambient activation statistics, or a `StatsSpec` estimated on the model at fit
            time. Required when any effective spec's method is `"lda"` or `"logreg"`.
        calibration_data: Optional mapping from probe name to calibration pairs; names absent
            from the mapping calibrate on their fit pairs.
    """

    data: Mapping[str, ContrastivePairs]
    spec: ProbeFitSpec | Mapping[str, ProbeFitSpec] = field(default_factory=ProbeFitSpec)
    stats: ActivationStats | StatsSpec | None = None
    calibration_data: Mapping[str, ContrastivePairs] | None = None

    def __post_init__(self):
        if not self.data:
            raise ValueError("data must map at least one probe name to contrastive pairs.")
        if isinstance(self.spec, Mapping):
            missing = [name for name in self.data if name not in self.spec]
            if missing:
                raise ValueError(f"spec mapping is missing entries for probe(s) {missing}.")
        if self.calibration_data is not None:
            unknown = [name for name in self.calibration_data if name not in self.data]
            if unknown:
                raise ValueError(
                    f"calibration_data names unknown probe(s) {unknown}; "
                    f"available: {sorted(self.data)}."
                )
        if self.stats is None:
            needy = [
                name for name in self.data
                if self._spec_for(name).method in ("lda", "logreg")
            ]
            if needy:
                raise ValueError(
                    f"probe(s) {needy} use a method that requires ambient activation "
                    "statistics; supply stats (an ActivationStats or a StatsSpec)."
                )

    def _spec_for(self, name: str) -> ProbeFitSpec:
        """The effective spec for one probe name."""
        return self.spec[name] if isinstance(self.spec, Mapping) else self.spec

    @property
    def names(self) -> tuple[str, ...]:
        """The probe names, available before fitting so routing rules validate at construction."""
        return tuple(self.data)

    def fit(self, model: PreTrainedModel | None, tokenizer: PreTrainedTokenizerBase, session=None) -> "ProbeSet":
        """Fit the recipe on `model`, resolving a `StatsSpec` on it first.

        Args:
            model: Model the probes are fitted on.
            tokenizer: Tokenizer for rendering and encoding the pairs.

        Returns:
            The fitted `ProbeSet`.
        """
        stats = self.stats
        if isinstance(stats, StatsSpec):
            stats = stats.estimate(model, tokenizer, session=session)
        return ProbeSet.fit(
            model,
            tokenizer,
            data=self.data,
            spec=self.spec,
            stats=stats,
            calibration_data=self.calibration_data,
            session=session,
        )

calibration_data = None class-attribute instance-attribute

data instance-attribute

names property

The probe names, available before fitting so routing rules validate at construction.

spec = field(default_factory=ProbeFitSpec) class-attribute instance-attribute

stats = None class-attribute instance-attribute

fit(model, tokenizer, session=None)

Fit the recipe on model, resolving a StatsSpec on it first.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model the probes are fitted on.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the pairs.	required

Returns:

Type	Description
ProbeSet	The fitted ProbeSet.

Source code in steerability/algorithms/core/internals/probes/probe_set.py

def fit(self, model: PreTrainedModel | None, tokenizer: PreTrainedTokenizerBase, session=None) -> "ProbeSet":
    """Fit the recipe on `model`, resolving a `StatsSpec` on it first.

    Args:
        model: Model the probes are fitted on.
        tokenizer: Tokenizer for rendering and encoding the pairs.

    Returns:
        The fitted `ProbeSet`.
    """
    stats = self.stats
    if isinstance(stats, StatsSpec):
        stats = stats.estimate(model, tokenizer, session=session)
    return ProbeSet.fit(
        model,
        tokenizer,
        data=self.data,
        spec=self.spec,
        stats=stats,
        calibration_data=self.calibration_data,
        session=session,
    )

calibrate_bias(pos_scores, neg_scores, calibration='max_f1')

Bias placing the operating point of raw_score + bias >= 0 at the chosen threshold.

The returned bias equals the negated threshold. "midpoint" places the threshold halfway between the class score means; "max_f1" picks the F1-maximizing threshold over the midpoints of consecutive distinct observed scores, breaking F1 ties toward the largest symmetric margin min(pos.min() - t, t - neg.max()) (predictions score >= threshold; a single distinct observed score is used as the threshold directly); ("target_fpr", x) places it at the (1 - x)-quantile of the negative scores. Calibration never repairs polarity; inverted inputs raise instead.

Parameters:

Name	Type	Description	Default
pos_scores	Tensor	Raw scores of positive examples, shape [N_pos].	required
neg_scores	Tensor	Raw scores of negative examples, shape [N_neg].	required
calibration	CalibrationSpec	The operating-point rule.	'max_f1'

Returns:

Type	Description
float	The bias b such that the decision is raw_score + b >= 0.

Raises:

Type	Description
ValueError	If either score tensor is empty, calibration is malformed, or mean(pos_scores) < mean(neg_scores) ("scores are inverted; negate your direction").

Source code in steerability/algorithms/core/internals/probes/fitting.py

def calibrate_bias(
    pos_scores: torch.Tensor,
    neg_scores: torch.Tensor,
    calibration: CalibrationSpec = "max_f1",
) -> float:
    """Bias placing the operating point of `raw_score + bias >= 0` at the chosen threshold.

    The returned bias equals the negated threshold. `"midpoint"` places the threshold halfway
    between the class score means; `"max_f1"` picks the F1-maximizing threshold over the midpoints
    of consecutive distinct observed scores, breaking F1 ties toward the largest symmetric margin
    `min(pos.min() - t, t - neg.max())` (predictions `score >= threshold`; a single distinct
    observed score is used as the threshold directly);
    `("target_fpr", x)` places it at the `(1 - x)`-quantile of the negative scores. Calibration
    never repairs polarity; inverted inputs raise instead.

    Args:
        pos_scores: Raw scores of positive examples, shape `[N_pos]`.
        neg_scores: Raw scores of negative examples, shape `[N_neg]`.
        calibration: The operating-point rule.

    Returns:
        The bias `b` such that the decision is `raw_score + b >= 0`.

    Raises:
        ValueError: If either score tensor is empty, `calibration` is malformed, or
            `mean(pos_scores) < mean(neg_scores)` ("scores are inverted; negate your direction").
    """
    _validate_calibration(calibration)
    pos = torch.as_tensor(pos_scores, dtype=torch.float32).reshape(-1)
    neg = torch.as_tensor(neg_scores, dtype=torch.float32).reshape(-1)
    if pos.numel() == 0 or neg.numel() == 0:
        raise ValueError("pos_scores and neg_scores must each contain at least one score.")
    if pos.mean() < neg.mean():
        raise ValueError("scores are inverted; negate your direction.")

    if calibration == "midpoint":
        threshold = float((pos.mean() + neg.mean()) / 2)
    elif calibration == "max_f1":
        values = torch.cat([pos, neg]).unique(sorted=True)
        if values.numel() == 1:
            threshold = float(values[0])
        else:
            pos_min, neg_max = float(pos.min()), float(neg.max())
            best = (-1.0, -float("inf"))
            threshold = float(values[:2].mean())
            for t in ((values[:-1] + values[1:]) / 2).tolist():
                f1 = _f1_at(pos, neg, t)
                margin = min(pos_min - t, t - neg_max)
                if (f1, margin) > best:
                    best, threshold = (f1, margin), t
    else:  # ("target_fpr", x)
        target = float(calibration[1])
        threshold = float(torch.quantile(neg, 1.0 - target))

    return -threshold

evaluate_probe(probe, model, tokenizer, data, *, prompt_format='raw', batch_size=8, max_length=None, session=None)

Score labeled data against a fitted probe in the probe's recorded space.

Renders data per prompt_format, captures at the probe's location, pools per the probe's pooling over the probe's layers, and applies probe.decision_function. The probe is not refitted or recalibrated; the decision uses the probe's calibrated bias, so score >= 0 is positive at the calibrated point.

Parameters:

Name	Type	Description	Default
probe	Probe	The fitted Probe to score against.	required
model	PreTrainedModel | None	Model whose activations are captured, or None when a capture-capable session is provided.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the data.	required
data	ContrastivePairs | LabeledExamples	Contrastive pairs or unpaired labeled examples to score. The positives are scored as the positive class and the negatives as the negative class.	required
prompt_format	PromptFormat	How the data is rendered before tokenization. Must match the format the probe was fitted under for the scores to be comparable.	'raw'
batch_size	int	Chunk size for feature extraction.	8
max_length	int | None	Tokenization truncation bound. None truncates to the tokenizer's model maximum.	None
session		Optional capture-capable session used in place of a live model.	None

Returns:

Type	Description
ProbeEvaluation	A ProbeEvaluation with the per-class scores, accuracy, and F1.

Source code in steerability/algorithms/core/internals/probes/fitting.py

def evaluate_probe(
    probe: Probe,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    data: ContrastivePairs | LabeledExamples,
    *,
    prompt_format: PromptFormat = "raw",
    batch_size: int = 8,
    max_length: int | None = None,
    session=None,
) -> ProbeEvaluation:
    """Score labeled data against a fitted probe in the probe's recorded space.

    Renders `data` per `prompt_format`, captures at the probe's `location`, pools per the probe's
    `pooling` over the probe's layers, and applies `probe.decision_function`. The probe is not
    refitted or recalibrated; the decision uses the probe's calibrated bias, so `score >= 0` is
    positive at the calibrated point.

    Args:
        probe: The fitted `Probe` to score against.
        model: Model whose activations are captured, or None when a capture-capable `session` is
            provided.
        tokenizer: Tokenizer for rendering and encoding the data.
        data: Contrastive pairs or unpaired labeled examples to score. The positives are scored as
            the positive class and the negatives as the negative class.
        prompt_format: How the data is rendered before tokenization. Must match the format the
            probe was fitted under for the scores to be comparable.
        batch_size: Chunk size for feature extraction.
        max_length: Tokenization truncation bound. None truncates to the tokenizer's model maximum.
        session: Optional capture-capable session used in place of a live model.

    Returns:
        A `ProbeEvaluation` with the per-class scores, accuracy, and F1.
    """
    spec = ProbeFitSpec(
        method="mean_diff",
        pooling=probe.pooling,
        location=probe.location,
        prompt_format=prompt_format,
        candidate_layers=list(probe.layer_ids),
        calibration="midpoint",
    )
    pos_features, neg_features = _pooled_features(
        model, tokenizer, data, spec, list(probe.layer_ids),
        batch_size=batch_size, max_length=max_length, session=session,
    )
    pos_scores = probe.decision_function(pos_features)
    neg_scores = probe.decision_function(neg_features)

    correct = int((pos_scores >= 0).sum()) + int((neg_scores < 0).sum())
    accuracy = correct / (pos_scores.numel() + neg_scores.numel())
    f1 = _f1_at(pos_scores, neg_scores, 0.0)
    return ProbeEvaluation(
        positive_scores=pos_scores, negative_scores=neg_scores, accuracy=accuracy, f1=f1
    )

fit_probe(model, tokenizer, *, data, spec=ProbeFitSpec(), stats=None, calibration_data=None, allow_model_mismatch=False, batch_size=8, max_length=None, session=None)

Fit a calibrated single-layer Probe from contrastive data.

The data is rendered per spec.prompt_format, tokenized, captured at spec.location (inside auxiliary_pass(aligned=True)), and pooled per spec.pooling. Extraction runs in chunks of batch_size, pooling each chunk immediately and retaining features only for the candidate layers, so memory holds pooled [N, H] features rather than per-token states. Candidate layers are swept; for each, a direction is fitted per spec.method, oriented on fit-set scores (weights are negated when the mean positive score falls below the mean negative score, which can occur only for "logreg", since for "mean_diff", "lda", and "fisher" the fit-set score gap is nonnegative by construction), and the bias is calibrated via calibrate_bias. The layer with the best F1 at the calibrated point is kept, with the calibration score gap in pooled within-class standard deviations as tie-break, and the full sweep is recorded in probe.meta["layer_sweep"].

Calibration scores come from calibration_data when supplied, else from data. The split lets the contrast be fitted on discriminative pairs while the operating point is calibrated on a covering set; orientation is never re-evaluated on calibration_data, and a calibration set whose scores contradict the fit is caught by the inversion raise in calibrate_bias.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model whose activations are captured.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the data.	required
data	ContrastivePairs | LabeledExamples	Contrastive pairs or unpaired labeled examples the direction is fitted on. Unpaired classes require spec.prompt_format of "raw" or "chat_prompt".	required
spec	ProbeFitSpec	Fitting configuration.	ProbeFitSpec()
stats	ActivationStats | None	Ambient activation statistics, required by "lda" and "logreg" (the standardization is folded into the stored weights). Ignored for the "mean_diff" and "fisher" directions.	None
calibration_data	ContrastivePairs | LabeledExamples | None	Optional data the operating point is calibrated on.	None
allow_model_mismatch	bool	When True, a stats artifact estimated on a different model is accepted instead of raising.	False
batch_size	int	Chunk size for feature extraction.	8
max_length	int | None	Tokenization truncation bound for feature extraction. None truncates to the tokenizer's model maximum length.	None

Returns:

Type	Description
Probe	A fitted Probe with canonical polarity, a single chosen layer, and a provenance
Probe	meta record.

Raises:

Type	Description
ValueError	If the method requires stats and none is supplied, a supplied stats was estimated on a different model (without allow_model_mismatch) or at a different location than spec.location, a requested candidate layer is out of range or has no recorded statistics, spec.prompt_format is "chat_completion" with unpaired data, or the calibration scores are inverted.

Source code in steerability/algorithms/core/internals/probes/fitting.py

def fit_probe(
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: ContrastivePairs | LabeledExamples,
    spec: ProbeFitSpec = ProbeFitSpec(),
    stats: ActivationStats | None = None,
    calibration_data: ContrastivePairs | LabeledExamples | None = None,
    allow_model_mismatch: bool = False,
    batch_size: int = 8,
    max_length: int | None = None,
    session=None,
) -> Probe:
    """Fit a calibrated single-layer `Probe` from contrastive data.

    The data is rendered per `spec.prompt_format`, tokenized, captured at `spec.location`
    (inside `auxiliary_pass(aligned=True)`), and pooled per `spec.pooling`. Extraction runs in
    chunks of `batch_size`, pooling each chunk immediately and retaining features only for the
    candidate layers, so memory holds pooled `[N, H]` features rather than per-token states.
    Candidate layers are swept; for each, a direction is fitted per `spec.method`, oriented on
    fit-set scores (weights are negated when the mean positive score falls below the mean
    negative score, which can occur only for `"logreg"`, since for `"mean_diff"`, `"lda"`, and
    `"fisher"` the fit-set score gap is nonnegative by construction), and the bias is
    calibrated via `calibrate_bias`. The layer with the best F1 at the calibrated point is
    kept, with the calibration score gap in pooled within-class standard deviations as
    tie-break, and the full sweep is recorded in `probe.meta["layer_sweep"]`.

    Calibration scores come from `calibration_data` when supplied, else from `data`. The split
    lets the contrast be fitted on discriminative pairs while the operating point is calibrated
    on a covering set; orientation is never re-evaluated on `calibration_data`, and a calibration
    set whose scores contradict the fit is caught by the inversion raise in `calibrate_bias`.

    Args:
        model: Model whose activations are captured.
        tokenizer: Tokenizer for rendering and encoding the data.
        data: Contrastive pairs or unpaired labeled examples the direction is fitted on.
            Unpaired classes require `spec.prompt_format` of `"raw"` or `"chat_prompt"`.
        spec: Fitting configuration.
        stats: Ambient activation statistics, required by `"lda"` and `"logreg"` (the
            standardization is folded into the stored weights). Ignored for the `"mean_diff"`
            and `"fisher"` directions.
        calibration_data: Optional data the operating point is calibrated on.
        allow_model_mismatch: When True, a `stats` artifact estimated on a different model is
            accepted instead of raising.
        batch_size: Chunk size for feature extraction.
        max_length: Tokenization truncation bound for feature extraction. None truncates to the
            tokenizer's model maximum length.

    Returns:
        A fitted `Probe` with canonical polarity, a single chosen layer, and a provenance
        `meta` record.

    Raises:
        ValueError: If the method requires `stats` and none is supplied, a supplied `stats` was
            estimated on a different model (without `allow_model_mismatch`) or at a different
            `location` than `spec.location`, a requested candidate layer is out of range or has
            no recorded statistics, `spec.prompt_format` is `"chat_completion"` with unpaired
            data, or the calibration scores are inverted.
    """
    if spec.method in ("lda", "logreg") and stats is None:
        raise ValueError(
            f"method={spec.method!r} requires ambient activation statistics; estimate "
            "ActivationStats once per model; see core.internals.stats."
        )

    if model is not None:
        fingerprint = model_fingerprint(model)
        fitted_model_type = getattr(model.config, "model_type", "unknown")
        session_meta: dict = {}
    else:
        fitted_model_type, session_meta = session_artifact_identity(session)
        fingerprint = session_meta.get("model_fingerprint")
    if (
        stats is not None
        and fingerprint is not None
        and stats.model_fingerprint != fingerprint
        and not allow_model_mismatch
    ):
        raise ValueError(
            "stats were estimated on a different model (fingerprint "
            f"{stats.model_fingerprint!r} vs {fingerprint!r}); whitening with another model's "
            "statistics produces a miscalibrated probe. Re-estimate ActivationStats on this "
            "model, or pass allow_model_mismatch=True."
        )
    if stats is not None and stats.location != spec.location:
        raise ValueError(
            f"stats were estimated at location {stats.location!r} but the fit captures at "
            f"{spec.location!r}; standardizing features with moments from a different boundary "
            "produces a miscalibrated probe. Re-estimate ActivationStats at the fit location."
        )

    num_layers = _resolve_num_layers(model, session)
    if spec.candidate_layers is not None:
        candidates = [int(lid) for lid in dict.fromkeys(spec.candidate_layers)]
        for lid in candidates:
            if not 0 <= lid < num_layers:
                raise ValueError(f"candidate layer {lid} out of range [0, {num_layers}).")
    elif spec.layer_range is not None:
        lo, hi = spec.layer_range
        candidates = [lid for lid in range(num_layers) if lo <= lid / max(num_layers - 1, 1) <= hi]
        if not candidates:
            raise ValueError(f"layer_range ({lo}, {hi}) selects no layers out of {num_layers}.")
    else:
        candidates = list(range(num_layers))

    if spec.method in ("lda", "logreg"):
        uncovered = [lid for lid in candidates if lid not in stats.mean]
        if uncovered:
            raise ValueError(
                f"stats record no statistics for candidate layer(s) {uncovered}; estimate "
                f"ActivationStats over these layers, or restrict candidate_layers to "
                f"{sorted(stats.mean)}."
            )

    pos_features, neg_features = _pooled_features(
        model, tokenizer, data, spec, candidates,
        batch_size=batch_size, max_length=max_length, session=session,
    )
    if calibration_data is not None:
        cal_pos_features, cal_neg_features = _pooled_features(
            model, tokenizer, calibration_data, spec, candidates,
            batch_size=batch_size, max_length=max_length, session=session,
        )
    else:
        cal_pos_features, cal_neg_features = pos_features, neg_features

    sweep: list[dict] = []
    best: dict | None = None
    for lid in candidates:
        w = _fit_direction(pos_features[lid], neg_features[lid], lid, spec, stats)

        fit_pos_scores = pos_features[lid] @ w
        fit_neg_scores = neg_features[lid] @ w
        fit_gap = float(fit_pos_scores.mean() - fit_neg_scores.mean())
        flipped = fit_gap < 0
        if flipped:
            assert spec.method == "logreg", (
                "fit-set orientation cannot invert for mean_diff, lda, or fisher; the score gap "
                "is a squared class-mean difference in the fit metric."
            )
            w = -w

        cal_pos_scores = cal_pos_features[lid] @ w
        cal_neg_scores = cal_neg_features[lid] @ w
        if float(cal_pos_scores.mean()) < float(cal_neg_scores.mean()):
            # the fit-set direction does not generalize to the calibration pool at
            # this layer; record the layer as unusable and keep sweeping rather than
            # aborting the whole fit (calibrate_bias would raise on these scores)
            sweep.append({
                "layer_id": lid,
                "f1": None,
                "margin": None,
                "gap": None,
                "fpr": None,
                "orientation_flipped": flipped,
                "calibration_inverted": True,
            })
            continue

        bias = calibrate_bias(cal_pos_scores, cal_neg_scores, spec.calibration)

        f1 = _f1_at(cal_pos_scores, cal_neg_scores, -bias)
        fpr = float((cal_neg_scores >= -bias).float().mean())
        gap = float(cal_pos_scores.min() - cal_neg_scores.max())
        margin = gap / _pooled_std(cal_pos_scores, cal_neg_scores)
        sweep.append({
            "layer_id": lid,
            "f1": f1,
            "margin": margin,
            "gap": gap,
            "fpr": fpr,
            "orientation_flipped": flipped,
            "calibration_inverted": False,
        })
        if best is None or (f1, margin) > (best["f1"], best["margin"]):
            best = {"layer_id": lid, "weights": w, "bias": bias, "f1": f1, "margin": margin,
                    "fpr": fpr, "orientation_flipped": flipped}

    if best is None:
        inverted = [entry["layer_id"] for entry in sweep if entry.get("calibration_inverted")]
        raise ValueError(
            f"Calibration scores are inverted at every candidate layer ({inverted}); the "
            "calibration pairs contradict the fit pairs for this concept. Revisit the "
            "pools, or the spec's pooling / layer_range."
    )

    calibration_record = {
        "kind": list(spec.calibration) if isinstance(spec.calibration, tuple) else spec.calibration,
        "on": "calibration_data" if calibration_data is not None else "data",
        "f1": best["f1"],
        "fpr": best["fpr"],
    }
    meta = {
        "method": spec.method,
        "pooling": spec.pooling,
        "location": spec.location,
        "prompt_format": spec.prompt_format,
        "n_pos": len(data.positives),
        "n_neg": len(data.negatives),
        "calibration": calibration_record,
        "layer_id": best["layer_id"],
        "layer_sweep": sweep,
        "orientation_flipped": best["orientation_flipped"],
        "stats_used": spec.method in ("lda", "logreg"),
        "model_fingerprint": fingerprint,
        "package_version": _PACKAGE_VERSION,
        "polarity": POLARITY_MARKER,
    }
    if model is not None:
        provenance = artifact_provenance_meta(model, tokenizer)
        for key in ("config_fingerprint", "chat_template_fingerprint"):
            if key in provenance:
                meta[key] = provenance[key]
    elif "model_ref" in session_meta:
        meta["model_ref"] = session_meta["model_ref"]
    if meta["stats_used"]:
        meta["stats_fingerprint"] = stats.fingerprint()

    return Probe(
        model_type=fitted_model_type,
        location=spec.location,
        pooling=spec.pooling,
        layer_ids=[best["layer_id"]],
        weights={best["layer_id"]: best["weights"]},
        bias=best["bias"],
        meta=meta,
    )

fitting

Probe fitting: contrastive feature extraction, direction estimation, and bias calibration.

CalibrationSpec = Literal['max_f1', 'midpoint'] | tuple[Literal['target_fpr'], float] module-attribute

logger = logging.getLogger(__name__) module-attribute

ProbeEvaluation dataclass

Held-out scores of a probe on labeled data.

Attributes:

Name	Type	Description
positive_scores	Tensor	Signed decision scores of the positive class, shape [N_pos], ordered by descending rendered text length rather than input order.
negative_scores	Tensor	Signed decision scores of the negative class, shape [N_neg], in the same order.
accuracy	float	Fraction correct at the calibrated point (score >= 0 is positive).
f1	float	F1 of the same decision.

Source code in steerability/algorithms/core/internals/probes/fitting.py

@dataclass(frozen=True, slots=True)
class ProbeEvaluation:
    """Held-out scores of a probe on labeled data.

    Attributes:
        positive_scores: Signed decision scores of the positive class, shape `[N_pos]`, ordered by
            descending rendered text length rather than input order.
        negative_scores: Signed decision scores of the negative class, shape `[N_neg]`, in the same
            order.
        accuracy: Fraction correct at the calibrated point (`score >= 0` is positive).
        f1: F1 of the same decision.
    """

    positive_scores: torch.Tensor
    negative_scores: torch.Tensor
    accuracy: float
    f1: float

accuracy instance-attribute

f1 instance-attribute

negative_scores instance-attribute

positive_scores instance-attribute

ProbeFitSpec dataclass

Configuration for fitting a Probe from contrastive pairs.

Attributes:

Name	Type	Description
method	Literal['lda', 'mean_diff', 'logreg', 'fisher']	Direction estimation method. "lda" (default) computes the difference in class means on standardized features (diagonal LDA) and requires stats; "logreg" is L2 logistic regression on standardized features and also requires stats; "mean_diff" computes the raw difference in class means and never consults stats; "fisher" is a full-covariance discriminant computed from the class features alone (the pooled within-class covariance, pseudo-inverted with singular values truncated at 1e-6, applied to the class-mean difference) and never consults stats. The truncated pseudo-inverse handles rank deficiency when sample counts are below the hidden size. "fisher" weights are unit-normalized; the scale is part of the method's contract, since downstream consumers score raw margins.
pooling	Literal['mean', 'last']	Token aggregation for feature extraction ("mean" or "last", mask-aware).
location	str	Residual-stream boundary features are captured at.
prompt_format	PromptFormat	How pairs are rendered into model-ready text before tokenization.
candidate_layers	Sequence[int] | None	Explicit layer ids to sweep. None sweeps all decoder layers, optionally windowed by layer_range.
layer_range	tuple[float, float] | None	Fractional-depth window (lo, hi) with 0 <= lo < hi <= 1, applied to the full layer count when candidate_layers is None.
calibration	CalibrationSpec	Operating-point rule folded into the bias (see calibrate_bias).
C	float	Inverse regularization strength for "logreg".
seed	int	Random seed forwarded to "logreg".

Source code in steerability/algorithms/core/internals/probes/fitting.py

@dataclass(frozen=True)
class ProbeFitSpec:
    """Configuration for fitting a `Probe` from contrastive pairs.

    Attributes:
        method: Direction estimation method. `"lda"` (default) computes the difference in class
            means on standardized features (diagonal LDA) and requires `stats`; `"logreg"` is L2
            logistic regression on standardized features and also requires `stats`;
            `"mean_diff"` computes the raw difference in class means and never consults `stats`;
            `"fisher"` is a full-covariance discriminant computed from the class features alone
            (the pooled within-class covariance, pseudo-inverted with singular values truncated
            at `1e-6`, applied to the class-mean difference) and never consults `stats`. The
            truncated pseudo-inverse handles rank deficiency when sample counts are below the
            hidden size. `"fisher"` weights are unit-normalized; the scale is part of the
            method's contract, since downstream consumers score raw margins.
        pooling: Token aggregation for feature extraction (`"mean"` or `"last"`, mask-aware).
        location: Residual-stream boundary features are captured at.
        prompt_format: How pairs are rendered into model-ready text before tokenization.
        candidate_layers: Explicit layer ids to sweep. None sweeps all decoder layers, optionally
            windowed by `layer_range`.
        layer_range: Fractional-depth window `(lo, hi)` with `0 <= lo < hi <= 1`, applied to the
            full layer count when `candidate_layers` is None.
        calibration: Operating-point rule folded into the bias (see `calibrate_bias`).
        C: Inverse regularization strength for `"logreg"`.
        seed: Random seed forwarded to `"logreg"`.
    """

    method: Literal["lda", "mean_diff", "logreg", "fisher"] = "lda"
    pooling: Literal["mean", "last"] = "last"
    location: str = "layer_input"
    prompt_format: PromptFormat = "chat_prompt"
    candidate_layers: Sequence[int] | None = None
    layer_range: tuple[float, float] | None = None
    calibration: CalibrationSpec = "max_f1"
    C: float = 1.0
    seed: int = 0

    def __post_init__(self):
        if self.method not in ("lda", "mean_diff", "logreg", "fisher"):
            raise ValueError(
                f"method must be 'lda', 'mean_diff', 'logreg', or 'fisher', got {self.method!r}."
            )
        if self.pooling not in ("mean", "last"):
            raise ValueError(f"pooling must be 'mean' or 'last', got {self.pooling!r}.")
        if self.location not in ("layer_input", "layer_output"):
            raise ValueError(
                f"location must be 'layer_input' or 'layer_output', got {self.location!r}."
            )
        if self.prompt_format not in ("raw", "chat_completion", "chat_prompt"):
            raise ValueError(
                f"prompt_format must be one of raw/chat_completion/chat_prompt, "
                f"got {self.prompt_format!r}."
            )
        if self.candidate_layers is not None and len(self.candidate_layers) == 0:
            raise ValueError("candidate_layers must name at least one layer when supplied.")
        if self.layer_range is not None:
            lo, hi = self.layer_range
            if not 0.0 <= lo < hi <= 1.0:
                raise ValueError(
                    f"layer_range must satisfy 0 <= lo < hi <= 1, got ({lo}, {hi})."
                )
        if self.C <= 0:
            raise ValueError(f"C must be positive, got {self.C}.")
        _validate_calibration(self.calibration)

C = 1.0 class-attribute instance-attribute

calibration = 'max_f1' class-attribute instance-attribute

candidate_layers = None class-attribute instance-attribute

layer_range = None class-attribute instance-attribute

location = 'layer_input' class-attribute instance-attribute

method = 'lda' class-attribute instance-attribute

pooling = 'last' class-attribute instance-attribute

prompt_format = 'chat_prompt' class-attribute instance-attribute

seed = 0 class-attribute instance-attribute

calibrate_bias(pos_scores, neg_scores, calibration='max_f1')

Bias placing the operating point of raw_score + bias >= 0 at the chosen threshold.

The returned bias equals the negated threshold. "midpoint" places the threshold halfway between the class score means; "max_f1" picks the F1-maximizing threshold over the midpoints of consecutive distinct observed scores, breaking F1 ties toward the largest symmetric margin min(pos.min() - t, t - neg.max()) (predictions score >= threshold; a single distinct observed score is used as the threshold directly); ("target_fpr", x) places it at the (1 - x)-quantile of the negative scores. Calibration never repairs polarity; inverted inputs raise instead.

Parameters:

Name	Type	Description	Default
pos_scores	Tensor	Raw scores of positive examples, shape [N_pos].	required
neg_scores	Tensor	Raw scores of negative examples, shape [N_neg].	required
calibration	CalibrationSpec	The operating-point rule.	'max_f1'

Returns:

Type	Description
float	The bias b such that the decision is raw_score + b >= 0.

Raises:

Type	Description
ValueError	If either score tensor is empty, calibration is malformed, or mean(pos_scores) < mean(neg_scores) ("scores are inverted; negate your direction").

Source code in steerability/algorithms/core/internals/probes/fitting.py

def calibrate_bias(
    pos_scores: torch.Tensor,
    neg_scores: torch.Tensor,
    calibration: CalibrationSpec = "max_f1",
) -> float:
    """Bias placing the operating point of `raw_score + bias >= 0` at the chosen threshold.

    The returned bias equals the negated threshold. `"midpoint"` places the threshold halfway
    between the class score means; `"max_f1"` picks the F1-maximizing threshold over the midpoints
    of consecutive distinct observed scores, breaking F1 ties toward the largest symmetric margin
    `min(pos.min() - t, t - neg.max())` (predictions `score >= threshold`; a single distinct
    observed score is used as the threshold directly);
    `("target_fpr", x)` places it at the `(1 - x)`-quantile of the negative scores. Calibration
    never repairs polarity; inverted inputs raise instead.

    Args:
        pos_scores: Raw scores of positive examples, shape `[N_pos]`.
        neg_scores: Raw scores of negative examples, shape `[N_neg]`.
        calibration: The operating-point rule.

    Returns:
        The bias `b` such that the decision is `raw_score + b >= 0`.

    Raises:
        ValueError: If either score tensor is empty, `calibration` is malformed, or
            `mean(pos_scores) < mean(neg_scores)` ("scores are inverted; negate your direction").
    """
    _validate_calibration(calibration)
    pos = torch.as_tensor(pos_scores, dtype=torch.float32).reshape(-1)
    neg = torch.as_tensor(neg_scores, dtype=torch.float32).reshape(-1)
    if pos.numel() == 0 or neg.numel() == 0:
        raise ValueError("pos_scores and neg_scores must each contain at least one score.")
    if pos.mean() < neg.mean():
        raise ValueError("scores are inverted; negate your direction.")

    if calibration == "midpoint":
        threshold = float((pos.mean() + neg.mean()) / 2)
    elif calibration == "max_f1":
        values = torch.cat([pos, neg]).unique(sorted=True)
        if values.numel() == 1:
            threshold = float(values[0])
        else:
            pos_min, neg_max = float(pos.min()), float(neg.max())
            best = (-1.0, -float("inf"))
            threshold = float(values[:2].mean())
            for t in ((values[:-1] + values[1:]) / 2).tolist():
                f1 = _f1_at(pos, neg, t)
                margin = min(pos_min - t, t - neg_max)
                if (f1, margin) > best:
                    best, threshold = (f1, margin), t
    else:  # ("target_fpr", x)
        target = float(calibration[1])
        threshold = float(torch.quantile(neg, 1.0 - target))

    return -threshold

evaluate_probe(probe, model, tokenizer, data, *, prompt_format='raw', batch_size=8, max_length=None, session=None)

Score labeled data against a fitted probe in the probe's recorded space.

Renders data per prompt_format, captures at the probe's location, pools per the probe's pooling over the probe's layers, and applies probe.decision_function. The probe is not refitted or recalibrated; the decision uses the probe's calibrated bias, so score >= 0 is positive at the calibrated point.

Parameters:

Name	Type	Description	Default
probe	Probe	The fitted Probe to score against.	required
model	PreTrainedModel | None	Model whose activations are captured, or None when a capture-capable session is provided.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the data.	required
data	ContrastivePairs | LabeledExamples	Contrastive pairs or unpaired labeled examples to score. The positives are scored as the positive class and the negatives as the negative class.	required
prompt_format	PromptFormat	How the data is rendered before tokenization. Must match the format the probe was fitted under for the scores to be comparable.	'raw'
batch_size	int	Chunk size for feature extraction.	8
max_length	int | None	Tokenization truncation bound. None truncates to the tokenizer's model maximum.	None
session		Optional capture-capable session used in place of a live model.	None

Returns:

Type	Description
ProbeEvaluation	A ProbeEvaluation with the per-class scores, accuracy, and F1.

Source code in steerability/algorithms/core/internals/probes/fitting.py

def evaluate_probe(
    probe: Probe,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    data: ContrastivePairs | LabeledExamples,
    *,
    prompt_format: PromptFormat = "raw",
    batch_size: int = 8,
    max_length: int | None = None,
    session=None,
) -> ProbeEvaluation:
    """Score labeled data against a fitted probe in the probe's recorded space.

    Renders `data` per `prompt_format`, captures at the probe's `location`, pools per the probe's
    `pooling` over the probe's layers, and applies `probe.decision_function`. The probe is not
    refitted or recalibrated; the decision uses the probe's calibrated bias, so `score >= 0` is
    positive at the calibrated point.

    Args:
        probe: The fitted `Probe` to score against.
        model: Model whose activations are captured, or None when a capture-capable `session` is
            provided.
        tokenizer: Tokenizer for rendering and encoding the data.
        data: Contrastive pairs or unpaired labeled examples to score. The positives are scored as
            the positive class and the negatives as the negative class.
        prompt_format: How the data is rendered before tokenization. Must match the format the
            probe was fitted under for the scores to be comparable.
        batch_size: Chunk size for feature extraction.
        max_length: Tokenization truncation bound. None truncates to the tokenizer's model maximum.
        session: Optional capture-capable session used in place of a live model.

    Returns:
        A `ProbeEvaluation` with the per-class scores, accuracy, and F1.
    """
    spec = ProbeFitSpec(
        method="mean_diff",
        pooling=probe.pooling,
        location=probe.location,
        prompt_format=prompt_format,
        candidate_layers=list(probe.layer_ids),
        calibration="midpoint",
    )
    pos_features, neg_features = _pooled_features(
        model, tokenizer, data, spec, list(probe.layer_ids),
        batch_size=batch_size, max_length=max_length, session=session,
    )
    pos_scores = probe.decision_function(pos_features)
    neg_scores = probe.decision_function(neg_features)

    correct = int((pos_scores >= 0).sum()) + int((neg_scores < 0).sum())
    accuracy = correct / (pos_scores.numel() + neg_scores.numel())
    f1 = _f1_at(pos_scores, neg_scores, 0.0)
    return ProbeEvaluation(
        positive_scores=pos_scores, negative_scores=neg_scores, accuracy=accuracy, f1=f1
    )

fit_probe(model, tokenizer, *, data, spec=ProbeFitSpec(), stats=None, calibration_data=None, allow_model_mismatch=False, batch_size=8, max_length=None, session=None)

Fit a calibrated single-layer Probe from contrastive data.

The data is rendered per spec.prompt_format, tokenized, captured at spec.location (inside auxiliary_pass(aligned=True)), and pooled per spec.pooling. Extraction runs in chunks of batch_size, pooling each chunk immediately and retaining features only for the candidate layers, so memory holds pooled [N, H] features rather than per-token states. Candidate layers are swept; for each, a direction is fitted per spec.method, oriented on fit-set scores (weights are negated when the mean positive score falls below the mean negative score, which can occur only for "logreg", since for "mean_diff", "lda", and "fisher" the fit-set score gap is nonnegative by construction), and the bias is calibrated via calibrate_bias. The layer with the best F1 at the calibrated point is kept, with the calibration score gap in pooled within-class standard deviations as tie-break, and the full sweep is recorded in probe.meta["layer_sweep"].

Calibration scores come from calibration_data when supplied, else from data. The split lets the contrast be fitted on discriminative pairs while the operating point is calibrated on a covering set; orientation is never re-evaluated on calibration_data, and a calibration set whose scores contradict the fit is caught by the inversion raise in calibrate_bias.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model whose activations are captured.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the data.	required
data	ContrastivePairs | LabeledExamples	Contrastive pairs or unpaired labeled examples the direction is fitted on. Unpaired classes require spec.prompt_format of "raw" or "chat_prompt".	required
spec	ProbeFitSpec	Fitting configuration.	ProbeFitSpec()
stats	ActivationStats | None	Ambient activation statistics, required by "lda" and "logreg" (the standardization is folded into the stored weights). Ignored for the "mean_diff" and "fisher" directions.	None
calibration_data	ContrastivePairs | LabeledExamples | None	Optional data the operating point is calibrated on.	None
allow_model_mismatch	bool	When True, a stats artifact estimated on a different model is accepted instead of raising.	False
batch_size	int	Chunk size for feature extraction.	8
max_length	int | None	Tokenization truncation bound for feature extraction. None truncates to the tokenizer's model maximum length.	None

Returns:

Type	Description
Probe	A fitted Probe with canonical polarity, a single chosen layer, and a provenance
Probe	meta record.

Raises:

Type	Description
ValueError	If the method requires stats and none is supplied, a supplied stats was estimated on a different model (without allow_model_mismatch) or at a different location than spec.location, a requested candidate layer is out of range or has no recorded statistics, spec.prompt_format is "chat_completion" with unpaired data, or the calibration scores are inverted.

Source code in steerability/algorithms/core/internals/probes/fitting.py

def fit_probe(
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: ContrastivePairs | LabeledExamples,
    spec: ProbeFitSpec = ProbeFitSpec(),
    stats: ActivationStats | None = None,
    calibration_data: ContrastivePairs | LabeledExamples | None = None,
    allow_model_mismatch: bool = False,
    batch_size: int = 8,
    max_length: int | None = None,
    session=None,
) -> Probe:
    """Fit a calibrated single-layer `Probe` from contrastive data.

    The data is rendered per `spec.prompt_format`, tokenized, captured at `spec.location`
    (inside `auxiliary_pass(aligned=True)`), and pooled per `spec.pooling`. Extraction runs in
    chunks of `batch_size`, pooling each chunk immediately and retaining features only for the
    candidate layers, so memory holds pooled `[N, H]` features rather than per-token states.
    Candidate layers are swept; for each, a direction is fitted per `spec.method`, oriented on
    fit-set scores (weights are negated when the mean positive score falls below the mean
    negative score, which can occur only for `"logreg"`, since for `"mean_diff"`, `"lda"`, and
    `"fisher"` the fit-set score gap is nonnegative by construction), and the bias is
    calibrated via `calibrate_bias`. The layer with the best F1 at the calibrated point is
    kept, with the calibration score gap in pooled within-class standard deviations as
    tie-break, and the full sweep is recorded in `probe.meta["layer_sweep"]`.

    Calibration scores come from `calibration_data` when supplied, else from `data`. The split
    lets the contrast be fitted on discriminative pairs while the operating point is calibrated
    on a covering set; orientation is never re-evaluated on `calibration_data`, and a calibration
    set whose scores contradict the fit is caught by the inversion raise in `calibrate_bias`.

    Args:
        model: Model whose activations are captured.
        tokenizer: Tokenizer for rendering and encoding the data.
        data: Contrastive pairs or unpaired labeled examples the direction is fitted on.
            Unpaired classes require `spec.prompt_format` of `"raw"` or `"chat_prompt"`.
        spec: Fitting configuration.
        stats: Ambient activation statistics, required by `"lda"` and `"logreg"` (the
            standardization is folded into the stored weights). Ignored for the `"mean_diff"`
            and `"fisher"` directions.
        calibration_data: Optional data the operating point is calibrated on.
        allow_model_mismatch: When True, a `stats` artifact estimated on a different model is
            accepted instead of raising.
        batch_size: Chunk size for feature extraction.
        max_length: Tokenization truncation bound for feature extraction. None truncates to the
            tokenizer's model maximum length.

    Returns:
        A fitted `Probe` with canonical polarity, a single chosen layer, and a provenance
        `meta` record.

    Raises:
        ValueError: If the method requires `stats` and none is supplied, a supplied `stats` was
            estimated on a different model (without `allow_model_mismatch`) or at a different
            `location` than `spec.location`, a requested candidate layer is out of range or has
            no recorded statistics, `spec.prompt_format` is `"chat_completion"` with unpaired
            data, or the calibration scores are inverted.
    """
    if spec.method in ("lda", "logreg") and stats is None:
        raise ValueError(
            f"method={spec.method!r} requires ambient activation statistics; estimate "
            "ActivationStats once per model; see core.internals.stats."
        )

    if model is not None:
        fingerprint = model_fingerprint(model)
        fitted_model_type = getattr(model.config, "model_type", "unknown")
        session_meta: dict = {}
    else:
        fitted_model_type, session_meta = session_artifact_identity(session)
        fingerprint = session_meta.get("model_fingerprint")
    if (
        stats is not None
        and fingerprint is not None
        and stats.model_fingerprint != fingerprint
        and not allow_model_mismatch
    ):
        raise ValueError(
            "stats were estimated on a different model (fingerprint "
            f"{stats.model_fingerprint!r} vs {fingerprint!r}); whitening with another model's "
            "statistics produces a miscalibrated probe. Re-estimate ActivationStats on this "
            "model, or pass allow_model_mismatch=True."
        )
    if stats is not None and stats.location != spec.location:
        raise ValueError(
            f"stats were estimated at location {stats.location!r} but the fit captures at "
            f"{spec.location!r}; standardizing features with moments from a different boundary "
            "produces a miscalibrated probe. Re-estimate ActivationStats at the fit location."
        )

    num_layers = _resolve_num_layers(model, session)
    if spec.candidate_layers is not None:
        candidates = [int(lid) for lid in dict.fromkeys(spec.candidate_layers)]
        for lid in candidates:
            if not 0 <= lid < num_layers:
                raise ValueError(f"candidate layer {lid} out of range [0, {num_layers}).")
    elif spec.layer_range is not None:
        lo, hi = spec.layer_range
        candidates = [lid for lid in range(num_layers) if lo <= lid / max(num_layers - 1, 1) <= hi]
        if not candidates:
            raise ValueError(f"layer_range ({lo}, {hi}) selects no layers out of {num_layers}.")
    else:
        candidates = list(range(num_layers))

    if spec.method in ("lda", "logreg"):
        uncovered = [lid for lid in candidates if lid not in stats.mean]
        if uncovered:
            raise ValueError(
                f"stats record no statistics for candidate layer(s) {uncovered}; estimate "
                f"ActivationStats over these layers, or restrict candidate_layers to "
                f"{sorted(stats.mean)}."
            )

    pos_features, neg_features = _pooled_features(
        model, tokenizer, data, spec, candidates,
        batch_size=batch_size, max_length=max_length, session=session,
    )
    if calibration_data is not None:
        cal_pos_features, cal_neg_features = _pooled_features(
            model, tokenizer, calibration_data, spec, candidates,
            batch_size=batch_size, max_length=max_length, session=session,
        )
    else:
        cal_pos_features, cal_neg_features = pos_features, neg_features

    sweep: list[dict] = []
    best: dict | None = None
    for lid in candidates:
        w = _fit_direction(pos_features[lid], neg_features[lid], lid, spec, stats)

        fit_pos_scores = pos_features[lid] @ w
        fit_neg_scores = neg_features[lid] @ w
        fit_gap = float(fit_pos_scores.mean() - fit_neg_scores.mean())
        flipped = fit_gap < 0
        if flipped:
            assert spec.method == "logreg", (
                "fit-set orientation cannot invert for mean_diff, lda, or fisher; the score gap "
                "is a squared class-mean difference in the fit metric."
            )
            w = -w

        cal_pos_scores = cal_pos_features[lid] @ w
        cal_neg_scores = cal_neg_features[lid] @ w
        if float(cal_pos_scores.mean()) < float(cal_neg_scores.mean()):
            # the fit-set direction does not generalize to the calibration pool at
            # this layer; record the layer as unusable and keep sweeping rather than
            # aborting the whole fit (calibrate_bias would raise on these scores)
            sweep.append({
                "layer_id": lid,
                "f1": None,
                "margin": None,
                "gap": None,
                "fpr": None,
                "orientation_flipped": flipped,
                "calibration_inverted": True,
            })
            continue

        bias = calibrate_bias(cal_pos_scores, cal_neg_scores, spec.calibration)

        f1 = _f1_at(cal_pos_scores, cal_neg_scores, -bias)
        fpr = float((cal_neg_scores >= -bias).float().mean())
        gap = float(cal_pos_scores.min() - cal_neg_scores.max())
        margin = gap / _pooled_std(cal_pos_scores, cal_neg_scores)
        sweep.append({
            "layer_id": lid,
            "f1": f1,
            "margin": margin,
            "gap": gap,
            "fpr": fpr,
            "orientation_flipped": flipped,
            "calibration_inverted": False,
        })
        if best is None or (f1, margin) > (best["f1"], best["margin"]):
            best = {"layer_id": lid, "weights": w, "bias": bias, "f1": f1, "margin": margin,
                    "fpr": fpr, "orientation_flipped": flipped}

    if best is None:
        inverted = [entry["layer_id"] for entry in sweep if entry.get("calibration_inverted")]
        raise ValueError(
            f"Calibration scores are inverted at every candidate layer ({inverted}); the "
            "calibration pairs contradict the fit pairs for this concept. Revisit the "
            "pools, or the spec's pooling / layer_range."
    )

    calibration_record = {
        "kind": list(spec.calibration) if isinstance(spec.calibration, tuple) else spec.calibration,
        "on": "calibration_data" if calibration_data is not None else "data",
        "f1": best["f1"],
        "fpr": best["fpr"],
    }
    meta = {
        "method": spec.method,
        "pooling": spec.pooling,
        "location": spec.location,
        "prompt_format": spec.prompt_format,
        "n_pos": len(data.positives),
        "n_neg": len(data.negatives),
        "calibration": calibration_record,
        "layer_id": best["layer_id"],
        "layer_sweep": sweep,
        "orientation_flipped": best["orientation_flipped"],
        "stats_used": spec.method in ("lda", "logreg"),
        "model_fingerprint": fingerprint,
        "package_version": _PACKAGE_VERSION,
        "polarity": POLARITY_MARKER,
    }
    if model is not None:
        provenance = artifact_provenance_meta(model, tokenizer)
        for key in ("config_fingerprint", "chat_template_fingerprint"):
            if key in provenance:
                meta[key] = provenance[key]
    elif "model_ref" in session_meta:
        meta["model_ref"] = session_meta["model_ref"]
    if meta["stats_used"]:
        meta["stats_fingerprint"] = stats.fingerprint()

    return Probe(
        model_type=fitted_model_type,
        location=spec.location,
        pooling=spec.pooling,
        layer_ids=[best["layer_id"]],
        weights={best["layer_id"]: best["weights"]},
        bias=best["bias"],
        meta=meta,
    )

probe

A calibrated linear probe on the pooled hidden states with scores oriented in a consistent direction.

POLARITY_MARKER = 'positives_high' module-attribute

PROBE_FORMAT_VERSION = 1 module-attribute

logger = logging.getLogger(__name__) module-attribute

Probe dataclass

Calibrated affine readout of a concept from pooled hidden states.

A probe scores pooled per-layer features as sum_l features[l] @ weights[l] + bias, and its decision is score >= 0 (ties open). Polarity is canonical, i.e., positives score at or above zero, and the operating threshold is folded into bias at calibration time, so no comparator or threshold field exists. Weights are stored in raw activation coordinates (any standardization used during fitting is folded into the weights), so scoring needs no activation statistics.

The class is model-free feature math. It runs no forwards and registers no hooks; batched reads over a live model go through ProbeSet.

Attributes:

Name	Type	Description
model_type	str	HuggingFace model_type of the model the probe was fitted on.
location	str	Residual-stream boundary the weights were fitted at ("layer_input" or "layer_output").
pooling	Literal['mean', 'last']	Token aggregation the probe expects ("mean" or "last", mask-aware).
layer_ids	list[int]	Layers whose features the decision function consumes.
weights	dict[int, Tensor]	Per-layer weight vectors, [H] float32, raw-activation coordinates.
bias	float	Calibrated offset; the decision is raw_score + bias >= 0.
meta	dict	Provenance record (fit method, sample counts, calibration record, layer sweep, fingerprints, package version, polarity marker). May be empty for hand-constructed probes, which disarms all fingerprint checks.

Source code in steerability/algorithms/core/internals/probes/probe.py

@dataclass
class Probe:
    """Calibrated affine readout of a concept from pooled hidden states.

    A probe scores pooled per-layer features as `sum_l features[l] @ weights[l] + bias`, and its
    decision is `score >= 0` (ties open). Polarity is canonical, i.e., positives score at or
    above zero, and the operating threshold is folded into `bias` at calibration time, so no
    comparator or threshold field exists. Weights are stored in raw activation coordinates (any
    standardization used during fitting is folded into the weights), so scoring needs no
    activation statistics.

    The class is model-free feature math. It runs no forwards and registers no hooks; batched
    reads over a live model go through `ProbeSet`.

    Attributes:
        model_type: HuggingFace `model_type` of the model the probe was fitted on.
        location: Residual-stream boundary the weights were fitted at (`"layer_input"` or
            `"layer_output"`).
        pooling: Token aggregation the probe expects (`"mean"` or `"last"`, mask-aware).
        layer_ids: Layers whose features the decision function consumes.
        weights: Per-layer weight vectors, `[H]` float32, raw-activation coordinates.
        bias: Calibrated offset; the decision is `raw_score + bias >= 0`.
        meta: Provenance record (fit method, sample counts, calibration record, layer sweep,
            fingerprints, package version, polarity marker). May be empty for hand-constructed
            probes, which disarms all fingerprint checks.
    """

    model_type: str
    location: str
    pooling: Literal["mean", "last"]
    layer_ids: list[int]
    weights: dict[int, torch.Tensor]
    bias: float
    meta: dict = field(default_factory=dict)

    def __post_init__(self):
        if self.pooling not in ("mean", "last"):
            raise ValueError(f"pooling must be 'mean' or 'last', got {self.pooling!r}.")
        if not self.layer_ids:
            raise ValueError("layer_ids must name at least one layer.")
        self.layer_ids = [int(lid) for lid in self.layer_ids]
        missing = [lid for lid in self.layer_ids if lid not in self.weights]
        if missing:
            raise ValueError(f"weights missing for layer(s) {missing}.")
        self.weights = {
            int(lid): torch.as_tensor(w).detach().reshape(-1).to(torch.float32)
            for lid, w in self.weights.items()
        }
        self.bias = float(self.bias)

    def decision_function(self, features: Mapping[int, torch.Tensor]) -> torch.Tensor:
        """Signed decision scores for pooled features.

        Args:
            features: Mapping from layer id to pooled features of shape `[B, H]` (a `[H]` tensor
                is treated as one row).

        Returns:
            Float32 tensor of shape `[B]`: the sum over `layer_ids` of `features[l] @ weights[l]`,
            plus `bias`.

        Raises:
            KeyError: If a layer in `layer_ids` is absent from `features`; the message names the
                missing layer and the layers supplied.
        """
        scores: torch.Tensor | None = None
        for lid in self.layer_ids:
            if lid not in features:
                raise KeyError(
                    f"Probe requires features for layer {lid}, but only layers "
                    f"{sorted(features)} were supplied; capture or pool every layer in "
                    f"probe.layer_ids ({self.layer_ids})."
                )
            feats = torch.as_tensor(features[lid]).to(torch.float32)
            if feats.ndim == 1:
                feats = feats.unsqueeze(0)
            contribution = feats @ self.weights[lid].to(feats.device)
            scores = contribution if scores is None else scores + contribution
        return (scores + self.bias).cpu()

    def predict(self, features: Mapping[int, torch.Tensor]) -> torch.Tensor:
        """Boolean decisions for pooled features: `decision_function(features) >= 0` (ties open).

        Args:
            features: Mapping from layer id to pooled features of shape `[B, H]`.

        Returns:
            Bool tensor of shape `[B]`.
        """
        return self.decision_function(features) >= 0

    def score_hidden(
        self,
        hidden: Mapping[int, torch.Tensor],
        prompt_mask: torch.Tensor | None = None,
    ) -> torch.Tensor:
        """Pool per-layer hidden states and apply the decision function.

        Each layer's `[B, T, H]` tensor is aggregated to `[B, H]` per `self.pooling` using the
        mask, so `"last"` selects the last real (non-pad) token per row.

        Args:
            hidden: Mapping from layer id to hidden states of shape `[B, T, H]`.
            prompt_mask: Attention mask of shape `[B, T]` (1 at real tokens), or None to treat
                every position as real.

        Returns:
            Float32 tensor of signed decision scores, shape `[B]`.

        Raises:
            KeyError: If a layer in `layer_ids` is absent from `hidden`.
        """
        features: dict[int, torch.Tensor] = {}
        for lid in self.layer_ids:
            if lid not in hidden:
                raise KeyError(
                    f"Probe requires hidden states for layer {lid}, but only layers "
                    f"{sorted(hidden)} were supplied; capture every layer in probe.layer_ids "
                    f"({self.layer_ids})."
                )
            features[lid] = aggregate_condition_hidden(
                hidden[lid], self.pooling, attention_mask=prompt_mask
            )
        return self.decision_function(features)

    def save(self, dir_path: str | Path) -> None:
        """Save the probe to a directory (safetensors for weights, JSON sidecar for metadata).

        Args:
            dir_path: Directory to write to; created if missing.
        """
        dir_path = Path(dir_path)
        dir_path.mkdir(parents=True, exist_ok=True)

        tensors = {
            f"weights.{lid}": self.weights[lid].to(torch.float32).contiguous()
            for lid in sorted(self.weights)
        }
        save_file(tensors, str(dir_path / _TENSOR_FILE))

        metadata = {
            "format_version": PROBE_FORMAT_VERSION,
            "polarity": POLARITY_MARKER,
            "model_type": self.model_type,
            "location": self.location,
            "pooling": self.pooling,
            "layer_ids": list(self.layer_ids),
            "weight_layer_ids": sorted(self.weights),
            "bias": self.bias,
            "meta": self.meta,
        }
        with open(dir_path / _META_FILE, "w") as f:
            json.dump(metadata, f, indent=2, sort_keys=True)
        logger.debug("Saved Probe to %s", dir_path)

    @classmethod
    def load(cls, dir_path: str | Path) -> "Probe":
        """Load a probe saved by `save`.

        Args:
            dir_path: Directory containing the safetensors file and JSON sidecar.

        Returns:
            The loaded `Probe`.

        Raises:
            ValueError: If the artifact-format version is unsupported or the recorded polarity
                marker is not canonical.
        """
        dir_path = Path(dir_path)
        with open(dir_path / _META_FILE) as f:
            metadata = json.load(f)

        version = metadata.get("format_version")
        if version != PROBE_FORMAT_VERSION:
            raise ValueError(
                f"Unsupported Probe format version {version!r}; expected {PROBE_FORMAT_VERSION}."
            )
        polarity = metadata.get("polarity")
        if polarity != POLARITY_MARKER:
            raise ValueError(
                f"Probe artifact records polarity {polarity!r}; expected {POLARITY_MARKER!r} "
                "(positives score at or above zero)."
            )

        tensors = load_file(str(dir_path / _TENSOR_FILE))
        weights = {int(lid): tensors[f"weights.{lid}"] for lid in metadata["weight_layer_ids"]}

        logger.debug("Loaded Probe from %s with layers %s", dir_path, metadata["layer_ids"])
        return cls(
            model_type=metadata["model_type"],
            location=metadata["location"],
            pooling=metadata["pooling"],
            layer_ids=[int(lid) for lid in metadata["layer_ids"]],
            weights=weights,
            bias=float(metadata["bias"]),
            meta=metadata["meta"],
        )

    def as_gate(self, *, allow_model_mismatch: bool = False) -> Gate:
        """A steering gate reproducing this probe's decision, for gated interventions.

        The gate reads the probe's layers at the probe's pooling, scores each pooled state
        against the probe's weights, and opens where the summed contributions plus the
        calibrated bias are at or above zero (ties open). The decision is evaluated on the
        prompt and holds for the whole generation.

        Args:
            allow_model_mismatch: When True, the intervention's model-identity check for this
                probe is disarmed.

        Returns:
            A `Gate` for an `Intervention`'s gate slot (e.g. `ActivationAdapter`'s `gate=`).
        """
        # function-local import to avoid a core/internals -> category package import at module load
        from steerability.algorithms.state_control.common.gating import gate_from_probe

        return gate_from_probe(self, allow_model_mismatch=allow_model_mismatch)

bias instance-attribute

layer_ids instance-attribute

location instance-attribute

meta = field(default_factory=dict) class-attribute instance-attribute

model_type instance-attribute

pooling instance-attribute

weights instance-attribute

as_gate(*, allow_model_mismatch=False)

A steering gate reproducing this probe's decision, for gated interventions.

The gate reads the probe's layers at the probe's pooling, scores each pooled state against the probe's weights, and opens where the summed contributions plus the calibrated bias are at or above zero (ties open). The decision is evaluated on the prompt and holds for the whole generation.

Parameters:

Name	Type	Description	Default
allow_model_mismatch	bool	When True, the intervention's model-identity check for this probe is disarmed.	False

Returns:

Type	Description
Gate	A Gate for an Intervention's gate slot (e.g. ActivationAdapter's gate=).

Source code in steerability/algorithms/core/internals/probes/probe.py

def as_gate(self, *, allow_model_mismatch: bool = False) -> Gate:
    """A steering gate reproducing this probe's decision, for gated interventions.

    The gate reads the probe's layers at the probe's pooling, scores each pooled state
    against the probe's weights, and opens where the summed contributions plus the
    calibrated bias are at or above zero (ties open). The decision is evaluated on the
    prompt and holds for the whole generation.

    Args:
        allow_model_mismatch: When True, the intervention's model-identity check for this
            probe is disarmed.

    Returns:
        A `Gate` for an `Intervention`'s gate slot (e.g. `ActivationAdapter`'s `gate=`).
    """
    # function-local import to avoid a core/internals -> category package import at module load
    from steerability.algorithms.state_control.common.gating import gate_from_probe

    return gate_from_probe(self, allow_model_mismatch=allow_model_mismatch)

decision_function(features)

Signed decision scores for pooled features.

Parameters:

Name	Type	Description	Default
features	Mapping[int, Tensor]	Mapping from layer id to pooled features of shape [B, H] (a [H] tensor is treated as one row).	required

Returns:

Type	Description
Tensor	Float32 tensor of shape [B]: the sum over layer_ids of features[l] @ weights[l],
Tensor	plus bias.

Raises:

Type	Description
KeyError	If a layer in layer_ids is absent from features; the message names the missing layer and the layers supplied.

Source code in steerability/algorithms/core/internals/probes/probe.py

def decision_function(self, features: Mapping[int, torch.Tensor]) -> torch.Tensor:
    """Signed decision scores for pooled features.

    Args:
        features: Mapping from layer id to pooled features of shape `[B, H]` (a `[H]` tensor
            is treated as one row).

    Returns:
        Float32 tensor of shape `[B]`: the sum over `layer_ids` of `features[l] @ weights[l]`,
        plus `bias`.

    Raises:
        KeyError: If a layer in `layer_ids` is absent from `features`; the message names the
            missing layer and the layers supplied.
    """
    scores: torch.Tensor | None = None
    for lid in self.layer_ids:
        if lid not in features:
            raise KeyError(
                f"Probe requires features for layer {lid}, but only layers "
                f"{sorted(features)} were supplied; capture or pool every layer in "
                f"probe.layer_ids ({self.layer_ids})."
            )
        feats = torch.as_tensor(features[lid]).to(torch.float32)
        if feats.ndim == 1:
            feats = feats.unsqueeze(0)
        contribution = feats @ self.weights[lid].to(feats.device)
        scores = contribution if scores is None else scores + contribution
    return (scores + self.bias).cpu()

load(dir_path) classmethod

Load a probe saved by save.

Parameters:

Name	Type	Description	Default
dir_path	str | Path	Directory containing the safetensors file and JSON sidecar.	required

Returns:

Type	Description
'Probe'	The loaded Probe.

Raises:

Type	Description
ValueError	If the artifact-format version is unsupported or the recorded polarity marker is not canonical.

Source code in steerability/algorithms/core/internals/probes/probe.py

@classmethod
def load(cls, dir_path: str | Path) -> "Probe":
    """Load a probe saved by `save`.

    Args:
        dir_path: Directory containing the safetensors file and JSON sidecar.

    Returns:
        The loaded `Probe`.

    Raises:
        ValueError: If the artifact-format version is unsupported or the recorded polarity
            marker is not canonical.
    """
    dir_path = Path(dir_path)
    with open(dir_path / _META_FILE) as f:
        metadata = json.load(f)

    version = metadata.get("format_version")
    if version != PROBE_FORMAT_VERSION:
        raise ValueError(
            f"Unsupported Probe format version {version!r}; expected {PROBE_FORMAT_VERSION}."
        )
    polarity = metadata.get("polarity")
    if polarity != POLARITY_MARKER:
        raise ValueError(
            f"Probe artifact records polarity {polarity!r}; expected {POLARITY_MARKER!r} "
            "(positives score at or above zero)."
        )

    tensors = load_file(str(dir_path / _TENSOR_FILE))
    weights = {int(lid): tensors[f"weights.{lid}"] for lid in metadata["weight_layer_ids"]}

    logger.debug("Loaded Probe from %s with layers %s", dir_path, metadata["layer_ids"])
    return cls(
        model_type=metadata["model_type"],
        location=metadata["location"],
        pooling=metadata["pooling"],
        layer_ids=[int(lid) for lid in metadata["layer_ids"]],
        weights=weights,
        bias=float(metadata["bias"]),
        meta=metadata["meta"],
    )

predict(features)

Boolean decisions for pooled features: decision_function(features) >= 0 (ties open).

Parameters:

Name	Type	Description	Default
features	Mapping[int, Tensor]	Mapping from layer id to pooled features of shape [B, H].	required

Returns:

Type	Description
Tensor	Bool tensor of shape [B].

Source code in steerability/algorithms/core/internals/probes/probe.py

def predict(self, features: Mapping[int, torch.Tensor]) -> torch.Tensor:
    """Boolean decisions for pooled features: `decision_function(features) >= 0` (ties open).

    Args:
        features: Mapping from layer id to pooled features of shape `[B, H]`.

    Returns:
        Bool tensor of shape `[B]`.
    """
    return self.decision_function(features) >= 0

save(dir_path)

Save the probe to a directory (safetensors for weights, JSON sidecar for metadata).

Parameters:

Name	Type	Description	Default
dir_path	str | Path	Directory to write to; created if missing.	required

Source code in steerability/algorithms/core/internals/probes/probe.py

def save(self, dir_path: str | Path) -> None:
    """Save the probe to a directory (safetensors for weights, JSON sidecar for metadata).

    Args:
        dir_path: Directory to write to; created if missing.
    """
    dir_path = Path(dir_path)
    dir_path.mkdir(parents=True, exist_ok=True)

    tensors = {
        f"weights.{lid}": self.weights[lid].to(torch.float32).contiguous()
        for lid in sorted(self.weights)
    }
    save_file(tensors, str(dir_path / _TENSOR_FILE))

    metadata = {
        "format_version": PROBE_FORMAT_VERSION,
        "polarity": POLARITY_MARKER,
        "model_type": self.model_type,
        "location": self.location,
        "pooling": self.pooling,
        "layer_ids": list(self.layer_ids),
        "weight_layer_ids": sorted(self.weights),
        "bias": self.bias,
        "meta": self.meta,
    }
    with open(dir_path / _META_FILE, "w") as f:
        json.dump(metadata, f, indent=2, sort_keys=True)
    logger.debug("Saved Probe to %s", dir_path)

score_hidden(hidden, prompt_mask=None)

Pool per-layer hidden states and apply the decision function.

Each layer's [B, T, H] tensor is aggregated to [B, H] per self.pooling using the mask, so "last" selects the last real (non-pad) token per row.

Parameters:

Name	Type	Description	Default
hidden	Mapping[int, Tensor]	Mapping from layer id to hidden states of shape [B, T, H].	required
prompt_mask	Tensor | None	Attention mask of shape [B, T] (1 at real tokens), or None to treat every position as real.	None

Returns:

Type	Description
Tensor	Float32 tensor of signed decision scores, shape [B].

Raises:

Type	Description
KeyError	If a layer in layer_ids is absent from hidden.

Source code in steerability/algorithms/core/internals/probes/probe.py

def score_hidden(
    self,
    hidden: Mapping[int, torch.Tensor],
    prompt_mask: torch.Tensor | None = None,
) -> torch.Tensor:
    """Pool per-layer hidden states and apply the decision function.

    Each layer's `[B, T, H]` tensor is aggregated to `[B, H]` per `self.pooling` using the
    mask, so `"last"` selects the last real (non-pad) token per row.

    Args:
        hidden: Mapping from layer id to hidden states of shape `[B, T, H]`.
        prompt_mask: Attention mask of shape `[B, T]` (1 at real tokens), or None to treat
            every position as real.

    Returns:
        Float32 tensor of signed decision scores, shape `[B]`.

    Raises:
        KeyError: If a layer in `layer_ids` is absent from `hidden`.
    """
    features: dict[int, torch.Tensor] = {}
    for lid in self.layer_ids:
        if lid not in hidden:
            raise KeyError(
                f"Probe requires hidden states for layer {lid}, but only layers "
                f"{sorted(hidden)} were supplied; capture every layer in probe.layer_ids "
                f"({self.layer_ids})."
            )
        features[lid] = aggregate_condition_hidden(
            hidden[lid], self.pooling, attention_mask=prompt_mask
        )
    return self.decision_function(features)

probe_set

Score all probes in a batch with a read-only forward pass.

logger = logging.getLogger(__name__) module-attribute

ProbeReadings dataclass

One ProbeSet.read() call's result.

Attributes:

Name	Type	Description
scores	dict[str, Tensor]	Mapping from probe name to signed decision scores, [B] per probe.
decisions	dict[str, Tensor]	Mapping from probe name to boolean decisions, [B] per probe (scores >= 0, ties open).

Source code in steerability/algorithms/core/internals/probes/probe_set.py

@dataclass
class ProbeReadings:
    """One `ProbeSet.read()` call's result.

    Attributes:
        scores: Mapping from probe name to signed decision scores, `[B]` per probe.
        decisions: Mapping from probe name to boolean decisions, `[B]` per probe
            (`scores >= 0`, ties open).
    """

    scores: dict[str, torch.Tensor]
    decisions: dict[str, torch.Tensor]

decisions instance-attribute

scores instance-attribute

ProbeSet

Named probes scored on a batch of prompts in one read-only forward.

All probes must share one model_type and one capture location, since a single forward supplies every probe's features. read() registers forward pre-hooks on the union of the probes' layers (the layer_input boundary), runs one forward, and returns per-probe signed scores and decisions. The call is self-contained, i.e., it registers no standing hooks and never edits hidden states.

Parameters:

Name	Type	Description	Default
probes	Mapping[str, Probe]	Mapping from probe name to a fitted Probe.	required

Attributes:

Name	Type	Description
probes	dict[str, Probe]	The probes, keyed by name.
names	tuple[str, ...]	The probe names, in mapping order.
layer_ids	list[int]	Sorted union of the probes' layers.
latest	ProbeReadings | None	The most recent ProbeReadings, overwritten by each read() call; None before the first. Diagnostics only.

Raises:

Type	Description
ValueError	If probes is empty, the probes disagree on model_type or location, or two probes record different model fingerprints (probes with no recorded fingerprint mix freely).

Source code in steerability/algorithms/core/internals/probes/probe_set.py

class ProbeSet:
    """Named probes scored on a batch of prompts in one read-only forward.

    All probes must share one `model_type` and one capture `location`, since a single forward
    supplies every probe's features. `read()` registers forward pre-hooks on the union of the
    probes' layers (the `layer_input` boundary), runs one forward, and returns per-probe signed
    scores and decisions. The call is self-contained, i.e., it registers no standing hooks and
    never edits hidden states.

    Args:
        probes: Mapping from probe name to a fitted `Probe`.

    Attributes:
        probes: The probes, keyed by name.
        names: The probe names, in mapping order.
        layer_ids: Sorted union of the probes' layers.
        latest: The most recent `ProbeReadings`, overwritten by each `read()` call; None
            before the first. Diagnostics only.

    Raises:
        ValueError: If `probes` is empty, the probes disagree on `model_type` or `location`, or
            two probes record different model fingerprints (probes with no recorded fingerprint
            mix freely).
    """

    def __init__(self, probes: Mapping[str, Probe]):
        if not probes:
            raise ValueError("ProbeSet requires at least one probe.")
        self.probes: dict[str, Probe] = dict(probes)

        model_types = {probe.model_type for probe in self.probes.values()}
        if len(model_types) > 1:
            raise ValueError(
                f"probes disagree on model_type ({sorted(model_types)}); one read() capture "
                "cannot serve probes fitted on different model families."
            )
        locations = {probe.location for probe in self.probes.values()}
        if len(locations) > 1:
            raise ValueError(
                f"probes disagree on location ({sorted(locations)}); one read() capture "
                "cannot serve probes fitted at different boundaries."
            )
        fingerprints = {
            probe.meta["model_fingerprint"]
            for probe in self.probes.values()
            if probe.meta.get("model_fingerprint") is not None
        }
        if len(fingerprints) > 1:
            raise ValueError(
                f"probes record different model fingerprints ({sorted(fingerprints)}); "
                "mixed-model probes cannot share one read() capture."
            )

        self.names: tuple[str, ...] = tuple(self.probes)
        self.layer_ids: list[int] = sorted({lid for probe in self.probes.values() for lid in probe.layer_ids})
        self.latest: ProbeReadings | None = None

    @property
    def model_type(self) -> str:
        """The shared `model_type` of the probes."""
        return next(iter(self.probes.values())).model_type

    @property
    def location(self) -> str:
        """The shared capture boundary of the probes."""
        return next(iter(self.probes.values())).location

    @classmethod
    def fit(
        cls,
        model: PreTrainedModel | None,
        tokenizer: PreTrainedTokenizerBase,
        *,
        data: Mapping[str, ContrastivePairs],
        spec: ProbeFitSpec | Mapping[str, ProbeFitSpec] | None = None,
        stats: ActivationStats | None = None,
        calibration_data: Mapping[str, ContrastivePairs] | None = None,
        session=None,
    ) -> "ProbeSet":
        """Fit one probe per name and return the set.

        Args:
            model: Model the probes are fitted on.
            tokenizer: Tokenizer for rendering and encoding the pairs.
            data: Mapping from probe name to the contrastive pairs it is fitted on.
            spec: A shared `ProbeFitSpec` (None uses the defaults), or a mapping from probe name
                to a per-name spec.
            stats: Ambient activation statistics for methods that require them.
            calibration_data: Optional mapping from probe name to calibration pairs.

        Returns:
            The fitted `ProbeSet`.

        Raises:
            ValueError: If `data` is empty, a spec mapping misses a name, or `fit_probe` raises
                for any probe.
        """
        if not data:
            raise ValueError("data must map at least one probe name to contrastive pairs.")
        if isinstance(spec, Mapping):
            missing = [name for name in data if name not in spec]
            if missing:
                raise ValueError(f"spec mapping is missing entries for probe(s) {missing}.")

        probes: dict[str, Probe] = {}
        for name, pairs in data.items():
            probe_spec = spec[name] if isinstance(spec, Mapping) else (spec or ProbeFitSpec())
            probes[name] = fit_probe(
                model,
                tokenizer,
                data=pairs,
                spec=probe_spec,
                stats=stats,
                calibration_data=(calibration_data or {}).get(name),
                session=session,
            )
        return cls(probes)

    def read(
        self,
        model: PreTrainedModel | None,
        input_ids: torch.Tensor,
        attention_mask: torch.Tensor | None = None,
        session=None,
    ) -> ProbeReadings:
        """Score a batch of prompts against every probe in one read-only forward.

        Registers one forward pre-hook per layer in the union (the `layer_input` boundary, the
        only boundary `read()` serves), runs a single forward under `torch.no_grad()` with
        `use_cache=False` wrapped in `auxiliary_pass(aligned=True)`, removes the hooks, then
        pools each probe's layers per its `pooling` with the attention mask and applies its
        decision function. The auxiliary marking means co-resident state controls' condition
        scorers, gates, and position counters ignore the pass, while `"all"`-scoped behavior
        transforms still apply, so scores measure the stream as deployed. A set fitted at
        `"layer_output"` raises, since a pre-hook capture would score features one block away
        from where those probes were fitted; score such probes offline via
        `layerwise_tokenwise_hidden` and `Probe.score_hidden`.

        Args:
            model: The model to read (must match the probes' `model_type`).
            input_ids: Prompt token ids of shape `[B, T]` (a 1-D tensor is treated as one row).
            attention_mask: Prompt attention mask matching `input_ids` (1 at real tokens). When
                None, every position is treated as real.

        Returns:
            A `ProbeReadings` with per-probe signed scores and decisions, also stashed on `latest`.

        Raises:
            ValueError: If the model's `model_type` does not match the probes', the set's
                `location` is not `"layer_input"`, a probe layer is out of range, or the
                architecture is unrecognized.
        """
        if model is not None:
            live_model_type = getattr(model.config, "model_type", "unknown")
            if live_model_type != self.model_type:
                raise ValueError(
                    f"ProbeSet was fitted on model_type {self.model_type!r} but read() received "
                    f"{live_model_type!r}."
                )
        elif session is None:
            raise ValueError("ProbeSet.read() requires a live model or a capture-capable session.")
        if self.location != "layer_input":
            raise ValueError(
                f"ProbeSet.read() serves the 'layer_input' boundary, but this set was fitted at "
                f"{self.location!r}; refit with location='layer_input', or score offline via "
                "layerwise_tokenwise_hidden and Probe.score_hidden."
            )

        device = next(model.parameters()).device if model is not None else torch.device("cpu")
        ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.as_tensor(input_ids, dtype=torch.long)
        if ids.ndim == 1:
            ids = ids.unsqueeze(0)
        ids = ids.to(device)

        if attention_mask is not None:
            mask = attention_mask if isinstance(attention_mask, torch.Tensor) else torch.as_tensor(attention_mask)
            if mask.ndim == 1:
                mask = mask.unsqueeze(0)
            mask = mask.to(device)
        else:
            mask = torch.ones_like(ids)

        if model is None:
            return self._read_via_session(session, ids, mask)

        layer_names = resolve_model_layout(model).layer_names
        for lid in self.layer_ids:
            if not 0 <= lid < len(layer_names):
                raise ValueError(f"probe layer {lid} out of range [0, {len(layer_names)}).")

        captured: dict[int, torch.Tensor] = {}

        def _pre_capture(layer_id: int):
            def _pre_hook(module, input_args, input_kwargs):
                hidden = input_args[0] if input_args else input_kwargs.get("hidden_states")
                if hidden is not None:
                    captured[layer_id] = hidden.detach().to("cpu", torch.float32)
                return input_args, input_kwargs

            return _pre_hook

        handles = []
        try:
            for lid in self.layer_ids:
                module = model.get_submodule(layer_names[lid])
                handles.append(module.register_forward_pre_hook(_pre_capture(lid), with_kwargs=True))
            with torch.no_grad(), auxiliary_pass(aligned=True):
                # explicit `cache_position` lets co-resident state hooks position this aligned
                # pass since transformers v5 threads the kwarg into decoder layers only when the
                # caller passes it
                model(
                    input_ids=ids,
                    attention_mask=mask,
                    use_cache=False,
                    cache_position=torch.arange(ids.size(1), device=ids.device),
                )
        finally:
            for handle in handles:
                handle.remove()

        cpu_mask = mask.to("cpu")
        scores: dict[str, torch.Tensor] = {}
        decisions: dict[str, torch.Tensor] = {}
        for name, probe in self.probes.items():
            features = {
                lid: aggregate_condition_hidden(captured[lid], probe.pooling, attention_mask=cpu_mask)
                for lid in probe.layer_ids
            }
            probe_scores = probe.decision_function(features)
            scores[name] = probe_scores
            decisions[name] = probe_scores >= 0

        readings = ProbeReadings(scores=scores, decisions=decisions)
        self.latest = readings
        return readings

    def _read_via_session(self, session, ids: torch.Tensor, mask: torch.Tensor) -> ProbeReadings:
        """Score through a capture-capable session's `capture` at the layer-input boundary."""
        from steerability.algorithms.core.execution.payloads import PreparedPrompt

        prompts = [
            PreparedPrompt.from_token_ids(ids[index:index + 1], mask[index:index + 1])
            for index in range(ids.size(0))
        ]
        result = session.capture(
            prompts, layers=list(self.layer_ids), mode="all_tokens", location="layer_input"
        )
        cpu_mask = result.attention_mask
        scores: dict[str, torch.Tensor] = {}
        decisions: dict[str, torch.Tensor] = {}
        for name, probe in self.probes.items():
            features = {
                lid: aggregate_condition_hidden(
                    result.hidden[lid].to(torch.float32), probe.pooling, attention_mask=cpu_mask
                )
                for lid in probe.layer_ids
            }
            probe_scores = probe.decision_function(features)
            scores[name] = probe_scores
            decisions[name] = probe_scores >= 0
        readings = ProbeReadings(scores=scores, decisions=decisions)
        self.latest = readings
        return readings

    def summary(self) -> dict[str, dict]:
        """Per-probe diagnostic table.

        Returns:
            Mapping from probe name to a dict with keys `"layer_ids"`, `"method"`, `"bias"`,
            `"calibration"`, `"f1"`, `"model_fingerprint"`, and `"stats_fingerprint"` (all but
            `"layer_ids"` and `"bias"` are None when the probe's `meta` does not record them).
        """
        return {
            name: {
                "layer_ids": list(probe.layer_ids),
                "method": probe.meta.get("method"),
                "bias": probe.bias,
                "calibration": probe.meta.get("calibration"),
                "f1": (probe.meta.get("calibration") or {}).get("f1"),
                "model_fingerprint": probe.meta.get("model_fingerprint"),
                "stats_fingerprint": probe.meta.get("stats_fingerprint"),
            }
            for name, probe in self.probes.items()
        }

latest = None instance-attribute

layer_ids = sorted({lid for probe in self.probes.values() for lid in probe.layer_ids}) instance-attribute

location property

The shared capture boundary of the probes.

model_type property

The shared model_type of the probes.

names = tuple(self.probes) instance-attribute

probes = dict(probes) instance-attribute

fit(model, tokenizer, *, data, spec=None, stats=None, calibration_data=None, session=None) classmethod

Fit one probe per name and return the set.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model the probes are fitted on.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the pairs.	required
data	Mapping[str, ContrastivePairs]	Mapping from probe name to the contrastive pairs it is fitted on.	required
spec	ProbeFitSpec | Mapping[str, ProbeFitSpec] | None	A shared ProbeFitSpec (None uses the defaults), or a mapping from probe name to a per-name spec.	None
stats	ActivationStats | None	Ambient activation statistics for methods that require them.	None
calibration_data	Mapping[str, ContrastivePairs] | None	Optional mapping from probe name to calibration pairs.	None

Returns:

Type	Description
ProbeSet	The fitted ProbeSet.

Raises:

Type	Description
ValueError	If data is empty, a spec mapping misses a name, or fit_probe raises for any probe.

Source code in steerability/algorithms/core/internals/probes/probe_set.py

@classmethod
def fit(
    cls,
    model: PreTrainedModel | None,
    tokenizer: PreTrainedTokenizerBase,
    *,
    data: Mapping[str, ContrastivePairs],
    spec: ProbeFitSpec | Mapping[str, ProbeFitSpec] | None = None,
    stats: ActivationStats | None = None,
    calibration_data: Mapping[str, ContrastivePairs] | None = None,
    session=None,
) -> "ProbeSet":
    """Fit one probe per name and return the set.

    Args:
        model: Model the probes are fitted on.
        tokenizer: Tokenizer for rendering and encoding the pairs.
        data: Mapping from probe name to the contrastive pairs it is fitted on.
        spec: A shared `ProbeFitSpec` (None uses the defaults), or a mapping from probe name
            to a per-name spec.
        stats: Ambient activation statistics for methods that require them.
        calibration_data: Optional mapping from probe name to calibration pairs.

    Returns:
        The fitted `ProbeSet`.

    Raises:
        ValueError: If `data` is empty, a spec mapping misses a name, or `fit_probe` raises
            for any probe.
    """
    if not data:
        raise ValueError("data must map at least one probe name to contrastive pairs.")
    if isinstance(spec, Mapping):
        missing = [name for name in data if name not in spec]
        if missing:
            raise ValueError(f"spec mapping is missing entries for probe(s) {missing}.")

    probes: dict[str, Probe] = {}
    for name, pairs in data.items():
        probe_spec = spec[name] if isinstance(spec, Mapping) else (spec or ProbeFitSpec())
        probes[name] = fit_probe(
            model,
            tokenizer,
            data=pairs,
            spec=probe_spec,
            stats=stats,
            calibration_data=(calibration_data or {}).get(name),
            session=session,
        )
    return cls(probes)

read(model, input_ids, attention_mask=None, session=None)

Score a batch of prompts against every probe in one read-only forward.

Registers one forward pre-hook per layer in the union (the layer_input boundary, the only boundary read() serves), runs a single forward under torch.no_grad() with use_cache=False wrapped in auxiliary_pass(aligned=True), removes the hooks, then pools each probe's layers per its pooling with the attention mask and applies its decision function. The auxiliary marking means co-resident state controls' condition scorers, gates, and position counters ignore the pass, while "all"-scoped behavior transforms still apply, so scores measure the stream as deployed. A set fitted at "layer_output" raises, since a pre-hook capture would score features one block away from where those probes were fitted; score such probes offline via layerwise_tokenwise_hidden and Probe.score_hidden.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	The model to read (must match the probes' model_type).	required
input_ids	Tensor	Prompt token ids of shape [B, T] (a 1-D tensor is treated as one row).	required
attention_mask	Tensor | None	Prompt attention mask matching input_ids (1 at real tokens). When None, every position is treated as real.	None

Returns:

Type	Description
ProbeReadings	A ProbeReadings with per-probe signed scores and decisions, also stashed on latest.

Raises:

Type	Description
ValueError	If the model's model_type does not match the probes', the set's location is not "layer_input", a probe layer is out of range, or the architecture is unrecognized.

Source code in steerability/algorithms/core/internals/probes/probe_set.py

def read(
    self,
    model: PreTrainedModel | None,
    input_ids: torch.Tensor,
    attention_mask: torch.Tensor | None = None,
    session=None,
) -> ProbeReadings:
    """Score a batch of prompts against every probe in one read-only forward.

    Registers one forward pre-hook per layer in the union (the `layer_input` boundary, the
    only boundary `read()` serves), runs a single forward under `torch.no_grad()` with
    `use_cache=False` wrapped in `auxiliary_pass(aligned=True)`, removes the hooks, then
    pools each probe's layers per its `pooling` with the attention mask and applies its
    decision function. The auxiliary marking means co-resident state controls' condition
    scorers, gates, and position counters ignore the pass, while `"all"`-scoped behavior
    transforms still apply, so scores measure the stream as deployed. A set fitted at
    `"layer_output"` raises, since a pre-hook capture would score features one block away
    from where those probes were fitted; score such probes offline via
    `layerwise_tokenwise_hidden` and `Probe.score_hidden`.

    Args:
        model: The model to read (must match the probes' `model_type`).
        input_ids: Prompt token ids of shape `[B, T]` (a 1-D tensor is treated as one row).
        attention_mask: Prompt attention mask matching `input_ids` (1 at real tokens). When
            None, every position is treated as real.

    Returns:
        A `ProbeReadings` with per-probe signed scores and decisions, also stashed on `latest`.

    Raises:
        ValueError: If the model's `model_type` does not match the probes', the set's
            `location` is not `"layer_input"`, a probe layer is out of range, or the
            architecture is unrecognized.
    """
    if model is not None:
        live_model_type = getattr(model.config, "model_type", "unknown")
        if live_model_type != self.model_type:
            raise ValueError(
                f"ProbeSet was fitted on model_type {self.model_type!r} but read() received "
                f"{live_model_type!r}."
            )
    elif session is None:
        raise ValueError("ProbeSet.read() requires a live model or a capture-capable session.")
    if self.location != "layer_input":
        raise ValueError(
            f"ProbeSet.read() serves the 'layer_input' boundary, but this set was fitted at "
            f"{self.location!r}; refit with location='layer_input', or score offline via "
            "layerwise_tokenwise_hidden and Probe.score_hidden."
        )

    device = next(model.parameters()).device if model is not None else torch.device("cpu")
    ids = input_ids if isinstance(input_ids, torch.Tensor) else torch.as_tensor(input_ids, dtype=torch.long)
    if ids.ndim == 1:
        ids = ids.unsqueeze(0)
    ids = ids.to(device)

    if attention_mask is not None:
        mask = attention_mask if isinstance(attention_mask, torch.Tensor) else torch.as_tensor(attention_mask)
        if mask.ndim == 1:
            mask = mask.unsqueeze(0)
        mask = mask.to(device)
    else:
        mask = torch.ones_like(ids)

    if model is None:
        return self._read_via_session(session, ids, mask)

    layer_names = resolve_model_layout(model).layer_names
    for lid in self.layer_ids:
        if not 0 <= lid < len(layer_names):
            raise ValueError(f"probe layer {lid} out of range [0, {len(layer_names)}).")

    captured: dict[int, torch.Tensor] = {}

    def _pre_capture(layer_id: int):
        def _pre_hook(module, input_args, input_kwargs):
            hidden = input_args[0] if input_args else input_kwargs.get("hidden_states")
            if hidden is not None:
                captured[layer_id] = hidden.detach().to("cpu", torch.float32)
            return input_args, input_kwargs

        return _pre_hook

    handles = []
    try:
        for lid in self.layer_ids:
            module = model.get_submodule(layer_names[lid])
            handles.append(module.register_forward_pre_hook(_pre_capture(lid), with_kwargs=True))
        with torch.no_grad(), auxiliary_pass(aligned=True):
            # explicit `cache_position` lets co-resident state hooks position this aligned
            # pass since transformers v5 threads the kwarg into decoder layers only when the
            # caller passes it
            model(
                input_ids=ids,
                attention_mask=mask,
                use_cache=False,
                cache_position=torch.arange(ids.size(1), device=ids.device),
            )
    finally:
        for handle in handles:
            handle.remove()

    cpu_mask = mask.to("cpu")
    scores: dict[str, torch.Tensor] = {}
    decisions: dict[str, torch.Tensor] = {}
    for name, probe in self.probes.items():
        features = {
            lid: aggregate_condition_hidden(captured[lid], probe.pooling, attention_mask=cpu_mask)
            for lid in probe.layer_ids
        }
        probe_scores = probe.decision_function(features)
        scores[name] = probe_scores
        decisions[name] = probe_scores >= 0

    readings = ProbeReadings(scores=scores, decisions=decisions)
    self.latest = readings
    return readings

summary()

Per-probe diagnostic table.

Returns:

Type	Description
dict[str, dict]	Mapping from probe name to a dict with keys "layer_ids", "method", "bias",
dict[str, dict]	"calibration", "f1", "model_fingerprint", and "stats_fingerprint" (all but
dict[str, dict]	"layer_ids" and "bias" are None when the probe's meta does not record them).

Source code in steerability/algorithms/core/internals/probes/probe_set.py

def summary(self) -> dict[str, dict]:
    """Per-probe diagnostic table.

    Returns:
        Mapping from probe name to a dict with keys `"layer_ids"`, `"method"`, `"bias"`,
        `"calibration"`, `"f1"`, `"model_fingerprint"`, and `"stats_fingerprint"` (all but
        `"layer_ids"` and `"bias"` are None when the probe's `meta` does not record them).
    """
    return {
        name: {
            "layer_ids": list(probe.layer_ids),
            "method": probe.meta.get("method"),
            "bias": probe.bias,
            "calibration": probe.meta.get("calibration"),
            "f1": (probe.meta.get("calibration") or {}).get("f1"),
            "model_fingerprint": probe.meta.get("model_fingerprint"),
            "stats_fingerprint": probe.meta.get("stats_fingerprint"),
        }
        for name, probe in self.probes.items()
    }

ProbeSetFit dataclass

Deferred fitting recipe; holds every fitting input except the model.

fit(model, tokenizer) estimates stats first when it is a StatsSpec, fits one probe per name, and returns a ProbeSet. RoutedDecoding.steer() calls fit on the model the pipeline provides, so pipelines whose structural controls produce the final weights fit on those weights, and ProbeFitSpec fields stay grid-searchable through ControlSpec.

Attributes:

Name	Type	Description
data	Mapping[str, ContrastivePairs]	Mapping from probe name to the contrastive pairs it is fitted on.
spec	ProbeFitSpec | Mapping[str, ProbeFitSpec]	A shared ProbeFitSpec, or a mapping from probe name to a per-name spec covering every name in data.
stats	ActivationStats | StatsSpec | None	Ambient activation statistics, or a StatsSpec estimated on the model at fit time. Required when any effective spec's method is "lda" or "logreg".
calibration_data	Mapping[str, ContrastivePairs] | None	Optional mapping from probe name to calibration pairs; names absent from the mapping calibrate on their fit pairs.

Source code in steerability/algorithms/core/internals/probes/probe_set.py

@dataclass
class ProbeSetFit:
    """Deferred fitting recipe; holds every fitting input except the model.

    `fit(model, tokenizer)` estimates `stats` first when it is a `StatsSpec`, fits one probe
    per name, and returns a `ProbeSet`. `RoutedDecoding.steer()` calls `fit` on the model the
    pipeline provides, so pipelines whose structural controls produce the final weights fit on
    those weights, and `ProbeFitSpec` fields stay grid-searchable through `ControlSpec`.

    Attributes:
        data: Mapping from probe name to the contrastive pairs it is fitted on.
        spec: A shared `ProbeFitSpec`, or a mapping from probe name to a per-name spec covering
            every name in `data`.
        stats: Ambient activation statistics, or a `StatsSpec` estimated on the model at fit
            time. Required when any effective spec's method is `"lda"` or `"logreg"`.
        calibration_data: Optional mapping from probe name to calibration pairs; names absent
            from the mapping calibrate on their fit pairs.
    """

    data: Mapping[str, ContrastivePairs]
    spec: ProbeFitSpec | Mapping[str, ProbeFitSpec] = field(default_factory=ProbeFitSpec)
    stats: ActivationStats | StatsSpec | None = None
    calibration_data: Mapping[str, ContrastivePairs] | None = None

    def __post_init__(self):
        if not self.data:
            raise ValueError("data must map at least one probe name to contrastive pairs.")
        if isinstance(self.spec, Mapping):
            missing = [name for name in self.data if name not in self.spec]
            if missing:
                raise ValueError(f"spec mapping is missing entries for probe(s) {missing}.")
        if self.calibration_data is not None:
            unknown = [name for name in self.calibration_data if name not in self.data]
            if unknown:
                raise ValueError(
                    f"calibration_data names unknown probe(s) {unknown}; "
                    f"available: {sorted(self.data)}."
                )
        if self.stats is None:
            needy = [
                name for name in self.data
                if self._spec_for(name).method in ("lda", "logreg")
            ]
            if needy:
                raise ValueError(
                    f"probe(s) {needy} use a method that requires ambient activation "
                    "statistics; supply stats (an ActivationStats or a StatsSpec)."
                )

    def _spec_for(self, name: str) -> ProbeFitSpec:
        """The effective spec for one probe name."""
        return self.spec[name] if isinstance(self.spec, Mapping) else self.spec

    @property
    def names(self) -> tuple[str, ...]:
        """The probe names, available before fitting so routing rules validate at construction."""
        return tuple(self.data)

    def fit(self, model: PreTrainedModel | None, tokenizer: PreTrainedTokenizerBase, session=None) -> "ProbeSet":
        """Fit the recipe on `model`, resolving a `StatsSpec` on it first.

        Args:
            model: Model the probes are fitted on.
            tokenizer: Tokenizer for rendering and encoding the pairs.

        Returns:
            The fitted `ProbeSet`.
        """
        stats = self.stats
        if isinstance(stats, StatsSpec):
            stats = stats.estimate(model, tokenizer, session=session)
        return ProbeSet.fit(
            model,
            tokenizer,
            data=self.data,
            spec=self.spec,
            stats=stats,
            calibration_data=self.calibration_data,
            session=session,
        )

calibration_data = None class-attribute instance-attribute

data instance-attribute

names property

The probe names, available before fitting so routing rules validate at construction.

spec = field(default_factory=ProbeFitSpec) class-attribute instance-attribute

stats = None class-attribute instance-attribute

fit(model, tokenizer, session=None)

Fit the recipe on model, resolving a StatsSpec on it first.

Parameters:

Name	Type	Description	Default
model	PreTrainedModel | None	Model the probes are fitted on.	required
tokenizer	PreTrainedTokenizerBase	Tokenizer for rendering and encoding the pairs.	required

Returns:

Type	Description
ProbeSet	The fitted ProbeSet.

Source code in steerability/algorithms/core/internals/probes/probe_set.py

def fit(self, model: PreTrainedModel | None, tokenizer: PreTrainedTokenizerBase, session=None) -> "ProbeSet":
    """Fit the recipe on `model`, resolving a `StatsSpec` on it first.

    Args:
        model: Model the probes are fitted on.
        tokenizer: Tokenizer for rendering and encoding the pairs.

    Returns:
        The fitted `ProbeSet`.
    """
    stats = self.stats
    if isinstance(stats, StatsSpec):
        stats = stats.estimate(model, tokenizer, session=session)
    return ProbeSet.fit(
        model,
        tokenizer,
        data=self.data,
        spec=self.spec,
        stats=stats,
        calibration_data=self.calibration_data,
        session=session,
    )