Skip to content

API Reference

The complete public API, generated from the source docstrings. Core symbols are importable from the top-level kamui package; the distributed-training helpers live in kamui.training.

Configuration & Model

kamui.model.config.ModelConfig dataclass

Configuration class for the KAMUI transformer model.

Holds all hyperparameters, performs validation, computes helper dimensions, and estimates the parameter count under explicit architectural assumptions.

Attributes:

Name Type Description
n_layers int

Number of transformer blocks.
d_model int

Residual stream / embedding dimension.
n_heads int

Number of attention heads.
d_ff int

Feed-forward network hidden layer dimension.
vocab_size int

Size of the token vocabulary.
context_length int

Maximum sequence length.
positional_encoding str

Type of positional encoding ('learned', 'sinusoidal', or 'rope'). 'rope' (rotary) is applied to Q/K inside attention and requires an even d_head.
dropout float

Dropout rate (must be in [0, 1]).
Source code in kamui/model/config.py 
 13
 14
 15
 16
 17
 18
 19
 20
 21
 22
 23
 24
 25
 26
 27
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
@dataclass
class ModelConfig:
    """Configuration class for the KAMUI transformer model.

    Holds all hyperparameters, performs validation, computes helper dimensions,
    and estimates the parameter count under explicit architectural assumptions.

    Attributes:
        n_layers: Number of transformer blocks.
        d_model: Residual stream / embedding dimension.
        n_heads: Number of attention heads.
        d_ff: Feed-forward network hidden layer dimension.
        vocab_size: Size of the token vocabulary.
        context_length: Maximum sequence length.
        positional_encoding: Type of positional encoding ('learned', 'sinusoidal',
            or 'rope'). 'rope' (rotary) is applied to Q/K inside attention and
            requires an even d_head.
        dropout: Dropout rate (must be in [0, 1]).
    """

    n_layers: int = 6
    d_model: int = 256
    n_heads: int = 8
    d_ff: int = 1024
    vocab_size: int = 8192
    context_length: int = 256
    positional_encoding: str = "learned"
    dropout: float = 0.0

    def __post_init__(self) -> None:  # noqa: C901 — flat list of independent guards
        """Enforces tensor dimension and hyperparameter sanity checks.

        Raises:
            ValueError: If any validation rule is violated.
        """
        if self.n_layers <= 0:
            raise ValueError(f"n_layers must be > 0, got {self.n_layers}")
        if self.d_model <= 0:
            raise ValueError(f"d_model must be > 0, got {self.d_model}")
        if self.n_heads <= 0:
            raise ValueError(f"n_heads must be > 0, got {self.n_heads}")
        if self.d_ff <= 0:
            raise ValueError(f"d_ff must be > 0, got {self.d_ff}")
        if self.vocab_size <= 0:
            raise ValueError(f"vocab_size must be > 0, got {self.vocab_size}")
        if self.context_length <= 0:
            raise ValueError(f"context_length must be > 0, got {self.context_length}")

        if self.d_model % self.n_heads != 0:
            raise ValueError(
                f"d_model ({self.d_model}) must be divisible by n_heads ({self.n_heads}) "
                f"to yield an integer d_head."
            )

        if not (0.0 <= self.dropout <= 1.0):
            raise ValueError(f"dropout must be in [0, 1], got {self.dropout}")

        if self.positional_encoding not in ("learned", "sinusoidal", "rope"):
            raise ValueError(
                f"positional_encoding must be 'learned', 'sinusoidal', or 'rope', "
                f"got '{self.positional_encoding}'"
            )

        # RoPE rotates the head dimension in 2-D planes, so d_head must be even.
        if self.positional_encoding == "rope" and self.d_head % 2 != 0:
            raise ValueError(
                f"positional_encoding='rope' requires an even d_head, got "
                f"d_head={self.d_head} (d_model={self.d_model}, n_heads={self.n_heads})"
            )

    @property
    def d_head(self) -> int:
        """Dimension of each attention head.

        Calculated as d_model // n_heads.
        """
        return self.d_model // self.n_heads

    @property
    def attention_parameters(self) -> int:
        """Estimates total attention parameters across all layers.

        Assumptions:
            - Multi-head causal self-attention is applied at each of the `n_layers`.
            - In each layer, the input (residual stream of shape d_model) is projected
              to Query, Key, and Value spaces. Each of these 3 projections is a linear
              layer of shape (d_model, d_model) with a bias vector of shape (d_model).
            - The output of the combined heads is projected back to the residual stream
              using an output projection of shape (d_model, d_model) with a bias of shape (d_model).
            - Formula: n_layers * 4 * d_model * (d_model + 1)
        """
        return self.n_layers * 4 * self.d_model * (self.d_model + 1)

    @property
    def feedforward_parameters(self) -> int:
        """Estimates total feed-forward network parameters across all layers.

        Assumptions:
            - A position-wise MLP is applied at each of the `n_layers`.
            - The MLP consists of two linear layers:
              1. Expansion layer: weight of shape (d_model, d_ff) and bias (d_ff).
              2. Projection layer: weight of shape (d_ff, d_model) and bias (d_model).
            - Formula: n_layers * (2 * d_model * d_ff + d_model + d_ff)
        """
        return self.n_layers * (2 * self.d_model * self.d_ff + self.d_model + self.d_ff)

    @property
    def embedding_parameters(self) -> int:
        """Estimates total embedding parameters.

        Assumptions:
            - Token embedding matrix of shape (vocab_size, d_model).
            - If `positional_encoding` is "learned", a learned positional embedding
              matrix of shape (context_length, d_model) is present.
            - If `positional_encoding` is "sinusoidal" or "rope", positional
              encodings are fixed and add 0 trainable parameters.
            - Formula: (vocab_size * d_model) + (context_length * d_model if learned else 0)
        """
        pos_params = (
            self.context_length * self.d_model if self.positional_encoding == "learned" else 0
        )
        return (self.vocab_size * self.d_model) + pos_params

    @property
    def estimated_total_parameters(self) -> int:
        """Estimates the total trainable parameter count of the model.

        Assumptions:
            - Includes attention, feedforward, and embedding parameters.
            - Includes normalization layer parameters:
              - 2 LayerNorms per block (attention input, FFN input). Each LayerNorm has
                learnable scale (gamma) and bias (beta) parameters of shape (d_model,).
              - 1 final LayerNorm before unembedding, also with scale and bias (d_model,).
              - LayerNorm parameters formula: n_layers * 4 * d_model + 2 * d_model.
            - Includes unembedding layer parameters:
                - The weight matrix is tied to the token embedding matrix transpose, so it
                  adds 0 trainable parameters.
                - The unembedding linear layer has NO learnable bias (bias=False), matching
                  standard GPT-2 weight-tying conventions.
        """
        ln_params = (self.n_layers * 4 * self.d_model) + (2 * self.d_model)
        return (
            self.embedding_parameters
            + self.attention_parameters
            + self.feedforward_parameters
            + ln_params
        )

    def __repr__(self) -> str:
        """Returns a helpful, formatted string representation of the model configuration."""
        return (
            f"ModelConfig(\n"
            f"  n_layers={self.n_layers},\n"
            f"  d_model={self.d_model},\n"
            f"  n_heads={self.n_heads} (d_head={self.d_head}),\n"
            f"  d_ff={self.d_ff},\n"
            f"  vocab_size={self.vocab_size},\n"
            f"  context_length={self.context_length},\n"
            f"  positional_encoding='{self.positional_encoding}',\n"
            f"  dropout={self.dropout},\n"
            f"  estimated_total_parameters={self.estimated_total_parameters:,}\n"
            f")"
        )

    @classmethod
    def from_yaml(cls, path: str | Path) -> "ModelConfig":
        """Loads a ModelConfig from a YAML file.

        The YAML file must contain a 'model' section.

        Args:
            path: Path to the YAML file.

        Returns:
            An instance of ModelConfig.
        """
        import yaml

        try:
            with open(path, encoding="utf-8") as f:
                data = yaml.safe_load(f)
        except Exception as e:
            raise ValueError(f"Failed to parse YAML file at {path}: {e}") from e

        if not isinstance(data, dict):
            raise ValueError(f"YAML file at {path} must contain a dictionary mapping.")

        if "model" not in data:
            raise ValueError("YAML configuration must contain a 'model' section.")

        model_data = data["model"]
        if not isinstance(model_data, dict):
            raise ValueError("Model configuration section in YAML must be a dictionary.")

        # Check for unknown configuration keys
        dataclass_fields = {field.name for field in fields(cls)}
        for key in model_data:
            if key not in dataclass_fields:
                raise ValueError(f"Unknown configuration key in 'model' section: '{key}'")

        return cls(**model_data)

    def to_yaml(self, path: str | Path) -> None:
        """Serializes the ModelConfig instance to a YAML file.

        If the file already exists, preserves other sections (e.g., 'training',
        'data', 'logging') and only updates the 'model' block.

        Args:
            path: Path where the YAML file should be written.
        """
        import yaml

        path_obj = Path(path)
        existing_data = {}
        if path_obj.exists() and path_obj.is_file():
            try:
                with open(path_obj, encoding="utf-8") as f:
                    content = yaml.safe_load(f)
                    if isinstance(content, dict):
                        existing_data = content
            except Exception:
                pass

        existing_data["model"] = asdict(self)

        with open(path_obj, "w", encoding="utf-8") as f:
            yaml.safe_dump(existing_data, f, default_flow_style=False, sort_keys=False)

d_head property

Dimension of each attention head.

Calculated as d_model // n_heads.

attention_parameters property

Estimates total attention parameters across all layers.

Assumptions
  • Multi-head causal self-attention is applied at each of the n_layers.
  • In each layer, the input (residual stream of shape d_model) is projected to Query, Key, and Value spaces. Each of these 3 projections is a linear layer of shape (d_model, d_model) with a bias vector of shape (d_model).
  • The output of the combined heads is projected back to the residual stream using an output projection of shape (d_model, d_model) with a bias of shape (d_model).
  • Formula: n_layers * 4 * d_model * (d_model + 1)

feedforward_parameters property

Estimates total feed-forward network parameters across all layers.

Assumptions
  • A position-wise MLP is applied at each of the n_layers.
  • The MLP consists of two linear layers:
  • Expansion layer: weight of shape (d_model, d_ff) and bias (d_ff).
  • Projection layer: weight of shape (d_ff, d_model) and bias (d_model).
  • Formula: n_layers * (2 * d_model * d_ff + d_model + d_ff)

embedding_parameters property

Estimates total embedding parameters.

Assumptions
  • Token embedding matrix of shape (vocab_size, d_model).
  • If positional_encoding is "learned", a learned positional embedding matrix of shape (context_length, d_model) is present.
  • If positional_encoding is "sinusoidal" or "rope", positional encodings are fixed and add 0 trainable parameters.
  • Formula: (vocab_size * d_model) + (context_length * d_model if learned else 0)

estimated_total_parameters property

Estimates the total trainable parameter count of the model.

Assumptions
  • Includes attention, feedforward, and embedding parameters.
  • Includes normalization layer parameters:
  • 2 LayerNorms per block (attention input, FFN input). Each LayerNorm has learnable scale (gamma) and bias (beta) parameters of shape (d_model,).
  • 1 final LayerNorm before unembedding, also with scale and bias (d_model,).
  • LayerNorm parameters formula: n_layers * 4 * d_model + 2 * d_model.
  • Includes unembedding layer parameters:
    • The weight matrix is tied to the token embedding matrix transpose, so it adds 0 trainable parameters.
    • The unembedding linear layer has NO learnable bias (bias=False), matching standard GPT-2 weight-tying conventions.

__post_init__()

Enforces tensor dimension and hyperparameter sanity checks.

Raises:

Type Description
ValueError

If any validation rule is violated.
Source code in kamui/model/config.py 
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
def __post_init__(self) -> None:  # noqa: C901 — flat list of independent guards
    """Enforces tensor dimension and hyperparameter sanity checks.

    Raises:
        ValueError: If any validation rule is violated.
    """
    if self.n_layers <= 0:
        raise ValueError(f"n_layers must be > 0, got {self.n_layers}")
    if self.d_model <= 0:
        raise ValueError(f"d_model must be > 0, got {self.d_model}")
    if self.n_heads <= 0:
        raise ValueError(f"n_heads must be > 0, got {self.n_heads}")
    if self.d_ff <= 0:
        raise ValueError(f"d_ff must be > 0, got {self.d_ff}")
    if self.vocab_size <= 0:
        raise ValueError(f"vocab_size must be > 0, got {self.vocab_size}")
    if self.context_length <= 0:
        raise ValueError(f"context_length must be > 0, got {self.context_length}")

    if self.d_model % self.n_heads != 0:
        raise ValueError(
            f"d_model ({self.d_model}) must be divisible by n_heads ({self.n_heads}) "
            f"to yield an integer d_head."
        )

    if not (0.0 <= self.dropout <= 1.0):
        raise ValueError(f"dropout must be in [0, 1], got {self.dropout}")

    if self.positional_encoding not in ("learned", "sinusoidal", "rope"):
        raise ValueError(
            f"positional_encoding must be 'learned', 'sinusoidal', or 'rope', "
            f"got '{self.positional_encoding}'"
        )

    # RoPE rotates the head dimension in 2-D planes, so d_head must be even.
    if self.positional_encoding == "rope" and self.d_head % 2 != 0:
        raise ValueError(
            f"positional_encoding='rope' requires an even d_head, got "
            f"d_head={self.d_head} (d_model={self.d_model}, n_heads={self.n_heads})"
        )

__repr__()

Returns a helpful, formatted string representation of the model configuration.

Source code in kamui/model/config.py 
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
def __repr__(self) -> str:
    """Returns a helpful, formatted string representation of the model configuration."""
    return (
        f"ModelConfig(\n"
        f"  n_layers={self.n_layers},\n"
        f"  d_model={self.d_model},\n"
        f"  n_heads={self.n_heads} (d_head={self.d_head}),\n"
        f"  d_ff={self.d_ff},\n"
        f"  vocab_size={self.vocab_size},\n"
        f"  context_length={self.context_length},\n"
        f"  positional_encoding='{self.positional_encoding}',\n"
        f"  dropout={self.dropout},\n"
        f"  estimated_total_parameters={self.estimated_total_parameters:,}\n"
        f")"
    )

from_yaml(path) classmethod

Loads a ModelConfig from a YAML file.

The YAML file must contain a 'model' section.

Parameters:

Name Type Description Default
path str | Path

Path to the YAML file.

 required

Returns:

Type Description
ModelConfig

An instance of ModelConfig.
Source code in kamui/model/config.py 
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
@classmethod
def from_yaml(cls, path: str | Path) -> "ModelConfig":
    """Loads a ModelConfig from a YAML file.

    The YAML file must contain a 'model' section.

    Args:
        path: Path to the YAML file.

    Returns:
        An instance of ModelConfig.
    """
    import yaml

    try:
        with open(path, encoding="utf-8") as f:
            data = yaml.safe_load(f)
    except Exception as e:
        raise ValueError(f"Failed to parse YAML file at {path}: {e}") from e

    if not isinstance(data, dict):
        raise ValueError(f"YAML file at {path} must contain a dictionary mapping.")

    if "model" not in data:
        raise ValueError("YAML configuration must contain a 'model' section.")

    model_data = data["model"]
    if not isinstance(model_data, dict):
        raise ValueError("Model configuration section in YAML must be a dictionary.")

    # Check for unknown configuration keys
    dataclass_fields = {field.name for field in fields(cls)}
    for key in model_data:
        if key not in dataclass_fields:
            raise ValueError(f"Unknown configuration key in 'model' section: '{key}'")

    return cls(**model_data)

to_yaml(path)

Serializes the ModelConfig instance to a YAML file.

If the file already exists, preserves other sections (e.g., 'training', 'data', 'logging') and only updates the 'model' block.

Parameters:

Name Type Description Default
path str | Path

Path where the YAML file should be written.

 required
Source code in kamui/model/config.py 
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
def to_yaml(self, path: str | Path) -> None:
    """Serializes the ModelConfig instance to a YAML file.

    If the file already exists, preserves other sections (e.g., 'training',
    'data', 'logging') and only updates the 'model' block.

    Args:
        path: Path where the YAML file should be written.
    """
    import yaml

    path_obj = Path(path)
    existing_data = {}
    if path_obj.exists() and path_obj.is_file():
        try:
            with open(path_obj, encoding="utf-8") as f:
                content = yaml.safe_load(f)
                if isinstance(content, dict):
                    existing_data = content
        except Exception:
            pass

    existing_data["model"] = asdict(self)

    with open(path_obj, "w", encoding="utf-8") as f:
        yaml.safe_dump(existing_data, f, default_flow_style=False, sort_keys=False)

kamui.model.transformer.KAMUITransformer

Bases: Module

The full decoder-only transformer language model.

Assembles the embedding, a stack of n_layers transformer blocks, a final LayerNorm, and a weight-tied linear unembedding. Given token IDs it returns next-token logits; given targets as well it returns the cross-entropy loss.

Attributes:

Name Type Description
embed

Combined token + positional embedding.
blocks

nn.ModuleList of TransformerBlock.
final_ln

LayerNorm applied before the unembedding.
unembed

Linear d_model → vocab_size (weight tied to the token embedding matrix; no bias).
Source code in kamui/model/transformer.py 
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
class KAMUITransformer(nn.Module):
    """The full decoder-only transformer language model.

    Assembles the embedding, a stack of ``n_layers`` transformer blocks, a
    final LayerNorm, and a weight-tied linear unembedding.  Given token IDs
    it returns next-token logits; given targets as well it returns the
    cross-entropy loss.

    Attributes:
        embed:    Combined token + positional embedding.
        blocks:   ``nn.ModuleList`` of ``TransformerBlock``.
        final_ln: LayerNorm applied before the unembedding.
        unembed:  Linear ``d_model → vocab_size`` (weight tied to the token
            embedding matrix; no bias).
    """

    def __init__(self, config: ModelConfig) -> None:
        """Construct the model from a configuration and initialise weights.

        Args:
            config: The model configuration.
        """
        super().__init__()
        self.config = config

        self.embed = Embedding(config)
        self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layers))
        self.final_ln = LayerNorm(config.d_model)
        self.unembed = nn.Linear(config.d_model, config.vocab_size, bias=False)

        # GPT-2-style scaled initialisation (before tying, so the shared
        # matrix ends up with the standard 0.02 embedding init).
        init_weights(self, config.n_layers)

        # Weight tying: the unembedding shares the token-embedding matrix.
        self.unembed.weight = self.embed.token_embedding.weight

    # ------------------------------------------------------------------
    # Constructors
    # ------------------------------------------------------------------

    @classmethod
    def from_config(cls, config: ModelConfig) -> KAMUITransformer:
        """Construct a model from a ``ModelConfig`` instance."""
        return cls(config)

    @classmethod
    def from_yaml(cls, path: str | Path) -> KAMUITransformer:
        """Load a ``ModelConfig`` from YAML and construct the model.

        Args:
            path: Path to a YAML file with a ``model`` section.
        """
        return cls(ModelConfig.from_yaml(path))

    # ------------------------------------------------------------------
    # Forward
    # ------------------------------------------------------------------

    def forward(self, token_ids: Tensor, targets: Tensor | None = None) -> Tensor:
        """Run the model.

        Args:
            token_ids: Integer tensor of shape ``(B, S)``.
            targets:   Optional next-token labels of shape ``(B, S)``.  When
                provided, the return value is the scalar cross-entropy loss.

        Returns:
            Logits of shape ``(B, S, vocab_size)`` if ``targets`` is None,
            otherwise the scalar cross-entropy loss.

        Raises:
            TypeError:  If ``targets`` is provided but is not a tensor.
            ValueError: If ``targets`` is provided but its shape differs from
                ``token_ids``.
        """
        x = self.embed(token_ids)  # (B, S, D)
        for block in self.blocks:
            x = block(x)  # (B, S, D)
        x = self.final_ln(x)  # (B, S, D)
        logits = self.unembed(x)  # (B, S, V)

        if targets is None:
            return logits

        if not isinstance(targets, Tensor):
            raise TypeError(f"targets must be a torch.Tensor, got {type(targets)}")
        if targets.shape != token_ids.shape:
            raise ValueError(
                f"targets shape {tuple(targets.shape)} must match token_ids "
                f"shape {tuple(token_ids.shape)}"
            )

        # Cross-entropy over all positions; flatten batch and sequence.
        loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), targets.reshape(-1))
        return loss

    # ------------------------------------------------------------------
    # Introspection
    # ------------------------------------------------------------------

    def num_parameters(self, trainable_only: bool = True) -> int:
        """Return the number of parameters.

        ``nn.Module.parameters()`` already yields each shared (weight-tied)
        parameter once, so the tied embedding/unembedding matrix is counted
        a single time.

        Args:
            trainable_only: If True (default), count only parameters with
                ``requires_grad=True``.

        Returns:
            The parameter count.
        """
        return sum(p.numel() for p in self.parameters() if p.requires_grad or not trainable_only)

    def __repr__(self) -> str:
        return (
            f"KAMUITransformer(n_layers={self.config.n_layers}, "
            f"d_model={self.config.d_model}, n_heads={self.config.n_heads}, "
            f"vocab_size={self.config.vocab_size}, "
            f"num_parameters={self.num_parameters():,})"
        )

__init__(config)

Construct the model from a configuration and initialise weights.

Parameters:

Name Type Description Default
config ModelConfig

The model configuration.

 required
Source code in kamui/model/transformer.py 
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
def __init__(self, config: ModelConfig) -> None:
    """Construct the model from a configuration and initialise weights.

    Args:
        config: The model configuration.
    """
    super().__init__()
    self.config = config

    self.embed = Embedding(config)
    self.blocks = nn.ModuleList(TransformerBlock(config) for _ in range(config.n_layers))
    self.final_ln = LayerNorm(config.d_model)
    self.unembed = nn.Linear(config.d_model, config.vocab_size, bias=False)

    # GPT-2-style scaled initialisation (before tying, so the shared
    # matrix ends up with the standard 0.02 embedding init).
    init_weights(self, config.n_layers)

    # Weight tying: the unembedding shares the token-embedding matrix.
    self.unembed.weight = self.embed.token_embedding.weight

from_config(config) classmethod

Construct a model from a ModelConfig instance.

Source code in kamui/model/transformer.py 
115
116
117
118
@classmethod
def from_config(cls, config: ModelConfig) -> KAMUITransformer:
    """Construct a model from a ``ModelConfig`` instance."""
    return cls(config)

from_yaml(path) classmethod

Load a ModelConfig from YAML and construct the model.

Parameters:

Name Type Description Default
path str | Path

Path to a YAML file with a model section.

 required
Source code in kamui/model/transformer.py 
120
121
122
123
124
125
126
127
@classmethod
def from_yaml(cls, path: str | Path) -> KAMUITransformer:
    """Load a ``ModelConfig`` from YAML and construct the model.

    Args:
        path: Path to a YAML file with a ``model`` section.
    """
    return cls(ModelConfig.from_yaml(path))

forward(token_ids, targets=None)

Run the model.

Parameters:

Name Type Description Default
token_ids Tensor

Integer tensor of shape (B, S).

 required
targets Tensor | None

Optional next-token labels of shape (B, S). When provided, the return value is the scalar cross-entropy loss.

 None

Returns:

Type Description
Tensor

Logits of shape (B, S, vocab_size) if targets is None,
Tensor

otherwise the scalar cross-entropy loss.

Raises:

Type Description
TypeError

If targets is provided but is not a tensor.
ValueError

If targets is provided but its shape differs from token_ids.
Source code in kamui/model/transformer.py 
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
def forward(self, token_ids: Tensor, targets: Tensor | None = None) -> Tensor:
    """Run the model.

    Args:
        token_ids: Integer tensor of shape ``(B, S)``.
        targets:   Optional next-token labels of shape ``(B, S)``.  When
            provided, the return value is the scalar cross-entropy loss.

    Returns:
        Logits of shape ``(B, S, vocab_size)`` if ``targets`` is None,
        otherwise the scalar cross-entropy loss.

    Raises:
        TypeError:  If ``targets`` is provided but is not a tensor.
        ValueError: If ``targets`` is provided but its shape differs from
            ``token_ids``.
    """
    x = self.embed(token_ids)  # (B, S, D)
    for block in self.blocks:
        x = block(x)  # (B, S, D)
    x = self.final_ln(x)  # (B, S, D)
    logits = self.unembed(x)  # (B, S, V)

    if targets is None:
        return logits

    if not isinstance(targets, Tensor):
        raise TypeError(f"targets must be a torch.Tensor, got {type(targets)}")
    if targets.shape != token_ids.shape:
        raise ValueError(
            f"targets shape {tuple(targets.shape)} must match token_ids "
            f"shape {tuple(token_ids.shape)}"
        )

    # Cross-entropy over all positions; flatten batch and sequence.
    loss = F.cross_entropy(logits.reshape(-1, logits.shape[-1]), targets.reshape(-1))
    return loss

num_parameters(trainable_only=True)

Return the number of parameters.

nn.Module.parameters() already yields each shared (weight-tied) parameter once, so the tied embedding/unembedding matrix is counted a single time.

Parameters:

Name Type Description Default
trainable_only bool

If True (default), count only parameters with requires_grad=True.

 True

Returns:

Type Description
int

The parameter count.
Source code in kamui/model/transformer.py 
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
def num_parameters(self, trainable_only: bool = True) -> int:
    """Return the number of parameters.

    ``nn.Module.parameters()`` already yields each shared (weight-tied)
    parameter once, so the tied embedding/unembedding matrix is counted
    a single time.

    Args:
        trainable_only: If True (default), count only parameters with
            ``requires_grad=True``.

    Returns:
        The parameter count.
    """
    return sum(p.numel() for p in self.parameters() if p.requires_grad or not trainable_only)

Tokenizer

kamui.tokenizer.bpe.BPETokenizer

Byte-pair encoding tokeniser, trained and applied from scratch.

The vocabulary is structured as follows
  • IDs 0–255: raw byte tokens (\x00 … \xff).
  • IDs 256 … 256+len(special_tokens)-1: special tokens in declaration order.
  • IDs 256+len(special_tokens) … vocab_size-1: merged subword tokens, one per BPE merge rule, in the order they were learned.

Special tokens are injected into the id→bytes mapping as unique sentinel byte-sequences that can never arise from UTF-8 text, so they cannot accidentally appear inside a regular token.

Attributes:

Name Type Description
vocab_size int

Total number of tokens (bytes + specials + merges).
special_tokens tuple[str, ...]

