entoscope.probe

Probing.

The operator, its input, configs, enums, collector, method base and factory are importable from entoscope.probe; the operator, its input, configs and enums from entoscope too.

class entoscope.probe.config.PoolingType(*values)

Bases: StrEnum

How ProbeConfig.pool reduces an axis. last keeps its final entry, the last token for a position axis.

class entoscope.probe.config.ProbeMetric(*values)

Bases: StrEnum

A score of a probe’s predictions on its examples.

auroc ranks the positive-class probabilities, and the others compare the 0/1 predictions with the labels.

class entoscope.probe.config.ProbeMethodConfig(*, type: str)

Bases: MethodConfig

The method config of a probe that takes no fields.

A probe with fields of its own names a dataclass extending this one.

classmethod factory() → type[ProbeMethodFactory]

Name the factory probe methods register with.

Returns:

ProbeMethodFactory.

class entoscope.probe.config.ProbeConfig(*, config_id: str | None = None, method: ~entoscope.probe.config.ProbeMethodConfig, pool: dict[~entoscope.model.config.ModelAxis, ~entoscope.probe.config.PoolingType] = <factory>, metric: list[~entoscope.probe.config.ProbeMetric], source: ~entoscope.model.config.ModelLocatorConfig)

Bases: AbstractConfig

What a probing run trains, where, and how it is scored.

Variables:

• method (entoscope.probe.config.ProbeMethodConfig) – The probe trained at every location, its type naming it. In a YAML file it is the mapping read strictly into the probe’s method config class, or just the probe’s type for its defaults.

• pool (dict[entoscope.model.config.ModelAxis, entoscope.probe.config.PoolingType]) – Each axis to reduce, mapped to the pooling that reduces it, matched by name: attn.pattern’s positions are pooled as dest_pos and src_pos. Axes left out are flattened into the feature vector. layer is never a tensor axis and neuron is the feature axis, so neither can be pooled.

• metric (list[entoscope.probe.config.ProbeMetric]) – The scores to report for every probe.

• source (entoscope.model.config.ModelLocatorConfig) – Where to read the features.

Raises:

ValueError – If pool names layer or neuron, or an axis no activation of source has.

class entoscope.probe.input.ProbeInput(config: ProbeConfig, positive_prompts: list[str], negative_prompts: list[str])

Bases: OperationInput

Positive and negative example prompts whose pooled activations become probe training data.

The location grid is parsed once. When every prompt tokenizes to the same length the grid uses that length, so a filtered position resolves to a concrete index. When prompt lengths differ the grid is parsed without one: explicit positions stay symbolic, a non-negative one counting from the start and a negative one from each prompt’s own end, and a grid that would enumerate every position (an unfiltered position, dest_pos or src_pos axis) is rejected because its size would depend on the prompt. Position axes left outside the grid stay full slices and adapt to each prompt.

model.to_tokens prepends BOS, so mean or sum pooling over positions includes it while last pooling (or a -1 position) reads the final content token.

Variables:

• config (entoscope.probe.config.ProbeConfig) – The probing config.

• positive_prompts (list[str]) – The examples labeled 1.

• negative_prompts (list[str]) – The examples labeled 0.

prepare(model: TransformerBridge | CachedModel, *, index_only: bool = False) → ProbeInput

Resolve the locations the probes read, then pool every prompt’s activation at each into its features.

Parameters:

• model – The model the run reads.

• index_only – Whether to stop after the locations, before any forward pass.

Returns:

The input itself, prepared in place.

Raises:

ValueError – If either set of prompts is empty, the model computes no transformer activation the locator names, an explicit position is out of range for the shortest prompt, or a location’s feature size varies across prompts.

activation_names() → list[str]

Name the transformer activations the probes read, which a run caches and a plan reports.

Returns:

The transformer activation names, sorted.

property labels: Int[Tensor, 'examples']

Every example’s label, positives first, once the input is prepared.

property features: dict[int, Float[Tensor, 'examples features']]

