activations.py

# Dataclasses and enums for storing neuron-indexed information about activations. Also, related
# helper functions.

import math
from dataclasses import dataclass, field
from typing import List, Optional, Union

import blobfile as bf
import boostedblob as bbb
from neuron_explainer.fast_dataclasses import FastDataclass, loads, register_dataclass


@register_dataclass
@dataclass
class ActivationRecord(FastDataclass):
    """Collated lists of tokens and their activations for a single neuron."""

    tokens: List[str]
    """Tokens in the text sequence, represented as strings."""
    activations: List[float]
    """Raw activation values for the neuron on each token in the text sequence."""


@register_dataclass
@dataclass
class NeuronId(FastDataclass):
    """Identifier for a neuron in an artificial neural network."""

    layer_index: int
    """The index of layer the neuron is in. The first layer used during inference has index 0."""
    neuron_index: int
    """The neuron's index within in its layer. Indices start from 0 in each layer."""


def _check_slices(
    slices_by_split: dict[str, slice],
    expected_num_values: int,
) -> None:
    """Assert that the slices are disjoint and fully cover the intended range."""
    indices = set()
    sum_of_slice_lengths = 0
    n_splits = len(slices_by_split.keys())
    for s in slices_by_split.values():
        subrange = range(expected_num_values)[s]
        sum_of_slice_lengths += len(subrange)
        indices |= set(subrange)
    assert (
        sum_of_slice_lengths == expected_num_values
    ), f"{sum_of_slice_lengths=} != {expected_num_values=}"
    stride = n_splits
    expected_indices = set.union(
        *[set(range(start_index, expected_num_values, stride)) for start_index in range(n_splits)]
    )
    assert indices == expected_indices, f"{indices=} != {expected_indices=}"


def get_slices_for_splits(
    splits: list[str],
    num_activation_records_per_split: int,
) -> dict[str, slice]:
    """
    Get equal-sized interleaved subsets for each of a list of splits, given the number of elements
    to include in each split.
    """

    stride = len(splits)
    num_activation_records_for_even_splits = num_activation_records_per_split * stride
    slices_by_split = {
        split: slice(split_index, num_activation_records_for_even_splits, stride)
        for split_index, split in enumerate(splits)
    }
    _check_slices(
        slices_by_split=slices_by_split,
        expected_num_values=num_activation_records_for_even_splits,
    )
    return slices_by_split


@dataclass
class ActivationRecordSliceParams:
    """How to select splits (train, valid, etc.) of activation records."""

    n_examples_per_split: Optional[int]
    """The number of examples to include in each split."""