Tuple of registered special token strings.
Source code in kamui/tokenizer/bpe.py 
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
467
468
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
class BPETokenizer:
    """Byte-pair encoding tokeniser, trained and applied from scratch.

    The vocabulary is structured as follows:
        - IDs 0–255: raw byte tokens (``\\x00`` … ``\\xff``).
        - IDs 256 … 256+len(special_tokens)-1: special tokens in declaration order.
        - IDs 256+len(special_tokens) … vocab_size-1: merged subword tokens,
          one per BPE merge rule, in the order they were learned.

    Special tokens are injected into the id→bytes mapping as unique sentinel
    byte-sequences that can never arise from UTF-8 text, so they cannot
    accidentally appear inside a regular token.

    Attributes:
        vocab_size: Total number of tokens (bytes + specials + merges).
        special_tokens: Tuple of registered special token strings.
    """

    # ------------------------------------------------------------------
    # Construction helpers
    # ------------------------------------------------------------------

    def __init__(
        self,
        merges: list[tuple[int, int]],
        vocab: dict[int, bytes],
        special_tokens: tuple[str, ...] = _DEFAULT_SPECIAL_TOKENS,
    ) -> None:
        """Internal constructor.  Use ``train`` or ``load`` instead.

        Args:
            merges:        Ordered list of (left_id, right_id) merge rules.
            vocab:         Mapping from integer ID to raw bytes representation.
            special_tokens: Tuple of special token strings in ID order.
        """
        self._merges: list[tuple[int, int]] = merges
        self._vocab: dict[int, bytes] = vocab
        self._special_tokens: tuple[str, ...] = special_tokens

        # Build reverse vocab: bytes → id (for encoding)
        self._bytes_to_id: dict[bytes, int] = {v: k for k, v in vocab.items()}

        # Build merge lookup: pair → new_id  (O(1) encode step)
        self._merge_map: dict[tuple[int, int], int] = {}
        for i, pair in enumerate(merges):
            new_id = _NUM_BYTE_TOKENS + len(special_tokens) + i
            self._merge_map[pair] = new_id

        # Build special-token lookup: string → id
        self._special_to_id: dict[str, int] = {
            tok: _NUM_BYTE_TOKENS + i for i, tok in enumerate(special_tokens)
        }
        self._id_to_special: dict[int, str] = {v: k for k, v in self._special_to_id.items()}

    # ------------------------------------------------------------------
    # Properties
    # ------------------------------------------------------------------

    @property
    def vocab_size(self) -> int:
        """Total number of tokens in the vocabulary."""
        return len(self._vocab)

    @property
    def special_tokens(self) -> tuple[str, ...]:
        """Tuple of registered special token strings, in ID order."""
        return self._special_tokens

    # ------------------------------------------------------------------
    # Training
    # ------------------------------------------------------------------

    @classmethod
    def train(  # noqa: C901 — inherently branchy: corpus loading + BPE merge loop
        cls,
        corpus: str | Path,
        vocab_size: int,
        special_tokens: tuple[str, ...] = _DEFAULT_SPECIAL_TOKENS,
        verbose: bool = False,
    ) -> BPETokenizer:
        """Train a BPE tokenizer on a text corpus.

        The training procedure:
        1. Read the full corpus as a UTF-8 string.
        2. Encode every character as its raw UTF-8 bytes (256 base tokens).
        3. Reserve IDs for special tokens immediately after the 256 byte tokens.
        4. Iteratively find the most-frequent adjacent pair of IDs across the
           entire corpus, assign it a new merged ID, and update the corpus.
           Repeat until ``vocab_size`` is reached.

        Args:
            corpus:        Path to the training corpus text file, or a raw string
                           corpus (if the value does not correspond to an existing
                           file, it is treated as the corpus text directly — useful
                           for tests).
            vocab_size:    Target vocabulary size.  Must be ≥ 256 + len(special_tokens).
            special_tokens: Special tokens to reserve, in ID order.  Defaults to
                            ``("<|endoftext|>", "<|pad|>")``.
            verbose:       If True, print progress every 100 merges.

        Returns:
            A trained ``BPETokenizer`` instance.

        Raises:
            ValueError: If ``vocab_size`` is too small to accommodate all byte
                tokens plus the requested special tokens.
            FileNotFoundError: If ``corpus`` is a path string and the file does
                not exist.
        """
        min_vocab = _NUM_BYTE_TOKENS + len(special_tokens)
        if vocab_size < min_vocab:
            raise ValueError(
                f"vocab_size ({vocab_size}) must be >= {min_vocab} "
                f"(256 byte tokens + {len(special_tokens)} special tokens)"
            )

        # Load corpus text.
        # If corpus is a Path object, always treat it as a file path.
        # If it is a plain str, only treat it as a path when it could plausibly
        # be one (≤ 4096 chars, no newlines) AND the file actually exists —
        # otherwise treat it as raw text directly (useful in tests).
        if isinstance(corpus, Path):
            text = corpus.read_text(encoding="utf-8")
        else:
            corpus_str = str(corpus)
            is_plausible_path = (
                len(corpus_str) <= 4096 and "\n" not in corpus_str and " " not in corpus_str
            )
            corpus_path = Path(corpus_str) if is_plausible_path else None
            if corpus_path is not None and corpus_path.exists() and corpus_path.is_file():
                text = corpus_path.read_text(encoding="utf-8")
            else:
                text = corpus_str

        # Split on special tokens so they are never merged with surrounding text.
        # We process each chunk between special tokens as an independent sequence.
        special_pattern = (
            "(" + "|".join(re.escape(s) for s in special_tokens) + ")" if special_tokens else None
        )
        chunks = re.split(special_pattern, text) if special_pattern else [text]

        # Build initial byte-level sequences (one per chunk, skipping special tokens)
        sequences: list[list[int]] = []
        for chunk in chunks:
            if chunk in special_tokens:
                continue
            if chunk:
                sequences.append(text_to_bytes(chunk))

        # Build base vocab: id → bytes
        vocab: dict[int, bytes] = {i: bytes([i]) for i in range(_NUM_BYTE_TOKENS)}

        # Reserve special token IDs
        for i, tok in enumerate(special_tokens):
            vocab[_NUM_BYTE_TOKENS + i] = tok.encode("utf-8")

        merges: list[tuple[int, int]] = []
        num_merges = vocab_size - min_vocab

        for merge_idx in range(num_merges):
            if not sequences:
                break

            stats = get_stats(sequences)
            if not stats:
                break  # corpus too short to find more pairs

            # Choose the most frequent pair; break ties deterministically by pair value
            best_pair = max(stats, key=lambda p: (stats[p], p))
            if stats[best_pair] < 2:
                # Only unique pairs remain.  Merging count-1 pairs would just
                # memorise the corpus (on a degenerate corpus it snowballs
                # into a single token spanning the whole text), so stop here;
                # the resulting vocabulary may be smaller than requested.
                break
            new_id = min_vocab + merge_idx

            # Perform the merge across all sequences
            sequences = [merge_pair(seq, best_pair, new_id) for seq in sequences]

            # Record merge rule and vocab entry
            merges.append(best_pair)
            vocab[new_id] = vocab[best_pair[0]] + vocab[best_pair[1]]

            if verbose and (merge_idx + 1) % 100 == 0:
                print(f"  merge {merge_idx + 1}/{num_merges}: {best_pair} → {new_id}")

        return cls(merges=merges, vocab=vocab, special_tokens=special_tokens)

    # ------------------------------------------------------------------
    # Encoding
    # ------------------------------------------------------------------

    def encode(self, text: str) -> list[int]:
        """Encode a string into a list of token IDs.

        Special tokens in the text are matched first (as whole strings) before
        the remainder is byte-encoded and merged.  Merge rules are applied
        greedily, left to right, until no more merges are possible.

        Args:
            text: Any Unicode string.

        Returns:
            A list of integer token IDs, possibly empty for an empty input.

        Raises:
            TypeError: If ``text`` is not a string.
        """
        if not isinstance(text, str):
            raise TypeError(f"text must be a str, got {type(text)}")
        if not text:
            return []

        # Split on special tokens first so they are never broken up by merges
        if self._special_tokens:
            pattern = "(" + "|".join(re.escape(s) for s in self._special_tokens) + ")"
            parts = re.split(pattern, text)
        else:
            parts = [text]

        result: list[int] = []
        for part in parts:
            if not part:
                continue
            if part in self._special_to_id:
                result.append(self._special_to_id[part])
            else:
                # Byte-encode then apply merges
                ids = text_to_bytes(part)
                ids = self._apply_merges(ids)
                result.extend(ids)
        return result

    def _apply_merges(self, ids: list[int]) -> list[int]:
        """Apply all learned merge rules to a byte-level token sequence.

        Merges are applied greedily in training order (lowest new_id first),
        which matches how GPT-2 and similar tokenisers work.

        Args:
            ids: Sequence of byte-level (0–255) token IDs.

        Returns:
            Merged token ID sequence.
        """
        # One pass per merge rule, in training order, replacing every
        # occurrence.  This is equivalent to repeatedly applying the highest-
        # priority applicable merge: a pair learned at training step k can only
        # involve tokens created by earlier merges, so by the time rule k is
        # applied every occurrence of its pair already exists.  O(merges × S)
        # instead of O(S) rescans per single replacement.
        for pair, new_id in self._merge_map.items():
            if len(ids) < 2:
                break
            ids = merge_pair(ids, pair, new_id)
        return ids

    # ------------------------------------------------------------------
    # Decoding
    # ------------------------------------------------------------------

    def decode(self, ids: list[int]) -> str:
        """Decode a list of token IDs back to a string.

        The decode is lossless for any sequence produced by ``encode``:
        ``decode(encode(text)) == text`` for all Unicode strings.

        Special-token IDs are decoded back to their string representation
        (e.g. token 256 → ``"<|endoftext|>"``).

        Args:
            ids: A list of integer token IDs.

        Returns:
            The decoded string.

        Raises:
            TypeError:  If ``ids`` is not a list or any element is not an int.
            ValueError: If any ID is outside [0, vocab_size).
        """
        if not isinstance(ids, list):
            raise TypeError(f"ids must be a list, got {type(ids)}")

        byte_buf: list[int] = []
        text_parts: list[str] = []

        for token_id in ids:
            if not isinstance(token_id, int) or isinstance(token_id, bool):
                raise TypeError(f"token ID must be an int, got {type(token_id)}")
            if token_id < 0 or token_id >= self.vocab_size:
                raise ValueError(f"token ID {token_id} is out of range [0, {self.vocab_size})")

            if token_id in self._id_to_special:
                # Flush any buffered bytes before emitting the special token
                if byte_buf:
                    text_parts.append(bytes_to_text(byte_buf))
                    byte_buf = []
                text_parts.append(self._id_to_special[token_id])
            else:
                byte_buf.extend(self._vocab[token_id])

        if byte_buf:
            text_parts.append(bytes_to_text(byte_buf))

        return "".join(text_parts)

    # ------------------------------------------------------------------
    # Serialisation
    # ------------------------------------------------------------------

    def save(self, path: str | Path) -> None:
        """Save the tokenizer to a JSON file.

        The JSON contains:
        - ``"special_tokens"``: list of special token strings in ID order.
        - ``"merges"``: list of [left_id, right_id] pairs in training order.
        - ``"vocab"``: dict mapping string token ID → base64-encoded bytes
          representation of the token.

        Args:
            path: Path to the output JSON file.  Parent directories are created
                  if they do not exist.

        Raises:
            IOError: If writing fails.
        """
        import base64

        path_obj = Path(path)

        # Encode vocab values as base64 strings so arbitrary bytes are JSON-safe
        vocab_serialized = {
            str(token_id): base64.b64encode(token_bytes).decode("ascii")
            for token_id, token_bytes in self._vocab.items()
        }

        data = {
            "special_tokens": list(self._special_tokens),
            "merges": [list(pair) for pair in self._merges],
            "vocab": vocab_serialized,
        }

        try:
            path_obj.parent.mkdir(parents=True, exist_ok=True)
            with open(path_obj, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=2, ensure_ascii=True, sort_keys=False)
        except Exception as e:
            raise OSError(f"Failed to save tokenizer to {path}: {e}") from e

    @classmethod
    def load(cls, path: str | Path) -> BPETokenizer:
        """Load a tokenizer from a JSON file created by ``save``.

        Args:
            path: Path to the JSON file.

        Returns:
            A ``BPETokenizer`` instance identical to the one that was saved.

        Raises:
            FileNotFoundError: If the file does not exist.
            ValueError: If the JSON structure is invalid.
            IOError: If reading fails.
        """
        import base64

        path_obj = Path(path)
        if not path_obj.exists():
            raise FileNotFoundError(f"Tokenizer file not found: {path}")

        try:
            with open(path_obj, encoding="utf-8") as f:
                data = json.load(f)
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON in tokenizer file {path}: {e}") from e
        except Exception as e:
            raise OSError(f"Failed to read tokenizer file {path}: {e}") from e

        if not isinstance(data, dict):
            raise ValueError("Tokenizer file must contain a JSON object")
        for key in ("special_tokens", "merges", "vocab"):
            if key not in data:
                raise ValueError(f"Tokenizer file missing required key: '{key}'")

        special_tokens = tuple(data["special_tokens"])

        merges: list[tuple[int, int]] = []
        for pair in data["merges"]:
            if not (isinstance(pair, list) and len(pair) == 2):
                raise ValueError(f"Each merge must be a 2-element list, got: {pair}")
            merges.append((int(pair[0]), int(pair[1])))

        vocab: dict[int, bytes] = {}
        for id_str, b64 in data["vocab"].items():
            vocab[int(id_str)] = base64.b64decode(b64)

        return cls(merges=merges, vocab=vocab, special_tokens=special_tokens)

    # ------------------------------------------------------------------
    # Special-token helpers
    # ------------------------------------------------------------------

    def token_to_id(self, token: str) -> int:
        """Return the integer ID for a special token string.

        Args:
            token: A special token string (must be in ``self.special_tokens``).

        Returns:
            The integer ID.

        Raises:
            KeyError: If the token is not a registered special token.
        """
        if token not in self._special_to_id:
            raise KeyError(f"'{token}' is not a registered special token")
        return self._special_to_id[token]

    # ------------------------------------------------------------------
    # Dunder helpers
    # ------------------------------------------------------------------

    def __repr__(self) -> str:
        return (
            f"BPETokenizer("
            f"vocab_size={self.vocab_size}, "
            f"merges={len(self._merges)}, "
            f"special_tokens={self._special_tokens}"
            f")"
        )

vocab_size property

Total number of tokens in the vocabulary.

special_tokens property

Tuple of registered special token strings, in ID order.

__init__(merges, vocab, special_tokens=_DEFAULT_SPECIAL_TOKENS)

Internal constructor. Use train or load instead.

Parameters:

Name Type Description Default
merges list[tuple[int, int]]

Ordered list of (left_id, right_id) merge rules.

 required
vocab dict[int, bytes]

Mapping from integer ID to raw bytes representation.

 required
special_tokens tuple[str, ...]

Tuple of special token strings in ID order.

 _DEFAULT_SPECIAL_TOKENS
Source code in kamui/tokenizer/bpe.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
def __init__(
    self,
    merges: list[tuple[int, int]],
    vocab: dict[int, bytes],
    special_tokens: tuple[str, ...] = _DEFAULT_SPECIAL_TOKENS,
) -> None:
    """Internal constructor.  Use ``train`` or ``load`` instead.

    Args:
        merges:        Ordered list of (left_id, right_id) merge rules.
        vocab:         Mapping from integer ID to raw bytes representation.
        special_tokens: Tuple of special token strings in ID order.
    """
    self._merges: list[tuple[int, int]] = merges
    self._vocab: dict[int, bytes] = vocab
    self._special_tokens: tuple[str, ...] = special_tokens

    # Build reverse vocab: bytes → id (for encoding)
    self._bytes_to_id: dict[bytes, int] = {v: k for k, v in vocab.items()}

    # Build merge lookup: pair → new_id  (O(1) encode step)
    self._merge_map: dict[tuple[int, int], int] = {}
    for i, pair in enumerate(merges):
        new_id = _NUM_BYTE_TOKENS + len(special_tokens) + i
        self._merge_map[pair] = new_id

    # Build special-token lookup: string → id
    self._special_to_id: dict[str, int] = {
        tok: _NUM_BYTE_TOKENS + i for i, tok in enumerate(special_tokens)
    }
    self._id_to_special: dict[int, str] = {v: k for k, v in self._special_to_id.items()}

train(corpus, vocab_size, special_tokens=_DEFAULT_SPECIAL_TOKENS, verbose=False) classmethod

Train a BPE tokenizer on a text corpus.

The training procedure: 1. Read the full corpus as a UTF-8 string. 2. Encode every character as its raw UTF-8 bytes (256 base tokens). 3. Reserve IDs for special tokens immediately after the 256 byte tokens. 4. Iteratively find the most-frequent adjacent pair of IDs across the entire corpus, assign it a new merged ID, and update the corpus. Repeat until vocab_size is reached.

Parameters:

Name Type Description Default
corpus str | Path

Path to the training corpus text file, or a raw string corpus (if the value does not correspond to an existing file, it is treated as the corpus text directly — useful for tests).

 required
vocab_size int

Target vocabulary size. Must be ≥ 256 + len(special_tokens).

 required
special_tokens tuple[str, ...]

Special tokens to reserve, in ID order. Defaults to ("<|endoftext|>", "<|pad|>").

 _DEFAULT_SPECIAL_TOKENS
verbose bool

If True, print progress every 100 merges.

 False

Returns:

Type Description
BPETokenizer

A trained BPETokenizer instance.

Raises:

Type Description
ValueError

If vocab_size is too small to accommodate all byte tokens plus the requested special tokens.
FileNotFoundError

If corpus is a path string and the file does not exist.
Source code in kamui/tokenizer/bpe.py 
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
@classmethod
def train(  # noqa: C901 — inherently branchy: corpus loading + BPE merge loop
    cls,
    corpus: str | Path,
    vocab_size: int,
    special_tokens: tuple[str, ...] = _DEFAULT_SPECIAL_TOKENS,
    verbose: bool = False,
) -> BPETokenizer:
    """Train a BPE tokenizer on a text corpus.

    The training procedure:
    1. Read the full corpus as a UTF-8 string.
    2. Encode every character as its raw UTF-8 bytes (256 base tokens).
    3. Reserve IDs for special tokens immediately after the 256 byte tokens.
    4. Iteratively find the most-frequent adjacent pair of IDs across the
       entire corpus, assign it a new merged ID, and update the corpus.
       Repeat until ``vocab_size`` is reached.

    Args:
        corpus:        Path to the training corpus text file, or a raw string
                       corpus (if the value does not correspond to an existing
                       file, it is treated as the corpus text directly — useful
                       for tests).
        vocab_size:    Target vocabulary size.  Must be ≥ 256 + len(special_tokens).
        special_tokens: Special tokens to reserve, in ID order.  Defaults to
                        ``("<|endoftext|>", "<|pad|>")``.
        verbose:       If True, print progress every 100 merges.

    Returns:
        A trained ``BPETokenizer`` instance.

    Raises:
        ValueError: If ``vocab_size`` is too small to accommodate all byte
            tokens plus the requested special tokens.
        FileNotFoundError: If ``corpus`` is a path string and the file does
            not exist.
    """
    min_vocab = _NUM_BYTE_TOKENS + len(special_tokens)
    if vocab_size < min_vocab:
        raise ValueError(
            f"vocab_size ({vocab_size}) must be >= {min_vocab} "
            f"(256 byte tokens + {len(special_tokens)} special tokens)"
        )

    # Load corpus text.
    # If corpus is a Path object, always treat it as a file path.
    # If it is a plain str, only treat it as a path when it could plausibly
    # be one (≤ 4096 chars, no newlines) AND the file actually exists —
    # otherwise treat it as raw text directly (useful in tests).
    if isinstance(corpus, Path):
        text = corpus.read_text(encoding="utf-8")
    else:
        corpus_str = str(corpus)
        is_plausible_path = (
            len(corpus_str) <= 4096 and "\n" not in corpus_str and " " not in corpus_str
        )
        corpus_path = Path(corpus_str) if is_plausible_path else None
        if corpus_path is not None and corpus_path.exists() and corpus_path.is_file():
            text = corpus_path.read_text(encoding="utf-8")
        else:
            text = corpus_str

    # Split on special tokens so they are never merged with surrounding text.
    # We process each chunk between special tokens as an independent sequence.
    special_pattern = (
        "(" + "|".join(re.escape(s) for s in special_tokens) + ")" if special_tokens else None
    )
    chunks = re.split(special_pattern, text) if special_pattern else [text]

    # Build initial byte-level sequences (one per chunk, skipping special tokens)
    sequences: list[list[int]] = []
    for chunk in chunks:
        if chunk in special_tokens:
            continue
        if chunk:
            sequences.append(text_to_bytes(chunk))

    # Build base vocab: id → bytes
    vocab: dict[int, bytes] = {i: bytes([i]) for i in range(_NUM_BYTE_TOKENS)}

    # Reserve special token IDs
    for i, tok in enumerate(special_tokens):
        vocab[_NUM_BYTE_TOKENS + i] = tok.encode("utf-8")

    merges: list[tuple[int, int]] = []
    num_merges = vocab_size - min_vocab

    for merge_idx in range(num_merges):
        if not sequences:
            break

        stats = get_stats(sequences)
        if not stats:
            break  # corpus too short to find more pairs

        # Choose the most frequent pair; break ties deterministically by pair value
        best_pair = max(stats, key=lambda p: (stats[p], p))
        if stats[best_pair] < 2:
            # Only unique pairs remain.  Merging count-1 pairs would just
            # memorise the corpus (on a degenerate corpus it snowballs
            # into a single token spanning the whole text), so stop here;
            # the resulting vocabulary may be smaller than requested.
            break
        new_id = min_vocab + merge_idx

        # Perform the merge across all sequences
        sequences = [merge_pair(seq, best_pair, new_id) for seq in sequences]

        # Record merge rule and vocab entry
        merges.append(best_pair)
        vocab[new_id] = vocab[best_pair[0]] + vocab[best_pair[1]]

        if verbose and (merge_idx + 1) % 100 == 0:
            print(f"  merge {merge_idx + 1}/{num_merges}: {best_pair} → {new_id}")

    return cls(merges=merges, vocab=vocab, special_tokens=special_tokens)

encode(text)

Encode a string into a list of token IDs.

Special tokens in the text are matched first (as whole strings) before the remainder is byte-encoded and merged. Merge rules are applied greedily, left to right, until no more merges are possible.

Parameters:

Name Type Description Default
text str

Any Unicode string.

 required

Returns:

Type Description
list[int]

A list of integer token IDs, possibly empty for an empty input.

Raises:

Type Description
TypeError

If text is not a string.
Source code in kamui/tokenizer/bpe.py 
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
def encode(self, text: str) -> list[int]:
    """Encode a string into a list of token IDs.

    Special tokens in the text are matched first (as whole strings) before
    the remainder is byte-encoded and merged.  Merge rules are applied
    greedily, left to right, until no more merges are possible.

    Args:
        text: Any Unicode string.

    Returns:
        A list of integer token IDs, possibly empty for an empty input.

    Raises:
        TypeError: If ``text`` is not a string.
    """
    if not isinstance(text, str):
        raise TypeError(f"text must be a str, got {type(text)}")
    if not text:
        return []

    # Split on special tokens first so they are never broken up by merges
    if self._special_tokens:
        pattern = "(" + "|".join(re.escape(s) for s in self._special_tokens) + ")"
        parts = re.split(pattern, text)
    else:
        parts = [text]

    result: list[int] = []
    for part in parts:
        if not part:
            continue
        if part in self._special_to_id:
            result.append(self._special_to_id[part])
        else:
            # Byte-encode then apply merges
            ids = text_to_bytes(part)
            ids = self._apply_merges(ids)
            result.extend(ids)
    return result

decode(ids)

Decode a list of token IDs back to a string.

The decode is lossless for any sequence produced by encode: decode(encode(text)) == text for all Unicode strings.

Special-token IDs are decoded back to their string representation (e.g. token 256 → "<|endoftext|>").

Parameters:

Name Type Description Default
ids list[int]

A list of integer token IDs.

 required

Returns:

Type Description
str

The decoded string.

Raises:

Type Description
TypeError

If ids is not a list or any element is not an int.
ValueError

If any ID is outside [0, vocab_size).
Source code in kamui/tokenizer/bpe.py 
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
def decode(self, ids: list[int]) -> str:
    """Decode a list of token IDs back to a string.

    The decode is lossless for any sequence produced by ``encode``:
    ``decode(encode(text)) == text`` for all Unicode strings.

    Special-token IDs are decoded back to their string representation
    (e.g. token 256 → ``"<|endoftext|>"``).

    Args:
        ids: A list of integer token IDs.

    Returns:
        The decoded string.

    Raises:
        TypeError:  If ``ids`` is not a list or any element is not an int.
        ValueError: If any ID is outside [0, vocab_size).
    """
    if not isinstance(ids, list):
        raise TypeError(f"ids must be a list, got {type(ids)}")

    byte_buf: list[int] = []
    text_parts: list[str] = []

    for token_id in ids:
        if not isinstance(token_id, int) or isinstance(token_id, bool):
            raise TypeError(f"token ID must be an int, got {type(token_id)}")
        if token_id < 0 or token_id >= self.vocab_size:
            raise ValueError(f"token ID {token_id} is out of range [0, {self.vocab_size})")

        if token_id in self._id_to_special:
            # Flush any buffered bytes before emitting the special token
            if byte_buf:
                text_parts.append(bytes_to_text(byte_buf))
                byte_buf = []
            text_parts.append(self._id_to_special[token_id])
        else:
            byte_buf.extend(self._vocab[token_id])

    if byte_buf:
        text_parts.append(bytes_to_text(byte_buf))

    return "".join(text_parts)

save(path)

Save the tokenizer to a JSON file.

The JSON contains: - "special_tokens": list of special token strings in ID order. - "merges": list of [left_id, right_id] pairs in training order. - "vocab": dict mapping string token ID → base64-encoded bytes representation of the token.

Parameters:

Name Type Description Default
path str | Path

Path to the output JSON file. Parent directories are created if they do not exist.

 required

Raises:

Type Description
IOError

If writing fails.
Source code in kamui/tokenizer/bpe.py 
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
def save(self, path: str | Path) -> None:
    """Save the tokenizer to a JSON file.

    The JSON contains:
    - ``"special_tokens"``: list of special token strings in ID order.
    - ``"merges"``: list of [left_id, right_id] pairs in training order.
    - ``"vocab"``: dict mapping string token ID → base64-encoded bytes
      representation of the token.

    Args:
        path: Path to the output JSON file.  Parent directories are created
              if they do not exist.

    Raises:
        IOError: If writing fails.
    """
    import base64

    path_obj = Path(path)

    # Encode vocab values as base64 strings so arbitrary bytes are JSON-safe
    vocab_serialized = {
        str(token_id): base64.b64encode(token_bytes).decode("ascii")
        for token_id, token_bytes in self._vocab.items()
    }

    data = {
        "special_tokens": list(self._special_tokens),
        "merges": [list(pair) for pair in self._merges],
        "vocab": vocab_serialized,
    }

    try:
        path_obj.parent.mkdir(parents=True, exist_ok=True)
        with open(path_obj, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=2, ensure_ascii=True, sort_keys=False)
    except Exception as e:
        raise OSError(f"Failed to save tokenizer to {path}: {e}") from e

load(path) classmethod

Load a tokenizer from a JSON file created by save.

Parameters:

Name Type Description Default
path str | Path

Path to the JSON file.

 required

Returns:

Type Description
BPETokenizer

A BPETokenizer instance identical to the one that was saved.

Raises:

Type Description
FileNotFoundError

If the file does not exist.
ValueError

If the JSON structure is invalid.
IOError

If reading fails.
Source code in kamui/tokenizer/bpe.py 
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
@classmethod
def load(cls, path: str | Path) -> BPETokenizer:
    """Load a tokenizer from a JSON file created by ``save``.

    Args:
        path: Path to the JSON file.

    Returns:
        A ``BPETokenizer`` instance identical to the one that was saved.

    Raises:
        FileNotFoundError: If the file does not exist.
        ValueError: If the JSON structure is invalid.
        IOError: If reading fails.
    """
    import base64

    path_obj = Path(path)
    if not path_obj.exists():
        raise FileNotFoundError(f"Tokenizer file not found: {path}")

    try:
        with open(path_obj, encoding="utf-8") as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON in tokenizer file {path}: {e}") from e
    except Exception as e:
        raise OSError(f"Failed to read tokenizer file {path}: {e}") from e

    if not isinstance(data, dict):
        raise ValueError("Tokenizer file must contain a JSON object")
    for key in ("special_tokens", "merges", "vocab"):
        if key not in data:
            raise ValueError(f"Tokenizer file missing required key: '{key}'")

    special_tokens = tuple(data["special_tokens"])

    merges: list[tuple[int, int]] = []
    for pair in data["merges"]:
        if not (isinstance(pair, list) and len(pair) == 2):
            raise ValueError(f"Each merge must be a 2-element list, got: {pair}")
        merges.append((int(pair[0]), int(pair[1])))

    vocab: dict[int, bytes] = {}
    for id_str, b64 in data["vocab"].items():
        vocab[int(id_str)] = base64.b64decode(b64)

    return cls(merges=merges, vocab=vocab, special_tokens=special_tokens)

token_to_id(token)

Return the integer ID for a special token string.

Parameters:

Name Type Description Default
token str

A special token string (must be in self.special_tokens).

 required

Returns:

Type Description
int

The integer ID.

Raises:

Type Description
KeyError

If the token is not a registered special token.
Source code in kamui/tokenizer/bpe.py 
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
def token_to_id(self, token: str) -> int:
    """Return the integer ID for a special token string.

    Args:
        token: A special token string (must be in ``self.special_tokens``).

    Returns:
        The integer ID.

    Raises:
        KeyError: If the token is not a registered special token.
    """
    if token not in self._special_to_id:
        raise KeyError(f"'{token}' is not a registered special token")
    return self._special_to_id[token]

kamui.tokenizer.vocab.Vocabulary

Vocabulary management for the KAMUI BPE tokenizer.

Maintains bidirectional mapping between string tokens and integer IDs, ensuring deterministic ID assignment, duplicate prevention, and stable serialization.