Each location’s features, one row per example, keyed by result row, once the input is prepared.

property index: Index[ModelLocatorConfig]

The locations the probes read, once the index is prepared.

property token_tensors: list[Int[Tensor, 'batch pos']]

Every example’s tokens, positives first, once the index is prepared.

to(device: str | device = 'cpu') → ProbeInput

Move the tensors the input holds to device.

Parameters:

device – The device to move to.

Returns:

The input itself, moved in place.

class entoscope.probe.prober.Prober(model: TransformerBridge | CachedModel, collector_class: type[ProbeCollector] = ProbeCollector)

Bases: Operator

Trains one probe per location and scores it on the training examples, with no held-out split.

Each example prompt is run once, caching only what the index reads.

Parameters:

• model – The model to read.

• collector_class – The collector to write the results with. Probing draws no page: its result is a table of scores and the fitted probes.

apply(probe_input: ProbeInput, save_dir: str | None = None, return_collector: bool = False, verbose: bool = True) → ProbeCollector | None

Train and score a probe at every row of the input’s index.

The method is built before any forward pass, so a config that names something the run cannot do fails before the model is touched. Where the results were saved is logged through the entoscope logger.

Parameters:

• probe_input – The config and example prompts to run, prepared in place.

• save_dir – The directory to write result.csv, index.csv, config.json and probes.pkl into, created if missing. None writes nothing.

• return_collector – Whether to return the collector.

• verbose – Whether to show a progress bar.

Returns:

The collector, whose result holds every location’s scores and whose probes the fitted probes, if return_collector is set, and None otherwise.

plan(probe_input: ProbeInput, save_dir: str | None = None, verbose: bool = True) → OperationPlan

Describe the run apply would make, without a forward pass and without touching save_dir.

The input’s index is resolved exactly as the run resolves it, and kept for apply.

Parameters:

• probe_input – The config and example prompts the run would use.

• save_dir – The directory the run would write into.

• verbose – Whether to also log the plan.

Returns:

The plan – one probe per index row and one forward pass per example prompt.

class entoscope.probe.collector.ProbeCollector(operation_input: ProbeInput, model: TransformerBridge | CachedModel, save_dir: str | None = None)

Bases: FrameCollector[ProbeInput]

One row per location, with a column per metric the config names.

Writes result.csv, index.csv, config.json and probes.pkl: the fitted probes, pickled as a dict from result row to estimator.

Parameters:

• operation_input – The input the run iterates.

• model – The model the run reads.

• save_dir – The directory save writes into, or None to write nothing.

property probes: dict[int, Any]

The fitted probes, keyed by result row.

collect(method: AbstractProbeMethod, metric_values: dict[ProbeMetric, float], index: int)

Take one location’s fitted probe and its scores.

Parameters:

• method – The fitted probe.

• metric_values – Its score under every metric the config names.

• index – The location’s row in the index.

output_files() → list[str]

List the file names save writes into save_dir, which a plan reports as the run’s output.

Subclasses extend this to mirror their own _save_result. It is deliberately not abstract: collector_class is an extension point, so a collector that declares no files keeps working and is reported as writing none. A plan asks a collector nothing was collected into, so an override may read the input but never the collected result.

Returns:

The file names, relative to save_dir.

class entoscope.probe.metric.AbstractProbeMetric

Bases: ABC

A score of a probe’s predictions against the labels.

abstractmethod classmethod compute(labels: ndarray, predictions: ndarray, probabilities: ndarray) → float

Score a probe’s predictions.

Parameters:

• labels – The 0/1 labels.

• predictions – The probe’s 0/1 predictions.

• probabilities – The probe’s positive-class probabilities.

Returns:

The score.

static get(metric: ProbeMetric) → type[AbstractProbeMetric]

Pick the implementation of a metric.

Parameters:

metric – The metric.

Returns:

The metric’s class.

Raises:

ValueError – If the metric is unknown.

class entoscope.probe.metric.AccuracyMetric