@register_dataclass
@dataclass
class NeuronRecord(FastDataclass):
    """Neuron-indexed activation data, including summary stats and notable activation records."""

    neuron_id: NeuronId
    """Identifier for the neuron."""

    random_sample: list[ActivationRecord] = field(default_factory=list)
    """
    Random activation records for this neuron. The random sample is independent from those used for
    other neurons.
    """
    random_sample_by_quantile: Optional[list[list[ActivationRecord]]] = None
    """
    Random samples of activation records in each of the specified quantiles. None if quantile
    tracking is disabled.
    """
    quantile_boundaries: Optional[list[float]] = None
    """Boundaries of the quantiles used to generate the random_sample_by_quantile field."""

    # Moments of activations
    mean: Optional[float] = math.nan
    variance: Optional[float] = math.nan
    skewness: Optional[float] = math.nan
    kurtosis: Optional[float] = math.nan

    most_positive_activation_records: list[ActivationRecord] = field(default_factory=list)
    """
    Activation records with the most positive figure of merit value for this neuron over all dataset
    examples.
    """

    @property
    def max_activation(self) -> float:
        """Return the maximum activation value over all top-activating activation records."""
        return max([max(ar.activations) for ar in self.most_positive_activation_records])

    def _get_top_activation_slices(
        self, activation_record_slice_params: ActivationRecordSliceParams
    ) -> dict[str, slice]:
        splits = ["train", "calibration", "valid", "test"]
        n_examples_per_split = activation_record_slice_params.n_examples_per_split
        if n_examples_per_split is None:
            n_examples_per_split = len(self.most_positive_activation_records) // len(splits)
        assert len(self.most_positive_activation_records) >= n_examples_per_split * len(splits)
        return get_slices_for_splits(splits, n_examples_per_split)

    def _get_random_activation_slices(
        self, activation_record_slice_params: ActivationRecordSliceParams
    ) -> dict[str, slice]:
        splits = ["calibration", "valid", "test"]
        n_examples_per_split = activation_record_slice_params.n_examples_per_split
        if n_examples_per_split is None:
            n_examples_per_split = len(self.random_sample) // len(splits)
        # NOTE: this assert could trigger on some old datasets with only 10 random samples, in which case you may have to remove "test" from the set of splits
        assert len(self.random_sample) >= n_examples_per_split * len(splits)
        return get_slices_for_splits(splits, n_examples_per_split)

    def train_activation_records(
        self,
        activation_record_slice_params: ActivationRecordSliceParams,
    ) -> list[ActivationRecord]:
        """
        Train split, typically used for generating explanations. Consists exclusively of
        top-activating records since context window limitations make it difficult to include
        random records.
        """
        return self.most_positive_activation_records[
            self._get_top_activation_slices(activation_record_slice_params)["train"]
        ]

    def calibration_activation_records(
        self,
        activation_record_slice_params: ActivationRecordSliceParams,
    ) -> list[ActivationRecord]:
        """
        Calibration split, typically used for calibrating neuron simulations. See
        http://go/neuron_explanation_methodology for an explanation of calibration. Consists of
        top-activating records and random records in a 1:1 ratio.
        """
        return (
            self.most_positive_activation_records[
                self._get_top_activation_slices(activation_record_slice_params)["calibration"]
            ]
            + self.random_sample[
                self._get_random_activation_slices(activation_record_slice_params)["calibration"]
            ]
        )

    def valid_activation_records(
        self,
        activation_record_slice_params: ActivationRecordSliceParams,
    ) -> list[ActivationRecord]:
        """
        Validation split, typically used for evaluating explanations, either automatically with
        simulation + correlation coefficient scoring, or manually by humans. Consists of
        top-activating records and random records in a 1:1 ratio.
        """
        return (
            self.most_positive_activation_records[
                self._get_top_activation_slices(activation_record_slice_params)["valid"]
            ]
            + self.random_sample[
                self._get_random_activation_slices(activation_record_slice_params)["valid"]
            ]
        )

    def test_activation_records(
        self,
        activation_record_slice_params: ActivationRecordSliceParams,
    ) -> list[ActivationRecord]:
        """
        Test split, typically used for explanation evaluations that can't use the validation split.
        Consists of top-activating records and random records in a 1:1 ratio.
        """
        return (
            self.most_positive_activation_records[
                self._get_top_activation_slices(activation_record_slice_params)["test"]
            ]
            + self.random_sample[
                self._get_random_activation_slices(activation_record_slice_params)["test"]
            ]
        )


def neuron_exists(
    dataset_path: str, layer_index: Union[str, int], neuron_index: Union[str, int]
) -> bool:
    """Return whether the specified neuron exists."""
    file = bf.join(dataset_path, "neurons", str(layer_index), f"{neuron_index}.json")
    return bf.exists(file)


def load_neuron(
    layer_index: Union[str, int], neuron_index: Union[str, int],
    dataset_path: str = "az://openaipublic/neuron-explainer/data/collated-activations",
) -> NeuronRecord:
    """Load the NeuronRecord for the specified neuron."""
    file = bf.join(dataset_path, str(layer_index), f"{neuron_index}.json")
    with bf.BlobFile(file, "r") as f:
        neuron_record = loads(f.read())
        if not isinstance(neuron_record, NeuronRecord):
            raise ValueError(
                f"Stored data incompatible with current version of NeuronRecord dataclass."
            )
        return neuron_record


@bbb.ensure_session
async def load_neuron_async(
    layer_index: Union[str, int], neuron_index: Union[str, int],
    dataset_path: str = "az://openaipublic/neuron-explainer/data/collated-activations",
) -> NeuronRecord:
    """Async version of load_neuron."""
    file = bf.join(dataset_path, str(layer_index), f"{neuron_index}.json")
    return await read_neuron_file(file)


@bbb.ensure_session
async def read_neuron_file(neuron_filename: str) -> NeuronRecord:
    """Like load_neuron_async, but takes a raw neuron filename."""
    raw_contents = await bbb.read.read_single(neuron_filename)
    neuron_record = loads(raw_contents.decode("utf-8"))
    if not isinstance(neuron_record, NeuronRecord):
        raise ValueError(
            f"Stored data incompatible with current version of NeuronRecord dataclass."
        )
    return neuron_record


def get_sorted_neuron_indices(dataset_path: str, layer_index: Union[str, int]) -> List[int]:
    """Returns the indices of all neurons in this layer, in ascending order."""
    layer_dir = bf.join(dataset_path, "neurons", str(layer_index))
    return sorted(
        [int(f.split(".")[0]) for f in bf.listdir(layer_dir) if f.split(".")[0].isnumeric()]
    )


def get_sorted_layers(dataset_path: str) -> List[str]:
    """
    Return the indices of all layers in this dataset, in ascending numerical order, as strings.
    """
    return [
        str(x)
        for x in sorted(
            [int(x) for x in bf.listdir(bf.join(dataset_path, "neurons")) if x.isnumeric()]
        )
    ]