Source code in kamui/tokenizer/vocab.py 
 28
 29
 30
 31
 32
 33
 34
 35
 36
 37
 38
 39
 40
 41
 42
 43
 44
 45
 46
 47
 48
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
class Vocabulary:
    """Vocabulary management for the KAMUI BPE tokenizer.

    Maintains bidirectional mapping between string tokens and integer IDs,
    ensuring deterministic ID assignment, duplicate prevention, and stable
    serialization.
    """

    def __init__(self, special_tokens: Iterable[str] | None = None) -> None:
        """Initialise a new Vocabulary.

        Args:
            special_tokens: Optional iterable of special tokens. If None,
                defaults to ("<pad>", "<bos>", "<eos>", "<unk>").

        Raises:
            ValueError: If empty strings are provided as special tokens,
                or if there are duplicate special tokens.
            TypeError: If special tokens contains non-string elements.
        """
        if special_tokens is None:
            specials: tuple[str, ...] = ("<pad>", "<bos>", "<eos>", "<unk>")
        else:
            specials = tuple(special_tokens)

        self._special_tokens: set[str] = set()
        self._token_to_id: dict[str, int] = {}
        self._id_to_token: list[str] = []

        # Register special tokens
        for token in specials:
            if not isinstance(token, str):
                raise TypeError(f"Special token must be a string, got type {type(token)}")
            if token == "":
                raise ValueError("Special token cannot be an empty string")
            if token in self._special_tokens:
                raise ValueError(f"Duplicate special token registered: '{token}'")

            token_id = len(self._id_to_token)
            self._token_to_id[token] = token_id
            self._id_to_token.append(token)
            self._special_tokens.add(token)

    @property
    def vocab_size(self) -> int:
        """Return the total number of tokens in the vocabulary."""
        return len(self._id_to_token)

    @property
    def special_tokens(self) -> set[str]:
        """Return a copy of the set of registered special tokens."""
        return self._special_tokens.copy()

    def add_token(self, token: str) -> int:
        """Add a token to the vocabulary and return its assigned integer ID.

        Args:
            token: The string token to add.

        Returns:
            The newly assigned integer ID of the token.

        Raises:
            ValueError: If the token is empty, already exists, or is a duplicate
                special token.
            TypeError: If the token is not a string.
        """
        if not isinstance(token, str):
            raise TypeError(f"Token must be a string, got type {type(token)}")
        if token == "":
            raise ValueError("Token cannot be an empty string")
        if token in self._token_to_id:
            raise ValueError(f"Token '{token}' already exists in the vocabulary")

        token_id = len(self._id_to_token)
        self._token_to_id[token] = token_id
        self._id_to_token.append(token)
        return token_id

    def add_tokens(self, tokens: Iterable[str]) -> None:
        """Add multiple tokens to the vocabulary.

        This operation is atomic. If any token in the input iterable is invalid
        (e.g., empty or duplicate), no tokens will be added.

        Args:
            tokens: Iterable of string tokens to add.

        Raises:
            ValueError: If any token is empty, already exists in the vocabulary,
                or is a duplicate within the input iterable.
            TypeError: If any token is not a string.
        """
        tokens_list = list(tokens)
        seen: set[str] = set()

        for token in tokens_list:
            if not isinstance(token, str):
                raise TypeError(f"Token must be a string, got type {type(token)}")
            if token == "":
                raise ValueError("Token cannot be an empty string")
            if token in self._token_to_id:
                raise ValueError(f"Token '{token}' already exists in the vocabulary")
            if token in seen:
                raise ValueError(f"Duplicate token in input iterable: '{token}'")
            seen.add(token)

        for token in tokens_list:
            self.add_token(token)

    def token_to_id(self, token: str) -> int:
        """Look up the integer ID of a string token.

        Args:
            token: The string token to look up.

        Returns:
            The integer ID.

        Raises:
            KeyError: If the token is not in the vocabulary.
            TypeError: If the token is not a string.
        """
        if not isinstance(token, str):
            raise TypeError(f"Token must be a string, got type {type(token)}")
        if token not in self._token_to_id:
            raise KeyError(f"Token '{token}' not found in vocabulary")
        return self._token_to_id[token]

    def id_to_token(self, token_id: int) -> str:
        """Look up the string representation of a token ID.

        Args:
            token_id: The integer ID to look up.

        Returns:
            The string token.

        Raises:
            ValueError: If the token ID is out of range.
            TypeError: If the token ID is not an integer.
        """
        if isinstance(token_id, bool) or not isinstance(token_id, int):
            raise TypeError(f"Token ID must be an integer, got type {type(token_id)}")
        if token_id < 0 or token_id >= len(self._id_to_token):
            raise ValueError(f"Token ID {token_id} is out of range [0, {len(self._id_to_token)})")
        return self._id_to_token[token_id]

    def __getitem__(self, key: str | int) -> int | str:
        """Look up a token string to get its ID, or a token ID to get its string.

        Args:
            key: Either a string token or an integer token ID.

        Returns:
            Integer ID if key is a string, or string token if key is an integer.

        Raises:
            KeyError: If a string key is not found.
            ValueError: If an integer key is out of range.
            TypeError: If key is neither a string nor an integer.
        """
        if isinstance(key, str):
            return self.token_to_id(key)
        elif isinstance(key, int) and not isinstance(key, bool):
            return self.id_to_token(key)
        else:
            raise TypeError(f"Key must be a string or an integer, got type {type(key)}")

    def __contains__(self, item: str | int) -> bool:
        """Check if a token or ID is in the vocabulary.

        Args:
            item: Either a string token or an integer token ID.

        Returns:
            True if the token or ID exists in the vocabulary, False otherwise.
        """
        if isinstance(item, str):
            return item in self._token_to_id
        elif isinstance(item, int) and not isinstance(item, bool):
            return 0 <= item < len(self._id_to_token)
        return False

    def contains(self, token: str) -> bool:
        """Return True if *token* is present in the vocabulary.

        This is the named method form of ``token in vocabulary``.  It only
        accepts string tokens; use the ``in`` operator to check by integer ID.

        Args:
            token: The string token to look up.

        Returns:
            True if the token exists in the vocabulary, False otherwise.

        Raises:
            TypeError: If *token* is not a string.
        """
        if not isinstance(token, str):
            raise TypeError(f"Token must be a string, got type {type(token)}")
        return token in self._token_to_id

    def __len__(self) -> int:
        """Return the total number of tokens (same as ``vocab_size``)."""
        return len(self._id_to_token)

    def __repr__(self) -> str:
        """Return a human-readable summary of the vocabulary."""
        return (
            f"Vocabulary("
            f"vocab_size={self.vocab_size}, "
            f"special_tokens={sorted(self._special_tokens, key=lambda t: self._token_to_id[t])}"
            f")"
        )

    def save(self, path: str | Path) -> None:
        """Save the vocabulary to a JSON file.

        Args:
            path: Path to the target JSON file.

        Raises:
            IOError: If writing to the file fails.
        """
        # Ensure special tokens are serialized in their ID order
        sorted_specials = sorted(list(self._special_tokens), key=lambda x: self._token_to_id[x])

        data = {
            "special_tokens": sorted_specials,
            "vocab": {token: token_id for token, token_id in self._token_to_id.items()},
        }

        path_obj = Path(path)
        try:
            path_obj.parent.mkdir(parents=True, exist_ok=True)
            with open(path_obj, "w", encoding="utf-8") as f:
                json.dump(data, f, indent=4, ensure_ascii=False)
        except Exception as e:
            raise OSError(f"Failed to save vocabulary to {path}: {e}") from e

    @classmethod
    def load(cls, path: str | Path) -> "Vocabulary":  # noqa: C901 — defensive JSON validation
        """Load a vocabulary from a JSON file.

        Args:
            path: Path to the JSON file to load.

        Returns:
            A new Vocabulary instance reconstructed from the file.

        Raises:
            ValueError: If the JSON format is invalid, IDs are not consecutive
                starting from 0, or special tokens are invalid or missing.
            FileNotFoundError: If the vocabulary file does not exist.
            IOError: If reading the file fails.
        """
        path_obj = Path(path)
        if not path_obj.exists():
            raise FileNotFoundError(f"Vocabulary file not found: {path}")

        try:
            with open(path_obj, encoding="utf-8") as f:
                data = json.load(f)
        except json.JSONDecodeError as e:
            raise ValueError(f"Invalid JSON format in vocabulary file {path}: {e}") from e
        except Exception as e:
            raise OSError(f"Failed to read vocabulary from {path}: {e}") from e

        if not isinstance(data, dict):
            raise ValueError("Vocabulary file must contain a JSON object at the top level")

        if "vocab" not in data or "special_tokens" not in data:
            raise ValueError("Vocabulary file must contain 'vocab' and 'special_tokens' keys")

        special_tokens_list = data["special_tokens"]
        vocab_dict = data["vocab"]

        if not isinstance(special_tokens_list, list):
            raise ValueError("'special_tokens' must be a list of strings")
        for t in special_tokens_list:
            if not isinstance(t, str):
                raise ValueError("All elements in 'special_tokens' must be strings")

        if not isinstance(vocab_dict, dict):
            raise ValueError("'vocab' must be a dictionary")

        # Sort vocab items by ID to verify they are consecutive and start from 0
        sorted_vocab: list[tuple[Any, Any]] = sorted(vocab_dict.items(), key=lambda item: item[1])

        for idx, (token, token_id) in enumerate(sorted_vocab):
            if not isinstance(token, str):
                raise ValueError(f"Vocabulary token must be a string, got type {type(token)}")
            if not isinstance(token_id, int) or isinstance(token_id, bool):
                raise ValueError(
                    f"Vocabulary token ID must be an integer, got type {type(token_id)}"
                )
            if token_id != idx:
                raise ValueError(
                    f"Vocabulary IDs must be consecutive starting from 0. "
                    f"Expected {idx}, got {token_id}"
                )

        # Verify all listed special tokens are actually in the vocabulary
        for token in special_tokens_list:
            if token not in vocab_dict:
                raise ValueError(f"Special token '{token}' not found in vocabulary mapping")

        # Reconstruct Vocabulary
        instance = cls(special_tokens=special_tokens_list)

        for token, token_id in sorted_vocab:
            if token in instance._special_tokens:
                # Ensure the ID assigned at constructor matches the serialized ID
                if instance.token_to_id(token) != token_id:
                    raise ValueError(
                        f"Special token '{token}' ID mismatch: expected {token_id}, "
                        f"got {instance.token_to_id(token)}"
                    )
            else:
                assigned_id = instance.add_token(token)
                if assigned_id != token_id:
                    raise ValueError(
                        f"Failed to load token '{token}' at expected ID {token_id}, "
                        f"assigned ID {assigned_id}"
                    )

        return instance

vocab_size property

Return the total number of tokens in the vocabulary.

special_tokens property

Return a copy of the set of registered special tokens.

__init__(special_tokens=None)

Initialise a new Vocabulary.

Parameters:

Name Type Description Default
special_tokens Iterable[str] | None

Optional iterable of special tokens. If None, defaults to ("", "", "", "").

 None

Raises:

Type Description
ValueError

If empty strings are provided as special tokens, or if there are duplicate special tokens.
TypeError

If special tokens contains non-string elements.
Source code in kamui/tokenizer/vocab.py 
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
def __init__(self, special_tokens: Iterable[str] | None = None) -> None:
    """Initialise a new Vocabulary.

    Args:
        special_tokens: Optional iterable of special tokens. If None,
            defaults to ("<pad>", "<bos>", "<eos>", "<unk>").

    Raises:
        ValueError: If empty strings are provided as special tokens,
            or if there are duplicate special tokens.
        TypeError: If special tokens contains non-string elements.
    """
    if special_tokens is None:
        specials: tuple[str, ...] = ("<pad>", "<bos>", "<eos>", "<unk>")
    else:
        specials = tuple(special_tokens)

    self._special_tokens: set[str] = set()
    self._token_to_id: dict[str, int] = {}
    self._id_to_token: list[str] = []

    # Register special tokens
    for token in specials:
        if not isinstance(token, str):
            raise TypeError(f"Special token must be a string, got type {type(token)}")
        if token == "":
            raise ValueError("Special token cannot be an empty string")
        if token in self._special_tokens:
            raise ValueError(f"Duplicate special token registered: '{token}'")

        token_id = len(self._id_to_token)
        self._token_to_id[token] = token_id
        self._id_to_token.append(token)
        self._special_tokens.add(token)

add_token(token)

Add a token to the vocabulary and return its assigned integer ID.

Parameters:

Name Type Description Default
token str

The string token to add.

 required

Returns:

Type Description
int

The newly assigned integer ID of the token.

Raises:

Type Description
ValueError

If the token is empty, already exists, or is a duplicate special token.
TypeError

If the token is not a string.
Source code in kamui/tokenizer/vocab.py 
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
def add_token(self, token: str) -> int:
    """Add a token to the vocabulary and return its assigned integer ID.

    Args:
        token: The string token to add.

    Returns:
        The newly assigned integer ID of the token.

    Raises:
        ValueError: If the token is empty, already exists, or is a duplicate
            special token.
        TypeError: If the token is not a string.
    """
    if not isinstance(token, str):
        raise TypeError(f"Token must be a string, got type {type(token)}")
    if token == "":
        raise ValueError("Token cannot be an empty string")
    if token in self._token_to_id:
        raise ValueError(f"Token '{token}' already exists in the vocabulary")

    token_id = len(self._id_to_token)
    self._token_to_id[token] = token_id
    self._id_to_token.append(token)
    return token_id

add_tokens(tokens)

Add multiple tokens to the vocabulary.

This operation is atomic. If any token in the input iterable is invalid (e.g., empty or duplicate), no tokens will be added.

Parameters:

Name Type Description Default
tokens Iterable[str]

Iterable of string tokens to add.

 required

Raises:

Type Description
ValueError

If any token is empty, already exists in the vocabulary, or is a duplicate within the input iterable.
TypeError

If any token is not a string.
Source code in kamui/tokenizer/vocab.py 
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
def add_tokens(self, tokens: Iterable[str]) -> None:
    """Add multiple tokens to the vocabulary.

    This operation is atomic. If any token in the input iterable is invalid
    (e.g., empty or duplicate), no tokens will be added.

    Args:
        tokens: Iterable of string tokens to add.

    Raises:
        ValueError: If any token is empty, already exists in the vocabulary,
            or is a duplicate within the input iterable.
        TypeError: If any token is not a string.
    """
    tokens_list = list(tokens)
    seen: set[str] = set()

    for token in tokens_list:
        if not isinstance(token, str):
            raise TypeError(f"Token must be a string, got type {type(token)}")
        if token == "":
            raise ValueError("Token cannot be an empty string")
        if token in self._token_to_id:
            raise ValueError(f"Token '{token}' already exists in the vocabulary")
        if token in seen:
            raise ValueError(f"Duplicate token in input iterable: '{token}'")
        seen.add(token)

    for token in tokens_list:
        self.add_token(token)

token_to_id(token)

Look up the integer ID of a string token.

Parameters:

Name Type Description Default
token str

The string token to look up.

 required

Returns:

Type Description
int

The integer ID.

Raises:

Type Description
KeyError

If the token is not in the vocabulary.
TypeError

If the token is not a string.
Source code in kamui/tokenizer/vocab.py 
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
def token_to_id(self, token: str) -> int:
    """Look up the integer ID of a string token.

    Args:
        token: The string token to look up.

    Returns:
        The integer ID.

    Raises:
        KeyError: If the token is not in the vocabulary.
        TypeError: If the token is not a string.
    """
    if not isinstance(token, str):
        raise TypeError(f"Token must be a string, got type {type(token)}")
    if token not in self._token_to_id:
        raise KeyError(f"Token '{token}' not found in vocabulary")
    return self._token_to_id[token]

id_to_token(token_id)

Look up the string representation of a token ID.

Parameters:

Name Type Description Default
token_id int

The integer ID to look up.

 required

Returns:

Type Description
str

The string token.

Raises:

Type Description
ValueError

If the token ID is out of range.
TypeError

If the token ID is not an integer.
Source code in kamui/tokenizer/vocab.py 
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
def id_to_token(self, token_id: int) -> str:
    """Look up the string representation of a token ID.

    Args:
        token_id: The integer ID to look up.

    Returns:
        The string token.

    Raises:
        ValueError: If the token ID is out of range.
        TypeError: If the token ID is not an integer.
    """
    if isinstance(token_id, bool) or not isinstance(token_id, int):
        raise TypeError(f"Token ID must be an integer, got type {type(token_id)}")
    if token_id < 0 or token_id >= len(self._id_to_token):
        raise ValueError(f"Token ID {token_id} is out of range [0, {len(self._id_to_token)})")
    return self._id_to_token[token_id]

__getitem__(key)

Look up a token string to get its ID, or a token ID to get its string.

Parameters:

Name Type Description Default
key str | int

Either a string token or an integer token ID.

 required

Returns:

Type Description
int | str

Integer ID if key is a string, or string token if key is an integer.

Raises:

Type Description
KeyError

If a string key is not found.
ValueError

If an integer key is out of range.
TypeError

If key is neither a string nor an integer.
Source code in kamui/tokenizer/vocab.py 
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
def __getitem__(self, key: str | int) -> int | str:
    """Look up a token string to get its ID, or a token ID to get its string.

    Args:
        key: Either a string token or an integer token ID.

    Returns:
        Integer ID if key is a string, or string token if key is an integer.

    Raises:
        KeyError: If a string key is not found.
        ValueError: If an integer key is out of range.
        TypeError: If key is neither a string nor an integer.
    """
    if isinstance(key, str):
        return self.token_to_id(key)
    elif isinstance(key, int) and not isinstance(key, bool):
        return self.id_to_token(key)
    else:
        raise TypeError(f"Key must be a string or an integer, got type {type(key)}")

__contains__(item)

Check if a token or ID is in the vocabulary.

Parameters:

Name Type Description Default
item str | int

Either a string token or an integer token ID.

 required

Returns:

Type Description
bool

True if the token or ID exists in the vocabulary, False otherwise.
Source code in kamui/tokenizer/vocab.py 
197
198
199
200
201
202
203
204
205
206
207
208
209
210
def __contains__(self, item: str | int) -> bool:
    """Check if a token or ID is in the vocabulary.

    Args:
        item: Either a string token or an integer token ID.

    Returns:
        True if the token or ID exists in the vocabulary, False otherwise.
    """
    if isinstance(item, str):
        return item in self._token_to_id
    elif isinstance(item, int) and not isinstance(item, bool):
        return 0 <= item < len(self._id_to_token)
    return False

contains(token)

Return True if token is present in the vocabulary.

This is the named method form of token in vocabulary. It only accepts string tokens; use the in operator to check by integer ID.

Parameters:

Name Type Description Default
token str

The string token to look up.

 required

Returns:

Type Description
bool

True if the token exists in the vocabulary, False otherwise.

Raises:

Type Description
TypeError

If token is not a string.
Source code in kamui/tokenizer/vocab.py 
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
def contains(self, token: str) -> bool:
    """Return True if *token* is present in the vocabulary.

    This is the named method form of ``token in vocabulary``.  It only
    accepts string tokens; use the ``in`` operator to check by integer ID.

    Args:
        token: The string token to look up.

    Returns:
        True if the token exists in the vocabulary, False otherwise.

    Raises:
        TypeError: If *token* is not a string.
    """
    if not isinstance(token, str):
        raise TypeError(f"Token must be a string, got type {type(token)}")
    return token in self._token_to_id

__len__()

Return the total number of tokens (same as vocab_size).

Source code in kamui/tokenizer/vocab.py 
231
232
233
def __len__(self) -> int:
    """Return the total number of tokens (same as ``vocab_size``)."""
    return len(self._id_to_token)

__repr__()

Return a human-readable summary of the vocabulary.

Source code in kamui/tokenizer/vocab.py 
235
236
237
238
239
240
241
242
def __repr__(self) -> str:
    """Return a human-readable summary of the vocabulary."""
    return (
        f"Vocabulary("
        f"vocab_size={self.vocab_size}, "
        f"special_tokens={sorted(self._special_tokens, key=lambda t: self._token_to_id[t])}"
        f")"
    )

save(path)

Save the vocabulary to a JSON file.

Parameters:

Name Type Description Default
path str | Path

Path to the target JSON file.

 required

Raises:

Type Description
IOError

If writing to the file fails.
Source code in kamui/tokenizer/vocab.py 
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
def save(self, path: str | Path) -> None:
    """Save the vocabulary to a JSON file.

    Args:
        path: Path to the target JSON file.

    Raises:
        IOError: If writing to the file fails.
    """
    # Ensure special tokens are serialized in their ID order
    sorted_specials = sorted(list(self._special_tokens), key=lambda x: self._token_to_id[x])

    data = {
        "special_tokens": sorted_specials,
        "vocab": {token: token_id for token, token_id in self._token_to_id.items()},
    }

    path_obj = Path(path)
    try:
        path_obj.parent.mkdir(parents=True, exist_ok=True)
        with open(path_obj, "w", encoding="utf-8") as f:
            json.dump(data, f, indent=4, ensure_ascii=False)
    except Exception as e:
        raise OSError(f"Failed to save vocabulary to {path}: {e}") from e

load(path) classmethod

Load a vocabulary from a JSON file.

Parameters:

Name Type Description Default
path str | Path

Path to the JSON file to load.

 required

Returns:

Type Description
Vocabulary

A new Vocabulary instance reconstructed from the file.

Raises:

Type Description
ValueError

If the JSON format is invalid, IDs are not consecutive starting from 0, or special tokens are invalid or missing.
FileNotFoundError

If the vocabulary file does not exist.
IOError

If reading the file fails.
Source code in kamui/tokenizer/vocab.py 
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
@classmethod
def load(cls, path: str | Path) -> "Vocabulary":  # noqa: C901 — defensive JSON validation
    """Load a vocabulary from a JSON file.

    Args:
        path: Path to the JSON file to load.

    Returns:
        A new Vocabulary instance reconstructed from the file.

    Raises:
        ValueError: If the JSON format is invalid, IDs are not consecutive
            starting from 0, or special tokens are invalid or missing.
        FileNotFoundError: If the vocabulary file does not exist.
        IOError: If reading the file fails.
    """
    path_obj = Path(path)
    if not path_obj.exists():
        raise FileNotFoundError(f"Vocabulary file not found: {path}")

    try:
        with open(path_obj, encoding="utf-8") as f:
            data = json.load(f)
    except json.JSONDecodeError as e:
        raise ValueError(f"Invalid JSON format in vocabulary file {path}: {e}") from e
    except Exception as e:
        raise OSError(f"Failed to read vocabulary from {path}: {e}") from e

    if not isinstance(data, dict):
        raise ValueError("Vocabulary file must contain a JSON object at the top level")

    if "vocab" not in data or "special_tokens" not in data:
        raise ValueError("Vocabulary file must contain 'vocab' and 'special_tokens' keys")

    special_tokens_list = data["special_tokens"]
    vocab_dict = data["vocab"]

    if not isinstance(special_tokens_list, list):
        raise ValueError("'special_tokens' must be a list of strings")
    for t in special_tokens_list:
        if not isinstance(t, str):
            raise ValueError("All elements in 'special_tokens' must be strings")

    if not isinstance(vocab_dict, dict):
        raise ValueError("'vocab' must be a dictionary")

    # Sort vocab items by ID to verify they are consecutive and start from 0
    sorted_vocab: list[tuple[Any, Any]] = sorted(vocab_dict.items(), key=lambda item: item[1])

    for idx, (token, token_id) in enumerate(sorted_vocab):
        if not isinstance(token, str):
            raise ValueError(f"Vocabulary token must be a string, got type {type(token)}")
        if not isinstance(token_id, int) or isinstance(token_id, bool):
            raise ValueError(
                f"Vocabulary token ID must be an integer, got type {type(token_id)}"
            )
        if token_id != idx:
            raise ValueError(
                f"Vocabulary IDs must be consecutive starting from 0. "
                f"Expected {idx}, got {token_id}"
            )

    # Verify all listed special tokens are actually in the vocabulary
    for token in special_tokens_list:
        if token not in vocab_dict:
            raise ValueError(f"Special token '{token}' not found in vocabulary mapping")

    # Reconstruct Vocabulary
    instance = cls(special_tokens=special_tokens_list)

    for token, token_id in sorted_vocab:
        if token in instance._special_tokens:
            # Ensure the ID assigned at constructor matches the serialized ID
            if instance.token_to_id(token) != token_id:
                raise ValueError(
                    f"Special token '{token}' ID mismatch: expected {token_id}, "
                    f"got {instance.token_to_id(token)}"
                )
        else:
            assigned_id = instance.add_token(token)
            if assigned_id != token_id:
                raise ValueError(
                    f"Failed to load token '{token}' at expected ID {token_id}, "
                    f"assigned ID {assigned_id}"
                )

    return instance

Training

kamui.training.trainer.TrainingConfig dataclass

All training hyperparameters.

Attributes:

Name Type Description
max_lr / min_lr

Peak / floor learning rates for the schedule.
warmup_steps int

Linear-warmup steps.
max_steps int

Cosine-decay horizon (LR floors after this).
grad_accum_steps int

Micro-batches accumulated per optimiser update.
max_grad_norm float

Gradient-clipping norm (<= 0 disables clipping).
weight_decay float

AdamW weight decay for weight matrices.
betas / eps

AdamW hyperparameters.
eval_interval int

Steps between validation evaluations (0 disables).
log bool

Whether to print structured log lines.
Source code in kamui/training/trainer.py 
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
@dataclass
class TrainingConfig:
    """All training hyperparameters.

    Attributes:
        max_lr / min_lr:  Peak / floor learning rates for the schedule.
        warmup_steps:     Linear-warmup steps.
        max_steps:        Cosine-decay horizon (LR floors after this).
        grad_accum_steps: Micro-batches accumulated per optimiser update.
        max_grad_norm:    Gradient-clipping norm (<= 0 disables clipping).
        weight_decay:     AdamW weight decay for weight matrices.
        betas / eps:      AdamW hyperparameters.
        eval_interval:    Steps between validation evaluations (0 disables).
        log:              Whether to print structured log lines.
    """

    max_lr: float = 3e-4
    min_lr: float = 3e-5
    warmup_steps: int = 100
    max_steps: int = 5000
    grad_accum_steps: int = 1
    max_grad_norm: float = 1.0
    weight_decay: float = 0.1
    betas: tuple[float, float] = (0.9, 0.95)
    eps: float = 1e-8
    eval_interval: int = 0
    log: bool = False

    def __post_init__(self) -> None:
        if self.grad_accum_steps < 1:
            raise ValueError(f"grad_accum_steps must be >= 1, got {self.grad_accum_steps}")
        if self.eval_interval < 0:
            raise ValueError(f"eval_interval must be >= 0, got {self.eval_interval}")

kamui.training.trainer.Trainer

An explicit training loop over a KAMUITransformer.

Attributes:

Name Type Description
model

The model being trained.
optimizer

AdamW with decayed / non-decayed parameter groups.
scheduler

The cosine-with-warmup LR schedule.
step

The global optimiser-step counter.
history list[dict[str, float]]

A list of per-step log dicts (step, train_loss, lr, grad_norm).
Source code in kamui/training/trainer.py 
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
class Trainer:
    """An explicit training loop over a ``KAMUITransformer``.

    Attributes:
        model:      The model being trained.
        optimizer:  AdamW with decayed / non-decayed parameter groups.
        scheduler:  The cosine-with-warmup LR schedule.
        step:       The global optimiser-step counter.
        history:    A list of per-step log dicts (step, train_loss, lr, grad_norm).
    """

    def __init__(
        self,
        model: nn.Module,
        train_loader: Iterable[Batch],
        val_loader: Iterable[Batch] | None = None,
        config: TrainingConfig | None = None,
    ) -> None:
        """Create a trainer.

        Args:
            model:        The model to train.
            train_loader: An iterable of ``(inputs, targets)`` batches.
            val_loader:   Optional iterable for validation.
            config:       Training hyperparameters (defaults used if None).
        """
        self.model = model
        self.train_loader = train_loader
        self.val_loader = val_loader
        self.config = config or TrainingConfig()

        self.optimizer = build_optimizer(
            model,
            lr=self.config.max_lr,
            weight_decay=self.config.weight_decay,
            betas=self.config.betas,
            eps=self.config.eps,
        )
        self.scheduler = CosineWithWarmup(
            max_lr=self.config.max_lr,
            min_lr=self.config.min_lr,
            warmup_steps=self.config.warmup_steps,
            max_steps=self.config.max_steps,
        )
        self.step = 0
        self.history: list[dict[str, float]] = []
        self._data_iter = _infinite(train_loader)

    def _set_lr(self, lr: float) -> None:
        for group in self.optimizer.param_groups:
            group["lr"] = lr

    def train(self, n_steps: int) -> list[dict[str, float]]:
        """Run ``n_steps`` optimiser updates.

        Each update accumulates ``grad_accum_steps`` micro-batches, clips the
        gradient, applies the scheduled LR, and steps the optimiser.

        Args:
            n_steps: Number of optimiser updates to run.

        Returns:
            The list of per-step log dicts recorded during this call.
        """
        self.model.train()
        records: list[dict[str, float]] = []

        for _ in range(n_steps):
            self.optimizer.zero_grad()
            accum_loss = 0.0
            for _ in range(self.config.grad_accum_steps):
                inputs, targets = next(self._data_iter)
                loss = self.model(inputs, targets=targets)
                (loss / self.config.grad_accum_steps).backward()
                accum_loss += loss.item() / self.config.grad_accum_steps

            if self.config.max_grad_norm > 0:
                grad_norm = torch.nn.utils.clip_grad_norm_(
                    self.model.parameters(), self.config.max_grad_norm
                ).item()
            else:
                grad_norm = _grad_norm(self.model)

            lr = self.scheduler.get_lr(self.step)
            self._set_lr(lr)
            self.optimizer.step()
            self.step += 1

            record: dict[str, float] = {
                "step": float(self.step),
                "train_loss": accum_loss,
                "lr": lr,
                "grad_norm": grad_norm,
            }
            if (
                self.config.eval_interval
                and self.val_loader is not None
                and self.step % self.config.eval_interval == 0
            ):
                record["val_loss"] = self.evaluate()

            if self.config.log:
                print(
                    f"step={self.step} train_loss={accum_loss:.4f} "
                    f"lr={lr:.2e} grad_norm={grad_norm:.3f}"
                )
            records.append(record)
            self.history.append(record)

        return records

    @torch.no_grad()
    def evaluate(self) -> float:
        """Return the mean validation loss over ``val_loader``.

        Returns:
            The token-weighted mean cross-entropy loss.

        Raises:
            ValueError: If no validation loader was provided, or it is empty.
        """
        if self.val_loader is None:
            raise ValueError("no val_loader was provided")

        was_training = self.model.training
        self.model.eval()
        total_loss = 0.0
        total_tokens = 0
        for inputs, targets in self.val_loader:
            loss = self.model(inputs, targets=targets)
            total_loss += loss.item() * targets.numel()
            total_tokens += targets.numel()
        if was_training:
            self.model.train()

        if total_tokens == 0:
            raise ValueError("val_loader produced no tokens")
        return total_loss / total_tokens

__init__(model, train_loader, val_loader=None, config=None)

Create a trainer.

Parameters:

Name Type Description Default
model Module

The model to train.

 required
train_loader Iterable[Batch]

An iterable of (inputs, targets) batches.

 required
val_loader Iterable[Batch] | None