Bases: AbstractProbeMetric

The share of predictions that match their label.

classmethod compute(labels: ndarray, predictions: ndarray, probabilities: ndarray) → float

Score a probe’s predictions.

Parameters:

• labels – The 0/1 labels.

• predictions – The probe’s 0/1 predictions.

• probabilities – The probe’s positive-class probabilities.

Returns:

The score.

class entoscope.probe.metric.F1Metric

Bases: AbstractProbeMetric

The harmonic mean of precision and recall, 0 when there are no positive predictions.

classmethod compute(labels: ndarray, predictions: ndarray, probabilities: ndarray) → float

Score a probe’s predictions.

Parameters:

• labels – The 0/1 labels.

• predictions – The probe’s 0/1 predictions.

• probabilities – The probe’s positive-class probabilities.

Returns:

The score.

class entoscope.probe.metric.PrecisionMetric

Bases: AbstractProbeMetric

The share of positive predictions that are positive, 0 when there are none.

classmethod compute(labels: ndarray, predictions: ndarray, probabilities: ndarray) → float

Score a probe’s predictions.

Parameters:

• labels – The 0/1 labels.

• predictions – The probe’s 0/1 predictions.

• probabilities – The probe’s positive-class probabilities.

Returns:

The score.

class entoscope.probe.metric.RecallMetric

Bases: AbstractProbeMetric

The share of positive examples predicted positive.

classmethod compute(labels: ndarray, predictions: ndarray, probabilities: ndarray) → float

Score a probe’s predictions.

Parameters:

• labels – The 0/1 labels.

• predictions – The probe’s 0/1 predictions.

• probabilities – The probe’s positive-class probabilities.

Returns:

The score.

class entoscope.probe.metric.AurocMetric

Bases: AbstractProbeMetric

The area under the ROC curve of the positive-class probabilities.

classmethod compute(labels: ndarray, predictions: ndarray, probabilities: ndarray) → float

Score a probe’s predictions.

Parameters:

• labels – The 0/1 labels.

• predictions – The probe’s 0/1 predictions.

• probabilities – The probe’s positive-class probabilities.

Returns:

The score.

class entoscope.probe.pool.AbstractPooling

Bases: ABC

Reduces one axis of an activation.

abstractmethod classmethod reduce(activation: Float[Tensor, '...'], dim: int) → Float[Tensor, '...']

Reduce one axis.

Parameters:

• activation – The activation.

• dim – The axis to reduce.

Returns:

The activation without that axis.

static get(pooling_type: PoolingType) → type[AbstractPooling]

Pick the implementation of a pooling type.

Parameters:

pooling_type – The pooling type.

Returns:

The pooling’s class.

Raises:

ValueError – If the pooling type is unknown.

class entoscope.probe.pool.MeanPooling

Bases: AbstractPooling

The mean along the axis.

classmethod reduce(activation: Float[Tensor, '...'], dim: int) → Float[Tensor, '...']

Reduce one axis.

Parameters:

• activation – The activation.

• dim – The axis to reduce.

Returns:

The activation without that axis.

class entoscope.probe.pool.MaxPooling

Bases: AbstractPooling

The maximum along the axis.

classmethod reduce(activation: Float[Tensor, '...'], dim: int) → Float[Tensor, '...']

Reduce one axis.

Parameters:

• activation – The activation.

• dim – The axis to reduce.

Returns:

The activation without that axis.

class entoscope.probe.pool.SumPooling

Bases: AbstractPooling

The sum along the axis.

classmethod reduce(activation: Float[Tensor, '...'], dim: int) → Float[Tensor, '...']

Reduce one axis.

Parameters:

• activation – The activation.

• dim – The axis to reduce.

Returns:

The activation without that axis.

class entoscope.probe.pool.LastPooling

Bases: AbstractPooling

The axis’s final entry.

classmethod reduce(activation: Float[Tensor, '...'], dim: int) → Float[Tensor, '...']

Reduce one axis.

Parameters:

• activation – The activation.

