Skip to content

API Reference

The complete public API, generated from the source docstrings. Everything below is importable from the top-level kamui package.

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' or 'sinusoidal').
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
@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' or 'sinusoidal').
        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:
        """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"):
            raise ValueError(
                f"positional_encoding must be 'learned' or 'sinusoidal', "
                f"got '{self.positional_encoding}'"
            )

    @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", 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", 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 
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
def __post_init__(self) -> None:
    """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"):
        raise ValueError(
            f"positional_encoding must be 'learned' or 'sinusoidal', "
            f"got '{self.positional_encoding}'"
        )

__repr__()

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

Source code in kamui/model/config.py 
152
153
154
155
156
157
158
159
160
161
162
163
164
165
166
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 
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
@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 
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 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

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))

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())