Optional iterable for validation.

 None
config TrainingConfig | None

Training hyperparameters (defaults used if None).

 None
Source code in kamui/training/trainer.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
def __init__(
    self,
    model: nn.Module,
    train_loader: Iterable[Batch],
    val_loader: Iterable[Batch] | None = None,
    config: TrainingConfig | None = None,
) -> None:
    """Create a trainer.

    Args:
        model:        The model to train.
        train_loader: An iterable of ``(inputs, targets)`` batches.
        val_loader:   Optional iterable for validation.
        config:       Training hyperparameters (defaults used if None).
    """
    self.model = model
    self.train_loader = train_loader
    self.val_loader = val_loader
    self.config = config or TrainingConfig()

    self.optimizer = build_optimizer(
        model,
        lr=self.config.max_lr,
        weight_decay=self.config.weight_decay,
        betas=self.config.betas,
        eps=self.config.eps,
    )
    self.scheduler = CosineWithWarmup(
        max_lr=self.config.max_lr,
        min_lr=self.config.min_lr,
        warmup_steps=self.config.warmup_steps,
        max_steps=self.config.max_steps,
    )
    self.step = 0
    self.history: list[dict[str, float]] = []
    self._data_iter = _infinite(train_loader)

train(n_steps)

Run n_steps optimiser updates.

Each update accumulates grad_accum_steps micro-batches, clips the gradient, applies the scheduled LR, and steps the optimiser.

Parameters:

Name Type Description Default
n_steps int

Number of optimiser updates to run.

 required

Returns:

Type Description
list[dict[str, float]]

The list of per-step log dicts recorded during this call.
Source code in kamui/training/trainer.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
def train(self, n_steps: int) -> list[dict[str, float]]:
    """Run ``n_steps`` optimiser updates.

    Each update accumulates ``grad_accum_steps`` micro-batches, clips the
    gradient, applies the scheduled LR, and steps the optimiser.

    Args:
        n_steps: Number of optimiser updates to run.

    Returns:
        The list of per-step log dicts recorded during this call.
    """
    self.model.train()
    records: list[dict[str, float]] = []

    for _ in range(n_steps):
        self.optimizer.zero_grad()
        accum_loss = 0.0
        for _ in range(self.config.grad_accum_steps):
            inputs, targets = next(self._data_iter)
            loss = self.model(inputs, targets=targets)
            (loss / self.config.grad_accum_steps).backward()
            accum_loss += loss.item() / self.config.grad_accum_steps

        if self.config.max_grad_norm > 0:
            grad_norm = torch.nn.utils.clip_grad_norm_(
                self.model.parameters(), self.config.max_grad_norm
            ).item()
        else:
            grad_norm = _grad_norm(self.model)

        lr = self.scheduler.get_lr(self.step)
        self._set_lr(lr)
        self.optimizer.step()
        self.step += 1

        record: dict[str, float] = {
            "step": float(self.step),
            "train_loss": accum_loss,
            "lr": lr,
            "grad_norm": grad_norm,
        }
        if (
            self.config.eval_interval
            and self.val_loader is not None
            and self.step % self.config.eval_interval == 0
        ):
            record["val_loss"] = self.evaluate()

        if self.config.log:
            print(
                f"step={self.step} train_loss={accum_loss:.4f} "
                f"lr={lr:.2e} grad_norm={grad_norm:.3f}"
            )
        records.append(record)
        self.history.append(record)

    return records

evaluate()

Return the mean validation loss over val_loader.

Returns:

Type Description
float

The token-weighted mean cross-entropy loss.

Raises:

Type Description
ValueError

If no validation loader was provided, or it is empty.
Source code in kamui/training/trainer.py 
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
@torch.no_grad()
def evaluate(self) -> float:
    """Return the mean validation loss over ``val_loader``.

    Returns:
        The token-weighted mean cross-entropy loss.

    Raises:
        ValueError: If no validation loader was provided, or it is empty.
    """
    if self.val_loader is None:
        raise ValueError("no val_loader was provided")

    was_training = self.model.training
    self.model.eval()
    total_loss = 0.0
    total_tokens = 0
    for inputs, targets in self.val_loader:
        loss = self.model(inputs, targets=targets)
        total_loss += loss.item() * targets.numel()
        total_tokens += targets.numel()
    if was_training:
        self.model.train()

    if total_tokens == 0:
        raise ValueError("val_loader produced no tokens")
    return total_loss / total_tokens

Distributed Training (v0.3)

kamui.training.distributed.wrap_ddp(model, device_ids=None, find_unused_parameters=False)

Wrap model in DistributedDataParallel.

On construction, DDP broadcasts rank 0's parameters to every rank, so all replicas start identical. The wrapper preserves model.parameters() and the forward signature, so build_optimizer and Trainer work unchanged.

Parameters:

Name Type Description Default
model Module

A model on this rank's device.

 required
device_ids list[int] | None

[local_gpu] for a single-GPU-per-process NCCL setup; None for CPU / gloo.

 None
find_unused_parameters bool

Enable only if some parameters do not receive gradients every step (it adds overhead).

 False

Returns:

Type Description
DistributedDataParallel

The DistributedDataParallel-wrapped model.

Raises:

Type Description
RuntimeError

If no process group is initialised.
Source code in kamui/training/distributed.py 
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
def wrap_ddp(
    model: nn.Module,
    device_ids: list[int] | None = None,
    find_unused_parameters: bool = False,
) -> DistributedDataParallel:
    """Wrap ``model`` in ``DistributedDataParallel``.

    On construction, DDP broadcasts rank 0's parameters to every rank, so all
    replicas start identical.  The wrapper preserves ``model.parameters()`` and
    the forward signature, so ``build_optimizer`` and ``Trainer`` work unchanged.

    Args:
        model:                  A model on this rank's device.
        device_ids:             ``[local_gpu]`` for a single-GPU-per-process
            NCCL setup; ``None`` for CPU / ``gloo``.
        find_unused_parameters: Enable only if some parameters do not receive
            gradients every step (it adds overhead).

    Returns:
        The ``DistributedDataParallel``-wrapped model.

    Raises:
        RuntimeError: If no process group is initialised.
    """
    if not is_dist_available_and_initialized():
        raise RuntimeError("call init_process_group before wrap_ddp")
    return DistributedDataParallel(
        model, device_ids=device_ids, find_unused_parameters=find_unused_parameters
    )

kamui.training.distributed.DistributedDataLoader

A batching loader that yields only this rank's shard of a TextDataset.

Each epoch, the dataset indices are (optionally) shuffled with an epoch-dependent seed shared by all ranks, truncated to a multiple of world_size, then split into contiguous per-rank blocks of equal size. Every rank therefore sees a disjoint shard and the same number of batches, keeping the replicas in lock-step.

Attributes:

Name Type Description
dataset

The wrapped dataset.
batch_size

Sequences per batch.
rank

This loader's rank.
world_size

Total number of ranks.
Source code in kamui/training/distributed.py 
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
class DistributedDataLoader:
    """A batching loader that yields only this rank's shard of a ``TextDataset``.

    Each epoch, the dataset indices are (optionally) shuffled with an
    epoch-dependent seed shared by all ranks, truncated to a multiple of
    ``world_size``, then split into contiguous per-rank blocks of equal size.
    Every rank therefore sees a disjoint shard and the same number of batches,
    keeping the replicas in lock-step.

    Attributes:
        dataset:    The wrapped dataset.
        batch_size: Sequences per batch.
        rank:       This loader's rank.
        world_size: Total number of ranks.
    """

    def __init__(
        self,
        dataset: TextDataset,
        batch_size: int,
        rank: int | None = None,
        world_size: int | None = None,
        shuffle: bool = True,
        seed: int = 0,
        drop_last: bool = True,
    ) -> None:
        """Create a distributed loader.

        Args:
            dataset:    A ``TextDataset``.
            batch_size: Sequences per batch (>= 1).
            rank:       Rank (defaults to the current rank).
            world_size: Process count (defaults to the current world size).
            shuffle:    Shuffle the global index order each epoch.
            seed:       Base RNG seed (the epoch index is added to it).
            drop_last:  Drop a trailing partial batch within the shard.

        Raises:
            ValueError: If ``batch_size < 1``, ``world_size < 1``, or ``rank``
                is out of range.
        """
        if batch_size < 1:
            raise ValueError(f"batch_size must be >= 1, got {batch_size}")
        self.dataset = dataset
        self.batch_size = batch_size
        self.rank = get_rank() if rank is None else rank
        self.world_size = get_world_size() if world_size is None else world_size
        if self.world_size < 1:
            raise ValueError(f"world_size must be >= 1, got {self.world_size}")
        if not (0 <= self.rank < self.world_size):
            raise ValueError(f"rank must be in [0, {self.world_size}), got {self.rank}")
        self.shuffle = shuffle
        self.seed = seed
        self.drop_last = drop_last
        self._epoch = 0
        self._per_rank = len(dataset) // self.world_size

    def set_epoch(self, epoch: int) -> None:
        """Set the epoch so shuffling differs across epochs but agrees across ranks."""
        self._epoch = epoch

    def _shard(self) -> list[int]:
        indices = list(range(len(self.dataset)))
        if self.shuffle:
            g = torch.Generator()
            g.manual_seed(self.seed + self._epoch)
            perm = torch.randperm(len(indices), generator=g).tolist()
            indices = [indices[i] for i in perm]
        usable = self._per_rank * self.world_size
        indices = indices[:usable]
        start = self.rank * self._per_rank
        return indices[start : start + self._per_rank]

    def __len__(self) -> int:
        """Number of batches this rank yields per epoch."""
        if self.drop_last:
            return self._per_rank // self.batch_size
        return (self._per_rank + self.batch_size - 1) // self.batch_size

    def __iter__(self) -> Iterator[tuple[Tensor, Tensor]]:
        """Yield ``(inputs, targets)`` batches from this rank's shard, then bump the epoch."""
        shard = self._shard()
        for start in range(0, len(shard), self.batch_size):
            batch_idx = shard[start : start + self.batch_size]
            if self.drop_last and len(batch_idx) < self.batch_size:
                break
            pairs = [self.dataset[i] for i in batch_idx]
            inputs = torch.stack([x for x, _ in pairs])
            targets = torch.stack([y for _, y in pairs])
            yield inputs, targets
        self._epoch += 1

__init__(dataset, batch_size, rank=None, world_size=None, shuffle=True, seed=0, drop_last=True)

Create a distributed loader.

Parameters:

Name Type Description Default
dataset TextDataset

A TextDataset.

 required
batch_size int

Sequences per batch (>= 1).

 required
rank int | None

Rank (defaults to the current rank).

 None
world_size int | None

Process count (defaults to the current world size).

 None
shuffle bool

Shuffle the global index order each epoch.

 True
seed int

Base RNG seed (the epoch index is added to it).

 0
drop_last bool

Drop a trailing partial batch within the shard.

 True

Raises:

Type Description
ValueError

If batch_size < 1, world_size < 1, or rank is out of range.
Source code in kamui/training/distributed.py 
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
def __init__(
    self,
    dataset: TextDataset,
    batch_size: int,
    rank: int | None = None,
    world_size: int | None = None,
    shuffle: bool = True,
    seed: int = 0,
    drop_last: bool = True,
) -> None:
    """Create a distributed loader.

    Args:
        dataset:    A ``TextDataset``.
        batch_size: Sequences per batch (>= 1).
        rank:       Rank (defaults to the current rank).
        world_size: Process count (defaults to the current world size).
        shuffle:    Shuffle the global index order each epoch.
        seed:       Base RNG seed (the epoch index is added to it).
        drop_last:  Drop a trailing partial batch within the shard.

    Raises:
        ValueError: If ``batch_size < 1``, ``world_size < 1``, or ``rank``
            is out of range.
    """
    if batch_size < 1:
        raise ValueError(f"batch_size must be >= 1, got {batch_size}")
    self.dataset = dataset
    self.batch_size = batch_size
    self.rank = get_rank() if rank is None else rank
    self.world_size = get_world_size() if world_size is None else world_size
    if self.world_size < 1:
        raise ValueError(f"world_size must be >= 1, got {self.world_size}")
    if not (0 <= self.rank < self.world_size):
        raise ValueError(f"rank must be in [0, {self.world_size}), got {self.rank}")
    self.shuffle = shuffle
    self.seed = seed
    self.drop_last = drop_last
    self._epoch = 0
    self._per_rank = len(dataset) // self.world_size

set_epoch(epoch)

Set the epoch so shuffling differs across epochs but agrees across ranks.

Source code in kamui/training/distributed.py 
342
343
344
def set_epoch(self, epoch: int) -> None:
    """Set the epoch so shuffling differs across epochs but agrees across ranks."""
    self._epoch = epoch

__len__()

Number of batches this rank yields per epoch.

Source code in kamui/training/distributed.py 
358
359
360
361
362
def __len__(self) -> int:
    """Number of batches this rank yields per epoch."""
    if self.drop_last:
        return self._per_rank // self.batch_size
    return (self._per_rank + self.batch_size - 1) // self.batch_size

__iter__()

Yield (inputs, targets) batches from this rank's shard, then bump the epoch.

Source code in kamui/training/distributed.py 
364
365
366
367
368
369
370
371
372
373
374
375
def __iter__(self) -> Iterator[tuple[Tensor, Tensor]]:
    """Yield ``(inputs, targets)`` batches from this rank's shard, then bump the epoch."""
    shard = self._shard()
    for start in range(0, len(shard), self.batch_size):
        batch_idx = shard[start : start + self.batch_size]
        if self.drop_last and len(batch_idx) < self.batch_size:
            break
        pairs = [self.dataset[i] for i in batch_idx]
        inputs = torch.stack([x for x, _ in pairs])
        targets = torch.stack([y for _, y in pairs])
        yield inputs, targets
    self._epoch += 1

kamui.training.distributed.spawn_workers(worker_fn, world_size, args=(), backend='gloo', init_method=None)

Launch worker_fn across world_size processes and wait for them.

Each process runs worker_fn(rank, world_size, *args) with the process group already initialised (and torn down on return). worker_fn and every element of args must be picklable (defined at module top level), because the spawn start method re-imports them in each child.

Parameters:

Name Type Description Default
worker_fn Callable[..., None]

worker_fn(rank, world_size, *args) -> None.

 required
world_size int

Number of processes to launch (>= 1).

 required
args tuple[object, ...]

Extra positional arguments forwarded to worker_fn.

 ()
backend str

"gloo" (CPU / portable) or "nccl" (CUDA).

 'gloo'
init_method str | None

Rendezvous URL. None allocates a private FileStore (portable, needs no free TCP port).

 None

Raises:

Type Description
ValueError

If world_size < 1.
Source code in kamui/training/distributed.py 
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
def spawn_workers(
    worker_fn: Callable[..., None],
    world_size: int,
    args: tuple[object, ...] = (),
    backend: str = "gloo",
    init_method: str | None = None,
) -> None:
    """Launch ``worker_fn`` across ``world_size`` processes and wait for them.

    Each process runs ``worker_fn(rank, world_size, *args)`` with the process
    group already initialised (and torn down on return).  ``worker_fn`` and
    every element of ``args`` must be picklable (defined at module top level),
    because the ``spawn`` start method re-imports them in each child.

    Args:
        worker_fn:   ``worker_fn(rank, world_size, *args) -> None``.
        world_size:  Number of processes to launch (>= 1).
        args:        Extra positional arguments forwarded to ``worker_fn``.
        backend:     ``"gloo"`` (CPU / portable) or ``"nccl"`` (CUDA).
        init_method: Rendezvous URL. ``None`` allocates a private ``FileStore``
            (portable, needs no free TCP port).

    Raises:
        ValueError: If ``world_size < 1``.
    """
    if world_size < 1:
        raise ValueError(f"world_size must be >= 1, got {world_size}")
    if init_method is None:
        store_path = os.path.join(tempfile.gettempdir(), f"kamui_ddp_{uuid.uuid4().hex}")
        init_method = "file:///" + store_path.replace("\\", "/")
    mp.spawn(
        _worker_entry,
        args=(world_size, init_method, backend, worker_fn, args),
        nprocs=world_size,
        join=True,
    )

kamui.training.distributed.all_reduce_mean(value)

Average a scalar / tensor across all ranks (identity when not initialised).

Parameters:

Name Type Description Default
value float | Tensor

A Python float or a tensor. A float returns a float; a tensor returns a new tensor (the input is not modified).

 required

Returns:

Type Description
float | Tensor

The mean of value over all ranks, in the same type as the input.
Source code in kamui/training/distributed.py 
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
def all_reduce_mean(value: float | Tensor) -> float | Tensor:
    """Average a scalar / tensor across all ranks (identity when not initialised).

    Args:
        value: A Python float or a tensor. A float returns a float; a tensor
            returns a new tensor (the input is not modified).

    Returns:
        The mean of ``value`` over all ranks, in the same type as the input.
    """
    if not is_dist_available_and_initialized():
        return value
    world_size = dist.get_world_size()
    if isinstance(value, Tensor):
        tensor = value.detach().clone()
        dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
        return tensor / world_size
    tensor = torch.tensor(float(value))
    dist.all_reduce(tensor, op=dist.ReduceOp.SUM)
    return (tensor / world_size).item()

Hooks

kamui.hooks.manager.HookManager

Context-managed activation capture over a model's named submodules.

Source code in kamui/hooks/manager.py 
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
class HookManager:
    """Context-managed activation capture over a model's named submodules."""

    def __init__(self, model: nn.Module) -> None:
        """Create a hook manager for ``model``.

        Args:
            model: The model to attach hooks to.  Must expose a ``config``
                attribute (a ``ModelConfig``) so hook points can be validated.
        """
        self.model = model
        self._cache: ActivationCache = {}
        self._handles: list[Any] = []
        self._wrapped: list[nn.Module] = []

    # ------------------------------------------------------------------
    # Context manager protocol
    # ------------------------------------------------------------------

    def __enter__(self) -> HookManager:
        return self

    def __exit__(self, *exc: object) -> None:
        self.remove()

    # ------------------------------------------------------------------
    # Attaching
    # ------------------------------------------------------------------

    def attach(self, module_path: str, point: str) -> HookManager:
        """Attach a capture hook at ``"<module_path>.<point>"``.

        Args:
            module_path: A ``named_modules()`` key, e.g. ``"blocks.3.attn"``.
            point:       One of ``"output"``, ``"input"``, ``"weights"``, ``"mid"``.

        Returns:
            ``self`` (so calls can be chained).

        Raises:
            ValueError: If the resulting hook point is not valid for this model.
        """
        hook_point = f"{module_path}.{point}"
        # HookManager works on any nn.Module exposing a ModelConfig at .config.
        config = cast(ModelConfig, self.model.config)
        if not HookRegistry.validate(hook_point, config):
            valid = HookRegistry.all_points(config)
            raise ValueError(f"'{hook_point}' is not a valid hook point. Valid points: {valid}")

        if point == "output":
            module = self._resolve(module_path)
            handle = module.register_forward_hook(self._output_saver(hook_point))
            self._handles.append(handle)
        elif point == "input":
            module = self._resolve(module_path)
            handle = module.register_forward_pre_hook(self._input_saver(hook_point))
            self._handles.append(handle)
        elif point == "mid":
            # FFN hidden activation = output of the GELU submodule.
            activation = self._resolve(f"{module_path}.activation")
            handle = activation.register_forward_hook(self._output_saver(hook_point))
            self._handles.append(handle)
        else:  # point == "weights" (registry guarantees no other value)
            module = self._resolve(module_path)
            self._wrap_for_weights(module, hook_point)

        return self

    # ------------------------------------------------------------------
    # Reading the cache
    # ------------------------------------------------------------------

    def get(self, hook_point: str) -> Tensor:
        """Return a cached activation.

        Args:
            hook_point: The full hook-point string, e.g. ``"embed.output"``.

        Returns:
            The captured tensor.

        Raises:
            KeyError: If nothing is cached at ``hook_point`` (not attached, or
                no forward pass has run since the last ``clear``).
        """
        if hook_point not in self._cache:
            raise KeyError(
                f"'{hook_point}' is not in the cache. Did you attach it and run "
                f"a forward pass? Cached points: {list(self._cache)}"
            )
        return self._cache[hook_point]

    def get_all(self) -> ActivationCache:
        """Return a shallow copy of all cached activations."""
        return dict(self._cache)

    def clear(self) -> None:
        """Clear the activation cache without removing hooks."""
        self._cache.clear()

    # ------------------------------------------------------------------
    # Teardown
    # ------------------------------------------------------------------

    def remove(self) -> None:
        """Remove all hooks and restore any wrapped forwards.

        The activation cache is left intact so captured tensors remain
        accessible after the ``with`` block exits.
        """
        for handle in self._handles:
            handle.remove()
        self._handles.clear()
        for module in self._wrapped:
            # Delete the instance-level override so the class ``forward`` returns.
            del module.forward
        self._wrapped.clear()

    # ------------------------------------------------------------------
    # Internals
    # ------------------------------------------------------------------

    def _resolve(self, module_path: str) -> nn.Module:
        # attach() validates the hook point against the registry first, so the
        # module path is always present here; a direct lookup suffices.
        return dict(self.model.named_modules())[module_path]

    def _output_saver(
        self, hook_point: str
    ) -> Callable[[nn.Module, tuple[Any, ...], Tensor], None]:
        # Every hookable point (embed, attn output, ffn output, ffn.activation)
        # yields a plain tensor: attention returns a bare tensor in the model's
        # forward, and the ``weights`` path is captured by the forward wrapper.
        def hook(module: nn.Module, inputs: tuple[Any, ...], output: Tensor) -> None:
            self._cache[hook_point] = output.detach()

        return hook

    def _input_saver(self, hook_point: str) -> Callable[[nn.Module, tuple[Any, ...]], None]:
        def pre_hook(module: nn.Module, inputs: tuple[Any, ...]) -> None:
            self._cache[hook_point] = inputs[0].detach()

        return pre_hook

    def _wrap_for_weights(self, module: nn.Module, hook_point: str) -> None:
        if module in self._wrapped:
            return
        original_forward = module.forward

        def wrapper(x: Tensor, *args: Any, **kwargs: Any) -> Tensor:
            output, weights = original_forward(x, return_weights=True)
            self._cache[hook_point] = weights.detach()
            return output

        module.forward = wrapper  # type: ignore[method-assign]
        self._wrapped.append(module)

__init__(model)

Create a hook manager for model.

Parameters:

Name Type Description Default
model Module

The model to attach hooks to. Must expose a config attribute (a ModelConfig) so hook points can be validated.

 required
Source code in kamui/hooks/manager.py 
80
81
82
83
84
85
86
87
88
89
90
def __init__(self, model: nn.Module) -> None:
    """Create a hook manager for ``model``.

    Args:
        model: The model to attach hooks to.  Must expose a ``config``
            attribute (a ``ModelConfig``) so hook points can be validated.
    """
    self.model = model
    self._cache: ActivationCache = {}
    self._handles: list[Any] = []
    self._wrapped: list[nn.Module] = []

attach(module_path, point)

Attach a capture hook at "<module_path>.<point>".

Parameters:

Name Type Description Default
module_path str

A named_modules() key, e.g. "blocks.3.attn".

 required
point str

One of "output", "input", "weights", "mid".

 required

Returns:

Type Description
HookManager

self (so calls can be chained).

Raises:

Type Description
ValueError

If the resulting hook point is not valid for this model.
Source code in kamui/hooks/manager.py 
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
def attach(self, module_path: str, point: str) -> HookManager:
    """Attach a capture hook at ``"<module_path>.<point>"``.

    Args:
        module_path: A ``named_modules()`` key, e.g. ``"blocks.3.attn"``.
        point:       One of ``"output"``, ``"input"``, ``"weights"``, ``"mid"``.

    Returns:
        ``self`` (so calls can be chained).

    Raises:
        ValueError: If the resulting hook point is not valid for this model.
    """
    hook_point = f"{module_path}.{point}"
    # HookManager works on any nn.Module exposing a ModelConfig at .config.
    config = cast(ModelConfig, self.model.config)
    if not HookRegistry.validate(hook_point, config):
        valid = HookRegistry.all_points(config)
        raise ValueError(f"'{hook_point}' is not a valid hook point. Valid points: {valid}")

    if point == "output":
        module = self._resolve(module_path)
        handle = module.register_forward_hook(self._output_saver(hook_point))
        self._handles.append(handle)
    elif point == "input":
        module = self._resolve(module_path)
        handle = module.register_forward_pre_hook(self._input_saver(hook_point))
        self._handles.append(handle)
    elif point == "mid":
        # FFN hidden activation = output of the GELU submodule.
        activation = self._resolve(f"{module_path}.activation")
        handle = activation.register_forward_hook(self._output_saver(hook_point))
        self._handles.append(handle)
    else:  # point == "weights" (registry guarantees no other value)
        module = self._resolve(module_path)
        self._wrap_for_weights(module, hook_point)

    return self

get(hook_point)

Return a cached activation.

Parameters:

Name Type Description Default
hook_point str

The full hook-point string, e.g. "embed.output".

 required

Returns:

Type Description
Tensor

The captured tensor.

Raises:

Type Description
KeyError

If nothing is cached at hook_point (not attached, or no forward pass has run since the last clear).
Source code in kamui/hooks/manager.py 
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
def get(self, hook_point: str) -> Tensor:
    """Return a cached activation.

    Args:
        hook_point: The full hook-point string, e.g. ``"embed.output"``.

    Returns:
        The captured tensor.

    Raises:
        KeyError: If nothing is cached at ``hook_point`` (not attached, or
            no forward pass has run since the last ``clear``).
    """
    if hook_point not in self._cache:
        raise KeyError(
            f"'{hook_point}' is not in the cache. Did you attach it and run "
            f"a forward pass? Cached points: {list(self._cache)}"
        )
    return self._cache[hook_point]

get_all()

Return a shallow copy of all cached activations.

Source code in kamui/hooks/manager.py 
169
170
171
def get_all(self) -> ActivationCache:
    """Return a shallow copy of all cached activations."""
    return dict(self._cache)

clear()

Clear the activation cache without removing hooks.

Source code in kamui/hooks/manager.py 
173
174
175
def clear(self) -> None:
    """Clear the activation cache without removing hooks."""
    self._cache.clear()

remove()

Remove all hooks and restore any wrapped forwards.

The activation cache is left intact so captured tensors remain accessible after the with block exits.

Source code in kamui/hooks/manager.py 
181
182
183
184
185
186
187
188
189
190
191
192
193
def remove(self) -> None:
    """Remove all hooks and restore any wrapped forwards.

    The activation cache is left intact so captured tensors remain
    accessible after the ``with`` block exits.
    """
    for handle in self._handles:
        handle.remove()
    self._handles.clear()
    for module in self._wrapped:
        # Delete the instance-level override so the class ``forward`` returns.
        del module.forward
    self._wrapped.clear()

kamui.hooks.registry.HookRegistry

Canonical registry of valid hook points for a KAMUITransformer.

Source code in kamui/hooks/registry.py 
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
class HookRegistry:
    """Canonical registry of valid hook points for a KAMUITransformer."""

    #: Per-layer hook point suffixes, in canonical order.
    _PER_LAYER_POINTS: tuple[str, ...] = (
        "attn.output",
        "attn.weights",
        "ffn.mid",
        "ffn.output",
    )

    @classmethod
    def all_points(cls, config: ModelConfig) -> list[str]:
        """Return every valid hook point string for ``config``.

        Args:
            config: The model configuration (only ``n_layers`` is used).

        Returns:
            An ordered list of hook-point strings: the embedding output, the
            per-layer attention/FFN points for each of ``n_layers`` blocks, and
            the unembedding input.
        """
        points: list[str] = ["embed.output"]
        for layer in range(config.n_layers):
            for suffix in cls._PER_LAYER_POINTS:
                points.append(f"blocks.{layer}.{suffix}")
        points.append("unembed.input")
        return points

    @classmethod
    def validate(cls, hook_point: str, config: ModelConfig) -> bool:
        """Return True if ``hook_point`` is valid for ``config``.

        Args:
            hook_point: A candidate hook-point string.
            config:     The model configuration.

        Returns:
            Whether ``hook_point`` is in ``all_points(config)``.
        """
        return hook_point in set(cls.all_points(config))

all_points(config) classmethod

Return every valid hook point string for config.

Parameters:

Name Type Description Default
config ModelConfig

The model configuration (only n_layers is used).

 required

Returns:

Type Description
list[str]

An ordered list of hook-point strings: the embedding output, the
list[str]

per-layer attention/FFN points for each of n_layers blocks, and
list[str]

the unembedding input.
Source code in kamui/hooks/registry.py 
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
@classmethod
def all_points(cls, config: ModelConfig) -> list[str]:
    """Return every valid hook point string for ``config``.

    Args:
        config: The model configuration (only ``n_layers`` is used).

    Returns:
        An ordered list of hook-point strings: the embedding output, the
        per-layer attention/FFN points for each of ``n_layers`` blocks, and
        the unembedding input.
    """
    points: list[str] = ["embed.output"]
    for layer in range(config.n_layers):
        for suffix in cls._PER_LAYER_POINTS:
            points.append(f"blocks.{layer}.{suffix}")
    points.append("unembed.input")
    return points

validate(hook_point, config) classmethod

Return True if hook_point is valid for config.

Parameters:

Name Type Description Default
hook_point str

A candidate hook-point string.

 required
config ModelConfig

The model configuration.

 required

Returns:

Type Description
bool