• dim – The axis to reduce.

Returns:

The activation without that axis.

entoscope.probe.pool.pool(activation: Float[Tensor, '...'], axes: Sequence[ModelAxis | str], config: dict[ModelAxis, PoolingType]) → Float[Tensor, 'features']

Reduce each axis the config names with its pooling type, and flatten the rest into features.

Configured axes are reduced left to right, which matters when different pooling types are mixed. Config entries that match no axis are ignored, such as an axis the location grid already pins.

Parameters:

• activation – One example’s activation at one location.

• axes – The activation’s axes, in order.

• config – Each axis to reduce, mapped to its pooling type.

Returns:

The feature vector, a single feature for a 0-D activation.

class entoscope.probe.method.probe.AbstractProbeMethod(model: TransformerBridge | CachedModel, config: ProbeConfig)

Bases: ABC, Generic[ProbeMethodConfigT]

The base of every probe: a classifier trained on the pooled activations of one location.

Built afresh for every index row, since each location trains its own probe, and scored on the examples it was fitted on. Its type argument names the dataclass of its method config, extending ProbeMethodConfig, which a config file’s method is read into. A probe naming none takes ProbeMethodConfig itself, which no field fits.

Parameters:

• model – The model the features are read from.

• config – The run’s config.

Variables:

method_config – The run’s method config, typed by the probe’s type argument.

abstractmethod fit(features: Float[Tensor, 'examples features'], labels: Int[Tensor, 'examples'])

Train the probe.

Parameters:

• features – One feature vector per example.

• labels – Every example’s 0/1 label.

abstractmethod predict(features: Float[Tensor, 'examples features']) → ndarray

Predict labels.

Parameters:

features – One feature vector per example.

Returns:

Every example’s predicted 0/1 label.

abstractmethod predict_proba(features: Float[Tensor, 'examples features']) → ndarray

Predict the positive class’s probability.

Parameters:

features – One feature vector per example.

Returns:

Every example’s positive-class probability.

abstract property probe: Any

The fitted estimator, which the collector pickles into probes.pkl, so it must pickle.

class entoscope.probe.method.linear.LinearProbeConfig(*, type: str, max_iter: int = 1000, seed: int = 42)

Bases: ProbeMethodConfig

The method config of linear, a logistic regression.

Variables:

• max_iter (int) – The most iterations its solver runs.

• seed (int) – The seed of its randomness.

class entoscope.probe.method.linear.LinearProbe(model: TransformerBridge | CachedModel, config: ProbeConfig)

Bases: AbstractProbeMethod[LinearProbeConfig]

A scikit-learn logistic regression.

Parameters:

• model – The model the features are read from.

• config – The run’s config.

fit(features: Float[Tensor, 'examples features'], labels: Int[Tensor, 'examples'])

Train the probe.

Parameters:

• features – One feature vector per example.

• labels – Every example’s 0/1 label.

predict(features: Float[Tensor, 'examples features']) → ndarray

Predict labels.

Parameters:

features – One feature vector per example.

Returns:

Every example’s predicted 0/1 label.

predict_proba(features: Float[Tensor, 'examples features']) → ndarray

Predict the positive class’s probability.

Parameters:

features – One feature vector per example.

Returns:

Every example’s positive-class probability.

property probe: LogisticRegression

The fitted estimator, which the collector pickles into probes.pkl, so it must pickle.

class entoscope.probe.method.factory.ProbeMethodFactory

Bases: RegistryFactory[AbstractProbeMethod]

Resolves the probe a config’s method.type names and builds one per location.

It holds linear and every probe registered with register.

classmethod create(model: TransformerBridge | CachedModel, config: ProbeConfig) → AbstractProbeMethod

Build a fresh, unfitted probe of the kind a config names, called once per index element.

Parameters:

• model – The model the features are read from.

• config – The run’s config, whose method.type names the probe.

Returns:

The probe.

Raises:

ValueError – If no probe is registered under the name, or the probe rejects the config.