Whether hook_point is in all_points(config).
Source code in kamui/hooks/registry.py 
82
83
84
85
86
87
88
89
90
91
92
93
@classmethod
def validate(cls, hook_point: str, config: ModelConfig) -> bool:
    """Return True if ``hook_point`` is valid for ``config``.

    Args:
        hook_point: A candidate hook-point string.
        config:     The model configuration.

    Returns:
        Whether ``hook_point`` is in ``all_points(config)``.
    """
    return hook_point in set(cls.all_points(config))

Interpretability Tools

kamui.mechinterp.attention_viz.AttentionVisualizer

Extract every head's attention pattern for a given input.

Attributes:

Name Type Description
model

A trained KAMUITransformer.
tokenizer

A tokenizer with decode (for axis labels).
Source code in kamui/mechinterp/attention_viz.py 
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
class AttentionVisualizer:
    """Extract every head's attention pattern for a given input.

    Attributes:
        model:     A trained ``KAMUITransformer``.
        tokenizer: A tokenizer with ``decode`` (for axis labels).
    """

    def __init__(self, model: KAMUITransformer, tokenizer: Any) -> None:
        """Create a visualizer over ``model``.

        Args:
            model:     A ``KAMUITransformer``.
            tokenizer: A tokenizer with ``decode(list[int]) -> str``.
        """
        self.model = model
        self.tokenizer = tokenizer

    @torch.no_grad()
    def run(self, token_ids: Tensor) -> AttentionResult:
        """Capture all attention weights for a single sequence.

        Args:
            token_ids: Token IDs of shape ``(S,)`` or ``(1, S)``.

        Returns:
            An ``AttentionResult``.

        Raises:
            ValueError: If ``token_ids`` is not a single sequence.
        """
        if token_ids.dim() == 1:
            token_ids = token_ids.unsqueeze(0)
        if token_ids.dim() != 2 or token_ids.shape[0] != 1:
            raise ValueError(
                f"token_ids must be a single sequence (S,) or (1, S), "
                f"got shape {tuple(token_ids.shape)}"
            )

        self.model.eval()
        n_layers = self.model.config.n_layers
        with HookManager(self.model) as hooks:
            for i in range(n_layers):
                hooks.attach(f"blocks.{i}.attn", "weights")
            self.model(token_ids)
            weights = torch.stack(
                [hooks.get(f"blocks.{i}.attn.weights")[0] for i in range(n_layers)]
            )  # (n_layers, n_heads, S, S)

        tokens = [_decode_token(self.tokenizer, int(t)) for t in token_ids[0].tolist()]
        return AttentionResult(weights=weights, tokens=tokens)

__init__(model, tokenizer)

Create a visualizer over model.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
tokenizer Any

A tokenizer with decode(list[int]) -> str.

 required
Source code in kamui/mechinterp/attention_viz.py 
198
199
200
201
202
203
204
205
206
def __init__(self, model: KAMUITransformer, tokenizer: Any) -> None:
    """Create a visualizer over ``model``.

    Args:
        model:     A ``KAMUITransformer``.
        tokenizer: A tokenizer with ``decode(list[int]) -> str``.
    """
    self.model = model
    self.tokenizer = tokenizer

run(token_ids)

Capture all attention weights for a single sequence.

Parameters:

Name Type Description Default
token_ids Tensor

Token IDs of shape (S,) or (1, S).

 required

Returns:

Type Description
AttentionResult

An AttentionResult.

Raises:

Type Description
ValueError

If token_ids is not a single sequence.
Source code in kamui/mechinterp/attention_viz.py 
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
@torch.no_grad()
def run(self, token_ids: Tensor) -> AttentionResult:
    """Capture all attention weights for a single sequence.

    Args:
        token_ids: Token IDs of shape ``(S,)`` or ``(1, S)``.

    Returns:
        An ``AttentionResult``.

    Raises:
        ValueError: If ``token_ids`` is not a single sequence.
    """
    if token_ids.dim() == 1:
        token_ids = token_ids.unsqueeze(0)
    if token_ids.dim() != 2 or token_ids.shape[0] != 1:
        raise ValueError(
            f"token_ids must be a single sequence (S,) or (1, S), "
            f"got shape {tuple(token_ids.shape)}"
        )

    self.model.eval()
    n_layers = self.model.config.n_layers
    with HookManager(self.model) as hooks:
        for i in range(n_layers):
            hooks.attach(f"blocks.{i}.attn", "weights")
        self.model(token_ids)
        weights = torch.stack(
            [hooks.get(f"blocks.{i}.attn.weights")[0] for i in range(n_layers)]
        )  # (n_layers, n_heads, S, S)

    tokens = [_decode_token(self.tokenizer, int(t)) for t in token_ids[0].tolist()]
    return AttentionResult(weights=weights, tokens=tokens)

kamui.mechinterp.logit_lens.LogitLens

Project the residual stream to vocabulary at every layer.

Attributes:

Name Type Description
model

A trained KAMUITransformer.
tokenizer

A tokenizer with decode (for token labels).
Source code in kamui/mechinterp/logit_lens.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
class LogitLens:
    """Project the residual stream to vocabulary at every layer.

    Attributes:
        model:     A trained ``KAMUITransformer``.
        tokenizer: A tokenizer with ``decode`` (for token labels).
    """

    def __init__(self, model: KAMUITransformer, tokenizer: Any) -> None:
        """Create a logit lens over ``model``.

        Args:
            model:     A ``KAMUITransformer``.
            tokenizer: A tokenizer with a ``decode(list[int]) -> str`` method.
        """
        self.model = model
        self.tokenizer = tokenizer

    @torch.no_grad()
    def run(self, token_ids: Tensor, top_k: int = 5) -> LogitLensResult:
        """Run the logit lens over a single sequence.

        Args:
            token_ids: Token IDs of shape ``(S,)`` or ``(1, S)``.
            top_k:     Number of top tokens to record per (layer, position).

        Returns:
            A ``LogitLensResult``.

        Raises:
            ValueError: If ``token_ids`` is not a single sequence, or ``top_k < 1``.
        """
        if top_k < 1:
            raise ValueError(f"top_k must be >= 1, got {top_k}")
        if token_ids.dim() == 1:
            token_ids = token_ids.unsqueeze(0)
        if token_ids.dim() != 2 or token_ids.shape[0] != 1:
            raise ValueError(
                f"token_ids must be a single sequence (S,) or (1, S), "
                f"got shape {tuple(token_ids.shape)}"
            )

        n_layers = self.model.config.n_layers
        self.model.eval()

        with HookManager(self.model) as hooks:
            hooks.attach("embed", "output")
            for i in range(n_layers):
                hooks.attach(f"blocks.{i}.attn", "output")
                hooks.attach(f"blocks.{i}.ffn", "output")
            self.model(token_ids)

            stream = hooks.get("embed.output")  # (1, S, D)
            streams = [stream]
            for i in range(n_layers):
                stream = (
                    stream
                    + hooks.get(f"blocks.{i}.attn.output")
                    + hooks.get(f"blocks.{i}.ffn.output")
                )
                streams.append(stream)

        # Project each residual stream through final_ln + unembed.
        probs_per_layer = []
        for stream in streams:
            logits = self.model.unembed(self.model.final_ln(stream))  # (1, S, V)
            probs_per_layer.append(torch.softmax(logits, dim=-1)[0])  # (S, V)
        probs = torch.stack(probs_per_layer)  # (L+1, S, V)

        vocab = probs.shape[-1]
        top_tokens = probs.topk(min(top_k, vocab), dim=-1).indices  # (L+1, S, k)
        labels = [_decode_token(self.tokenizer, int(t)) for t in token_ids[0].tolist()]

        return LogitLensResult(
            probs=probs, top_tokens=top_tokens, tokens=labels, tokenizer=self.tokenizer
        )

__init__(model, tokenizer)

Create a logit lens over model.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
tokenizer Any

A tokenizer with a decode(list[int]) -> str method.

 required
Source code in kamui/mechinterp/logit_lens.py 
173
174
175
176
177
178
179
180
181
def __init__(self, model: KAMUITransformer, tokenizer: Any) -> None:
    """Create a logit lens over ``model``.

    Args:
        model:     A ``KAMUITransformer``.
        tokenizer: A tokenizer with a ``decode(list[int]) -> str`` method.
    """
    self.model = model
    self.tokenizer = tokenizer

run(token_ids, top_k=5)

Run the logit lens over a single sequence.

Parameters:

Name Type Description Default
token_ids Tensor

Token IDs of shape (S,) or (1, S).

 required
top_k int

Number of top tokens to record per (layer, position).

 5

Returns:

Type Description
LogitLensResult

A LogitLensResult.

Raises:

Type Description
ValueError

If token_ids is not a single sequence, or top_k < 1.
Source code in kamui/mechinterp/logit_lens.py 
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
@torch.no_grad()
def run(self, token_ids: Tensor, top_k: int = 5) -> LogitLensResult:
    """Run the logit lens over a single sequence.

    Args:
        token_ids: Token IDs of shape ``(S,)`` or ``(1, S)``.
        top_k:     Number of top tokens to record per (layer, position).

    Returns:
        A ``LogitLensResult``.

    Raises:
        ValueError: If ``token_ids`` is not a single sequence, or ``top_k < 1``.
    """
    if top_k < 1:
        raise ValueError(f"top_k must be >= 1, got {top_k}")
    if token_ids.dim() == 1:
        token_ids = token_ids.unsqueeze(0)
    if token_ids.dim() != 2 or token_ids.shape[0] != 1:
        raise ValueError(
            f"token_ids must be a single sequence (S,) or (1, S), "
            f"got shape {tuple(token_ids.shape)}"
        )

    n_layers = self.model.config.n_layers
    self.model.eval()

    with HookManager(self.model) as hooks:
        hooks.attach("embed", "output")
        for i in range(n_layers):
            hooks.attach(f"blocks.{i}.attn", "output")
            hooks.attach(f"blocks.{i}.ffn", "output")
        self.model(token_ids)

        stream = hooks.get("embed.output")  # (1, S, D)
        streams = [stream]
        for i in range(n_layers):
            stream = (
                stream
                + hooks.get(f"blocks.{i}.attn.output")
                + hooks.get(f"blocks.{i}.ffn.output")
            )
            streams.append(stream)

    # Project each residual stream through final_ln + unembed.
    probs_per_layer = []
    for stream in streams:
        logits = self.model.unembed(self.model.final_ln(stream))  # (1, S, V)
        probs_per_layer.append(torch.softmax(logits, dim=-1)[0])  # (S, V)
    probs = torch.stack(probs_per_layer)  # (L+1, S, V)

    vocab = probs.shape[-1]
    top_tokens = probs.topk(min(top_k, vocab), dim=-1).indices  # (L+1, S, k)
    labels = [_decode_token(self.tokenizer, int(t)) for t in token_ids[0].tolist()]

    return LogitLensResult(
        probs=probs, top_tokens=top_tokens, tokens=labels, tokenizer=self.tokenizer
    )

kamui.mechinterp.probing.LinearProbe

Train linear classifiers on cached activations.

Attributes:

Name Type Description
model

A trained KAMUITransformer.
Source code in kamui/mechinterp/probing.py 
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
class LinearProbe:
    """Train linear classifiers on cached activations.

    Attributes:
        model: A trained ``KAMUITransformer``.
    """

    def __init__(self, model: KAMUITransformer) -> None:
        """Create a prober over ``model``.

        Args:
            model: A ``KAMUITransformer``.
        """
        self.model = model

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _validate(dataset: list[Tensor], labels: list[int]) -> Tensor:
        if not dataset:
            raise ValueError("dataset must be non-empty")
        if len(dataset) != len(labels):
            raise ValueError(
                f"dataset ({len(dataset)}) and labels ({len(labels)}) differ in length"
            )
        label_tensor = torch.tensor(labels, dtype=torch.long)
        if label_tensor.unique().numel() < 2:
            raise ValueError("labels must contain at least 2 distinct classes")
        return label_tensor

    @staticmethod
    def _split(n: int) -> tuple[slice, slice]:
        n_val = max(1, int(n * _VAL_FRACTION))
        return slice(0, n - n_val), slice(n - n_val, n)

    @torch.no_grad()
    def _capture(self, hook_point: str, dataset: list[Tensor]) -> Tensor:
        """Last-position activation at ``hook_point`` for each sequence."""
        module_path, point = hook_point.rsplit(".", 1)
        feats = []
        for ids in dataset:
            batch = ids.unsqueeze(0) if ids.dim() == 1 else ids
            with HookManager(self.model) as hooks:
                hooks.attach(module_path, point)
                self.model(batch)
                feats.append(hooks.get(hook_point)[0, -1])
        return torch.stack(feats)  # (N, D)

    @torch.no_grad()
    def _capture_streams(self, dataset: list[Tensor]) -> list[Tensor]:
        """Last-position residual stream at every depth, per sequence.

        Returns one ``(N, D)`` matrix per depth (0 = embedding, i = after
        block i-1), reconstructed as ``embed + Σ attn_out + ffn_out``.
        """
        n_layers = self.model.config.n_layers
        per_depth: list[list[Tensor]] = [[] for _ in range(n_layers + 1)]
        for ids in dataset:
            batch = ids.unsqueeze(0) if ids.dim() == 1 else ids
            with HookManager(self.model) as hooks:
                hooks.attach("embed", "output")
                for i in range(n_layers):
                    hooks.attach(f"blocks.{i}.attn", "output")
                    hooks.attach(f"blocks.{i}.ffn", "output")
                self.model(batch)
                stream = hooks.get("embed.output")
                per_depth[0].append(stream[0, -1])
                for i in range(n_layers):
                    stream = (
                        stream
                        + hooks.get(f"blocks.{i}.attn.output")
                        + hooks.get(f"blocks.{i}.ffn.output")
                    )
                    per_depth[i + 1].append(stream[0, -1])
        return [torch.stack(feats) for feats in per_depth]

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    def train(
        self,
        hook_point: str,
        dataset: list[Tensor],
        labels: list[int],
        epochs: int = 200,
        lr: float = 0.05,
        seed: int = 0,
    ) -> ProbeResult:
        """Train one probe on activations captured at ``hook_point``.

        Args:
            hook_point: A registry hook point (e.g. ``"blocks.1.ffn.output"``).
            dataset:    One 1-D token-ID tensor per example.
            labels:     One integer class label per example.
            epochs:     Full-batch gradient steps for the probe.
            lr:         Probe learning rate.
            seed:       RNG seed for probe init.

        Returns:
            A ``ProbeResult``.

        Raises:
            ValueError: If the dataset/labels are empty, mismatched, or contain
                fewer than 2 classes.
        """
        label_tensor = self._validate(dataset, labels)
        self.model.eval()
        features = self._capture(hook_point, dataset)

        train_slice, val_slice = self._split(len(dataset))
        n_classes = int(label_tensor.max().item()) + 1
        train_acc, val_acc, weights = _fit_logistic(
            features[train_slice],
            label_tensor[train_slice],
            features[val_slice],
            label_tensor[val_slice],
            n_classes,
            epochs,
            lr,
            seed,
        )
        return ProbeResult(
            hook_point=hook_point, train_acc=train_acc, val_acc=val_acc, weights=weights
        )

    def probe_all_layers(
        self,
        dataset: list[Tensor],
        labels: list[int],
        epochs: int = 200,
        lr: float = 0.05,
        seed: int = 0,
    ) -> LayerProbeResult:
        """Train one probe on the residual stream at every depth.

        Args:
            dataset: One 1-D token-ID tensor per example.
            labels:  One integer class label per example.
            epochs:  Full-batch gradient steps per probe.
            lr:      Probe learning rate.
            seed:    RNG seed for probe init.

        Returns:
            A ``LayerProbeResult`` with accuracies at each depth.

        Raises:
            ValueError: If the dataset/labels are invalid.
        """
        label_tensor = self._validate(dataset, labels)
        self.model.eval()
        streams = self._capture_streams(dataset)

        train_slice, val_slice = self._split(len(dataset))
        n_classes = int(label_tensor.max().item()) + 1
        train_accs: list[float] = []
        val_accs: list[float] = []
        for features in streams:
            train_acc, val_acc, _ = _fit_logistic(
                features[train_slice],
                label_tensor[train_slice],
                features[val_slice],
                label_tensor[val_slice],
                n_classes,
                epochs,
                lr,
                seed,
            )
            train_accs.append(train_acc)
            val_accs.append(val_acc)
        return LayerProbeResult(val_accs=val_accs, train_accs=train_accs)

__init__(model)

Create a prober over model.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
Source code in kamui/mechinterp/probing.py 
136
137
138
139
140
141
142
def __init__(self, model: KAMUITransformer) -> None:
    """Create a prober over ``model``.

    Args:
        model: A ``KAMUITransformer``.
    """
    self.model = model

train(hook_point, dataset, labels, epochs=200, lr=0.05, seed=0)

Train one probe on activations captured at hook_point.

Parameters:

Name Type Description Default
hook_point str

A registry hook point (e.g. "blocks.1.ffn.output").

 required
dataset list[Tensor]

One 1-D token-ID tensor per example.

 required
labels list[int]

One integer class label per example.

 required
epochs int

Full-batch gradient steps for the probe.

 200
lr float

Probe learning rate.

 0.05
seed int

RNG seed for probe init.

 0

Returns:

Type Description
ProbeResult

A ProbeResult.

Raises:

Type Description
ValueError

If the dataset/labels are empty, mismatched, or contain fewer than 2 classes.
Source code in kamui/mechinterp/probing.py 
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
def train(
    self,
    hook_point: str,
    dataset: list[Tensor],
    labels: list[int],
    epochs: int = 200,
    lr: float = 0.05,
    seed: int = 0,
) -> ProbeResult:
    """Train one probe on activations captured at ``hook_point``.

    Args:
        hook_point: A registry hook point (e.g. ``"blocks.1.ffn.output"``).
        dataset:    One 1-D token-ID tensor per example.
        labels:     One integer class label per example.
        epochs:     Full-batch gradient steps for the probe.
        lr:         Probe learning rate.
        seed:       RNG seed for probe init.

    Returns:
        A ``ProbeResult``.

    Raises:
        ValueError: If the dataset/labels are empty, mismatched, or contain
            fewer than 2 classes.
    """
    label_tensor = self._validate(dataset, labels)
    self.model.eval()
    features = self._capture(hook_point, dataset)

    train_slice, val_slice = self._split(len(dataset))
    n_classes = int(label_tensor.max().item()) + 1
    train_acc, val_acc, weights = _fit_logistic(
        features[train_slice],
        label_tensor[train_slice],
        features[val_slice],
        label_tensor[val_slice],
        n_classes,
        epochs,
        lr,
        seed,
    )
    return ProbeResult(
        hook_point=hook_point, train_acc=train_acc, val_acc=val_acc, weights=weights
    )

probe_all_layers(dataset, labels, epochs=200, lr=0.05, seed=0)

Train one probe on the residual stream at every depth.

Parameters:

Name Type Description Default
dataset list[Tensor]

One 1-D token-ID tensor per example.

 required
labels list[int]

One integer class label per example.

 required
epochs int

Full-batch gradient steps per probe.

 200
lr float

Probe learning rate.

 0.05
seed int

RNG seed for probe init.

 0

Returns:

Type Description
LayerProbeResult

A LayerProbeResult with accuracies at each depth.

Raises:

Type Description
ValueError

If the dataset/labels are invalid.
Source code in kamui/mechinterp/probing.py 
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
def probe_all_layers(
    self,
    dataset: list[Tensor],
    labels: list[int],
    epochs: int = 200,
    lr: float = 0.05,
    seed: int = 0,
) -> LayerProbeResult:
    """Train one probe on the residual stream at every depth.

    Args:
        dataset: One 1-D token-ID tensor per example.
        labels:  One integer class label per example.
        epochs:  Full-batch gradient steps per probe.
        lr:      Probe learning rate.
        seed:    RNG seed for probe init.

    Returns:
        A ``LayerProbeResult`` with accuracies at each depth.

    Raises:
        ValueError: If the dataset/labels are invalid.
    """
    label_tensor = self._validate(dataset, labels)
    self.model.eval()
    streams = self._capture_streams(dataset)

    train_slice, val_slice = self._split(len(dataset))
    n_classes = int(label_tensor.max().item()) + 1
    train_accs: list[float] = []
    val_accs: list[float] = []
    for features in streams:
        train_acc, val_acc, _ = _fit_logistic(
            features[train_slice],
            label_tensor[train_slice],
            features[val_slice],
            label_tensor[val_slice],
            n_classes,
            epochs,
            lr,
            seed,
        )
        train_accs.append(train_acc)
        val_accs.append(val_acc)
    return LayerProbeResult(val_accs=val_accs, train_accs=train_accs)

kamui.mechinterp.activation_patch.ActivationPatcher

Causal localisation via activation patching.

Attributes:

Name Type Description
model

A trained KAMUITransformer.
Source code in kamui/mechinterp/activation_patch.py 
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
388
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
class ActivationPatcher:
    """Causal localisation via activation patching.

    Attributes:
        model: A trained ``KAMUITransformer``.
    """

    def __init__(self, model: KAMUITransformer) -> None:
        """Create a patcher over ``model``.

        Args:
            model: A ``KAMUITransformer``.
        """
        self.model = model

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    @staticmethod
    def _as_batch(ids: Tensor) -> Tensor:
        if ids.dim() == 1:
            ids = ids.unsqueeze(0)
        if ids.dim() != 2 or ids.shape[0] != 1:
            raise ValueError(
                f"ids must be a single sequence (S,) or (1, S), got shape {tuple(ids.shape)}"
            )
        return ids

    def _resolve(self, module_path: str) -> nn.Module:
        return dict(self.model.named_modules())[module_path]

    @staticmethod
    def _metric(logits: Tensor, answer: int, corrupt: int) -> float:
        return (logits[0, -1, answer] - logits[0, -1, corrupt]).item()

    def _run_patched(self, ids: Tensor, module: nn.Module, replacement: Tensor) -> Tensor:
        """Run the model with ``module``'s output replaced by ``replacement``."""

        def hook(_m: nn.Module, _inp: tuple, _out: Tensor) -> Tensor:
            return replacement

        handle = module.register_forward_hook(hook)
        try:
            return self.model(ids)
        finally:
            handle.remove()

    def _recovery(
        self, patched_logits: Tensor, answer: int, corrupt: int, l_clean: float, l_corrupt: float
    ) -> float:
        denom = l_clean - l_corrupt
        if abs(denom) < _EPS:
            return 0.0
        return (self._metric(patched_logits, answer, corrupt) - l_corrupt) / denom

    @staticmethod
    def _answers(
        clean_logits: Tensor,
        corrupt_logits: Tensor,
        answer_token: int | None,
        corrupt_token: int | None,
    ) -> tuple[int, int]:
        answer = answer_token if answer_token is not None else int(clean_logits[0, -1].argmax())
        corrupt = (
            corrupt_token if corrupt_token is not None else int(corrupt_logits[0, -1].argmax())
        )
        return answer, corrupt

    # ------------------------------------------------------------------
    # Public API
    # ------------------------------------------------------------------

    @torch.no_grad()
    def patch_single(
        self,
        clean_ids: Tensor,
        corrupted_ids: Tensor,
        hook_point: str,
        answer_token: int | None = None,
        corrupt_token: int | None = None,
    ) -> float:
        """Patch a single output hook point and return the recovery.

        Args:
            clean_ids:     Clean token IDs ``(S,)`` or ``(1, S)``.
            corrupted_ids: Corrupted token IDs of the same shape.
            hook_point:    A patchable output point: ``"embed.output"``,
                ``"blocks.{i}.attn.output"``, or ``"blocks.{i}.ffn.output"``.
            answer_token:  Optional metric answer token (defaults to the clean
                run's top prediction).
            corrupt_token: Optional metric corrupt token (defaults to the
                corrupted run's top prediction).

        Returns:
            The logit-difference recovery in roughly ``[0, 1]``.

        Raises:
            ValueError: If shapes differ or ``hook_point`` is not patchable.
        """
        clean_ids = self._as_batch(clean_ids)
        corrupted_ids = self._as_batch(corrupted_ids)
        if clean_ids.shape != corrupted_ids.shape:
            raise ValueError("clean_ids and corrupted_ids must have the same shape")
        if not self._is_patchable(hook_point):
            raise ValueError(
                f"'{hook_point}' is not a patchable output point "
                f"(use embed.output / blocks.i.attn.output / blocks.i.ffn.output)"
            )

        self.model.eval()
        module_path = hook_point.rsplit(".", 1)[0]

        with HookManager(self.model) as hooks:
            hooks.attach(module_path, "output")
            clean_logits = self.model(clean_ids)
            clean_act = hooks.get(hook_point)
        corrupt_logits = self.model(corrupted_ids)

        answer, corrupt = self._answers(clean_logits, corrupt_logits, answer_token, corrupt_token)
        l_clean = self._metric(clean_logits, answer, corrupt)
        l_corrupt = self._metric(corrupt_logits, answer, corrupt)

        patched = self._run_patched(corrupted_ids, self._resolve(module_path), clean_act)
        return self._recovery(patched, answer, corrupt, l_clean, l_corrupt)

    @torch.no_grad()
    def patch_all_layers(
        self,
        clean_ids: Tensor,
        corrupted_ids: Tensor,
        component: str = "attn",
        answer_token: int | None = None,
        corrupt_token: int | None = None,
    ) -> PatchingResult:
        """Patch each layer's ``attn`` or ``ffn`` output, one at a time.

        Args:
            clean_ids:     Clean token IDs.
            corrupted_ids: Corrupted token IDs of the same shape.
            component:     ``"attn"`` or ``"ffn"``.
            answer_token:  Optional metric answer token.
            corrupt_token: Optional metric corrupt token.

        Returns:
            A ``PatchingResult`` with one recovery per layer.

        Raises:
            ValueError: If ``component`` is not ``"attn"``/``"ffn"`` or shapes differ.
        """
        if component not in ("attn", "ffn"):
            raise ValueError(f"component must be 'attn' or 'ffn', got '{component}'")
        clean_ids = self._as_batch(clean_ids)
        corrupted_ids = self._as_batch(corrupted_ids)
        if clean_ids.shape != corrupted_ids.shape:
            raise ValueError("clean_ids and corrupted_ids must have the same shape")

        self.model.eval()
        n_layers = self.model.config.n_layers
        module_paths = [f"blocks.{i}.{component}" for i in range(n_layers)]

        with HookManager(self.model) as hooks:
            for path in module_paths:
                hooks.attach(path, "output")
            clean_logits = self.model(clean_ids)
            clean_acts = {path: hooks.get(f"{path}.output") for path in module_paths}
        corrupt_logits = self.model(corrupted_ids)

        answer, corrupt = self._answers(clean_logits, corrupt_logits, answer_token, corrupt_token)
        l_clean = self._metric(clean_logits, answer, corrupt)
        l_corrupt = self._metric(corrupt_logits, answer, corrupt)

        effects = []
        for path in module_paths:
            patched = self._run_patched(corrupted_ids, self._resolve(path), clean_acts[path])
            effects.append(self._recovery(patched, answer, corrupt, l_clean, l_corrupt))
        return PatchingResult(torch.tensor(effects), component)

    @torch.no_grad()
    def patch_all_heads(
        self,
        clean_ids: Tensor,
        corrupted_ids: Tensor,
        answer_token: int | None = None,
        corrupt_token: int | None = None,
    ) -> HeadPatchingResult:
        """Patch each attention head's contribution, one at a time.

        A head is patched by replacing its slice of the ``out_proj`` input (the
        concatenated per-head outputs) with the clean run's slice.

        Args:
            clean_ids:     Clean token IDs.
            corrupted_ids: Corrupted token IDs of the same shape.
            answer_token:  Optional metric answer token.
            corrupt_token: Optional metric corrupt token.

        Returns:
            A ``HeadPatchingResult`` with an ``(n_layers, n_heads)`` effect matrix.

        Raises:
            ValueError: If shapes differ.
        """
        clean_ids = self._as_batch(clean_ids)
        corrupted_ids = self._as_batch(corrupted_ids)
        if clean_ids.shape != corrupted_ids.shape:
            raise ValueError("clean_ids and corrupted_ids must have the same shape")

        self.model.eval()
        n_layers = self.model.config.n_layers
        n_heads = self.model.config.n_heads
        d_head = self.model.config.d_head

        # Capture the clean out_proj inputs (concatenated head outputs) per layer.
        clean_head_inputs: dict[int, Tensor] = {}
        handles = []
        for i in range(n_layers):
            out_proj = self._resolve(f"blocks.{i}.attn.out_proj")

            def _capture(layer: int) -> Callable[[nn.Module, tuple[Tensor, ...]], None]:
                def pre_hook(_m: nn.Module, inp: tuple[Tensor, ...]) -> None:
                    clean_head_inputs[layer] = inp[0].detach()

                return pre_hook

            handles.append(out_proj.register_forward_pre_hook(_capture(i)))
        clean_logits = self.model(clean_ids)
        for handle in handles:
            handle.remove()

        corrupt_logits = self.model(corrupted_ids)
        answer, corrupt = self._answers(clean_logits, corrupt_logits, answer_token, corrupt_token)
        l_clean = self._metric(clean_logits, answer, corrupt)
        l_corrupt = self._metric(corrupt_logits, answer, corrupt)

        effects = torch.zeros(n_layers, n_heads)
        for i in range(n_layers):
            out_proj = self._resolve(f"blocks.{i}.attn.out_proj")
            clean_in = clean_head_inputs[i]
            for head in range(n_heads):
                patched = self._run_head_patched(
                    corrupted_ids, out_proj, clean_in, head, n_heads, d_head
                )
                effects[i, head] = self._recovery(patched, answer, corrupt, l_clean, l_corrupt)
        return HeadPatchingResult(effects)

    def _run_head_patched(
        self,
        ids: Tensor,
        out_proj: nn.Module,
        clean_in: Tensor,
        head: int,
        n_heads: int,
        d_head: int,
    ) -> Tensor:
        """Run the model, replacing one head's slice of ``out_proj``'s input."""

        def pre_hook(_m: nn.Module, inp: tuple) -> tuple:
            current = inp[0]
            b, s, _ = current.shape
            patched = current.clone().view(b, s, n_heads, d_head)
            clean_reshaped = clean_in.view(b, s, n_heads, d_head)
            patched[:, :, head, :] = clean_reshaped[:, :, head, :]
            return (patched.view(b, s, n_heads * d_head),)

        handle = out_proj.register_forward_pre_hook(pre_hook)
        try:
            return self.model(ids)
        finally:
            handle.remove()

    @staticmethod
    def _is_patchable(hook_point: str) -> bool:
        return (
            hook_point == "embed.output"
            or hook_point.endswith(".attn.output")
            or hook_point.endswith(".ffn.output")
        )

__init__(model)

Create a patcher over model.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
Source code in kamui/mechinterp/activation_patch.py 
134
135
136
137
138
139
140
def __init__(self, model: KAMUITransformer) -> None:
    """Create a patcher over ``model``.

    Args:
        model: A ``KAMUITransformer``.
    """
    self.model = model

patch_single(clean_ids, corrupted_ids, hook_point, answer_token=None, corrupt_token=None)

Patch a single output hook point and return the recovery.

Parameters:

Name Type Description Default
clean_ids Tensor

Clean token IDs (S,) or (1, S).

 required
corrupted_ids Tensor

Corrupted token IDs of the same shape.

 required
hook_point str

A patchable output point: "embed.output", "blocks.{i}.attn.output", or "blocks.{i}.ffn.output".

 required
answer_token int | None

Optional metric answer token (defaults to the clean run's top prediction).

 None
corrupt_token int | None

Optional metric corrupt token (defaults to the corrupted run's top prediction).

 None

Returns:

Type Description
float

The logit-difference recovery in roughly [0, 1].

Raises:

Type Description
ValueError

If shapes differ or hook_point is not patchable.
Source code in kamui/mechinterp/activation_patch.py 
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
@torch.no_grad()
def patch_single(
    self,
    clean_ids: Tensor,
    corrupted_ids: Tensor,
    hook_point: str,
    answer_token: int | None = None,
    corrupt_token: int | None = None,
) -> float:
    """Patch a single output hook point and return the recovery.

    Args:
        clean_ids:     Clean token IDs ``(S,)`` or ``(1, S)``.
        corrupted_ids: Corrupted token IDs of the same shape.
        hook_point:    A patchable output point: ``"embed.output"``,
            ``"blocks.{i}.attn.output"``, or ``"blocks.{i}.ffn.output"``.
        answer_token:  Optional metric answer token (defaults to the clean
            run's top prediction).
        corrupt_token: Optional metric corrupt token (defaults to the
            corrupted run's top prediction).

    Returns:
        The logit-difference recovery in roughly ``[0, 1]``.

    Raises:
        ValueError: If shapes differ or ``hook_point`` is not patchable.
    """
    clean_ids = self._as_batch(clean_ids)
    corrupted_ids = self._as_batch(corrupted_ids)
    if clean_ids.shape != corrupted_ids.shape:
        raise ValueError("clean_ids and corrupted_ids must have the same shape")
    if not self._is_patchable(hook_point):
        raise ValueError(
            f"'{hook_point}' is not a patchable output point "
            f"(use embed.output / blocks.i.attn.output / blocks.i.ffn.output)"
        )

    self.model.eval()
    module_path = hook_point.rsplit(".", 1)[0]

    with HookManager(self.model) as hooks:
        hooks.attach(module_path, "output")
        clean_logits = self.model(clean_ids)
        clean_act = hooks.get(hook_point)
    corrupt_logits = self.model(corrupted_ids)

    answer, corrupt = self._answers(clean_logits, corrupt_logits, answer_token, corrupt_token)
    l_clean = self._metric(clean_logits, answer, corrupt)
    l_corrupt = self._metric(corrupt_logits, answer, corrupt)

    patched = self._run_patched(corrupted_ids, self._resolve(module_path), clean_act)
    return self._recovery(patched, answer, corrupt, l_clean, l_corrupt)

patch_all_layers(clean_ids, corrupted_ids, component='attn', answer_token=None, corrupt_token=None)

Patch each layer's attn or ffn output, one at a time.

Parameters:

Name Type Description Default
clean_ids Tensor

Clean token IDs.

 required
corrupted_ids Tensor

Corrupted token IDs of the same shape.

 required
component str

"attn" or "ffn".

 'attn'
answer_token int | None

Optional metric answer token.

 None
corrupt_token int | None

Optional metric corrupt token.

 None

Returns:

Type Description
PatchingResult

A PatchingResult with one recovery per layer.

Raises:

Type Description
ValueError

If component is not "attn"/"ffn" or shapes differ.
Source code in kamui/mechinterp/activation_patch.py 
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
@torch.no_grad()
def patch_all_layers(
    self,
    clean_ids: Tensor,
    corrupted_ids: Tensor,
    component: str = "attn",
    answer_token: int | None = None,
    corrupt_token: int | None = None,
) -> PatchingResult:
    """Patch each layer's ``attn`` or ``ffn`` output, one at a time.

    Args:
        clean_ids:     Clean token IDs.
        corrupted_ids: Corrupted token IDs of the same shape.
        component:     ``"attn"`` or ``"ffn"``.
        answer_token:  Optional metric answer token.
        corrupt_token: Optional metric corrupt token.

    Returns:
        A ``PatchingResult`` with one recovery per layer.

    Raises:
        ValueError: If ``component`` is not ``"attn"``/``"ffn"`` or shapes differ.
    """
    if component not in ("attn", "ffn"):
        raise ValueError(f"component must be 'attn' or 'ffn', got '{component}'")
    clean_ids = self._as_batch(clean_ids)
    corrupted_ids = self._as_batch(corrupted_ids)
    if clean_ids.shape != corrupted_ids.shape:
        raise ValueError("clean_ids and corrupted_ids must have the same shape")

    self.model.eval()
    n_layers = self.model.config.n_layers
    module_paths = [f"blocks.{i}.{component}" for i in range(n_layers)]

    with HookManager(self.model) as hooks:
        for path in module_paths:
            hooks.attach(path, "output")
        clean_logits = self.model(clean_ids)
        clean_acts = {path: hooks.get(f"{path}.output") for path in module_paths}
    corrupt_logits = self.model(corrupted_ids)

    answer, corrupt = self._answers(clean_logits, corrupt_logits, answer_token, corrupt_token)
    l_clean = self._metric(clean_logits, answer, corrupt)
    l_corrupt = self._metric(corrupt_logits, answer, corrupt)

    effects = []
    for path in module_paths:
        patched = self._run_patched(corrupted_ids, self._resolve(path), clean_acts[path])
        effects.append(self._recovery(patched, answer, corrupt, l_clean, l_corrupt))
    return PatchingResult(torch.tensor(effects), component)

patch_all_heads(clean_ids, corrupted_ids, answer_token=None, corrupt_token=None)

Patch each attention head's contribution, one at a time.

A head is patched by replacing its slice of the out_proj input (the concatenated per-head outputs) with the clean run's slice.

Parameters:

Name Type Description Default
clean_ids Tensor

Clean token IDs.

 required
corrupted_ids Tensor

Corrupted token IDs of the same shape.

 required
answer_token int | None

Optional metric answer token.

 None
corrupt_token int | None

Optional metric corrupt token.

 None

Returns:

Type Description
HeadPatchingResult

A HeadPatchingResult with an (n_layers, n_heads) effect matrix.

Raises:

Type Description
ValueError

If shapes differ.
Source code in kamui/mechinterp/activation_patch.py 
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
346
347
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
@torch.no_grad()
def patch_all_heads(
    self,
    clean_ids: Tensor,
    corrupted_ids: Tensor,
    answer_token: int | None = None,
    corrupt_token: int | None = None,
) -> HeadPatchingResult:
    """Patch each attention head's contribution, one at a time.

    A head is patched by replacing its slice of the ``out_proj`` input (the
    concatenated per-head outputs) with the clean run's slice.

    Args:
        clean_ids:     Clean token IDs.
        corrupted_ids: Corrupted token IDs of the same shape.
        answer_token:  Optional metric answer token.
        corrupt_token: Optional metric corrupt token.

    Returns:
        A ``HeadPatchingResult`` with an ``(n_layers, n_heads)`` effect matrix.

    Raises:
        ValueError: If shapes differ.
    """
    clean_ids = self._as_batch(clean_ids)
    corrupted_ids = self._as_batch(corrupted_ids)
    if clean_ids.shape != corrupted_ids.shape:
        raise ValueError("clean_ids and corrupted_ids must have the same shape")

    self.model.eval()
    n_layers = self.model.config.n_layers
    n_heads = self.model.config.n_heads
    d_head = self.model.config.d_head

    # Capture the clean out_proj inputs (concatenated head outputs) per layer.
    clean_head_inputs: dict[int, Tensor] = {}
    handles = []
    for i in range(n_layers):
        out_proj = self._resolve(f"blocks.{i}.attn.out_proj")

        def _capture(layer: int) -> Callable[[nn.Module, tuple[Tensor, ...]], None]:
            def pre_hook(_m: nn.Module, inp: tuple[Tensor, ...]) -> None:
                clean_head_inputs[layer] = inp[0].detach()

            return pre_hook

        handles.append(out_proj.register_forward_pre_hook(_capture(i)))
    clean_logits = self.model(clean_ids)
    for handle in handles:
        handle.remove()

    corrupt_logits = self.model(corrupted_ids)
    answer, corrupt = self._answers(clean_logits, corrupt_logits, answer_token, corrupt_token)
    l_clean = self._metric(clean_logits, answer, corrupt)
    l_corrupt = self._metric(corrupt_logits, answer, corrupt)

    effects = torch.zeros(n_layers, n_heads)
    for i in range(n_layers):
        out_proj = self._resolve(f"blocks.{i}.attn.out_proj")
        clean_in = clean_head_inputs[i]
        for head in range(n_heads):
            patched = self._run_head_patched(
                corrupted_ids, out_proj, clean_in, head, n_heads, d_head
            )
            effects[i, head] = self._recovery(patched, answer, corrupt, l_clean, l_corrupt)
    return HeadPatchingResult(effects)

kamui.mechinterp.induction.InductionHeadDetector

Score and causally test every attention head for induction behaviour.

Attributes:

Name Type Description
model

A trained KAMUITransformer.
Source code in kamui/mechinterp/induction.py 
 49
 50
 51
 52
 53
 54
 55
 56
 57
 58
 59
 60
 61
 62
 63
 64
 65
 66
 67
 68
 69
 70
 71
 72
 73
 74
 75
 76
 77
 78
 79
 80
 81
 82
 83
 84
 85
 86
 87
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
class InductionHeadDetector:
    """Score and causally test every attention head for induction behaviour.

    Attributes:
        model: A trained ``KAMUITransformer``.
    """

    def __init__(self, model: KAMUITransformer) -> None:
        """Create a detector over ``model``.

        Args:
            model: A ``KAMUITransformer``.
        """
        self.model = model

    # ------------------------------------------------------------------
    # Helpers
    # ------------------------------------------------------------------

    def _repeated_sequence(self, prefix_len: int, seed: int) -> Tensor:
        """Return a ``(1, 2 * prefix_len)`` random sequence repeated twice."""
        generator = torch.Generator().manual_seed(seed)
        vocab = self.model.config.vocab_size
        prefix = torch.randint(0, vocab, (prefix_len,), generator=generator)
        return torch.cat([prefix, prefix]).unsqueeze(0)

    def _resolve_prefix_len(self, prefix_len: int | None) -> int:
        ctx = self.model.config.context_length
        if prefix_len is None:
            prefix_len = ctx // 2
        if not (2 <= prefix_len <= ctx // 2):
            raise ValueError(
                f"prefix_len must be in [2, context_length // 2 = {ctx // 2}], got {prefix_len}"
            )
        return prefix_len

    # ------------------------------------------------------------------
    # Scoring
    # ------------------------------------------------------------------

    @torch.no_grad()
    def score_all_heads(
        self, prefix_len: int | None = None, n_samples: int = 4, seed: int = 0
    ) -> dict[tuple[int, int], float]:
        """Compute the induction score for every (layer, head) pair.

        Args:
            prefix_len: Length of the repeated prefix (defaults to
                ``context_length // 2``).
            n_samples:  Number of random repeated sequences to average over.
            seed:       Base RNG seed (sample ``i`` uses ``seed + i``).

        Returns:
            ``{(layer, head): score}`` with scores in ``[0, 1]``.

        Raises:
            ValueError: If ``prefix_len`` is out of range or ``n_samples < 1``.
        """
        if n_samples < 1:
            raise ValueError(f"n_samples must be >= 1, got {n_samples}")
        prefix_len = self._resolve_prefix_len(prefix_len)

        self.model.eval()
        n_layers = self.model.config.n_layers
        n_heads = self.model.config.n_heads
        totals = torch.zeros(n_layers, n_heads)

        for sample in range(n_samples):
            ids = self._repeated_sequence(prefix_len, seed + sample)
            with HookManager(self.model) as hooks:
                for i in range(n_layers):
                    hooks.attach(f"blocks.{i}.attn", "weights")
                self.model(ids)
                # Mean attention on the induction offset t -> t - L + 1 for
                # every query position t in the second half.
                queries = torch.arange(prefix_len, 2 * prefix_len)
                keys = queries - prefix_len + 1
                for i in range(n_layers):
                    w = hooks.get(f"blocks.{i}.attn.weights")[0]  # (H, S, S)
                    totals[i] += w[:, queries, keys].mean(dim=-1)

        totals /= n_samples
        return {
            (layer, head): totals[layer, head].item()
            for layer in range(n_layers)
            for head in range(n_heads)
        }

    def plot_scores(
        self,
        scores: dict[tuple[int, int], float],
        figsize: tuple[float, float] | None = None,
    ) -> Figure:
        """Heatmap of induction scores by (layer, head).

        Args:
            scores:  The dict returned by ``score_all_heads``.
            figsize: Optional figure size.

        Returns:
            The matplotlib ``Figure``.
        """
        import matplotlib.pyplot as plt

        n_layers = self.model.config.n_layers
        n_heads = self.model.config.n_heads
        grid = torch.zeros(n_layers, n_heads)
        for (layer, head), score in scores.items():
            grid[layer, head] = score

        fig, ax = plt.subplots(figsize=figsize or (max(5, n_heads), max(4, n_layers)))
        im = ax.imshow(grid.numpy(), aspect="auto", cmap="Blues", vmin=0.0, vmax=1.0)
        ax.set_xlabel("head")
        ax.set_ylabel("layer")
        ax.set_title("Induction scores")
        ax.set_xticks(range(n_heads))
        ax.set_yticks(range(n_layers))
        fig.colorbar(im, ax=ax, label="induction score")
        fig.tight_layout()
        return fig

    # ------------------------------------------------------------------
    # Causal confirmation
    # ------------------------------------------------------------------

    @torch.no_grad()
    def ablate_and_measure(
        self,
        heads: list[tuple[int, int]],
        prefix_len: int | None = None,
        seed: int = 0,
    ) -> float:
        """Zero-ablate ``heads`` and measure the rise in in-context loss.

        In-context performance is the mean next-token loss over the second
        half of a repeated random sequence (predictable purely from context).
        A positive return value means ablation hurt in-context learning.

        Args:
            heads:      ``(layer, head)`` pairs to zero-ablate.
            prefix_len: Repeated-prefix length (defaults to ``context_length // 2``).
            seed:       RNG seed for the sequence.

        Returns:
            ``ablated_loss - baseline_loss`` on second-half tokens.

        Raises:
            ValueError: If ``heads`` is empty or any index is out of range.
        """
        if not heads:
            raise ValueError("heads must be a non-empty list of (layer, head) pairs")
        n_layers = self.model.config.n_layers
        n_heads = self.model.config.n_heads
        for layer, head in heads:
            if not (0 <= layer < n_layers and 0 <= head < n_heads):
                raise ValueError(f"(layer={layer}, head={head}) out of range")

        prefix_len = self._resolve_prefix_len(prefix_len)
        self.model.eval()
        ids = self._repeated_sequence(prefix_len, seed)

        baseline = self._second_half_loss(ids, prefix_len)

        # Zero the ablated heads' slices of each affected layer's out_proj input.
        d_head = self.model.config.d_head
        by_layer: dict[int, list[int]] = {}
        for layer, head in heads:
            by_layer.setdefault(layer, []).append(head)

        handles = []
        modules = dict(self.model.named_modules())
        for layer, layer_heads in by_layer.items():
            out_proj = modules[f"blocks.{layer}.attn.out_proj"]

            def _make(
                layer_heads: list[int],
            ) -> Callable[[nn.Module, tuple[Tensor, ...]], tuple[Tensor]]:
                def pre_hook(_m: nn.Module, inp: tuple[Tensor, ...]) -> tuple[Tensor]:
                    x = inp[0]
                    b, s, _ = x.shape
                    x = x.clone().view(b, s, n_heads, d_head)
                    for h in layer_heads:
                        x[:, :, h, :] = 0.0
                    return (x.view(b, s, n_heads * d_head),)

                return pre_hook

            handles.append(out_proj.register_forward_pre_hook(_make(layer_heads)))

        try:
            ablated = self._second_half_loss(ids, prefix_len)
        finally:
            for handle in handles:
                handle.remove()

        return ablated - baseline

    def _second_half_loss(self, ids: Tensor, prefix_len: int) -> float:
        """Mean next-token loss over the second half of ``ids``."""
        logits = self.model(ids)  # (1, 2L, V)
        # Predict tokens at positions prefix_len..2L-1 from their predecessors.
        preds = logits[0, prefix_len - 1 : -1]  # (L, V)
        targets = ids[0, prefix_len:]  # (L,)
        return F.cross_entropy(preds, targets).item()

__init__(model)

Create a detector over model.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
Source code in kamui/mechinterp/induction.py 
56
57
58
59
60
61
62
def __init__(self, model: KAMUITransformer) -> None:
    """Create a detector over ``model``.

    Args:
        model: A ``KAMUITransformer``.
    """
    self.model = model

score_all_heads(prefix_len=None, n_samples=4, seed=0)

Compute the induction score for every (layer, head) pair.

Parameters:

Name Type Description Default
prefix_len int | None

Length of the repeated prefix (defaults to context_length // 2).

 None
n_samples int

Number of random repeated sequences to average over.

 4
seed int

Base RNG seed (sample i uses seed + i).

 0

Returns:

Type Description
dict[tuple[int, int], float]

{(layer, head): score} with scores in [0, 1].

Raises:

Type Description
ValueError

If prefix_len is out of range or n_samples < 1.
Source code in kamui/mechinterp/induction.py 
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
@torch.no_grad()
def score_all_heads(
    self, prefix_len: int | None = None, n_samples: int = 4, seed: int = 0
) -> dict[tuple[int, int], float]:
    """Compute the induction score for every (layer, head) pair.

    Args:
        prefix_len: Length of the repeated prefix (defaults to
            ``context_length // 2``).
        n_samples:  Number of random repeated sequences to average over.
        seed:       Base RNG seed (sample ``i`` uses ``seed + i``).

    Returns:
        ``{(layer, head): score}`` with scores in ``[0, 1]``.

    Raises:
        ValueError: If ``prefix_len`` is out of range or ``n_samples < 1``.
    """
    if n_samples < 1:
        raise ValueError(f"n_samples must be >= 1, got {n_samples}")
    prefix_len = self._resolve_prefix_len(prefix_len)

    self.model.eval()
    n_layers = self.model.config.n_layers
    n_heads = self.model.config.n_heads
    totals = torch.zeros(n_layers, n_heads)

    for sample in range(n_samples):
        ids = self._repeated_sequence(prefix_len, seed + sample)
        with HookManager(self.model) as hooks:
            for i in range(n_layers):
                hooks.attach(f"blocks.{i}.attn", "weights")
            self.model(ids)
            # Mean attention on the induction offset t -> t - L + 1 for
            # every query position t in the second half.
            queries = torch.arange(prefix_len, 2 * prefix_len)
            keys = queries - prefix_len + 1
            for i in range(n_layers):
                w = hooks.get(f"blocks.{i}.attn.weights")[0]  # (H, S, S)
                totals[i] += w[:, queries, keys].mean(dim=-1)

    totals /= n_samples
    return {
        (layer, head): totals[layer, head].item()
        for layer in range(n_layers)
        for head in range(n_heads)
    }

plot_scores(scores, figsize=None)

Heatmap of induction scores by (layer, head).

Parameters:

Name Type Description Default
scores dict[tuple[int, int], float]

The dict returned by score_all_heads.

 required
figsize tuple[float, float] | None

Optional figure size.

 None

Returns:

Type Description
Figure

The matplotlib Figure.
Source code in kamui/mechinterp/induction.py 
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
def plot_scores(
    self,
    scores: dict[tuple[int, int], float],
    figsize: tuple[float, float] | None = None,
) -> Figure:
    """Heatmap of induction scores by (layer, head).

    Args:
        scores:  The dict returned by ``score_all_heads``.
        figsize: Optional figure size.

    Returns:
        The matplotlib ``Figure``.
    """
    import matplotlib.pyplot as plt

    n_layers = self.model.config.n_layers
    n_heads = self.model.config.n_heads
    grid = torch.zeros(n_layers, n_heads)
    for (layer, head), score in scores.items():
        grid[layer, head] = score

    fig, ax = plt.subplots(figsize=figsize or (max(5, n_heads), max(4, n_layers)))
    im = ax.imshow(grid.numpy(), aspect="auto", cmap="Blues", vmin=0.0, vmax=1.0)
    ax.set_xlabel("head")
    ax.set_ylabel("layer")
    ax.set_title("Induction scores")
    ax.set_xticks(range(n_heads))
    ax.set_yticks(range(n_layers))
    fig.colorbar(im, ax=ax, label="induction score")
    fig.tight_layout()
    return fig

ablate_and_measure(heads, prefix_len=None, seed=0)

Zero-ablate heads and measure the rise in in-context loss.

In-context performance is the mean next-token loss over the second half of a repeated random sequence (predictable purely from context). A positive return value means ablation hurt in-context learning.

Parameters:

Name Type Description Default
heads list[tuple[int, int]]

(layer, head) pairs to zero-ablate.

 required
prefix_len int | None

Repeated-prefix length (defaults to context_length // 2).

 None
seed int

RNG seed for the sequence.

 0

Returns:

Type Description
float

ablated_loss - baseline_loss on second-half tokens.

Raises:

Type Description
ValueError

If heads is empty or any index is out of range.
Source code in kamui/mechinterp/induction.py 
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
@torch.no_grad()
def ablate_and_measure(
    self,
    heads: list[tuple[int, int]],
    prefix_len: int | None = None,
    seed: int = 0,
) -> float:
    """Zero-ablate ``heads`` and measure the rise in in-context loss.

    In-context performance is the mean next-token loss over the second
    half of a repeated random sequence (predictable purely from context).
    A positive return value means ablation hurt in-context learning.

    Args:
        heads:      ``(layer, head)`` pairs to zero-ablate.
        prefix_len: Repeated-prefix length (defaults to ``context_length // 2``).
        seed:       RNG seed for the sequence.

    Returns:
        ``ablated_loss - baseline_loss`` on second-half tokens.

    Raises:
        ValueError: If ``heads`` is empty or any index is out of range.
    """
    if not heads:
        raise ValueError("heads must be a non-empty list of (layer, head) pairs")
    n_layers = self.model.config.n_layers
    n_heads = self.model.config.n_heads
    for layer, head in heads:
        if not (0 <= layer < n_layers and 0 <= head < n_heads):
            raise ValueError(f"(layer={layer}, head={head}) out of range")

    prefix_len = self._resolve_prefix_len(prefix_len)
    self.model.eval()
    ids = self._repeated_sequence(prefix_len, seed)

    baseline = self._second_half_loss(ids, prefix_len)

    # Zero the ablated heads' slices of each affected layer's out_proj input.
    d_head = self.model.config.d_head
    by_layer: dict[int, list[int]] = {}
    for layer, head in heads:
        by_layer.setdefault(layer, []).append(head)

    handles = []
    modules = dict(self.model.named_modules())
    for layer, layer_heads in by_layer.items():
        out_proj = modules[f"blocks.{layer}.attn.out_proj"]

        def _make(
            layer_heads: list[int],
        ) -> Callable[[nn.Module, tuple[Tensor, ...]], tuple[Tensor]]:
            def pre_hook(_m: nn.Module, inp: tuple[Tensor, ...]) -> tuple[Tensor]:
                x = inp[0]
                b, s, _ = x.shape
                x = x.clone().view(b, s, n_heads, d_head)
                for h in layer_heads:
                    x[:, :, h, :] = 0.0
                return (x.view(b, s, n_heads * d_head),)

            return pre_hook

        handles.append(out_proj.register_forward_pre_hook(_make(layer_heads)))

    try:
        ablated = self._second_half_loss(ids, prefix_len)
    finally:
        for handle in handles:
            handle.remove()

    return ablated - baseline

kamui.mechinterp.circuits.CircuitAblator

Zero- and mean-ablation of model components.

Attributes:

Name Type Description
model

A trained KAMUITransformer.
Source code in kamui/mechinterp/circuits.py 
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
class CircuitAblator:
    """Zero- and mean-ablation of model components.

    Attributes:
        model: A trained ``KAMUITransformer``.
    """

    def __init__(self, model: KAMUITransformer) -> None:
        """Create an ablator over ``model``.

        Args:
            model: A ``KAMUITransformer``.
        """
        self.model = model

    def _run_ablated(
        self, ids: Tensor, module_paths: list[str], replacements: dict[str, Tensor] | None
    ) -> Tensor:
        """Run the model with each listed module's output replaced.

        ``replacements`` maps module path → replacement tensor (broadcastable
        to the output).  When None, outputs are zeroed.
        """
        modules = dict(self.model.named_modules())
        handles = []
        for path in module_paths:

            def _make(path: str) -> Callable[[nn.Module, tuple, Tensor], Tensor]:
                def hook(_m: nn.Module, _inp: tuple, out: Tensor) -> Tensor:
                    if replacements is None:
                        return torch.zeros_like(out)
                    return replacements[path].expand_as(out)

                return hook

            handles.append(modules[path].register_forward_hook(_make(path)))
        try:
            return self.model(ids)
        finally:
            for handle in handles:
                handle.remove()

    @torch.no_grad()
    def ablate(self, components: list[str], token_ids: Tensor, metric: Metric) -> AblationResult:
        """Zero-ablate ``components`` and measure the metric change.

        Args:
            components: Output hook points to ablate
                (``"blocks.{i}.attn.output"`` / ``"blocks.{i}.ffn.output"``).
            token_ids:  Input token IDs ``(S,)`` or ``(B, S)``.
            metric:     Maps logits to a scalar (higher = better).

        Returns:
            An ``AblationResult``.

        Raises:
            ValueError: If ``components`` is empty or contains an invalid point.
        """
        paths = _validate_components(self.model, components)
        ids = _as_batch(token_ids)
        self.model.eval()

        baseline = metric(self.model(ids))
        value = metric(self._run_ablated(ids, paths, replacements=None))
        return AblationResult(value=value, baseline=baseline, components=list(components))

    @torch.no_grad()
    def mean_ablate(
        self,
        components: list[str],
        token_ids: Tensor,
        baseline_ids: Tensor,
        metric: Metric,
    ) -> AblationResult:
        """Replace components' activations with their baseline-batch means.

        Args:
            components:   Output hook points to ablate.
            token_ids:    Task input ``(S,)`` or ``(B, S)``.
            baseline_ids: Baseline batch ``(S,)`` or ``(B, S)`` used to compute
                each component's mean activation (averaged over batch and
                positions, broadcast back at ablation time).
            metric:       Maps logits to a scalar (higher = better).

        Returns:
            An ``AblationResult``.

        Raises:
            ValueError: If ``components`` is empty or contains an invalid point.
        """
        paths = _validate_components(self.model, components)
        ids = _as_batch(token_ids)
        base = _as_batch(baseline_ids)
        self.model.eval()

        # Cache mean activations over the baseline batch.
        means: dict[str, Tensor] = {}
        modules = dict(self.model.named_modules())
        handles = []
        for path in paths:

            def _make(path: str) -> Callable[[nn.Module, tuple, Tensor], None]:
                def hook(_m: nn.Module, _inp: tuple, out: Tensor) -> None:
                    means[path] = out.detach().mean(dim=(0, 1), keepdim=True)  # (1, 1, D)

                return hook

            handles.append(modules[path].register_forward_hook(_make(path)))
        try:
            self.model(base)
        finally:
            for handle in handles:
                handle.remove()

        baseline_metric = metric(self.model(ids))
        value = metric(self._run_ablated(ids, paths, replacements=means))
        return AblationResult(value=value, baseline=baseline_metric, components=list(components))

__init__(model)

Create an ablator over model.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
Source code in kamui/mechinterp/circuits.py 
109
110
111
112
113
114
115
def __init__(self, model: KAMUITransformer) -> None:
    """Create an ablator over ``model``.

    Args:
        model: A ``KAMUITransformer``.
    """
    self.model = model

ablate(components, token_ids, metric)

Zero-ablate components and measure the metric change.

Parameters:

Name Type Description Default
components list[str]

Output hook points to ablate ("blocks.{i}.attn.output" / "blocks.{i}.ffn.output").

 required
token_ids Tensor

Input token IDs (S,) or (B, S).

 required
metric Metric

Maps logits to a scalar (higher = better).

 required

Returns:

Type Description
AblationResult

An AblationResult.

Raises:

Type Description
ValueError

If components is empty or contains an invalid point.
Source code in kamui/mechinterp/circuits.py 
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
@torch.no_grad()
def ablate(self, components: list[str], token_ids: Tensor, metric: Metric) -> AblationResult:
    """Zero-ablate ``components`` and measure the metric change.

    Args:
        components: Output hook points to ablate
            (``"blocks.{i}.attn.output"`` / ``"blocks.{i}.ffn.output"``).
        token_ids:  Input token IDs ``(S,)`` or ``(B, S)``.
        metric:     Maps logits to a scalar (higher = better).

    Returns:
        An ``AblationResult``.

    Raises:
        ValueError: If ``components`` is empty or contains an invalid point.
    """
    paths = _validate_components(self.model, components)
    ids = _as_batch(token_ids)
    self.model.eval()

    baseline = metric(self.model(ids))
    value = metric(self._run_ablated(ids, paths, replacements=None))
    return AblationResult(value=value, baseline=baseline, components=list(components))

mean_ablate(components, token_ids, baseline_ids, metric)

Replace components' activations with their baseline-batch means.

Parameters:

Name Type Description Default
components list[str]

Output hook points to ablate.

 required
token_ids Tensor

Task input (S,) or (B, S).

 required
baseline_ids Tensor

Baseline batch (S,) or (B, S) used to compute each component's mean activation (averaged over batch and positions, broadcast back at ablation time).

 required
metric Metric

Maps logits to a scalar (higher = better).

 required

Returns:

Type Description
AblationResult

An AblationResult.

Raises:

Type Description
ValueError

If components is empty or contains an invalid point.
Source code in kamui/mechinterp/circuits.py 
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
@torch.no_grad()
def mean_ablate(
    self,
    components: list[str],
    token_ids: Tensor,
    baseline_ids: Tensor,
    metric: Metric,
) -> AblationResult:
    """Replace components' activations with their baseline-batch means.

    Args:
        components:   Output hook points to ablate.
        token_ids:    Task input ``(S,)`` or ``(B, S)``.
        baseline_ids: Baseline batch ``(S,)`` or ``(B, S)`` used to compute
            each component's mean activation (averaged over batch and
            positions, broadcast back at ablation time).
        metric:       Maps logits to a scalar (higher = better).

    Returns:
        An ``AblationResult``.

    Raises:
        ValueError: If ``components`` is empty or contains an invalid point.
    """
    paths = _validate_components(self.model, components)
    ids = _as_batch(token_ids)
    base = _as_batch(baseline_ids)
    self.model.eval()

    # Cache mean activations over the baseline batch.
    means: dict[str, Tensor] = {}
    modules = dict(self.model.named_modules())
    handles = []
    for path in paths:

        def _make(path: str) -> Callable[[nn.Module, tuple, Tensor], None]:
            def hook(_m: nn.Module, _inp: tuple, out: Tensor) -> None:
                means[path] = out.detach().mean(dim=(0, 1), keepdim=True)  # (1, 1, D)

            return hook

        handles.append(modules[path].register_forward_hook(_make(path)))
    try:
        self.model(base)
    finally:
        for handle in handles:
            handle.remove()

    baseline_metric = metric(self.model(ids))
    value = metric(self._run_ablated(ids, paths, replacements=means))
    return AblationResult(value=value, baseline=baseline_metric, components=list(components))

Interpretability Tools (v0.2)

kamui.mechinterp.attribution.GradientAttribution

Attribute a target prediction to input tokens via gradients.

Attributes:

Name Type Description
model

A trained KAMUITransformer.
tokenizer

Optional tokenizer with decode (for token labels).
Source code in kamui/mechinterp/attribution.py 
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
class GradientAttribution:
    """Attribute a target prediction to input tokens via gradients.

    Attributes:
        model:     A trained ``KAMUITransformer``.
        tokenizer: Optional tokenizer with ``decode`` (for token labels).
    """

    def __init__(self, model: KAMUITransformer, tokenizer: Any = None) -> None:
        """Create an attributor over ``model``.

        Args:
            model:     A ``KAMUITransformer``.
            tokenizer: Optional tokenizer with ``decode(list[int]) -> str``.
        """
        self.model = model
        self.tokenizer = tokenizer

    @staticmethod
    def _as_batch(ids: Tensor) -> Tensor:
        if ids.dim() == 1:
            ids = ids.unsqueeze(0)
        if ids.dim() != 2 or ids.shape[0] != 1:
            raise ValueError(
                f"token_ids must be a single sequence (S,) or (1, S), got shape {tuple(ids.shape)}"
            )
        return ids

    def _scaled_embed_hook(
        self, scale: float, store: dict[str, Tensor]
    ) -> Callable[[nn.Module, tuple[Any, ...], Tensor], Tensor]:
        """Forward hook that scales the embedding output and retains its grad."""

        def hook(_m: nn.Module, _inp: tuple[Any, ...], out: Tensor) -> Tensor:
            scaled = out * scale
            scaled.retain_grad()
            store["scaled"] = scaled
            store["unscaled"] = out
            return scaled

        return hook

    def _embedding_gradient(
        self, ids: Tensor, target_token: int, scale: float
    ) -> tuple[Tensor, Tensor]:
        """Return ``(embedding, d metric / d scaled_embedding)`` after one backward.

        The embedding output is scaled by ``scale`` (for the IG path); the
        returned embedding is the *unscaled* value used for the ``e · grad``
        product.
        """
        store: dict[str, Tensor] = {}
        handle = self.model.embed.register_forward_hook(self._scaled_embed_hook(scale, store))
        self.model.zero_grad(set_to_none=True)
        try:
            logits = self.model(ids)
            metric = logits[0, -1, target_token]
            metric.backward()
        finally:
            handle.remove()
        grad = store["scaled"].grad
        assert grad is not None  # populated by backward() on the retained tensor
        return store["unscaled"].detach(), grad.detach()

    @torch.no_grad()
    def _metric_at_scale(self, ids: Tensor, target_token: int, scale: float) -> float:
        """Return the target logit with the embedding scaled by ``scale`` (no grad)."""

        def hook(_m: nn.Module, _inp: tuple[Any, ...], out: Tensor) -> Tensor:
            return out * scale

        handle = self.model.embed.register_forward_hook(hook)
        try:
            return float(self.model(ids)[0, -1, target_token])
        finally:
            handle.remove()

    def token_attribution(
        self,
        token_ids: Tensor,
        target_token: int | None = None,
        method: str = "input_x_grad",
        steps: int = 32,
    ) -> AttributionResult:
        """Attribute the target-token logit to each input token.

        Args:
            token_ids:    Token IDs of shape ``(S,)`` or ``(1, S)``.
            target_token: Vocab ID to attribute (defaults to the model's top
                prediction at the final position).
            method:       ``"input_x_grad"`` or ``"integrated_gradients"``.
            steps:        Riemann steps for integrated gradients (>= 1).

        Returns:
            An ``AttributionResult`` with one signed score per input position.

        Raises:
            ValueError: If ``method`` is unknown, ``steps < 1``, or ``token_ids``
                is not a single sequence.
        """
        if method not in _METHODS:
            raise ValueError(f"unknown method '{method}'; expected one of {_METHODS}")
        ids = self._as_batch(token_ids)
        self.model.eval()

        if target_token is None:
            with torch.no_grad():
                target_token = int(self.model(ids)[0, -1].argmax())

        input_metric: float | None = None
        baseline_metric: float | None = None

        if method == "input_x_grad":
            embedding, grad = self._embedding_gradient(ids, target_token, scale=1.0)
            scores = (embedding * grad).sum(dim=-1)[0]
        else:  # integrated_gradients
            if steps < 1:
                raise ValueError(f"steps must be >= 1, got {steps}")
            embedding = torch.zeros(1)
            total_grad: Tensor | None = None
            for k in range(1, steps + 1):
                alpha = k / steps
                embedding, grad = self._embedding_gradient(ids, target_token, scale=alpha)
                total_grad = grad if total_grad is None else total_grad + grad
            avg_grad = total_grad / steps  # type: ignore[operator]
            scores = (embedding * avg_grad).sum(dim=-1)[0]
            input_metric = self._metric_at_scale(ids, target_token, 1.0)
            baseline_metric = self._metric_at_scale(ids, target_token, 0.0)

        tokens = [_decode_token(self.tokenizer, int(t)) for t in ids[0].tolist()]
        return AttributionResult(
            scores=scores.detach(),
            tokens=tokens,
            target_token=target_token,
            method=method,
            input_metric=input_metric,
            baseline_metric=baseline_metric,
        )

__init__(model, tokenizer=None)

Create an attributor over model.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
tokenizer Any

Optional tokenizer with decode(list[int]) -> str.

 None
Source code in kamui/mechinterp/attribution.py 
102
103
104
105
106
107
108
109
110
def __init__(self, model: KAMUITransformer, tokenizer: Any = None) -> None:
    """Create an attributor over ``model``.

    Args:
        model:     A ``KAMUITransformer``.
        tokenizer: Optional tokenizer with ``decode(list[int]) -> str``.
    """
    self.model = model
    self.tokenizer = tokenizer

token_attribution(token_ids, target_token=None, method='input_x_grad', steps=32)

Attribute the target-token logit to each input token.

Parameters:

Name Type Description Default
token_ids Tensor

Token IDs of shape (S,) or (1, S).

 required
target_token int | None

Vocab ID to attribute (defaults to the model's top prediction at the final position).

 None
method str

"input_x_grad" or "integrated_gradients".

 'input_x_grad'
steps int

Riemann steps for integrated gradients (>= 1).

 32

Returns:

Type Description
AttributionResult

An AttributionResult with one signed score per input position.

Raises:

Type Description
ValueError

If method is unknown, steps < 1, or token_ids is not a single sequence.
Source code in kamui/mechinterp/attribution.py 
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
def token_attribution(
    self,
    token_ids: Tensor,
    target_token: int | None = None,
    method: str = "input_x_grad",
    steps: int = 32,
) -> AttributionResult:
    """Attribute the target-token logit to each input token.

    Args:
        token_ids:    Token IDs of shape ``(S,)`` or ``(1, S)``.
        target_token: Vocab ID to attribute (defaults to the model's top
            prediction at the final position).
        method:       ``"input_x_grad"`` or ``"integrated_gradients"``.
        steps:        Riemann steps for integrated gradients (>= 1).

    Returns:
        An ``AttributionResult`` with one signed score per input position.

    Raises:
        ValueError: If ``method`` is unknown, ``steps < 1``, or ``token_ids``
            is not a single sequence.
    """
    if method not in _METHODS:
        raise ValueError(f"unknown method '{method}'; expected one of {_METHODS}")
    ids = self._as_batch(token_ids)
    self.model.eval()

    if target_token is None:
        with torch.no_grad():
            target_token = int(self.model(ids)[0, -1].argmax())

    input_metric: float | None = None
    baseline_metric: float | None = None

    if method == "input_x_grad":
        embedding, grad = self._embedding_gradient(ids, target_token, scale=1.0)
        scores = (embedding * grad).sum(dim=-1)[0]
    else:  # integrated_gradients
        if steps < 1:
            raise ValueError(f"steps must be >= 1, got {steps}")
        embedding = torch.zeros(1)
        total_grad: Tensor | None = None
        for k in range(1, steps + 1):
            alpha = k / steps
            embedding, grad = self._embedding_gradient(ids, target_token, scale=alpha)
            total_grad = grad if total_grad is None else total_grad + grad
        avg_grad = total_grad / steps  # type: ignore[operator]
        scores = (embedding * avg_grad).sum(dim=-1)[0]
        input_metric = self._metric_at_scale(ids, target_token, 1.0)
        baseline_metric = self._metric_at_scale(ids, target_token, 0.0)

    tokens = [_decode_token(self.tokenizer, int(t)) for t in ids[0].tolist()]
    return AttributionResult(
        scores=scores.detach(),
        tokens=tokens,
        target_token=target_token,
        method=method,
        input_metric=input_metric,
        baseline_metric=baseline_metric,
    )

kamui.mechinterp.superposition.SparseAutoencoder

Bases: Module

An overcomplete sparse autoencoder over model activations.

Attributes:

Name Type Description
d_model

Dimension of the activations being decomposed.
n_features

Dictionary size (k * d_model for an overcomplete SAE).
l1_coeff

Weight of the L1 sparsity penalty.
W_enc

(d_model, n_features) encoder weights.
b_enc

(n_features,) encoder bias.
W_dec

(n_features, d_model) decoder dictionary (unit-norm rows).
b_dec

(d_model,) decoder pre-bias.
Source code in kamui/mechinterp/superposition.py 
 88
 89
 90
 91
 92
 93
 94
 95
 96
 97
 98
 99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
class SparseAutoencoder(nn.Module):
    """An overcomplete sparse autoencoder over model activations.

    Attributes:
        d_model:    Dimension of the activations being decomposed.
        n_features: Dictionary size (``k * d_model`` for an overcomplete SAE).
        l1_coeff:   Weight of the L1 sparsity penalty.
        W_enc:      ``(d_model, n_features)`` encoder weights.
        b_enc:      ``(n_features,)`` encoder bias.
        W_dec:      ``(n_features, d_model)`` decoder dictionary (unit-norm rows).
        b_dec:      ``(d_model,)`` decoder pre-bias.
    """

    def __init__(self, d_model: int, n_features: int, l1_coeff: float = _DEFAULT_L1) -> None:
        """Create a sparse autoencoder.

        Args:
            d_model:    Activation dimension.  Must be > 0.
            n_features: Dictionary size.  Must be > 0 (typically ``k * d_model``).
            l1_coeff:   L1 sparsity coefficient.  Must be >= 0.

        Raises:
            ValueError: If any argument is out of range.
        """
        super().__init__()
        if d_model <= 0:
            raise ValueError(f"d_model must be > 0, got {d_model}")
        if n_features <= 0:
            raise ValueError(f"n_features must be > 0, got {n_features}")
        if l1_coeff < 0:
            raise ValueError(f"l1_coeff must be >= 0, got {l1_coeff}")

        self.d_model = d_model
        self.n_features = n_features
        self.l1_coeff = l1_coeff

        self.b_dec = nn.Parameter(torch.zeros(d_model))
        self.b_enc = nn.Parameter(torch.zeros(n_features))
        self.W_enc = nn.Parameter(torch.empty(d_model, n_features))
        self.W_dec = nn.Parameter(torch.empty(n_features, d_model))
        nn.init.normal_(self.W_dec, std=0.1)
        with torch.no_grad():
            self.normalize_decoder()
            self.W_enc.copy_(self.W_dec.t())  # tied init from the normalised dictionary

    def normalize_decoder(self) -> None:
        """Rescale each feature's decoder direction (row of ``W_dec``) to unit norm."""
        with torch.no_grad():
            norms = self.W_dec.norm(dim=1, keepdim=True).clamp_min(_NORM_EPS)
            self.W_dec.div_(norms)

    def encode(self, x: Tensor) -> Tensor:
        """Encode activations to sparse, non-negative feature coefficients.

        Args:
            x: Activations of shape ``(..., d_model)``.

        Returns:
            Feature activations of shape ``(..., n_features)`` (>= 0 via ReLU).
        """
        return torch.relu((x - self.b_dec) @ self.W_enc + self.b_enc)

    def decode(self, features: Tensor) -> Tensor:
        """Reconstruct activations from feature coefficients.

        Args:
            features: Feature activations of shape ``(..., n_features)``.

        Returns:
            Reconstructed activations of shape ``(..., d_model)``.
        """
        return features @ self.W_dec + self.b_dec

    def forward(self, x: Tensor) -> tuple[Tensor, Tensor]:
        """Return ``(reconstruction, features)`` for input activations ``x``."""
        features = self.encode(x)
        return self.decode(features), features

    def loss(self, x: Tensor) -> SAELoss:
        """Compute the SAE objective for a batch of activations.

        Args:
            x: Activations of shape ``(N, d_model)``.

        Returns:
            An ``SAELoss`` with ``total``, ``reconstruction``, and ``sparsity``.
        """
        reconstruction, features = self.forward(x)
        recon = F.mse_loss(reconstruction, x)
        sparsity = features.abs().mean()
        return SAELoss(
            total=recon + self.l1_coeff * sparsity, reconstruction=recon, sparsity=sparsity
        )

    def __repr__(self) -> str:
        return (
            f"SparseAutoencoder(d_model={self.d_model}, "
            f"n_features={self.n_features}, l1_coeff={self.l1_coeff})"
        )

    def save(self, path: str | Path) -> None:
        """Save the SAE's config and weights to a ``.pt`` file.

        Both the architecture (``d_model`` / ``n_features`` / ``l1_coeff``) and
        the learned weights are stored, so ``SparseAutoencoder.load`` can rebuild
        the model without the caller re-specifying its shape.  Parent directories
        are created as needed.

        Args:
            path: Destination ``.pt`` file.
        """
        path_obj = Path(path)
        path_obj.parent.mkdir(parents=True, exist_ok=True)
        torch.save(
            {
                "d_model": self.d_model,
                "n_features": self.n_features,
                "l1_coeff": self.l1_coeff,
                "state_dict": self.state_dict(),
            },
            path_obj,
        )

    @classmethod
    def load(cls, path: str | Path) -> SparseAutoencoder:
        """Rebuild an SAE saved by :meth:`save`.

        Loads with ``weights_only=True`` (no arbitrary-code execution — the file
        holds only the config scalars and weight tensors).

        Args:
            path: A ``.pt`` file written by :meth:`save`.

        Returns:
            A ``SparseAutoencoder`` with the saved architecture and weights.
        """
        checkpoint = torch.load(Path(path), map_location="cpu", weights_only=True)
        sae = cls(
            d_model=int(checkpoint["d_model"]),
            n_features=int(checkpoint["n_features"]),
            l1_coeff=float(checkpoint["l1_coeff"]),
        )
        sae.load_state_dict(checkpoint["state_dict"])
        return sae

__init__(d_model, n_features, l1_coeff=_DEFAULT_L1)

Create a sparse autoencoder.

Parameters:

Name Type Description Default
d_model int

Activation dimension. Must be > 0.

 required
n_features int

Dictionary size. Must be > 0 (typically k * d_model).

 required
l1_coeff float

L1 sparsity coefficient. Must be >= 0.

 _DEFAULT_L1

Raises:

Type Description
ValueError

If any argument is out of range.
Source code in kamui/mechinterp/superposition.py 
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
def __init__(self, d_model: int, n_features: int, l1_coeff: float = _DEFAULT_L1) -> None:
    """Create a sparse autoencoder.

    Args:
        d_model:    Activation dimension.  Must be > 0.
        n_features: Dictionary size.  Must be > 0 (typically ``k * d_model``).
        l1_coeff:   L1 sparsity coefficient.  Must be >= 0.

    Raises:
        ValueError: If any argument is out of range.
    """
    super().__init__()
    if d_model <= 0:
        raise ValueError(f"d_model must be > 0, got {d_model}")
    if n_features <= 0:
        raise ValueError(f"n_features must be > 0, got {n_features}")
    if l1_coeff < 0:
        raise ValueError(f"l1_coeff must be >= 0, got {l1_coeff}")

    self.d_model = d_model
    self.n_features = n_features
    self.l1_coeff = l1_coeff

    self.b_dec = nn.Parameter(torch.zeros(d_model))
    self.b_enc = nn.Parameter(torch.zeros(n_features))
    self.W_enc = nn.Parameter(torch.empty(d_model, n_features))
    self.W_dec = nn.Parameter(torch.empty(n_features, d_model))
    nn.init.normal_(self.W_dec, std=0.1)
    with torch.no_grad():
        self.normalize_decoder()
        self.W_enc.copy_(self.W_dec.t())  # tied init from the normalised dictionary

normalize_decoder()

Rescale each feature's decoder direction (row of W_dec) to unit norm.

Source code in kamui/mechinterp/superposition.py 
133
134
135
136
137
def normalize_decoder(self) -> None:
    """Rescale each feature's decoder direction (row of ``W_dec``) to unit norm."""
    with torch.no_grad():
        norms = self.W_dec.norm(dim=1, keepdim=True).clamp_min(_NORM_EPS)
        self.W_dec.div_(norms)

encode(x)

Encode activations to sparse, non-negative feature coefficients.

Parameters:

Name Type Description Default
x Tensor

Activations of shape (..., d_model).

 required

Returns:

Type Description
Tensor

Feature activations of shape (..., n_features) (>= 0 via ReLU).
Source code in kamui/mechinterp/superposition.py 
139
140
141
142
143
144
145
146
147
148
def encode(self, x: Tensor) -> Tensor:
    """Encode activations to sparse, non-negative feature coefficients.

    Args:
        x: Activations of shape ``(..., d_model)``.

    Returns:
        Feature activations of shape ``(..., n_features)`` (>= 0 via ReLU).
    """
    return torch.relu((x - self.b_dec) @ self.W_enc + self.b_enc)

decode(features)

Reconstruct activations from feature coefficients.

Parameters:

Name Type Description Default
features Tensor

Feature activations of shape (..., n_features).

 required

Returns:

Type Description
Tensor

Reconstructed activations of shape (..., d_model).
Source code in kamui/mechinterp/superposition.py 
150
151
152
153
154
155
156
157
158
159
def decode(self, features: Tensor) -> Tensor:
    """Reconstruct activations from feature coefficients.

    Args:
        features: Feature activations of shape ``(..., n_features)``.

    Returns:
        Reconstructed activations of shape ``(..., d_model)``.
    """
    return features @ self.W_dec + self.b_dec

forward(x)

Return (reconstruction, features) for input activations x.

Source code in kamui/mechinterp/superposition.py 
161
162
163
164
def forward(self, x: Tensor) -> tuple[Tensor, Tensor]:
    """Return ``(reconstruction, features)`` for input activations ``x``."""
    features = self.encode(x)
    return self.decode(features), features

loss(x)

Compute the SAE objective for a batch of activations.

Parameters:

Name Type Description Default
x Tensor

Activations of shape (N, d_model).

 required

Returns:

Type Description
SAELoss

An SAELoss with total, reconstruction, and sparsity.
Source code in kamui/mechinterp/superposition.py 
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
def loss(self, x: Tensor) -> SAELoss:
    """Compute the SAE objective for a batch of activations.

    Args:
        x: Activations of shape ``(N, d_model)``.

    Returns:
        An ``SAELoss`` with ``total``, ``reconstruction``, and ``sparsity``.
    """
    reconstruction, features = self.forward(x)
    recon = F.mse_loss(reconstruction, x)
    sparsity = features.abs().mean()
    return SAELoss(
        total=recon + self.l1_coeff * sparsity, reconstruction=recon, sparsity=sparsity
    )

save(path)

Save the SAE's config and weights to a .pt file.

Both the architecture (d_model / n_features / l1_coeff) and the learned weights are stored, so SparseAutoencoder.load can rebuild the model without the caller re-specifying its shape. Parent directories are created as needed.

Parameters:

Name Type Description Default
path str | Path

Destination .pt file.

 required
Source code in kamui/mechinterp/superposition.py 
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
def save(self, path: str | Path) -> None:
    """Save the SAE's config and weights to a ``.pt`` file.

    Both the architecture (``d_model`` / ``n_features`` / ``l1_coeff``) and
    the learned weights are stored, so ``SparseAutoencoder.load`` can rebuild
    the model without the caller re-specifying its shape.  Parent directories
    are created as needed.

    Args:
        path: Destination ``.pt`` file.
    """
    path_obj = Path(path)
    path_obj.parent.mkdir(parents=True, exist_ok=True)
    torch.save(
        {
            "d_model": self.d_model,
            "n_features": self.n_features,
            "l1_coeff": self.l1_coeff,
            "state_dict": self.state_dict(),
        },
        path_obj,
    )

load(path) classmethod

Rebuild an SAE saved by :meth:save.

Loads with weights_only=True (no arbitrary-code execution — the file holds only the config scalars and weight tensors).

Parameters:

Name Type Description Default
path str | Path

A .pt file written by :meth:save.

 required

Returns:

Type Description
SparseAutoencoder

A SparseAutoencoder with the saved architecture and weights.
Source code in kamui/mechinterp/superposition.py 
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
@classmethod
def load(cls, path: str | Path) -> SparseAutoencoder:
    """Rebuild an SAE saved by :meth:`save`.

    Loads with ``weights_only=True`` (no arbitrary-code execution — the file
    holds only the config scalars and weight tensors).

    Args:
        path: A ``.pt`` file written by :meth:`save`.

    Returns:
        A ``SparseAutoencoder`` with the saved architecture and weights.
    """
    checkpoint = torch.load(Path(path), map_location="cpu", weights_only=True)
    sae = cls(
        d_model=int(checkpoint["d_model"]),
        n_features=int(checkpoint["n_features"]),
        l1_coeff=float(checkpoint["l1_coeff"]),
    )
    sae.load_state_dict(checkpoint["state_dict"])
    return sae

kamui.mechinterp.superposition.interpret_features(sae, activations, token_ids, top_k=10, features=None)

Profile SAE features by the tokens that most strongly activate them.

A sparse autoencoder learns a dictionary of feature directions, but a feature is only interpretable once you see what makes it fire. For each requested feature this returns the tokens whose encoded activation at that feature is largest, along with how often the feature fires (density) and how strongly (mean activation) — the token-level form of Anthropic's "what does this feature detect?" analysis. Pair top_token_ids with a tokenizer's decode to read the feature's meaning.

Parameters:

Name Type Description Default
sae SparseAutoencoder

A (typically trained) SparseAutoencoder.

 required
activations Tensor

(N, d_model) activations, one row per token.

 required
token_ids Tensor

(N,) token ID for each activation row (parallel to activations — usually the tokens passed to collect_activations).

 required
top_k int

Number of top-activating tokens to keep per feature.

 10
features Iterable[int] | None

Feature indices to profile (defaults to every feature).

 None

Returns:

Type Description
list[FeatureProfile]

One FeatureProfile per requested feature, in the requested order.

Raises:

Type Description
ValueError

If shapes are inconsistent, top_k < 1, or a requested feature index is out of range.
Source code in kamui/mechinterp/superposition.py 
389
390
391
392
393
394
395
396
397
398
399
400
401
402
403
404
405
406
407
408
409
410
411
412
413
414
415
416
417
418
419
420
421
422
423
424
425
426
427
428
429
430
431
432
433
434
435
436
437
438
439
440
441
442
443
444
445
446
447
448
449
450
451
452
453
454
455
456
457
458
459
460
461
462
463
464
465
466
@torch.no_grad()
def interpret_features(
    sae: SparseAutoencoder,
    activations: Tensor,
    token_ids: Tensor,
    top_k: int = 10,
    features: Iterable[int] | None = None,
) -> list[FeatureProfile]:
    """Profile SAE features by the tokens that most strongly activate them.

    A sparse autoencoder learns a dictionary of feature directions, but a
    feature is only *interpretable* once you see what makes it fire.  For each
    requested feature this returns the tokens whose encoded activation at that
    feature is largest, along with how often the feature fires (density) and how
    strongly (mean activation) — the token-level form of Anthropic's "what does
    this feature detect?" analysis.  Pair ``top_token_ids`` with a tokenizer's
    ``decode`` to read the feature's meaning.

    Args:
        sae:         A (typically trained) ``SparseAutoencoder``.
        activations: ``(N, d_model)`` activations, one row per token.
        token_ids:   ``(N,)`` token ID for each activation row (parallel to
            ``activations`` — usually the tokens passed to ``collect_activations``).
        top_k:       Number of top-activating tokens to keep per feature.
        features:    Feature indices to profile (defaults to every feature).

    Returns:
        One ``FeatureProfile`` per requested feature, in the requested order.

    Raises:
        ValueError: If shapes are inconsistent, ``top_k < 1``, or a requested
            feature index is out of range.
    """
    if activations.dim() != 2 or activations.shape[1] != sae.d_model:
        raise ValueError(
            f"activations must be (N, d_model={sae.d_model}), got {tuple(activations.shape)}"
        )
    if token_ids.dim() != 1 or token_ids.shape[0] != activations.shape[0]:
        raise ValueError(
            f"token_ids must be 1-D with length {activations.shape[0]}, "
            f"got shape {tuple(token_ids.shape)}"
        )
    if top_k < 1:
        raise ValueError(f"top_k must be >= 1, got {top_k}")

    feature_list = list(range(sae.n_features)) if features is None else list(features)
    coded = sae.encode(activations)  # (N, n_features)
    n_rows = coded.shape[0]
    ids = token_ids.tolist()

    profiles: list[FeatureProfile] = []
    for feature in feature_list:
        if not (0 <= feature < sae.n_features):
            raise ValueError(f"feature must be in [0, {sae.n_features}), got {feature}")
        column = coded[:, feature]
        active = column > 0
        n_active = int(active.sum().item())
        density = n_active / n_rows
        mean_activation = float(column[active].mean().item()) if n_active > 0 else 0.0

        k = min(top_k, n_active)  # never rank in inactive (zero) tokens
        if k > 0:
            values, indices = torch.topk(column, k)
            top_token_ids = [ids[i] for i in indices.tolist()]
            top_activations = [float(v) for v in values.tolist()]
        else:
            top_token_ids = []
            top_activations = []
        profiles.append(
            FeatureProfile(
                feature=feature,
                density=density,
                mean_activation=mean_activation,
                top_token_ids=top_token_ids,
                top_activations=top_activations,
            )
        )
    return profiles

kamui.mechinterp.superposition.feature_cooccurrence(sae, activations, features=None)

Compute how often pairs of SAE features co-activate.

Monosemantic features rarely fire in isolation — related features (e.g. the products of feature splitting) tend to co-occur. This returns the co-activation density matrix M where::

M[i, j] = fraction of tokens on which features i and j are both active

M is symmetric and its diagonal M[i, i] is feature i's own firing density (as in sae_feature_metrics). A dead feature yields a zero row and column. Read M alongside interpret_features: the profiles say what each feature detects, M says which detections travel together.

Parameters:

Name Type Description Default
sae SparseAutoencoder

A (typically trained) SparseAutoencoder.

 required
activations Tensor

(N, d_model) activations, one row per token.

 required
features Iterable[int] | None

Feature indices to include (defaults to every feature). The returned matrix is ordered to match this list.

 None

Returns:

Type Description
Tensor

A (k, k) co-activation density matrix for the k requested
Tensor

features (k == n_features by default).

Raises:

Type Description
ValueError

If activations is not (N, d_model) matching the SAE, or a requested feature index is out of range.
Source code in kamui/mechinterp/superposition.py 
469
470
471
472
473
474
475
476
477
478
479
480
481
482
483
484
485
486
487
488
489
490
491
492
493
494
495
496
497
498
499
500
501
502
503
504
505
506
507
508
509
510
511
512
513
514
@torch.no_grad()
def feature_cooccurrence(
    sae: SparseAutoencoder,
    activations: Tensor,
    features: Iterable[int] | None = None,
) -> Tensor:
    """Compute how often pairs of SAE features co-activate.

    Monosemantic features rarely fire in isolation — related features (e.g. the
    products of feature splitting) tend to co-occur.  This returns the
    **co-activation density** matrix ``M`` where::

        M[i, j] = fraction of tokens on which features i and j are both active

    ``M`` is symmetric and its diagonal ``M[i, i]`` is feature ``i``'s own firing
    density (as in ``sae_feature_metrics``).  A dead feature yields a zero row and
    column.  Read ``M`` alongside ``interpret_features``: the profiles say what
    each feature detects, ``M`` says which detections travel together.

    Args:
        sae:         A (typically trained) ``SparseAutoencoder``.
        activations: ``(N, d_model)`` activations, one row per token.
        features:    Feature indices to include (defaults to every feature). The
            returned matrix is ordered to match this list.

    Returns:
        A ``(k, k)`` co-activation density matrix for the ``k`` requested
        features (``k == n_features`` by default).

    Raises:
        ValueError: If ``activations`` is not ``(N, d_model)`` matching the SAE,
            or a requested feature index is out of range.
    """
    if activations.dim() != 2 or activations.shape[1] != sae.d_model:
        raise ValueError(
            f"activations must be (N, d_model={sae.d_model}), got {tuple(activations.shape)}"
        )
    feature_list = list(range(sae.n_features)) if features is None else list(features)
    for feature in feature_list:
        if not (0 <= feature < sae.n_features):
            raise ValueError(f"feature must be in [0, {sae.n_features}), got {feature}")

    active = (sae.encode(activations) > 0).float()  # (N, n_features)
    selected = active[:, feature_list]  # (N, k)
    n_rows = selected.shape[0]
    return (selected.t() @ selected) / n_rows

kamui.mechinterp.steering.FeatureSteerer

Causally steer a model along activation-space directions.

Attributes:

Name Type Description
model

A trained KAMUITransformer.
sae

An optional SparseAutoencoder whose decoder rows supply feature directions for steer_with_feature.
Source code in kamui/mechinterp/steering.py 
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
203
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
247
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
273
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
316
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
class FeatureSteerer:
    """Causally steer a model along activation-space directions.

    Attributes:
        model: A trained ``KAMUITransformer``.
        sae:   An optional ``SparseAutoencoder`` whose decoder rows supply
            feature directions for ``steer_with_feature``.
    """

    def __init__(self, model: KAMUITransformer, sae: SparseAutoencoder | None = None) -> None:
        """Create a steerer.

        Args:
            model: A ``KAMUITransformer``.
            sae:   Optional SAE providing feature directions (``W_dec``).
        """
        self.model = model
        self.sae = sae

    @staticmethod
    def _as_batch(ids: Tensor) -> Tensor:
        if ids.dim() == 1:
            ids = ids.unsqueeze(0)
        if ids.dim() != 2 or ids.shape[0] != 1:
            raise ValueError(
                f"input_ids must be a single sequence (S,) or (1, S), got shape {tuple(ids.shape)}"
            )
        return ids

    @staticmethod
    def _is_steerable(hook_point: str) -> bool:
        return (
            hook_point == "embed.output"
            or hook_point.endswith(".attn.output")
            or hook_point.endswith(".ffn.output")
        )

    def _resolve(self, module_path: str) -> nn.Module:
        return dict(self.model.named_modules())[module_path]

    def _validate(self, hook_point: str, direction: Tensor) -> None:
        if not self._is_steerable(hook_point):
            raise ValueError(
                f"'{hook_point}' is not a steerable point "
                f"(use embed.output / blocks.i.attn.output / blocks.i.ffn.output)"
            )
        d_model = self.model.config.d_model
        if direction.shape != (d_model,):
            raise ValueError(
                f"direction must have shape (d_model={d_model},), got {tuple(direction.shape)}"
            )

    def _feature_direction(self, feature: int) -> Tensor:
        if self.sae is None:
            raise ValueError(
                "steering with a feature requires an SAE; pass one to FeatureSteerer(...)"
            )
        if not (0 <= feature < self.sae.n_features):
            raise ValueError(f"feature must be in [0, {self.sae.n_features}), got {feature}")
        return self.sae.W_dec[feature].detach()

    @torch.no_grad()
    def steer(
        self,
        input_ids: Tensor,
        hook_point: str,
        direction: Tensor,
        coefficient: float = 1.0,
        position: int | None = None,
    ) -> SteeringResult:
        """Add ``coefficient * direction`` to a residual-stream point and rerun.

        Args:
            input_ids:   Token IDs ``(S,)`` or ``(1, S)``.
            hook_point:  A steerable output point: ``"embed.output"``,
                ``"blocks.{i}.attn.output"``, or ``"blocks.{i}.ffn.output"``.
            direction:   A ``(d_model,)`` direction to add.
            coefficient: The scalar the direction is multiplied by (may be
                negative to suppress).
            position:    If given, steer only this sequence position; otherwise
                steer every position.

        Returns:
            A ``SteeringResult`` with baseline and steered logits.

        Raises:
            ValueError: If ``hook_point`` is not steerable, ``input_ids`` is not
                a single sequence, or ``direction`` is not shape ``(d_model,)``.
        """
        self._validate(hook_point, direction)

        ids = self._as_batch(input_ids)
        self.model.eval()
        baseline_logits = self.model(ids)

        add = coefficient * direction.to(baseline_logits.dtype)
        module = self._resolve(hook_point.rsplit(".", 1)[0])
        handle = module.register_forward_hook(_steering_hook(add, position))
        try:
            steered_logits = self.model(ids)
        finally:
            handle.remove()

        return SteeringResult(baseline_logits, steered_logits, hook_point, coefficient)

    def steer_with_feature(
        self,
        input_ids: Tensor,
        hook_point: str,
        feature: int,
        coefficient: float = 1.0,
        position: int | None = None,
    ) -> SteeringResult:
        """Steer along an SAE feature's decoder direction (``W_dec[feature]``).

        Args:
            input_ids:   Token IDs ``(S,)`` or ``(1, S)``.
            hook_point:  A steerable residual-stream point.
            feature:     The SAE feature index whose direction to add.
            coefficient: The scalar the direction is multiplied by.
            position:    If given, steer only this sequence position.

        Returns:
            A ``SteeringResult``.

        Raises:
            ValueError: If no SAE was provided or ``feature`` is out of range.
        """
        direction = self._feature_direction(feature)
        return self.steer(input_ids, hook_point, direction, coefficient, position)

    @torch.no_grad()
    def generate_steered(
        self,
        tokenizer: TokenizerLike,
        prompt: str,
        hook_point: str,
        direction: Tensor,
        coefficient: float = 1.0,
        **generate_kwargs: Any,
    ) -> str:
        """Generate text while steering along ``direction`` at every step.

        The direction is added at every position of every forward pass during
        autoregressive generation, so the whole continuation is produced under
        the intervention — the tangible payoff of steering (compare against a
        plain ``generate`` call to see the effect on the text).

        Args:
            tokenizer:       An object with ``encode`` / ``decode``.
            prompt:          The (non-empty) prompt string.
            hook_point:      A steerable residual-stream point.
            direction:       A ``(d_model,)`` direction to add.
            coefficient:     The scalar the direction is multiplied by.
            **generate_kwargs: Forwarded to ``kamui.evaluate.generate``
                (``max_new_tokens``, ``strategy``, ``temperature``, ``top_k``,
                ``top_p``, ``seed``).

        Returns:
            The full decoded string (prompt + steered continuation).

        Raises:
            ValueError: If ``hook_point`` is not steerable or ``direction`` is
                not shape ``(d_model,)``.
        """
        self._validate(hook_point, direction)
        add = coefficient * direction.detach().to(next(self.model.parameters()).dtype)
        module = self._resolve(hook_point.rsplit(".", 1)[0])
        handle = module.register_forward_hook(_steering_hook(add, position=None))
        try:
            return generate(self.model, tokenizer, prompt, **generate_kwargs)
        finally:
            handle.remove()

    def generate_steered_with_feature(
        self,
        tokenizer: TokenizerLike,
        prompt: str,
        hook_point: str,
        feature: int,
        coefficient: float = 1.0,
        **generate_kwargs: Any,
    ) -> str:
        """Generate text steered along an SAE feature's decoder direction.

        Args:
            tokenizer:       An object with ``encode`` / ``decode``.
            prompt:          The (non-empty) prompt string.
            hook_point:      A steerable residual-stream point.
            feature:         The SAE feature index whose direction to add.
            coefficient:     The scalar the direction is multiplied by.
            **generate_kwargs: Forwarded to ``kamui.evaluate.generate``.

        Returns:
            The full decoded string (prompt + steered continuation).

        Raises:
            ValueError: If no SAE was provided or ``feature`` is out of range.
        """
        direction = self._feature_direction(feature)
        return self.generate_steered(
            tokenizer, prompt, hook_point, direction, coefficient, **generate_kwargs
        )

__init__(model, sae=None)

Create a steerer.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
sae SparseAutoencoder | None

Optional SAE providing feature directions (W_dec).

 None
Source code in kamui/mechinterp/steering.py 
152
153
154
155
156
157
158
159
160
def __init__(self, model: KAMUITransformer, sae: SparseAutoencoder | None = None) -> None:
    """Create a steerer.

    Args:
        model: A ``KAMUITransformer``.
        sae:   Optional SAE providing feature directions (``W_dec``).
    """
    self.model = model
    self.sae = sae

steer(input_ids, hook_point, direction, coefficient=1.0, position=None)

Add coefficient * direction to a residual-stream point and rerun.

Parameters:

Name Type Description Default
input_ids Tensor

Token IDs (S,) or (1, S).

 required
hook_point str

A steerable output point: "embed.output", "blocks.{i}.attn.output", or "blocks.{i}.ffn.output".

 required
direction Tensor

A (d_model,) direction to add.

 required
coefficient float

The scalar the direction is multiplied by (may be negative to suppress).

 1.0
position int | None

If given, steer only this sequence position; otherwise steer every position.

 None

Returns:

Type Description
SteeringResult

A SteeringResult with baseline and steered logits.

Raises:

Type Description
ValueError

If hook_point is not steerable, input_ids is not a single sequence, or direction is not shape (d_model,).
Source code in kamui/mechinterp/steering.py 
204
205
206
207
208
209
210
211
212
213
214
215
216
217
218
219
220
221
222
223
224
225
226
227
228
229
230
231
232
233
234
235
236
237
238
239
240
241
242
243
244
245
246
@torch.no_grad()
def steer(
    self,
    input_ids: Tensor,
    hook_point: str,
    direction: Tensor,
    coefficient: float = 1.0,
    position: int | None = None,
) -> SteeringResult:
    """Add ``coefficient * direction`` to a residual-stream point and rerun.

    Args:
        input_ids:   Token IDs ``(S,)`` or ``(1, S)``.
        hook_point:  A steerable output point: ``"embed.output"``,
            ``"blocks.{i}.attn.output"``, or ``"blocks.{i}.ffn.output"``.
        direction:   A ``(d_model,)`` direction to add.
        coefficient: The scalar the direction is multiplied by (may be
            negative to suppress).
        position:    If given, steer only this sequence position; otherwise
            steer every position.

    Returns:
        A ``SteeringResult`` with baseline and steered logits.

    Raises:
        ValueError: If ``hook_point`` is not steerable, ``input_ids`` is not
            a single sequence, or ``direction`` is not shape ``(d_model,)``.
    """
    self._validate(hook_point, direction)

    ids = self._as_batch(input_ids)
    self.model.eval()
    baseline_logits = self.model(ids)

    add = coefficient * direction.to(baseline_logits.dtype)
    module = self._resolve(hook_point.rsplit(".", 1)[0])
    handle = module.register_forward_hook(_steering_hook(add, position))
    try:
        steered_logits = self.model(ids)
    finally:
        handle.remove()

    return SteeringResult(baseline_logits, steered_logits, hook_point, coefficient)

steer_with_feature(input_ids, hook_point, feature, coefficient=1.0, position=None)

Steer along an SAE feature's decoder direction (W_dec[feature]).

Parameters:

Name Type Description Default
input_ids Tensor

Token IDs (S,) or (1, S).

 required
hook_point str

A steerable residual-stream point.

 required
feature int

The SAE feature index whose direction to add.

 required
coefficient float

The scalar the direction is multiplied by.

 1.0
position int | None

If given, steer only this sequence position.

 None

Returns:

Type Description
SteeringResult

A SteeringResult.

Raises:

Type Description
ValueError

If no SAE was provided or feature is out of range.
Source code in kamui/mechinterp/steering.py 
248
249
250
251
252
253
254
255
256
257
258
259
260
261
262
263
264
265
266
267
268
269
270
271
272
def steer_with_feature(
    self,
    input_ids: Tensor,
    hook_point: str,
    feature: int,
    coefficient: float = 1.0,
    position: int | None = None,
) -> SteeringResult:
    """Steer along an SAE feature's decoder direction (``W_dec[feature]``).

    Args:
        input_ids:   Token IDs ``(S,)`` or ``(1, S)``.
        hook_point:  A steerable residual-stream point.
        feature:     The SAE feature index whose direction to add.
        coefficient: The scalar the direction is multiplied by.
        position:    If given, steer only this sequence position.

    Returns:
        A ``SteeringResult``.

    Raises:
        ValueError: If no SAE was provided or ``feature`` is out of range.
    """
    direction = self._feature_direction(feature)
    return self.steer(input_ids, hook_point, direction, coefficient, position)

generate_steered(tokenizer, prompt, hook_point, direction, coefficient=1.0, **generate_kwargs)

Generate text while steering along direction at every step.

The direction is added at every position of every forward pass during autoregressive generation, so the whole continuation is produced under the intervention — the tangible payoff of steering (compare against a plain generate call to see the effect on the text).

Parameters:

Name Type Description Default
tokenizer TokenizerLike

An object with encode / decode.

 required
prompt str

The (non-empty) prompt string.

 required
hook_point str

A steerable residual-stream point.

 required
direction Tensor

A (d_model,) direction to add.

 required
coefficient float

The scalar the direction is multiplied by.

 1.0
**generate_kwargs Any

Forwarded to kamui.evaluate.generate (max_new_tokens, strategy, temperature, top_k, top_p, seed).

 {}

Returns:

Type Description
str

The full decoded string (prompt + steered continuation).

Raises:

Type Description
ValueError

If hook_point is not steerable or direction is not shape (d_model,).
Source code in kamui/mechinterp/steering.py 
274
275
276
277
278
279
280
281
282
283
284
285
286
287
288
289
290
291
292
293
294
295
296
297
298
299
300
301
302
303
304
305
306
307
308
309
310
311
312
313
314
315
@torch.no_grad()
def generate_steered(
    self,
    tokenizer: TokenizerLike,
    prompt: str,
    hook_point: str,
    direction: Tensor,
    coefficient: float = 1.0,
    **generate_kwargs: Any,
) -> str:
    """Generate text while steering along ``direction`` at every step.

    The direction is added at every position of every forward pass during
    autoregressive generation, so the whole continuation is produced under
    the intervention — the tangible payoff of steering (compare against a
    plain ``generate`` call to see the effect on the text).

    Args:
        tokenizer:       An object with ``encode`` / ``decode``.
        prompt:          The (non-empty) prompt string.
        hook_point:      A steerable residual-stream point.
        direction:       A ``(d_model,)`` direction to add.
        coefficient:     The scalar the direction is multiplied by.
        **generate_kwargs: Forwarded to ``kamui.evaluate.generate``
            (``max_new_tokens``, ``strategy``, ``temperature``, ``top_k``,
            ``top_p``, ``seed``).

    Returns:
        The full decoded string (prompt + steered continuation).

    Raises:
        ValueError: If ``hook_point`` is not steerable or ``direction`` is
            not shape ``(d_model,)``.
    """
    self._validate(hook_point, direction)
    add = coefficient * direction.detach().to(next(self.model.parameters()).dtype)
    module = self._resolve(hook_point.rsplit(".", 1)[0])
    handle = module.register_forward_hook(_steering_hook(add, position=None))
    try:
        return generate(self.model, tokenizer, prompt, **generate_kwargs)
    finally:
        handle.remove()

generate_steered_with_feature(tokenizer, prompt, hook_point, feature, coefficient=1.0, **generate_kwargs)

Generate text steered along an SAE feature's decoder direction.

Parameters:

Name Type Description Default
tokenizer TokenizerLike

An object with encode / decode.

 required
prompt str

The (non-empty) prompt string.

 required
hook_point str

A steerable residual-stream point.

 required
feature int

The SAE feature index whose direction to add.

 required
coefficient float

The scalar the direction is multiplied by.

 1.0
**generate_kwargs Any

Forwarded to kamui.evaluate.generate.

 {}

Returns:

Type Description
str

The full decoded string (prompt + steered continuation).

Raises:

Type Description
ValueError

If no SAE was provided or feature is out of range.
Source code in kamui/mechinterp/steering.py 
317
318
319
320
321
322
323
324
325
326
327
328
329
330
331
332
333
334
335
336
337
338
339
340
341
342
343
344
345
def generate_steered_with_feature(
    self,
    tokenizer: TokenizerLike,
    prompt: str,
    hook_point: str,
    feature: int,
    coefficient: float = 1.0,
    **generate_kwargs: Any,
) -> str:
    """Generate text steered along an SAE feature's decoder direction.

    Args:
        tokenizer:       An object with ``encode`` / ``decode``.
        prompt:          The (non-empty) prompt string.
        hook_point:      A steerable residual-stream point.
        feature:         The SAE feature index whose direction to add.
        coefficient:     The scalar the direction is multiplied by.
        **generate_kwargs: Forwarded to ``kamui.evaluate.generate``.

    Returns:
        The full decoded string (prompt + steered continuation).

    Raises:
        ValueError: If no SAE was provided or ``feature`` is out of range.
    """
    direction = self._feature_direction(feature)
    return self.generate_steered(
        tokenizer, prompt, hook_point, direction, coefficient, **generate_kwargs
    )

kamui.mechinterp.steering.build_steering_vector(model, hook_point, positive, negative, normalize=False)

Build a contrastive steering direction (ActAdd) from example sequences.

The direction is the difference of mean activations at hook_point::

d = mean_token(activation | positive) - mean_token(activation | negative)

where the mean is taken over every token of every sequence. Added to the residual stream (via FeatureSteerer.steer), d pushes the model toward whatever distinguishes the positive prompts from the negative ones — the data-driven alternative to steer_with_feature when you have contrasting examples rather than a trained SAE (Turner et al., 2023).

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
hook_point str

A residual-stream hook point, e.g. "blocks.3.ffn.output".

 required
positive Iterable[Tensor]

Sequences exhibiting the target behaviour (each (S,) or (B, S)).

 required
negative Iterable[Tensor]

Contrasting sequences.

 required
normalize bool

If True, return a unit-norm direction.

 False

Returns:

Type Description
Tensor

A (d_model,) steering direction.

Raises:

Type Description
ValueError

If positive or negative yields no activations.
Source code in kamui/mechinterp/steering.py 
348
349
350
351
352
353
354
355
356
357
358
359
360
361
362
363
364
365
366
367
368
369
370
371
372
373
374
375
376
377
378
379
380
381
382
383
384
385
386
387
@torch.no_grad()
def build_steering_vector(
    model: KAMUITransformer,
    hook_point: str,
    positive: Iterable[Tensor],
    negative: Iterable[Tensor],
    normalize: bool = False,
) -> Tensor:
    """Build a contrastive steering direction (ActAdd) from example sequences.

    The direction is the difference of mean activations at ``hook_point``::

        d = mean_token(activation | positive) - mean_token(activation | negative)

    where the mean is taken over every token of every sequence.  Added to the
    residual stream (via ``FeatureSteerer.steer``), ``d`` pushes the model toward
    whatever distinguishes the positive prompts from the negative ones — the
    data-driven alternative to ``steer_with_feature`` when you have contrasting
    examples rather than a trained SAE (Turner et al., 2023).

    Args:
        model:      A ``KAMUITransformer``.
        hook_point: A residual-stream hook point, e.g. ``"blocks.3.ffn.output"``.
        positive:   Sequences exhibiting the target behaviour (each ``(S,)`` or
            ``(B, S)``).
        negative:   Contrasting sequences.
        normalize:  If True, return a unit-norm direction.

    Returns:
        A ``(d_model,)`` steering direction.

    Raises:
        ValueError: If ``positive`` or ``negative`` yields no activations.
    """
    positive_mean = collect_activations(model, hook_point, positive).mean(dim=0)
    negative_mean = collect_activations(model, hook_point, negative).mean(dim=0)
    direction = positive_mean - negative_mean
    if normalize:
        direction = direction / direction.norm().clamp_min(_NORM_EPS)
    return direction

Evaluation

kamui.evaluate.perplexity.compute_perplexity(model, dataloader)

Compute token-level perplexity over a dataset.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
dataloader Iterable[tuple[Tensor, Tensor] | Tensor]

An iterable of batches. Each batch is either a (inputs, targets) pair (both (B, S)) or a single (B, S) token tensor (split into next-token inputs/targets internally).

 required

Returns:

Type Description
float

The corpus perplexity as a float.

Raises:

Type Description
ValueError

If the dataloader yields no tokens.
Source code in kamui/evaluate/perplexity.py 
165
166
167
168
169
170
171
172
173
174
175
176
177
178
179
180
181
182
183
184
185
186
187
188
189
190
191
192
193
194
195
196
197
198
199
200
201
202
@torch.no_grad()
def compute_perplexity(
    model: KAMUITransformer, dataloader: Iterable[tuple[Tensor, Tensor] | Tensor]
) -> float:
    """Compute token-level perplexity over a dataset.

    Args:
        model:      A ``KAMUITransformer``.
        dataloader: An iterable of batches.  Each batch is either a
            ``(inputs, targets)`` pair (both ``(B, S)``) or a single ``(B, S)``
            token tensor (split into next-token inputs/targets internally).

    Returns:
        The corpus perplexity as a float.

    Raises:
        ValueError: If the dataloader yields no tokens.
    """
    was_training = model.training
    model.eval()

    total_loss = 0.0
    total_tokens = 0
    for batch in dataloader:
        inputs, targets = _split_batch(batch)
        logits = model(inputs)
        total_loss += F.cross_entropy(
            logits.reshape(-1, logits.shape[-1]),
            targets.reshape(-1),
            reduction="sum",
        ).item()
        total_tokens += targets.numel()

    if was_training:
        model.train()
    if total_tokens == 0:
        raise ValueError("dataloader produced no tokens")
    return math.exp(total_loss / total_tokens)

kamui.evaluate.generation.generate(model, tokenizer, prompt, max_new_tokens=50, strategy='greedy', temperature=1.0, top_k=50, top_p=0.9, seed=None)

Generate a continuation of prompt.

Parameters:

Name Type Description Default
model KAMUITransformer

A KAMUITransformer.

 required
tokenizer TokenizerLike

An object with encode(str) -> list[int] and decode(list[int]) -> str (e.g. BPETokenizer).

 required
prompt str

The (non-empty) prompt string.

 required
max_new_tokens int

Number of tokens to generate.

 50
strategy str

"greedy", "top_k", "nucleus", or "temperature".

 'greedy'
temperature float

Logit temperature (> 0), applied before filtering.

 1.0
top_k int

k for top_k sampling (>= 1).

 50
top_p float

p for nucleus sampling (in (0, 1]).

 0.9
seed int | None

Optional RNG seed for reproducible sampling.

 None

Returns:

Type Description
str

The full decoded string (prompt + continuation).

Raises:

Type Description
ValueError

If prompt is empty, strategy is unknown, or a hyperparameter is out of range.
Source code in kamui/evaluate/generation.py 
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128
129
130
131
132
133
134
135
136
137
138
139
140
141
142
143
144
145
146
147
148
149
150
151
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
167
168
169
170
171
172
173
174
175
@torch.no_grad()
def generate(
    model: KAMUITransformer,
    tokenizer: TokenizerLike,
    prompt: str,
    max_new_tokens: int = 50,
    strategy: str = "greedy",
    temperature: float = 1.0,
    top_k: int = 50,
    top_p: float = 0.9,
    seed: int | None = None,
) -> str:
    """Generate a continuation of ``prompt``.

    Args:
        model:          A ``KAMUITransformer``.
        tokenizer:      An object with ``encode(str) -> list[int]`` and
            ``decode(list[int]) -> str`` (e.g. ``BPETokenizer``).
        prompt:         The (non-empty) prompt string.
        max_new_tokens: Number of tokens to generate.
        strategy:       ``"greedy"``, ``"top_k"``, ``"nucleus"``, or
            ``"temperature"``.
        temperature:    Logit temperature (> 0), applied before filtering.
        top_k:          k for ``top_k`` sampling (>= 1).
        top_p:          p for ``nucleus`` sampling (in (0, 1]).
        seed:           Optional RNG seed for reproducible sampling.

    Returns:
        The full decoded string (prompt + continuation).

    Raises:
        ValueError: If ``prompt`` is empty, ``strategy`` is unknown, or a
            hyperparameter is out of range.
    """
    if temperature <= 0:
        raise ValueError(f"temperature must be > 0, got {temperature}")
    if top_k < 1:
        raise ValueError(f"top_k must be >= 1, got {top_k}")
    if not (0.0 < top_p <= 1.0):
        raise ValueError(f"top_p must be in (0, 1], got {top_p}")

    if seed is not None:
        torch.manual_seed(seed)

    ids = tokenizer.encode(prompt)
    if not ids:
        raise ValueError("prompt must be non-empty (encoded to zero tokens)")

    was_training = model.training
    model.eval()

    ctx = model.config.context_length
    tokens = torch.tensor([ids], dtype=torch.long)  # (1, S)
    for _ in range(max_new_tokens):
        window = tokens[:, -ctx:]
        logits = model(window)  # (1, s, V)
        next_id = _sample_next(logits[0, -1], strategy, temperature, top_k, top_p)
        tokens = torch.cat([tokens, torch.tensor([[next_id]], dtype=torch.long)], dim=1)

    if was_training:
        model.train()
    return tokenizer.decode(tokens[0].tolist())

kamui.evaluate.calibration.expected_calibration_error(probs, labels, n_bins=15)

Compute the expected calibration error of top-1 predictions.

Parameters:

Name Type Description Default
probs Tensor

(N, C) probability distributions.

 required
labels Tensor

(N,) integer ground-truth labels.

 required
n_bins int

Number of equal-width confidence bins.

 15

Returns:

Type Description
float

The ECE in [0, 1] (0 = perfectly calibrated).

Raises:

Type Description
ValueError

If shapes are inconsistent or n_bins < 1.
Source code in kamui/evaluate/calibration.py 
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
def expected_calibration_error(probs: Tensor, labels: Tensor, n_bins: int = 15) -> float:
    """Compute the expected calibration error of top-1 predictions.

    Args:
        probs:  ``(N, C)`` probability distributions.
        labels: ``(N,)`` integer ground-truth labels.
        n_bins: Number of equal-width confidence bins.

    Returns:
        The ECE in ``[0, 1]`` (0 = perfectly calibrated).

    Raises:
        ValueError: If shapes are inconsistent or ``n_bins < 1``.
    """
    _validate_probs_labels(probs, labels, n_bins)
    conf_mean, acc_mean, counts = _bin_stats(probs, labels, n_bins)
    total = counts.sum()
    return float((counts / total * (conf_mean - acc_mean).abs()).sum().item())