Apply coding standards from 1cb5747 (files only)

README.md

@@ -102,12 +102,14 @@ Residue labels have shape `(b, l)` and use `-100` outside biological positions.
 
 ```python
 import torch
+
 from transformers import AutoTokenizer
 from transformers import (
     AutoModelForSequenceClassification,
     AutoModelForTokenClassification,
 )
 
+
 model_id = "Synthyra/ESM3_small"
 sequence_model = AutoModelForSequenceClassification.from_pretrained(
     model_id, num_labels=2, trust_remote_code=True
@@ -118,13 +120,13 @@ token_model = AutoModelForTokenClassification.from_pretrained(
 tokenizer = AutoTokenizer.from_pretrained(model_id, trust_remote_code=True)
 sequences = ["MSTNPKPQRKTKRNT", "MKTIIALSYIFCLVFA"]
 batch = tokenizer(sequences, padding=True, return_tensors="pt")
-biological = batch["attention_mask"].bool()
+biological = batch["attention_mask"].bool()  # (b, l)
 for special_id in tokenizer.all_special_ids:
-    biological &= batch["input_ids"].ne(special_id)
+    biological &= batch["input_ids"].ne(special_id)  # (b, l)
 
-sequence_labels = torch.zeros(len(sequences), dtype=torch.long)
-token_labels = torch.full_like(batch["input_ids"], -100)
-token_labels[biological] = 0
+sequence_labels = torch.zeros(len(sequences), dtype=torch.long)  # (b,)
+token_labels = torch.full_like(batch["input_ids"], -100)  # (b, l)
+token_labels[biological] = 0  # selected biological positions; labels stay (b, l)
 
 with torch.inference_mode():
     sequence_output = sequence_model(**batch, labels=sequence_labels)
@@ -144,6 +146,7 @@ python -m pip install "datasets>=4.8,<5" "peft>=0.19,<0.20"
 ```python
 from peft import LoraConfig, TaskType, get_peft_model
 
+
 peft_model = get_peft_model(
     sequence_model,
     LoraConfig(
@@ -171,6 +174,7 @@ adapters. Base checkpoint weights stay frozen:
 ```python
 from transformers import AutoModel
 
+
 ttt_model = AutoModel.from_pretrained(
     "Synthyra/ESM3_small",
     trust_remote_code=True,
@@ -196,6 +200,7 @@ multimodal helpers:
 ```python
 import torch
 
+
 batch = model.tokenize_sequences(
     ["MKTAYIAKQ", "GGGG"],
     device=model.device,
@@ -218,6 +223,7 @@ Generate masked sequence positions with an explicit seed:
 ```python
 from fastplms.models.esm3.modeling_esm3 import FastESM3GenerationConfig
 
+
 config = FastESM3GenerationConfig(
     num_steps=8,
     temperature=1.0,

fastplms/attention/_core.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import warnings
 import torch
+
 from collections import OrderedDict
 from collections.abc import Callable
 from dataclasses import dataclass

fastplms/attention/_kernel_lock.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import json
 import os
+
 from pathlib import Path
 from typing import Any
 

fastplms/attention/interfaces.py

@@ -3,6 +3,7 @@
 from __future__ import annotations
 
 import torch
+
 from collections.abc import Mapping
 from functools import partial
 from typing import Any

fastplms/embeddings/batches.py

@@ -1,378 +1,379 @@
-"""Execute model-specific batches and return ordered residue-aware CPU tensors."""
-
-from __future__ import annotations
-
-import torch
-from collections.abc import Callable, Iterator, Sequence
-from contextlib import contextmanager
-from dataclasses import dataclass, field
-from typing import Any
-from torch import Tensor
-
-from .identity import _model_device
-from .inputs import _planned_batches
-from .pooling import Pooler
-from .types import EmbeddingBatch, EmbeddingInput, EmbeddingRecord
-
-
-_MAX_PARTI_RESIDUES = 2_048
-
-
-def _validate_parti_length(M: Tensor) -> None:
-    """Reject an oversized attention graph before model inference."""
-
-    # M: (b, l)
-    n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item())
-    if n_residues > _MAX_PARTI_RESIDUES:
-        raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.")
-
-
-def select_hidden_state_embeddings(
-    last_hidden_state: Tensor,
-    hidden_states: tuple[Tensor, ...] | None,
-    *,
-    hidden_state_index: int = -1,
-    store_all_hidden_states: bool = False,
-) -> Tensor:
-    """Select one hidden state or stack every state without changing values."""
-    # last_hidden_state and each hidden_states entry: (b, l, d)
-    if store_all_hidden_states:
-        if not hidden_states:
-            raise ValueError("store_all_hidden_states requires model hidden states.")
-        # H has shape (b, n, l, d), where n follows the model's output order.
-        return torch.stack(hidden_states, dim=1)  # (b, n, l, d)
-    if hidden_state_index == -1:
-        return last_hidden_state  # (b, l, d)
-    if not hidden_states:
-        raise ValueError("hidden_state_index requires model hidden states.")
-    return hidden_states[hidden_state_index]  # (b, l, d)
-
-
-def _residue_embeddings(X: Tensor, M: Tensor) -> list[Tensor]:
-    """Copy every sample's biological residues to the host in one transfer.
-
-    Boolean indexing packs the selected rows in batch order, so splitting the
-    packed rows by residue count gives the values that indexing each sample
-    would. Each returned tensor owns its storage, as a per-sample copy does.
-    """
-    # X: (b, l, d); M: (b, l)
-    residue_counts = M.sum(dim=1).tolist()  # b counts r_i
-    packed = X[M].detach().cpu()  # (sum of r_i, d)
-    return [sample.clone() for sample in torch.split(packed, residue_counts)]  # each: (r_i, d)
-
-
-@contextmanager
-def _temporary_eval(model: Any) -> Iterator[None]:
-    was_training = getattr(model, "training", None)
-    eval_method = getattr(model, "eval", None)
-    train_method = getattr(model, "train", None)
-    if (
-        not isinstance(was_training, bool)
-        or not callable(eval_method)
-        or not callable(train_method)
-    ):
-        yield
-        return
-    eval_method()
-    try:
-        yield
-    finally:
-        train_method(was_training)
-
-
-def _biological_residue_mask(
-    input_ids: Tensor,
-    attention_mask: Tensor,
-    tokenizer: Any,
-) -> Tensor:
-    """Remove padding and tokenizer-declared special tokens from M."""
-
-    # input_ids, attention_mask: (b, l)
-    M = attention_mask.to(dtype=torch.bool)  # (b, l)
-    special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ()))
-    if special_ids:
-        specials = torch.tensor(  # (n_special,)
-            special_ids,
-            device=input_ids.device,
-            dtype=input_ids.dtype,
-        )
-        M = M & ~torch.isin(input_ids, specials)  # (b, l)
-    return M  # (b, l)
-
-
-def _generic_embedding_batch(
-    model: Any,
-    sequences: list[str],
-    *,
-    tokenizer: Any | None,
-    max_length: int | None,
-    truncate: bool,
-    need_attentions: bool,
-    model_kwargs: dict[str, Any],
-) -> EmbeddingBatch:
-    config = getattr(model, "config", None)
-    model_type = str(getattr(config, "model_type", "")).lower()
-    if tokenizer is None:
-        tokenizer = getattr(model, "tokenizer", None)
-
-    if tokenizer is None and model_type == "e1":
-        output = model._embed(sequences, return_attention_mask=True, **model_kwargs)
-        if not isinstance(output, tuple) or len(output) != 2:
-            raise TypeError("E1 _embed must return (X, residue_mask).")
-        X, M = output  # (b, l, d), (b, l)
-        preparer = getattr(model, "prep_tokens", None)
-        if preparer is not None and hasattr(preparer, "get_batch_kwargs"):
-            prepared = preparer.get_batch_kwargs(sequences, device=X.device)
-            input_ids = prepared["input_ids"]  # (b, l)
-            boundary_ids = preparer.boundary_token_ids.to(  # (n_boundary,)
-                device=input_ids.device, dtype=input_ids.dtype
-            )
-            # E1 wraps each raw sequence in BOS, context-label, terminal-label,
-            # and EOS tokens. Only amino-acid rows are biological residues.
-            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)  # (b, l)
-        if need_attentions:
-            raise ValueError("parti is not available for tokenizer-free E1 embedding.")
-        return EmbeddingBatch(  # X: (b, l, d); residue_mask: (b, l)
-            X=X,
-            residue_mask=M.to(dtype=torch.bool),
-        )
-    if tokenizer is None:
-        raise ValueError("A tokenizer is required for this model's embedding path.")
-
-    tokenize_kwargs: dict[str, Any] = {
-        "return_tensors": "pt",
-        "padding": True,
-        "truncation": truncate,
-    }
-    if max_length is not None and truncate:
-        # ``max_length`` is a biological-residue limit. Tokenizer limits include
-        # boundary tokens, so reserve their declared width instead of dropping
-        # residues at the exact boundary.
-        special_token_count = 0
-        num_special_tokens_to_add = getattr(tokenizer, "num_special_tokens_to_add", None)
-        if callable(num_special_tokens_to_add):
-            special_token_count = int(num_special_tokens_to_add(pair=False))
-        tokenize_kwargs["max_length"] = max_length + special_token_count
-    sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None)
-    if callable(sequence_tokenizer):
-        encoded = sequence_tokenizer(sequences, tokenizer=tokenizer, **tokenize_kwargs)
-    else:
-        encoded = tokenizer(sequences, **tokenize_kwargs)
-    device = _model_device(model)
-    input_ids = encoded["input_ids"].to(device)  # (b, l)
-    attention_mask = encoded.get(  # (b, l)
-        "attention_mask",
-        input_ids.new_ones(input_ids.shape),
-    ).to(device)
-    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)  # (b, l)
-    if need_attentions:
-        # Validate l before either the backbone or its quadratic attention graph
-        # is materialized. M has shape (b, l).
-        _validate_parti_length(M)
-    X = model._embed(input_ids, attention_mask, **model_kwargs)  # (b, l, d)
-    attentions = None
-    if need_attentions:
-        output = model(
-            input_ids=input_ids,
-            attention_mask=attention_mask,
-            output_attentions=True,
-            return_dict=True,
-        )
-        attentions = getattr(output, "attentions", None)  # each: (b, h, l, l)
-        if attentions is None:
-            raise ValueError("The model did not return attentions required by parti.")
-    return EmbeddingBatch(  # X: (b, l, d); M: (b, l)
-        X=X,
-        residue_mask=M,
-        attentions=attentions,
-    )
-
-
-@dataclass(eq=False)
-class BatchExecutor:
-    """Model and batch policy for one bounded embedding window at a time."""
-
-    model: Any
-    batch_size: int
-    max_tokens_per_batch: int | None
-    max_length: int | None
-    truncate: bool
-    model_kwargs: dict[str, Any]
-    hidden_state_source: str
-    normalized_decoder_inputs: tuple[str, ...] | None
-    decoder_input_ids: Tensor | None
-    decoder_attention_mask: Tensor | None
-    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None
-    tokenizer: Any | None
-    store_all_hidden_states: bool
-    full_embeddings: bool
-    dtype: torch.dtype | None
-    pooler: Pooler | None
-    attention_backend: str | None
-    need_attentions: bool
-    model_type: str = field(init=False)
-    resolved_tokenizer: Any = field(init=False)
-
-    def __post_init__(self) -> None:
-        config = getattr(self.model, "config", None)
-        self.model_type = str(getattr(config, "model_type", "")).lower()
-        self.resolved_tokenizer = (
-            self.tokenizer if self.tokenizer is not None else getattr(self.model, "tokenizer", None)
-        )
-
-    def run_window(
-        self,
-        window_records: Sequence[EmbeddingInput],
-        *,
-        window_start: int,
-    ) -> tuple[list[EmbeddingRecord], dict[str, tuple[int, int]]]:
-        """Restore source order after length-bucketed inference and pooling."""
-
-        pool_slices: dict[str, tuple[int, int]] = {}
-        window_results: dict[int, EmbeddingRecord] = {}
-        for local_positions in _planned_batches(
-            window_records,
-            range(len(window_records)),
-            batch_size=self.batch_size,
-            max_tokens_per_batch=self.max_tokens_per_batch,
-            max_length=self.max_length,
-            truncate=self.truncate,
-        ):
-            batch_positions = [window_start + position for position in local_positions]
-            batch_records = [window_records[position] for position in local_positions]
-            sequences = [
-                record.sequence[: self.max_length]
-                if self.truncate and self.max_length is not None
-                else record.sequence
-                for record in batch_records
-            ]
-            batch_model_kwargs = dict(self.model_kwargs)
-            if self.model_type == "fast_ankh" or self.hidden_state_source == "decoder":
-                batch_model_kwargs["hidden_state_source"] = self.hidden_state_source
-                if self.normalized_decoder_inputs is not None:
-                    batch_model_kwargs["decoder_inputs"] = [
-                        self.normalized_decoder_inputs[position] for position in batch_positions
-                    ]
-                if self.decoder_input_ids is not None:
-                    # decoder_input_ids: (n_records, l_decoder)
-                    indices = torch.tensor(  # (b,)
-                        batch_positions,
-                        device=self.decoder_input_ids.device,
-                        dtype=torch.long,
-                    )
-                    batch_model_kwargs["decoder_input_ids"] = (  # (b, l_decoder)
-                        self.decoder_input_ids.index_select(0, indices)
-                    )
-                if self.decoder_attention_mask is not None:
-                    # decoder_attention_mask: (n_records, l_decoder)
-                    indices = torch.tensor(  # (b,)
-                        batch_positions,
-                        device=self.decoder_attention_mask.device,
-                        dtype=torch.long,
-                    )
-                    batch_model_kwargs["decoder_attention_mask"] = (
-                        self.decoder_attention_mask.index_select(0, indices)  # (b, l_decoder)
-                    )
-            custom_batch = self._embedding_batch_fn or getattr(self.model, "_embedding_batch", None)
-            if custom_batch is not None:
-                if self.model_type == "fast_ankh":
-                    batch = custom_batch(
-                        sequences,
-                        tokenizer=self.resolved_tokenizer,
-                        max_length=self.max_length,
-                        truncate=self.truncate,
-                        need_attentions=self.need_attentions,
-                        **batch_model_kwargs,
-                    )
-                else:
-                    batch = custom_batch(sequences, **batch_model_kwargs)
-                if not isinstance(batch, EmbeddingBatch):
-                    raise TypeError("_embedding_batch must return EmbeddingBatch.")
-            else:
-                batch = _generic_embedding_batch(
-                    self.model,
-                    sequences,
-                    tokenizer=self.tokenizer,
-                    max_length=self.max_length,
-                    truncate=self.truncate,
-                    need_attentions=self.need_attentions,
-                    model_kwargs=batch_model_kwargs,
-                )
-            X = batch.X  # (b, l, d) or (b, n_states, l, d)
-            raw_mask = batch.residue_mask  # (b, l)
-            if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor):
-                raise TypeError("Embedding batches must provide Tensor X and residue_mask.")
-            if X.is_meta or raw_mask.is_meta:
-                raise ValueError("Embedding batches cannot contain meta tensors.")
-            if not X.is_floating_point():
-                raise TypeError("Embedding batches must use a floating-point X dtype.")
-            if raw_mask.is_complex() or not bool(torch.isfinite(raw_mask).all()):
-                raise ValueError("Embedding residue_mask must contain finite binary values.")
-            if not bool(((raw_mask == 0) | (raw_mask == 1)).all()):
-                raise ValueError("Embedding residue_mask must contain finite binary values.")
-            M = raw_mask.to(device=X.device, dtype=torch.bool)  # (b, l)
-            valid_X_shape = (
-                X.ndim == 3
-                and X.shape[0] == len(batch_records)
-                and X.shape[-1] > 0
-                and M.shape == X.shape[:2]
-            )
-            valid_all_states_shape = (
-                X.ndim == 4
-                and self.store_all_hidden_states
-                and self.full_embeddings
-                and X.shape[0] == len(batch_records)
-                and X.shape[1] > 0
-                and X.shape[-1] > 0
-                and M.shape == (X.shape[0], X.shape[2])
-            )
-            if not (valid_X_shape or valid_all_states_shape):
-                raise ValueError(
-                    "Embedding batches must provide X with shape (b, l, d), or "
-                    "(b, states, l, d) when storing all hidden states, and "
-                    "residue_mask with shape (b, l)."
-                )
-            if not bool(M.any(dim=1).all()):
-                raise ValueError("Every embedding sample must contain a biological residue.")
-            finite_selected = (  # X.shape
-                torch.isfinite(X) | ~M.unsqueeze(-1)
-                if X.ndim == 3
-                else torch.isfinite(X) | ~M[:, None, :, None]
-            )
-            if not bool(finite_selected.all()):
-                raise ValueError("Biological residue embeddings produced non-finite output.")
-            if self.need_attentions:
-                # Validate the biological graph only after mask integrity is established.
-                _validate_parti_length(M)
-            if self.dtype is not None:
-                X = X.to(dtype=self.dtype)  # unchanged shape
-
-            if self.full_embeddings:
-                if X.ndim == 4:
-                    values = [
-                        X_i[:, M_i, :].detach().cpu()  # (n_states, r_i, d)
-                        for X_i, M_i in zip(X, M, strict=True)
-                    ]
-                else:
-                    values = _residue_embeddings(X, M)  # each: (r_i, d)
-            else:
-                if self.pooler is None:
-                    raise RuntimeError(
-                        "Pooled embedding output was requested without an initialized pooler."
-                    )
-                Y = self.pooler(  # (b, n_poolers * d)
-                    X,
-                    M,
-                    attentions=batch.attentions,
-                    attention_backend=self.attention_backend,
-                )
-                pool_slices = self.pooler.output_slices(X.shape[-1])
-                values = list(Y.detach().cpu().unbind(0))  # each: (n_poolers * d,)
-            for position, record, value in zip(batch_positions, batch_records, values, strict=True):
-                window_results[position] = EmbeddingRecord(record.id, record.sequence, value)
-
-        new_records = [
-            window_results[position]
-            for position in range(window_start, window_start + len(window_records))
-        ]
-        return new_records, pool_slices
+"""Execute model-specific batches and return ordered residue-aware CPU tensors."""
+
+from __future__ import annotations
+
+import torch
+
+from collections.abc import Callable, Iterator, Sequence
+from contextlib import contextmanager
+from dataclasses import dataclass, field
+from typing import Any
+from torch import Tensor
+
+from .identity import _model_device
+from .inputs import _planned_batches
+from .pooling import Pooler
+from .types import EmbeddingBatch, EmbeddingInput, EmbeddingRecord
+
+
+_MAX_PARTI_RESIDUES = 2_048
+
+
+def _validate_parti_length(M: Tensor) -> None:
+    """Reject an oversized attention graph before model inference."""
+
+    # M: (b, l)
+    n_residues = int(M.to(dtype=torch.int64).sum(dim=1).max().item())
+    if n_residues > _MAX_PARTI_RESIDUES:
+        raise ValueError(f"parti supports at most {_MAX_PARTI_RESIDUES:,} biological residues.")
+
+
+def select_hidden_state_embeddings(
+    last_hidden_state: Tensor,
+    hidden_states: tuple[Tensor, ...] | None,
+    *,
+    hidden_state_index: int = -1,
+    store_all_hidden_states: bool = False,
+) -> Tensor:
+    """Select one hidden state or stack every state without changing values."""
+    # last_hidden_state and each hidden_states entry: (b, l, d)
+    if store_all_hidden_states:
+        if not hidden_states:
+            raise ValueError("store_all_hidden_states requires model hidden states.")
+        # H has shape (b, n, l, d), where n follows the model's output order.
+        return torch.stack(hidden_states, dim=1)  # (b, n, l, d)
+    if hidden_state_index == -1:
+        return last_hidden_state  # (b, l, d)
+    if not hidden_states:
+        raise ValueError("hidden_state_index requires model hidden states.")
+    return hidden_states[hidden_state_index]  # (b, l, d)
+
+
+def _residue_embeddings(X: Tensor, M: Tensor) -> list[Tensor]:
+    """Copy every sample's biological residues to the host in one transfer.
+
+    Boolean indexing packs the selected rows in batch order, so splitting the
+    packed rows by residue count gives the values that indexing each sample
+    would. Each returned tensor owns its storage, as a per-sample copy does.
+    """
+    # X: (b, l, d); M: (b, l)
+    residue_counts = M.sum(dim=1).tolist()  # b counts r_i
+    packed = X[M].detach().cpu()  # (sum of r_i, d)
+    return [sample.clone() for sample in torch.split(packed, residue_counts)]  # each: (r_i, d)
+
+
+@contextmanager
+def _temporary_eval(model: Any) -> Iterator[None]:
+    was_training = getattr(model, "training", None)
+    eval_method = getattr(model, "eval", None)
+    train_method = getattr(model, "train", None)
+    if (
+        not isinstance(was_training, bool)
+        or not callable(eval_method)
+        or not callable(train_method)
+    ):
+        yield
+        return
+    eval_method()
+    try:
+        yield
+    finally:
+        train_method(was_training)
+
+
+def _biological_residue_mask(
+    input_ids: Tensor,
+    attention_mask: Tensor,
+    tokenizer: Any,
+) -> Tensor:
+    """Remove padding and tokenizer-declared special tokens from M."""
+
+    # input_ids, attention_mask: (b, l)
+    M = attention_mask.to(dtype=torch.bool)  # (b, l)
+    special_ids = tuple(int(token_id) for token_id in getattr(tokenizer, "all_special_ids", ()))
+    if special_ids:
+        specials = torch.tensor(  # (n_special,)
+            special_ids,
+            device=input_ids.device,
+            dtype=input_ids.dtype,
+        )
+        M = M & ~torch.isin(input_ids, specials)  # (b, l)
+    return M  # (b, l)
+
+
+def _generic_embedding_batch(
+    model: Any,
+    sequences: list[str],
+    *,
+    tokenizer: Any | None,
+    max_length: int | None,
+    truncate: bool,
+    need_attentions: bool,
+    model_kwargs: dict[str, Any],
+) -> EmbeddingBatch:
+    config = getattr(model, "config", None)
+    model_type = str(getattr(config, "model_type", "")).lower()
+    if tokenizer is None:
+        tokenizer = getattr(model, "tokenizer", None)
+
+    if tokenizer is None and model_type == "e1":
+        output = model._embed(sequences, return_attention_mask=True, **model_kwargs)
+        if not isinstance(output, tuple) or len(output) != 2:
+            raise TypeError("E1 _embed must return (X, residue_mask).")
+        X, M = output  # (b, l, d), (b, l)
+        preparer = getattr(model, "prep_tokens", None)
+        if preparer is not None and hasattr(preparer, "get_batch_kwargs"):
+            prepared = preparer.get_batch_kwargs(sequences, device=X.device)
+            input_ids = prepared["input_ids"]  # (b, l)
+            boundary_ids = preparer.boundary_token_ids.to(  # (n_boundary,)
+                device=input_ids.device, dtype=input_ids.dtype
+            )
+            # E1 wraps each raw sequence in BOS, context-label, terminal-label,
+            # and EOS tokens. Only amino-acid rows are biological residues.
+            M = M.to(dtype=torch.bool) & ~torch.isin(input_ids, boundary_ids)  # (b, l)
+        if need_attentions:
+            raise ValueError("parti is not available for tokenizer-free E1 embedding.")
+        return EmbeddingBatch(  # X: (b, l, d); residue_mask: (b, l)
+            X=X,
+            residue_mask=M.to(dtype=torch.bool),
+        )
+    if tokenizer is None:
+        raise ValueError("A tokenizer is required for this model's embedding path.")
+
+    tokenize_kwargs: dict[str, Any] = {
+        "return_tensors": "pt",
+        "padding": True,
+        "truncation": truncate,
+    }
+    if max_length is not None and truncate:
+        # ``max_length`` is a biological-residue limit. Tokenizer limits include
+        # boundary tokens, so reserve their declared width instead of dropping
+        # residues at the exact boundary.
+        special_token_count = 0
+        num_special_tokens_to_add = getattr(tokenizer, "num_special_tokens_to_add", None)
+        if callable(num_special_tokens_to_add):
+            special_token_count = int(num_special_tokens_to_add(pair=False))
+        tokenize_kwargs["max_length"] = max_length + special_token_count
+    sequence_tokenizer = getattr(model, "_tokenize_sequence_batch", None)
+    if callable(sequence_tokenizer):
+        encoded = sequence_tokenizer(sequences, tokenizer=tokenizer, **tokenize_kwargs)
+    else:
+        encoded = tokenizer(sequences, **tokenize_kwargs)
+    device = _model_device(model)
+    input_ids = encoded["input_ids"].to(device)  # (b, l)
+    attention_mask = encoded.get(  # (b, l)
+        "attention_mask",
+        input_ids.new_ones(input_ids.shape),
+    ).to(device)
+    M = _biological_residue_mask(input_ids, attention_mask, tokenizer)  # (b, l)
+    if need_attentions:
+        # Validate l before either the backbone or its quadratic attention graph
+        # is materialized. M has shape (b, l).
+        _validate_parti_length(M)
+    X = model._embed(input_ids, attention_mask, **model_kwargs)  # (b, l, d)
+    attentions = None
+    if need_attentions:
+        output = model(
+            input_ids=input_ids,
+            attention_mask=attention_mask,
+            output_attentions=True,
+            return_dict=True,
+        )
+        attentions = getattr(output, "attentions", None)  # each: (b, h, l, l)
+        if attentions is None:
+            raise ValueError("The model did not return attentions required by parti.")
+    return EmbeddingBatch(  # X: (b, l, d); M: (b, l)
+        X=X,
+        residue_mask=M,
+        attentions=attentions,
+    )
+
+
+@dataclass(eq=False)
+class BatchExecutor:
+    """Model and batch policy for one bounded embedding window at a time."""
+
+    model: Any
+    batch_size: int
+    max_tokens_per_batch: int | None
+    max_length: int | None
+    truncate: bool
+    model_kwargs: dict[str, Any]
+    hidden_state_source: str
+    normalized_decoder_inputs: tuple[str, ...] | None
+    decoder_input_ids: Tensor | None
+    decoder_attention_mask: Tensor | None
+    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None
+    tokenizer: Any | None
+    store_all_hidden_states: bool
+    full_embeddings: bool
+    dtype: torch.dtype | None
+    pooler: Pooler | None
+    attention_backend: str | None
+    need_attentions: bool
+    model_type: str = field(init=False)
+    resolved_tokenizer: Any = field(init=False)
+
+    def __post_init__(self) -> None:
+        config = getattr(self.model, "config", None)
+        self.model_type = str(getattr(config, "model_type", "")).lower()
+        self.resolved_tokenizer = (
+            self.tokenizer if self.tokenizer is not None else getattr(self.model, "tokenizer", None)
+        )
+
+    def run_window(
+        self,
+        window_records: Sequence[EmbeddingInput],
+        *,
+        window_start: int,
+    ) -> tuple[list[EmbeddingRecord], dict[str, tuple[int, int]]]:
+        """Restore source order after length-bucketed inference and pooling."""
+
+        pool_slices: dict[str, tuple[int, int]] = {}
+        window_results: dict[int, EmbeddingRecord] = {}
+        for local_positions in _planned_batches(
+            window_records,
+            range(len(window_records)),
+            batch_size=self.batch_size,
+            max_tokens_per_batch=self.max_tokens_per_batch,
+            max_length=self.max_length,
+            truncate=self.truncate,
+        ):
+            batch_positions = [window_start + position for position in local_positions]
+            batch_records = [window_records[position] for position in local_positions]
+            sequences = [
+                record.sequence[: self.max_length]
+                if self.truncate and self.max_length is not None
+                else record.sequence
+                for record in batch_records
+            ]
+            batch_model_kwargs = dict(self.model_kwargs)
+            if self.model_type == "fast_ankh" or self.hidden_state_source == "decoder":
+                batch_model_kwargs["hidden_state_source"] = self.hidden_state_source
+                if self.normalized_decoder_inputs is not None:
+                    batch_model_kwargs["decoder_inputs"] = [
+                        self.normalized_decoder_inputs[position] for position in batch_positions
+                    ]
+                if self.decoder_input_ids is not None:
+                    # decoder_input_ids: (n_records, l_decoder)
+                    indices = torch.tensor(  # (b,)
+                        batch_positions,
+                        device=self.decoder_input_ids.device,
+                        dtype=torch.long,
+                    )
+                    batch_model_kwargs["decoder_input_ids"] = (  # (b, l_decoder)
+                        self.decoder_input_ids.index_select(0, indices)
+                    )
+                if self.decoder_attention_mask is not None:
+                    # decoder_attention_mask: (n_records, l_decoder)
+                    indices = torch.tensor(  # (b,)
+                        batch_positions,
+                        device=self.decoder_attention_mask.device,
+                        dtype=torch.long,
+                    )
+                    batch_model_kwargs["decoder_attention_mask"] = (
+                        self.decoder_attention_mask.index_select(0, indices)  # (b, l_decoder)
+                    )
+            custom_batch = self._embedding_batch_fn or getattr(self.model, "_embedding_batch", None)
+            if custom_batch is not None:
+                if self.model_type == "fast_ankh":
+                    batch = custom_batch(
+                        sequences,
+                        tokenizer=self.resolved_tokenizer,
+                        max_length=self.max_length,
+                        truncate=self.truncate,
+                        need_attentions=self.need_attentions,
+                        **batch_model_kwargs,
+                    )
+                else:
+                    batch = custom_batch(sequences, **batch_model_kwargs)
+                if not isinstance(batch, EmbeddingBatch):
+                    raise TypeError("_embedding_batch must return EmbeddingBatch.")
+            else:
+                batch = _generic_embedding_batch(
+                    self.model,
+                    sequences,
+                    tokenizer=self.tokenizer,
+                    max_length=self.max_length,
+                    truncate=self.truncate,
+                    need_attentions=self.need_attentions,
+                    model_kwargs=batch_model_kwargs,
+                )
+            X = batch.X  # (b, l, d) or (b, n_states, l, d)
+            raw_mask = batch.residue_mask  # (b, l)
+            if not isinstance(X, Tensor) or not isinstance(raw_mask, Tensor):
+                raise TypeError("Embedding batches must provide Tensor X and residue_mask.")
+            if X.is_meta or raw_mask.is_meta:
+                raise ValueError("Embedding batches cannot contain meta tensors.")
+            if not X.is_floating_point():
+                raise TypeError("Embedding batches must use a floating-point X dtype.")
+            if raw_mask.is_complex() or not bool(torch.isfinite(raw_mask).all()):
+                raise ValueError("Embedding residue_mask must contain finite binary values.")
+            if not bool(((raw_mask == 0) | (raw_mask == 1)).all()):
+                raise ValueError("Embedding residue_mask must contain finite binary values.")
+            M = raw_mask.to(device=X.device, dtype=torch.bool)  # (b, l)
+            valid_X_shape = (
+                X.ndim == 3
+                and X.shape[0] == len(batch_records)
+                and X.shape[-1] > 0
+                and M.shape == X.shape[:2]
+            )
+            valid_all_states_shape = (
+                X.ndim == 4
+                and self.store_all_hidden_states
+                and self.full_embeddings
+                and X.shape[0] == len(batch_records)
+                and X.shape[1] > 0
+                and X.shape[-1] > 0
+                and M.shape == (X.shape[0], X.shape[2])
+            )
+            if not (valid_X_shape or valid_all_states_shape):
+                raise ValueError(
+                    "Embedding batches must provide X with shape (b, l, d), or "
+                    "(b, states, l, d) when storing all hidden states, and "
+                    "residue_mask with shape (b, l)."
+                )
+            if not bool(M.any(dim=1).all()):
+                raise ValueError("Every embedding sample must contain a biological residue.")
+            finite_selected = (  # X.shape
+                torch.isfinite(X) | ~M.unsqueeze(-1)
+                if X.ndim == 3
+                else torch.isfinite(X) | ~M[:, None, :, None]
+            )
+            if not bool(finite_selected.all()):
+                raise ValueError("Biological residue embeddings produced non-finite output.")
+            if self.need_attentions:
+                # Validate the biological graph only after mask integrity is established.
+                _validate_parti_length(M)
+            if self.dtype is not None:
+                X = X.to(dtype=self.dtype)  # unchanged shape
+
+            if self.full_embeddings:
+                if X.ndim == 4:
+                    values = [
+                        X_i[:, M_i, :].detach().cpu()  # (n_states, r_i, d)
+                        for X_i, M_i in zip(X, M, strict=True)
+                    ]
+                else:
+                    values = _residue_embeddings(X, M)  # each: (r_i, d)
+            else:
+                if self.pooler is None:
+                    raise RuntimeError(
+                        "Pooled embedding output was requested without an initialized pooler."
+                    )
+                Y = self.pooler(  # (b, n_poolers * d)
+                    X,
+                    M,
+                    attentions=batch.attentions,
+                    attention_backend=self.attention_backend,
+                )
+                pool_slices = self.pooler.output_slices(X.shape[-1])
+                values = list(Y.detach().cpu().unbind(0))  # each: (n_poolers * d,)
+            for position, record, value in zip(batch_positions, batch_records, values, strict=True):
+                window_results[position] = EmbeddingRecord(record.id, record.sequence, value)
+
+        new_records = [
+            window_results[position]
+            for position in range(window_start, window_start + len(window_records))
+        ]
+        return new_records, pool_slices

fastplms/embeddings/identity.py

@@ -1,510 +1,511 @@
-"""Deterministic identity for embedding inputs, models, tokenizers, and execution."""
-
-from __future__ import annotations
-
-import hashlib
-import json
-import platform
-import torch
-from collections.abc import Iterable, Mapping, Sequence
-from pathlib import Path
-from typing import Any
-from torch import Tensor
-
+"""Deterministic identity for embedding inputs, models, tokenizers, and execution."""
+
+from __future__ import annotations
+
+import hashlib
+import json
+import platform
+import torch
+
+from collections.abc import Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
+from torch import Tensor
+
 from .inputs import _InputSpool
 from .storage import tensor_sha256
 from .types import EmbeddingInput
-
-
-_RUN_FINGERPRINT_SCHEMA_VERSION = 3
-_MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
-
-
-def _model_device(model: Any) -> torch.device:
-    try:
-        return torch.device(next(model.parameters()).device)
-    except (AttributeError, StopIteration):
-        return torch.device("cpu")
-
-
-def _attention_backend(model: Any) -> str | None:
-    config = getattr(model, "config", None)
-    for name in ("_attn_implementation", "attn_implementation", "attn_backend"):
-        value = getattr(config, name, None)
-        if value:
-            return str(value)
-    return None
-
-
-def _attention_kernel_metadata(backend: str | None) -> dict[str, Any] | None:
-    if backend not in {"flash_attention_2", "flash_attention_3"}:
-        return None
-    from fastplms.registry import get_model_registry
-
-    spec = get_model_registry().attention_kernels[backend]
-    return {
-        "repository": spec.repository,
-        "revision": spec.revision,
-        "version": spec.version,
-        "expected_variant": spec.expected_variant,
-        "dtypes": list(spec.dtypes),
-    }
-
-
-def _fingerprint_jsonable(value: Any) -> Any:
-    if isinstance(value, Mapping):
-        return {str(key): _fingerprint_jsonable(item) for key, item in value.items()}
-    if isinstance(value, (list, tuple)):
-        return [_fingerprint_jsonable(item) for item in value]
-    if isinstance(value, (set, frozenset)):
-        return sorted((_fingerprint_jsonable(item) for item in value), key=repr)
-    if isinstance(value, Path):
-        return str(value)
-    if isinstance(value, Tensor):
-        return {
-            "dtype": str(value.dtype).removeprefix("torch."),
-            "shape": list(value.shape),
-            "sha256": tensor_sha256(value),
-        }
-    if isinstance(value, torch.dtype):
-        return str(value).removeprefix("torch.")
-    if isinstance(value, torch.device):
-        return str(value)
-    if value is None or isinstance(value, (str, int, float, bool)):
-        return value
-    return {
-        "class": f"{value.__class__.__module__}.{value.__class__.__qualname__}",
-        "value": str(value),
-    }
-
-
-def _tokenizer_content_sha256(tokenizer: Any) -> str:
-    content: dict[str, Any] = {
-        "init_kwargs": getattr(tokenizer, "init_kwargs", None),
-        "special_tokens_map": getattr(tokenizer, "special_tokens_map", None),
-        "model_max_length": getattr(tokenizer, "model_max_length", None),
-        "padding_side": getattr(tokenizer, "padding_side", None),
-        "truncation_side": getattr(tokenizer, "truncation_side", None),
-    }
-    get_vocab = getattr(tokenizer, "get_vocab", None)
-    if callable(get_vocab):
-        content["vocabulary"] = get_vocab()
-    get_added_vocab = getattr(tokenizer, "get_added_vocab", None)
-    if callable(get_added_vocab):
-        content["added_vocabulary"] = get_added_vocab()
-    backend = getattr(tokenizer, "backend_tokenizer", None)
-    backend_to_str = getattr(backend, "to_str", None)
-    if callable(backend_to_str):
-        content["backend"] = backend_to_str()
-    serialized = json.dumps(
-        _fingerprint_jsonable(content),
-        sort_keys=True,
-        separators=(",", ":"),
-        ensure_ascii=False,
-    ).encode()
-    return hashlib.sha256(serialized).hexdigest()
-
-
-def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]:
-    resolved = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
-    if resolved is None:
-        # Raw-sequence families such as E1 retain their loader context on the
-        # model/encoder rather than exposing a Transformers tokenizer. Bind the
-        # non-secret source policy to resume identity without serializing a Hub
-        # token or forcing lazy tokenizer initialization.
-        for candidate in (model, getattr(model, "model", None)):
-            settings = getattr(candidate, "__dict__", {}).get("_fastplms_tokenizer_kwargs")
-            if isinstance(settings, Mapping):
-                token_value = settings.get("token")
-                return {
-                    "mode": "native-sequence",
-                    "source": (
-                        str(settings.get("tokenizer_source"))
-                        if settings.get("tokenizer_source") is not None
-                        else None
-                    ),
-                    "revision": settings.get("revision"),
-                    "cache_dir": (
-                        str(settings.get("cache_dir"))
-                        if settings.get("cache_dir") is not None
-                        else None
-                    ),
-                    "local_files_only": bool(settings.get("local_files_only", False)),
-                    "token_policy": (
-                        "disabled"
-                        if token_value is False
-                        else "provided"
-                        if token_value is not None
-                        else "default"
-                    ),
-                }
-        return {"mode": "native-sequence"}
-    return {
-        "mode": "tokenizer",
-        "class": f"{resolved.__class__.__module__}.{resolved.__class__.__qualname__}",
-        "name_or_path": getattr(resolved, "name_or_path", None),
-        "vocab_size": getattr(resolved, "vocab_size", None),
-        "special_token_ids": list(getattr(resolved, "all_special_ids", ())),
-        "content_sha256": _tokenizer_content_sha256(resolved),
-    }
-
-
-def _software_versions() -> dict[str, str | None]:
-    try:
-        import fastplms
-
-        fastplms_version = fastplms.__version__
-    except (AttributeError, ImportError):
-        fastplms_version = None
-    try:
-        import safetensors
-
-        safetensors_version = safetensors.__version__
-    except ImportError:
-        safetensors_version = None
-    try:
-        import transformers
-
-        transformers_version = transformers.__version__
-    except ImportError:
-        transformers_version = None
-    return {
-        "fastplms": fastplms_version,
-        "python": platform.python_version(),
-        "safetensors": safetensors_version,
-        "torch": torch.__version__,
-        "torch_cuda": torch.version.cuda,
-        "transformers": transformers_version,
-    }
-
-
-def _adapter_identity_metadata(model: Any) -> dict[str, Any] | None:
-    """Return deterministic PEFT/adapter identity without tensor payloads."""
-
-    peft_config = getattr(model, "peft_config", None)
-    if not isinstance(peft_config, Mapping) or not peft_config:
-        return None
-    configurations: dict[str, Any] = {}
-    for name, config in sorted(peft_config.items(), key=lambda item: str(item[0])):
-        to_dict = getattr(config, "to_dict", None)
-        if callable(to_dict):
-            value = to_dict()
-        else:
-            try:
-                value = vars(config)
-            except TypeError:
-                value = config
-        configurations[str(name)] = _fingerprint_jsonable(value)
-    active_adapters = getattr(model, "active_adapters", None)
-    if callable(active_adapters):
-        active_adapters = active_adapters()
-    return {
-        "active": _fingerprint_jsonable(active_adapters),
-        "configurations": configurations,
-    }
-
-
-def _execution_identity_metadata(model: Any) -> dict[str, Any]:
-    """Capture runtime policy that can change persisted numerical results."""
-
-    parameter_dtypes = sorted(
-        {
-            str(parameter.dtype).removeprefix("torch.")
-            for parameter in getattr(model, "parameters", lambda: ())()
-        }
-    )
-    return {
-        "device": _model_device(model).type,
-        "hf_device_map": _fingerprint_jsonable(getattr(model, "hf_device_map", None)),
-        "parameter_dtypes": parameter_dtypes,
-        "software": _software_versions(),
-    }
-
-
-def _first_metadata_value(*values: Any) -> Any:
-    for value in values:
-        if isinstance(value, str):
-            if value.strip():
-                return value
-        elif value is not None:
-            return value
-    return None
-
-
-def _model_identity_metadata(model: Any) -> dict[str, Any]:
-    """Resolve model and checkpoint identity, including local artifact fallbacks."""
-
-    config = getattr(model, "config", None)
-    checkpoint_revision = _first_metadata_value(
-        getattr(config, "fastplms_checkpoint_revision", None),
-        getattr(config, "_commit_hash", None),
-    )
-    return {
-        "model_id": _first_metadata_value(
-            getattr(config, "fastplms_model_id", None),
-            getattr(config, "_name_or_path", None),
-        ),
-        "model_revision": _first_metadata_value(
-            getattr(config, "_commit_hash", None),
-            checkpoint_revision,
-        ),
-        "checkpoint_repo_id": getattr(config, "fastplms_checkpoint_repo_id", None),
-        "checkpoint_revision": checkpoint_revision,
-        "checkpoint_hash": _first_metadata_value(
-            getattr(model, "checkpoint_hash", None),
-            getattr(config, "checkpoint_hash", None),
-            getattr(config, "fastplms_checkpoint_hash", None),
-        ),
-        "weights_revision": getattr(config, "fastplms_weights_revision", None),
-        "runtime_revision": getattr(config, "fastplms_runtime_revision", None),
-        "source_tree_sha256": getattr(config, "fastplms_source_tree_sha256", None),
-        "runtime_bundle_sha256": getattr(config, "fastplms_runtime_bundle_sha256", None),
-    }
-
-
-def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
-    """Yield X in logical row-major order without materializing a full copy."""
-
-    # X: (...)
-    if X.numel() == 0:
-        return
-    if X.ndim == 0:
-        yield X
-        return
-    trailing_elements = 1
-    for size in X.shape[1:]:
-        trailing_elements *= int(size)
-    if trailing_elements <= max_elements:
-        rows_per_chunk = max(1, max_elements // trailing_elements)
-        for start in range(0, X.shape[0], rows_per_chunk):
-            yield X[start : start + rows_per_chunk]  # (chunk_rows, ...)
-        return
-    for row in X:
-        yield from _bounded_tensor_chunks(row, max_elements)
-
-
-def _model_state_sha256(model: Any) -> str:
-    """Hash named parameters and persistent buffers using bounded CPU copies."""
-
-    # Never cache this digest from tensor identity or ``Tensor._version``.
-    # ``Parameter.data`` and independent tensor aliases can mutate shared storage
-    # without changing either signal, while persisted resume identity must bind
-    # the authoritative bytes visible at the start of this run.
-    state = model.state_dict(keep_vars=True)
-    digest = hashlib.sha256()
-    for name, value in sorted(state.items()):
-        if not isinstance(value, Tensor):
-            raise TypeError(f"Model state entry {name!r} is not a tensor.")
-        if value.is_meta:
-            raise ValueError(
-                f"Cannot fingerprint meta-device model state entry {name!r}; pass "
-                "model_state_fingerprint with a caller-owned state identity."
-            )
-        header = json.dumps(
-            {
-                "name": name,
-                "dtype": str(value.dtype).removeprefix("torch."),
-                "shape": list(value.shape),
-            },
-            sort_keys=True,
-            separators=(",", ":"),
-        ).encode()
-        digest.update(len(header).to_bytes(8, "big"))
-        digest.update(header)
-        max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size())
-        for chunk in _bounded_tensor_chunks(value.detach(), max_elements):
-            cpu_chunk = chunk.to(device="cpu").contiguous()  # chunk.shape
-            digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes())
-    return digest.hexdigest()
-
-
-def _input_sha256(records: Iterable[EmbeddingInput]) -> str:
-    """Hash an ordered input stream without constructing a duplicate JSON payload."""
-
-    precomputed = getattr(records, "input_fingerprint", None)
-    if isinstance(precomputed, str):
-        return precomputed
-    digest = hashlib.sha256()
-    count = 0
-    for record in records:
-        count += 1
-        for value in (record.id, record.sequence):
-            encoded = value.encode("utf-8")
-            digest.update(len(encoded).to_bytes(8, "big"))
-            digest.update(encoded)
-    digest.update(count.to_bytes(8, "big"))
-    return digest.hexdigest()
-
-
-def _run_fingerprint(
-    model: Any,
-    records: Sequence[EmbeddingInput],
-    *,
-    pooling: Sequence[str],
-    full_embeddings: bool,
-    max_length: int | None,
-    truncate: bool,
-    dtype: torch.dtype | None,
-    model_kwargs: dict[str, Any],
-    tokenizer_metadata: dict[str, Any],
-    model_state_fingerprint: str | None,
-    persist_output: bool,
-    embedding_context: Mapping[str, Any],
-    batch_size: int,
-    batch_window_size: int,
-    max_tokens_per_batch: int | None,
-) -> tuple[str, str, str | None, str]:
-    input_fingerprint = _input_sha256(records)
-    attention_backend = _attention_backend(model)
-    model_identity = _model_identity_metadata(model)
-    if model_state_fingerprint is None and persist_output:
-        resolved_model_state_fingerprint = _model_state_sha256(model)
-        model_state_fingerprint_source = "computed"
-    elif model_state_fingerprint is not None:
-        resolved_model_state_fingerprint = model_state_fingerprint.strip()
-        if not resolved_model_state_fingerprint:
-            raise ValueError("model_state_fingerprint must not be empty.")
-        model_state_fingerprint_source = "caller"
-    else:
-        resolved_model_state_fingerprint = None
-        model_state_fingerprint_source = "not-computed"
-    payload = {
-        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-        "input_fingerprint": input_fingerprint,
-        "model_state_fingerprint": resolved_model_state_fingerprint,
-        "model_state_fingerprint_source": model_state_fingerprint_source,
-        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
-        **model_identity,
-        "attention_backend": attention_backend,
-        "attention_kernel": _attention_kernel_metadata(attention_backend),
-        "layer": repr(
-            getattr(model, "embedding_layer", model_kwargs.get("hidden_state_index", -1))
-        ),
-        "projection": getattr(model, "embedding_projection", None),
-        "esmc_source": getattr(model, "_esmc_source", None),
-        "esmc_revision": getattr(model, "_esmc_source_revision", None),
-        "esmc_files": getattr(model, "_esmc_source_files", None),
-        "token_policy": getattr(model, "embedding_token_policy", None),
-        "tokenizer": tokenizer_metadata,
-        "adapter": _adapter_identity_metadata(model),
-        "execution": _execution_identity_metadata(model),
-        "embedding_context": _fingerprint_jsonable(embedding_context),
-        "pooling": list(pooling),
-        "full_embeddings": full_embeddings,
-        "max_length": max_length,
-        "truncate": truncate,
-        "dtype": str(dtype) if dtype is not None else None,
-        "batching": {
-            "batch_size": batch_size,
-            "batch_window_size": batch_window_size,
-            "max_tokens_per_batch": max_tokens_per_batch,
-            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
-        },
-        "model_kwargs": {
-            key: _fingerprint_jsonable(value) for key, value in sorted(model_kwargs.items())
-        },
-        "residue_mask_policy": "attention-mask-minus-special-tokens",
-    }
-    run_fingerprint = hashlib.sha256(
-        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
-    ).hexdigest()
-    return (
-        input_fingerprint,
-        run_fingerprint,
-        resolved_model_state_fingerprint,
-        model_state_fingerprint_source,
-    )
-
-
-def _ordered_string_sha256(values: Sequence[str]) -> str:
-    digest = hashlib.sha256()
-    for value in values:
-        encoded = value.encode("utf-8")
-        digest.update(len(encoded).to_bytes(8, "big"))
-        digest.update(encoded)
-    digest.update(len(values).to_bytes(8, "big"))
-    return digest.hexdigest()
-
-
-def _embedding_context(
-    model: Any,
-    records: Sequence[EmbeddingInput],
-    *,
-    hidden_state_source: str,
-    decoder_inputs: Sequence[str] | None,
-    decoder_input_ids: Tensor | None,
-    decoder_attention_mask: Tensor | None,
-    model_kwargs: Mapping[str, Any],
-) -> tuple[dict[str, Any], tuple[str, ...] | None]:
-    if hidden_state_source not in {"encoder", "decoder"}:
-        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
-    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
-    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
-        raise TypeError("hidden_state_index must be an integer.")
-    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
-    if not isinstance(store_all_hidden_states, bool):
-        raise TypeError("store_all_hidden_states must be a boolean.")
-    normalized_decoder_inputs: tuple[str, ...] | None = None
-    has_decoder_inputs = decoder_inputs is not None
-    has_decoder_ids = decoder_input_ids is not None
-    if hidden_state_source == "encoder":
-        if has_decoder_inputs or has_decoder_ids or decoder_attention_mask is not None:
-            raise ValueError("Decoder inputs are only valid when hidden_state_source='decoder'.")
-    else:
-        if has_decoder_inputs == has_decoder_ids:
-            raise ValueError(
-                "Decoder embedding requires exactly one of decoder_inputs or decoder_input_ids."
-            )
-    decoder_input_fingerprint: str | None = None
-    if decoder_inputs is not None:
-        if isinstance(decoder_inputs, (str, bytes)) or not isinstance(decoder_inputs, Sequence):
-            raise TypeError("decoder_inputs must be an aligned sequence of strings.")
-        normalized_decoder_inputs = tuple(decoder_inputs)
-        if not all(isinstance(value, str) and value for value in normalized_decoder_inputs):
-            raise ValueError("decoder_inputs must contain non-empty strings.")
-        if len(normalized_decoder_inputs) != len(records):
-            raise ValueError("decoder_inputs must align one-to-one with embedding inputs.")
-        decoder_input_fingerprint = _ordered_string_sha256(normalized_decoder_inputs)
-        if decoder_attention_mask is not None:
-            raise ValueError("decoder_attention_mask requires decoder_input_ids.")
-    if decoder_input_ids is not None:
-        if not isinstance(decoder_input_ids, Tensor) or decoder_input_ids.ndim != 2:
-            raise ValueError("decoder_input_ids must have shape (batch, sequence).")
-        if decoder_input_ids.shape[0] != len(records):
-            raise ValueError("decoder_input_ids must align one-to-one with embedding inputs.")
-        if decoder_input_ids.dtype == torch.bool or decoder_input_ids.is_floating_point():
-            raise TypeError("decoder_input_ids must use an integer token dtype.")
-        decoder_input_fingerprint = tensor_sha256(decoder_input_ids)
-    decoder_mask_fingerprint: str | None = None
-    if decoder_attention_mask is not None:
-        if not isinstance(decoder_attention_mask, Tensor):
-            raise TypeError("decoder_attention_mask must be a tensor.")
-        if decoder_input_ids is None or decoder_attention_mask.shape != decoder_input_ids.shape:
-            raise ValueError("decoder_attention_mask must match decoder_input_ids shape.")
-        decoder_mask_fingerprint = tensor_sha256(decoder_attention_mask)
-
-    context: dict[str, Any] = {
-        "hidden_state_source": hidden_state_source,
-        "hidden_state_index": hidden_state_index,
-        "store_all_hidden_states": store_all_hidden_states,
-        "decoder_input_fingerprint": decoder_input_fingerprint,
-        "decoder_attention_mask_fingerprint": decoder_mask_fingerprint,
-        "decoder_alignment": "input-position" if hidden_state_source == "decoder" else None,
-    }
-    metadata_hook = getattr(model, "_embedding_metadata", None)
-    model_metadata: Mapping[str, Any] | None = None
-    if callable(metadata_hook):
-        model_metadata = metadata_hook(**context)
-        if not isinstance(model_metadata, Mapping):
-            raise TypeError("_embedding_metadata must return a mapping.")
-        context["model_embedding"] = _fingerprint_jsonable(model_metadata)
-    if hidden_state_source == "decoder":
-        has_decoder_batch = callable(getattr(model, "_embedding_batch", None))
-        declares_decoder_stack = (
-            model_metadata is not None and model_metadata.get("hidden_state_stack") == "decoder"
-        )
-        if not has_decoder_batch or not declares_decoder_stack:
-            raise ValueError(
-                f"{model.__class__.__name__} does not declare decoder embedding support."
-            )
-    return context, normalized_decoder_inputs
+
+
+_RUN_FINGERPRINT_SCHEMA_VERSION = 3
+_MODEL_STATE_HASH_CHUNK_BYTES = 16 * 1024**2
+
+
+def _model_device(model: Any) -> torch.device:
+    try:
+        return torch.device(next(model.parameters()).device)
+    except (AttributeError, StopIteration):
+        return torch.device("cpu")
+
+
+def _attention_backend(model: Any) -> str | None:
+    config = getattr(model, "config", None)
+    for name in ("_attn_implementation", "attn_implementation", "attn_backend"):
+        value = getattr(config, name, None)
+        if value:
+            return str(value)
+    return None
+
+
+def _attention_kernel_metadata(backend: str | None) -> dict[str, Any] | None:
+    if backend not in {"flash_attention_2", "flash_attention_3"}:
+        return None
+    from fastplms.registry import get_model_registry
+
+    spec = get_model_registry().attention_kernels[backend]
+    return {
+        "repository": spec.repository,
+        "revision": spec.revision,
+        "version": spec.version,
+        "expected_variant": spec.expected_variant,
+        "dtypes": list(spec.dtypes),
+    }
+
+
+def _fingerprint_jsonable(value: Any) -> Any:
+    if isinstance(value, Mapping):
+        return {str(key): _fingerprint_jsonable(item) for key, item in value.items()}
+    if isinstance(value, (list, tuple)):
+        return [_fingerprint_jsonable(item) for item in value]
+    if isinstance(value, (set, frozenset)):
+        return sorted((_fingerprint_jsonable(item) for item in value), key=repr)
+    if isinstance(value, Path):
+        return str(value)
+    if isinstance(value, Tensor):
+        return {
+            "dtype": str(value.dtype).removeprefix("torch."),
+            "shape": list(value.shape),
+            "sha256": tensor_sha256(value),
+        }
+    if isinstance(value, torch.dtype):
+        return str(value).removeprefix("torch.")
+    if isinstance(value, torch.device):
+        return str(value)
+    if value is None or isinstance(value, (str, int, float, bool)):
+        return value
+    return {
+        "class": f"{value.__class__.__module__}.{value.__class__.__qualname__}",
+        "value": str(value),
+    }
+
+
+def _tokenizer_content_sha256(tokenizer: Any) -> str:
+    content: dict[str, Any] = {
+        "init_kwargs": getattr(tokenizer, "init_kwargs", None),
+        "special_tokens_map": getattr(tokenizer, "special_tokens_map", None),
+        "model_max_length": getattr(tokenizer, "model_max_length", None),
+        "padding_side": getattr(tokenizer, "padding_side", None),
+        "truncation_side": getattr(tokenizer, "truncation_side", None),
+    }
+    get_vocab = getattr(tokenizer, "get_vocab", None)
+    if callable(get_vocab):
+        content["vocabulary"] = get_vocab()
+    get_added_vocab = getattr(tokenizer, "get_added_vocab", None)
+    if callable(get_added_vocab):
+        content["added_vocabulary"] = get_added_vocab()
+    backend = getattr(tokenizer, "backend_tokenizer", None)
+    backend_to_str = getattr(backend, "to_str", None)
+    if callable(backend_to_str):
+        content["backend"] = backend_to_str()
+    serialized = json.dumps(
+        _fingerprint_jsonable(content),
+        sort_keys=True,
+        separators=(",", ":"),
+        ensure_ascii=False,
+    ).encode()
+    return hashlib.sha256(serialized).hexdigest()
+
+
+def _tokenizer_metadata(model: Any, tokenizer: Any | None) -> dict[str, Any]:
+    resolved = tokenizer if tokenizer is not None else getattr(model, "tokenizer", None)
+    if resolved is None:
+        # Raw-sequence families such as E1 retain their loader context on the
+        # model/encoder rather than exposing a Transformers tokenizer. Bind the
+        # non-secret source policy to resume identity without serializing a Hub
+        # token or forcing lazy tokenizer initialization.
+        for candidate in (model, getattr(model, "model", None)):
+            settings = getattr(candidate, "__dict__", {}).get("_fastplms_tokenizer_kwargs")
+            if isinstance(settings, Mapping):
+                token_value = settings.get("token")
+                return {
+                    "mode": "native-sequence",
+                    "source": (
+                        str(settings.get("tokenizer_source"))
+                        if settings.get("tokenizer_source") is not None
+                        else None
+                    ),
+                    "revision": settings.get("revision"),
+                    "cache_dir": (
+                        str(settings.get("cache_dir"))
+                        if settings.get("cache_dir") is not None
+                        else None
+                    ),
+                    "local_files_only": bool(settings.get("local_files_only", False)),
+                    "token_policy": (
+                        "disabled"
+                        if token_value is False
+                        else "provided"
+                        if token_value is not None
+                        else "default"
+                    ),
+                }
+        return {"mode": "native-sequence"}
+    return {
+        "mode": "tokenizer",
+        "class": f"{resolved.__class__.__module__}.{resolved.__class__.__qualname__}",
+        "name_or_path": getattr(resolved, "name_or_path", None),
+        "vocab_size": getattr(resolved, "vocab_size", None),
+        "special_token_ids": list(getattr(resolved, "all_special_ids", ())),
+        "content_sha256": _tokenizer_content_sha256(resolved),
+    }
+
+
+def _software_versions() -> dict[str, str | None]:
+    try:
+        import fastplms
+
+        fastplms_version = fastplms.__version__
+    except (AttributeError, ImportError):
+        fastplms_version = None
+    try:
+        import safetensors
+
+        safetensors_version = safetensors.__version__
+    except ImportError:
+        safetensors_version = None
+    try:
+        import transformers
+
+        transformers_version = transformers.__version__
+    except ImportError:
+        transformers_version = None
+    return {
+        "fastplms": fastplms_version,
+        "python": platform.python_version(),
+        "safetensors": safetensors_version,
+        "torch": torch.__version__,
+        "torch_cuda": torch.version.cuda,
+        "transformers": transformers_version,
+    }
+
+
+def _adapter_identity_metadata(model: Any) -> dict[str, Any] | None:
+    """Return deterministic PEFT/adapter identity without tensor payloads."""
+
+    peft_config = getattr(model, "peft_config", None)
+    if not isinstance(peft_config, Mapping) or not peft_config:
+        return None
+    configurations: dict[str, Any] = {}
+    for name, config in sorted(peft_config.items(), key=lambda item: str(item[0])):
+        to_dict = getattr(config, "to_dict", None)
+        if callable(to_dict):
+            value = to_dict()
+        else:
+            try:
+                value = vars(config)
+            except TypeError:
+                value = config
+        configurations[str(name)] = _fingerprint_jsonable(value)
+    active_adapters = getattr(model, "active_adapters", None)
+    if callable(active_adapters):
+        active_adapters = active_adapters()
+    return {
+        "active": _fingerprint_jsonable(active_adapters),
+        "configurations": configurations,
+    }
+
+
+def _execution_identity_metadata(model: Any) -> dict[str, Any]:
+    """Capture runtime policy that can change persisted numerical results."""
+
+    parameter_dtypes = sorted(
+        {
+            str(parameter.dtype).removeprefix("torch.")
+            for parameter in getattr(model, "parameters", lambda: ())()
+        }
+    )
+    return {
+        "device": _model_device(model).type,
+        "hf_device_map": _fingerprint_jsonable(getattr(model, "hf_device_map", None)),
+        "parameter_dtypes": parameter_dtypes,
+        "software": _software_versions(),
+    }
+
+
+def _first_metadata_value(*values: Any) -> Any:
+    for value in values:
+        if isinstance(value, str):
+            if value.strip():
+                return value
+        elif value is not None:
+            return value
+    return None
+
+
+def _model_identity_metadata(model: Any) -> dict[str, Any]:
+    """Resolve model and checkpoint identity, including local artifact fallbacks."""
+
+    config = getattr(model, "config", None)
+    checkpoint_revision = _first_metadata_value(
+        getattr(config, "fastplms_checkpoint_revision", None),
+        getattr(config, "_commit_hash", None),
+    )
+    return {
+        "model_id": _first_metadata_value(
+            getattr(config, "fastplms_model_id", None),
+            getattr(config, "_name_or_path", None),
+        ),
+        "model_revision": _first_metadata_value(
+            getattr(config, "_commit_hash", None),
+            checkpoint_revision,
+        ),
+        "checkpoint_repo_id": getattr(config, "fastplms_checkpoint_repo_id", None),
+        "checkpoint_revision": checkpoint_revision,
+        "checkpoint_hash": _first_metadata_value(
+            getattr(model, "checkpoint_hash", None),
+            getattr(config, "checkpoint_hash", None),
+            getattr(config, "fastplms_checkpoint_hash", None),
+        ),
+        "weights_revision": getattr(config, "fastplms_weights_revision", None),
+        "runtime_revision": getattr(config, "fastplms_runtime_revision", None),
+        "source_tree_sha256": getattr(config, "fastplms_source_tree_sha256", None),
+        "runtime_bundle_sha256": getattr(config, "fastplms_runtime_bundle_sha256", None),
+    }
+
+
+def _bounded_tensor_chunks(X: Tensor, max_elements: int) -> Iterable[Tensor]:
+    """Yield X in logical row-major order without materializing a full copy."""
+
+    # X: (...)
+    if X.numel() == 0:
+        return
+    if X.ndim == 0:
+        yield X
+        return
+    trailing_elements = 1
+    for size in X.shape[1:]:
+        trailing_elements *= int(size)
+    if trailing_elements <= max_elements:
+        rows_per_chunk = max(1, max_elements // trailing_elements)
+        for start in range(0, X.shape[0], rows_per_chunk):
+            yield X[start : start + rows_per_chunk]  # (chunk_rows, ...)
+        return
+    for row in X:
+        yield from _bounded_tensor_chunks(row, max_elements)
+
+
+def _model_state_sha256(model: Any) -> str:
+    """Hash named parameters and persistent buffers using bounded CPU copies."""
+
+    # Never cache this digest from tensor identity or ``Tensor._version``.
+    # ``Parameter.data`` and independent tensor aliases can mutate shared storage
+    # without changing either signal, while persisted resume identity must bind
+    # the authoritative bytes visible at the start of this run.
+    state = model.state_dict(keep_vars=True)
+    digest = hashlib.sha256()
+    for name, value in sorted(state.items()):
+        if not isinstance(value, Tensor):
+            raise TypeError(f"Model state entry {name!r} is not a tensor.")
+        if value.is_meta:
+            raise ValueError(
+                f"Cannot fingerprint meta-device model state entry {name!r}; pass "
+                "model_state_fingerprint with a caller-owned state identity."
+            )
+        header = json.dumps(
+            {
+                "name": name,
+                "dtype": str(value.dtype).removeprefix("torch."),
+                "shape": list(value.shape),
+            },
+            sort_keys=True,
+            separators=(",", ":"),
+        ).encode()
+        digest.update(len(header).to_bytes(8, "big"))
+        digest.update(header)
+        max_elements = max(1, _MODEL_STATE_HASH_CHUNK_BYTES // value.element_size())
+        for chunk in _bounded_tensor_chunks(value.detach(), max_elements):
+            cpu_chunk = chunk.to(device="cpu").contiguous()  # chunk.shape
+            digest.update(cpu_chunk.reshape(-1).view(torch.uint8).numpy().tobytes())
+    return digest.hexdigest()
+
+
+def _input_sha256(records: Iterable[EmbeddingInput]) -> str:
+    """Hash an ordered input stream without constructing a duplicate JSON payload."""
+
+    precomputed = getattr(records, "input_fingerprint", None)
+    if isinstance(precomputed, str):
+        return precomputed
+    digest = hashlib.sha256()
+    count = 0
+    for record in records:
+        count += 1
+        for value in (record.id, record.sequence):
+            encoded = value.encode("utf-8")
+            digest.update(len(encoded).to_bytes(8, "big"))
+            digest.update(encoded)
+    digest.update(count.to_bytes(8, "big"))
+    return digest.hexdigest()
+
+
+def _run_fingerprint(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    pooling: Sequence[str],
+    full_embeddings: bool,
+    max_length: int | None,
+    truncate: bool,
+    dtype: torch.dtype | None,
+    model_kwargs: dict[str, Any],
+    tokenizer_metadata: dict[str, Any],
+    model_state_fingerprint: str | None,
+    persist_output: bool,
+    embedding_context: Mapping[str, Any],
+    batch_size: int,
+    batch_window_size: int,
+    max_tokens_per_batch: int | None,
+) -> tuple[str, str, str | None, str]:
+    input_fingerprint = _input_sha256(records)
+    attention_backend = _attention_backend(model)
+    model_identity = _model_identity_metadata(model)
+    if model_state_fingerprint is None and persist_output:
+        resolved_model_state_fingerprint = _model_state_sha256(model)
+        model_state_fingerprint_source = "computed"
+    elif model_state_fingerprint is not None:
+        resolved_model_state_fingerprint = model_state_fingerprint.strip()
+        if not resolved_model_state_fingerprint:
+            raise ValueError("model_state_fingerprint must not be empty.")
+        model_state_fingerprint_source = "caller"
+    else:
+        resolved_model_state_fingerprint = None
+        model_state_fingerprint_source = "not-computed"
+    payload = {
+        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+        "input_fingerprint": input_fingerprint,
+        "model_state_fingerprint": resolved_model_state_fingerprint,
+        "model_state_fingerprint_source": model_state_fingerprint_source,
+        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
+        **model_identity,
+        "attention_backend": attention_backend,
+        "attention_kernel": _attention_kernel_metadata(attention_backend),
+        "layer": repr(
+            getattr(model, "embedding_layer", model_kwargs.get("hidden_state_index", -1))
+        ),
+        "projection": getattr(model, "embedding_projection", None),
+        "esmc_source": getattr(model, "_esmc_source", None),
+        "esmc_revision": getattr(model, "_esmc_source_revision", None),
+        "esmc_files": getattr(model, "_esmc_source_files", None),
+        "token_policy": getattr(model, "embedding_token_policy", None),
+        "tokenizer": tokenizer_metadata,
+        "adapter": _adapter_identity_metadata(model),
+        "execution": _execution_identity_metadata(model),
+        "embedding_context": _fingerprint_jsonable(embedding_context),
+        "pooling": list(pooling),
+        "full_embeddings": full_embeddings,
+        "max_length": max_length,
+        "truncate": truncate,
+        "dtype": str(dtype) if dtype is not None else None,
+        "batching": {
+            "batch_size": batch_size,
+            "batch_window_size": batch_window_size,
+            "max_tokens_per_batch": max_tokens_per_batch,
+            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
+        },
+        "model_kwargs": {
+            key: _fingerprint_jsonable(value) for key, value in sorted(model_kwargs.items())
+        },
+        "residue_mask_policy": "attention-mask-minus-special-tokens",
+    }
+    run_fingerprint = hashlib.sha256(
+        json.dumps(payload, sort_keys=True, separators=(",", ":")).encode()
+    ).hexdigest()
+    return (
+        input_fingerprint,
+        run_fingerprint,
+        resolved_model_state_fingerprint,
+        model_state_fingerprint_source,
+    )
+
+
+def _ordered_string_sha256(values: Sequence[str]) -> str:
+    digest = hashlib.sha256()
+    for value in values:
+        encoded = value.encode("utf-8")
+        digest.update(len(encoded).to_bytes(8, "big"))
+        digest.update(encoded)
+    digest.update(len(values).to_bytes(8, "big"))
+    return digest.hexdigest()
+
+
+def _embedding_context(
+    model: Any,
+    records: Sequence[EmbeddingInput],
+    *,
+    hidden_state_source: str,
+    decoder_inputs: Sequence[str] | None,
+    decoder_input_ids: Tensor | None,
+    decoder_attention_mask: Tensor | None,
+    model_kwargs: Mapping[str, Any],
+) -> tuple[dict[str, Any], tuple[str, ...] | None]:
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
+    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
+    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
+        raise TypeError("hidden_state_index must be an integer.")
+    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
+    if not isinstance(store_all_hidden_states, bool):
+        raise TypeError("store_all_hidden_states must be a boolean.")
+    normalized_decoder_inputs: tuple[str, ...] | None = None
+    has_decoder_inputs = decoder_inputs is not None
+    has_decoder_ids = decoder_input_ids is not None
+    if hidden_state_source == "encoder":
+        if has_decoder_inputs or has_decoder_ids or decoder_attention_mask is not None:
+            raise ValueError("Decoder inputs are only valid when hidden_state_source='decoder'.")
+    else:
+        if has_decoder_inputs == has_decoder_ids:
+            raise ValueError(
+                "Decoder embedding requires exactly one of decoder_inputs or decoder_input_ids."
+            )
+    decoder_input_fingerprint: str | None = None
+    if decoder_inputs is not None:
+        if isinstance(decoder_inputs, (str, bytes)) or not isinstance(decoder_inputs, Sequence):
+            raise TypeError("decoder_inputs must be an aligned sequence of strings.")
+        normalized_decoder_inputs = tuple(decoder_inputs)
+        if not all(isinstance(value, str) and value for value in normalized_decoder_inputs):
+            raise ValueError("decoder_inputs must contain non-empty strings.")
+        if len(normalized_decoder_inputs) != len(records):
+            raise ValueError("decoder_inputs must align one-to-one with embedding inputs.")
+        decoder_input_fingerprint = _ordered_string_sha256(normalized_decoder_inputs)
+        if decoder_attention_mask is not None:
+            raise ValueError("decoder_attention_mask requires decoder_input_ids.")
+    if decoder_input_ids is not None:
+        if not isinstance(decoder_input_ids, Tensor) or decoder_input_ids.ndim != 2:
+            raise ValueError("decoder_input_ids must have shape (batch, sequence).")
+        if decoder_input_ids.shape[0] != len(records):
+            raise ValueError("decoder_input_ids must align one-to-one with embedding inputs.")
+        if decoder_input_ids.dtype == torch.bool or decoder_input_ids.is_floating_point():
+            raise TypeError("decoder_input_ids must use an integer token dtype.")
+        decoder_input_fingerprint = tensor_sha256(decoder_input_ids)
+    decoder_mask_fingerprint: str | None = None
+    if decoder_attention_mask is not None:
+        if not isinstance(decoder_attention_mask, Tensor):
+            raise TypeError("decoder_attention_mask must be a tensor.")
+        if decoder_input_ids is None or decoder_attention_mask.shape != decoder_input_ids.shape:
+            raise ValueError("decoder_attention_mask must match decoder_input_ids shape.")
+        decoder_mask_fingerprint = tensor_sha256(decoder_attention_mask)
+
+    context: dict[str, Any] = {
+        "hidden_state_source": hidden_state_source,
+        "hidden_state_index": hidden_state_index,
+        "store_all_hidden_states": store_all_hidden_states,
+        "decoder_input_fingerprint": decoder_input_fingerprint,
+        "decoder_attention_mask_fingerprint": decoder_mask_fingerprint,
+        "decoder_alignment": "input-position" if hidden_state_source == "decoder" else None,
+    }
+    metadata_hook = getattr(model, "_embedding_metadata", None)
+    model_metadata: Mapping[str, Any] | None = None
+    if callable(metadata_hook):
+        model_metadata = metadata_hook(**context)
+        if not isinstance(model_metadata, Mapping):
+            raise TypeError("_embedding_metadata must return a mapping.")
+        context["model_embedding"] = _fingerprint_jsonable(model_metadata)
+    if hidden_state_source == "decoder":
+        has_decoder_batch = callable(getattr(model, "_embedding_batch", None))
+        declares_decoder_stack = (
+            model_metadata is not None and model_metadata.get("hidden_state_stack") == "decoder"
+        )
+        if not has_decoder_batch or not declares_decoder_stack:
+            raise ValueError(
+                f"{model.__class__.__name__} does not declare decoder embedding support."
+            )
+    return context, normalized_decoder_inputs

fastplms/embeddings/inputs.py

@@ -1,263 +1,264 @@
-"""Normalize ordered inputs and plan bounded windows without retaining a full stream."""
-
-from __future__ import annotations
-
-import hashlib
-import sqlite3
-import tempfile
-from collections.abc import Iterable, Iterator, Mapping, Sequence
-from pathlib import Path
-from typing import overload
-
-from .types import EmbeddingInput
-
-
-def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
-    """Yield FASTA records in source order without reading the file into memory."""
-
-    identifier: str | None = None
-    sequence_parts: list[str] = []
-    found_record = False
-    with Path(path).open("r", encoding="utf-8") as handle:
-        for line_number, raw_line in enumerate(handle, start=1):
-            line = raw_line.strip()
-            if not line:
-                continue
-            if line.startswith(">"):
-                if identifier is not None:
-                    found_record = True
-                    yield EmbeddingInput(identifier, "".join(sequence_parts))
-                identifier = line[1:].strip().split(maxsplit=1)[0]
-                if not identifier:
-                    raise ValueError(f"Missing FASTA identifier on line {line_number}.")
-                sequence_parts = []
-            else:
-                if identifier is None:
-                    raise ValueError(
-                        f"Sequence data precedes the first FASTA header on line {line_number}."
-                    )
-                sequence_parts.append("".join(line.split()))
-    if identifier is not None:
-        found_record = True
-        yield EmbeddingInput(identifier, "".join(sequence_parts))
-    if not found_record:
-        raise ValueError(f"No FASTA records found in {path}.")
-
-
-def parse_fasta(path: str | Path) -> list[EmbeddingInput]:
-    """Parse FASTA records while preserving identifiers, order, and duplicates."""
-
-    return list(iter_fasta(path))
-
-
-def _normalize_input_item(
-    position: int,
-    item: str | EmbeddingInput | tuple[str, str],
-) -> EmbeddingInput:
-    if isinstance(item, EmbeddingInput):
-        return item
-    if isinstance(item, str):
-        return EmbeddingInput(str(position), item)
-    if isinstance(item, tuple) and len(item) == 2:
-        return EmbeddingInput(str(item[0]), str(item[1]))
-    raise TypeError(
-        "inputs must contain sequences, EmbeddingInput values, or (id, sequence) tuples."
-    )
-
-
-class _InputSpool(Sequence[EmbeddingInput]):
-    """Immutable disk-backed normalized inputs with an incremental digest."""
-
-    def __init__(
-        self,
-        values: Iterable[str | EmbeddingInput | tuple[str, str]],
-    ) -> None:
-        self._temporary: tempfile.TemporaryDirectory[str] | None = tempfile.TemporaryDirectory(
-            prefix="fastplms-inputs-"
-        )
-        self.path = Path(self._temporary.name) / "inputs.sqlite"
-        self._connection: sqlite3.Connection | None = sqlite3.connect(self.path)
-        self._connection.execute(
-            "CREATE TABLE inputs ("
-            "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)"
-        )
-        digest = hashlib.sha256()
-        count = 0
-        pending: list[tuple[int, str, str]] = []
-        try:
-            for position, item in enumerate(values):
-                record = _normalize_input_item(position, item)
-                for value in (record.id, record.sequence):
-                    encoded = value.encode("utf-8")
-                    digest.update(len(encoded).to_bytes(8, "big"))
-                    digest.update(encoded)
-                pending.append((position, record.id, record.sequence))
-                count += 1
-                if len(pending) == 1_024:
-                    self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
-                    pending.clear()
-            if pending:
-                self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
-            if count == 0:
-                raise ValueError("inputs must contain at least one sequence.")
-            self._connection.commit()
-            self._connection.close()
-            self._connection = sqlite3.connect(
-                f"{self.path.resolve().as_uri()}?mode=ro",
-                uri=True,
-            )
-        except BaseException:
-            self.close()
-            raise
-        digest.update(count.to_bytes(8, "big"))
-        self.input_fingerprint = digest.hexdigest()
-        self._count = count
-
-    def _require_connection(self) -> sqlite3.Connection:
-        if self._connection is None:
-            raise RuntimeError("Input spool is closed.")
-        return self._connection
-
-    def __len__(self) -> int:
-        return self._count
-
-    def __iter__(self) -> Iterator[EmbeddingInput]:
-        cursor = self._require_connection().execute(
-            "SELECT input_id, sequence FROM inputs ORDER BY position"
-        )
-        while rows := cursor.fetchmany(1_024):
-            for input_id, sequence in rows:
-                yield EmbeddingInput(input_id, sequence)
-
-    @overload
-    def __getitem__(self, index: int, /) -> EmbeddingInput: ...
-
-    @overload
-    def __getitem__(self, index: slice, /) -> list[EmbeddingInput]: ...
-
-    def __getitem__(self, index: int | slice) -> EmbeddingInput | list[EmbeddingInput]:
-        connection = self._require_connection()
-
-        if isinstance(index, slice):
-            start, stop, step = index.indices(self._count)
-            if step != 1:
-                return [self[position] for position in range(start, stop, step)]
-            rows = connection.execute(
-                "SELECT input_id, sequence FROM inputs "
-                "WHERE position >= ? AND position < ? ORDER BY position",
-                (start, stop),
-            ).fetchall()
-            return [EmbeddingInput(input_id, sequence) for input_id, sequence in rows]
-        position = index + self._count if index < 0 else index
-        if position < 0 or position >= self._count:
-            raise IndexError(index)
-        row = connection.execute(
-            "SELECT input_id, sequence FROM inputs WHERE position = ?", (position,)
-        ).fetchone()
-        if row is None:
-            raise IndexError(index)
-        return EmbeddingInput(row[0], row[1])
-
-    def close(self) -> None:
-        connection = getattr(self, "_connection", None)
-        if connection is not None:
-            connection.close()
-            self._connection = None
-        temporary = getattr(self, "_temporary", None)
-        if temporary is not None:
-            temporary.cleanup()
-            self._temporary = None
-
-    def __del__(self) -> None:
-        self.close()
-
-
-def _normalize_inputs(
-    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
-    *,
-    disk_backed: bool,
-) -> Sequence[EmbeddingInput]:
-    is_fasta_path = isinstance(inputs, Path)
-    if isinstance(inputs, str):
-        try:
-            is_fasta_path = Path(inputs).is_file()
-        except OSError:
-            is_fasta_path = False
-    should_spool = disk_backed or is_fasta_path or not isinstance(inputs, (str, Sequence, Mapping))
-    values: Iterable[str | EmbeddingInput | tuple[str, str]]
-    if isinstance(inputs, Path):
-        values = iter_fasta(inputs)
-    elif isinstance(inputs, str):
-        values = iter_fasta(inputs) if is_fasta_path else [inputs]
-    elif isinstance(inputs, Mapping):
-        values = inputs.items()
-    else:
-        values = inputs
-    if should_spool:
-        return _InputSpool(values)
-    records: list[EmbeddingInput] = []
-    for position, item in enumerate(values):
-        records.append(_normalize_input_item(position, item))
-    if not records:
-        raise ValueError("inputs must contain at least one sequence.")
-    return records
-
-
-def _validate_untruncated_lengths(
-    records: Sequence[EmbeddingInput],
-    *,
-    max_length: int | None,
-    truncate: bool,
-) -> None:
-    """Fail before inference when a biological-residue limit would be exceeded."""
-
-    if max_length is None or truncate:
-        return
-    for position, record in enumerate(records):
-        residue_count = len(record.sequence)
-        if residue_count > max_length:
-            raise ValueError(
-                f"Input at position {position} with id {record.id!r} has "
-                f"{residue_count} biological residues, exceeding max_length={max_length} "
-                "while truncate=False."
-            )
-
-
-def _planned_batches(
-    records: Sequence[EmbeddingInput],
-    positions: range,
-    *,
-    batch_size: int,
-    max_tokens_per_batch: int | None,
-    max_length: int | None,
-    truncate: bool,
-) -> Iterator[list[int]]:
-    """Length-bucket one bounded window while retaining stable output positions."""
-
-    def effective_length(position: int) -> int:
-        length = len(records[position].sequence)
-        return min(length, max_length) if truncate and max_length is not None else length
-
-    ordered = sorted(positions, key=lambda position: (-effective_length(position), position))
-    batch: list[int] = []
-    longest = 0
-    for position in ordered:
-        length = effective_length(position)
-        if max_tokens_per_batch is not None and length > max_tokens_per_batch:
-            raise ValueError(
-                f"Input at position {position} has {length} residues, exceeding "
-                f"max_tokens_per_batch={max_tokens_per_batch}."
-            )
-        candidate_longest = max(longest, length)
-        exceeds_tokens = (
-            max_tokens_per_batch is not None
-            and candidate_longest * (len(batch) + 1) > max_tokens_per_batch
-        )
-        if batch and (len(batch) >= batch_size or exceeds_tokens):
-            yield batch
-            batch = []
-            longest = 0
-        batch.append(position)
-        longest = max(longest, length)
-    if batch:
-        yield batch
+"""Normalize ordered inputs and plan bounded windows without retaining a full stream."""
+
+from __future__ import annotations
+
+import hashlib
+import sqlite3
+import tempfile
+
+from collections.abc import Iterable, Iterator, Mapping, Sequence
+from pathlib import Path
+from typing import overload
+
+from .types import EmbeddingInput
+
+
+def iter_fasta(path: str | Path) -> Iterator[EmbeddingInput]:
+    """Yield FASTA records in source order without reading the file into memory."""
+
+    identifier: str | None = None
+    sequence_parts: list[str] = []
+    found_record = False
+    with Path(path).open("r", encoding="utf-8") as handle:
+        for line_number, raw_line in enumerate(handle, start=1):
+            line = raw_line.strip()
+            if not line:
+                continue
+            if line.startswith(">"):
+                if identifier is not None:
+                    found_record = True
+                    yield EmbeddingInput(identifier, "".join(sequence_parts))
+                identifier = line[1:].strip().split(maxsplit=1)[0]
+                if not identifier:
+                    raise ValueError(f"Missing FASTA identifier on line {line_number}.")
+                sequence_parts = []
+            else:
+                if identifier is None:
+                    raise ValueError(
+                        f"Sequence data precedes the first FASTA header on line {line_number}."
+                    )
+                sequence_parts.append("".join(line.split()))
+    if identifier is not None:
+        found_record = True
+        yield EmbeddingInput(identifier, "".join(sequence_parts))
+    if not found_record:
+        raise ValueError(f"No FASTA records found in {path}.")
+
+
+def parse_fasta(path: str | Path) -> list[EmbeddingInput]:
+    """Parse FASTA records while preserving identifiers, order, and duplicates."""
+
+    return list(iter_fasta(path))
+
+
+def _normalize_input_item(
+    position: int,
+    item: str | EmbeddingInput | tuple[str, str],
+) -> EmbeddingInput:
+    if isinstance(item, EmbeddingInput):
+        return item
+    if isinstance(item, str):
+        return EmbeddingInput(str(position), item)
+    if isinstance(item, tuple) and len(item) == 2:
+        return EmbeddingInput(str(item[0]), str(item[1]))
+    raise TypeError(
+        "inputs must contain sequences, EmbeddingInput values, or (id, sequence) tuples."
+    )
+
+
+class _InputSpool(Sequence[EmbeddingInput]):
+    """Immutable disk-backed normalized inputs with an incremental digest."""
+
+    def __init__(
+        self,
+        values: Iterable[str | EmbeddingInput | tuple[str, str]],
+    ) -> None:
+        self._temporary: tempfile.TemporaryDirectory[str] | None = tempfile.TemporaryDirectory(
+            prefix="fastplms-inputs-"
+        )
+        self.path = Path(self._temporary.name) / "inputs.sqlite"
+        self._connection: sqlite3.Connection | None = sqlite3.connect(self.path)
+        self._connection.execute(
+            "CREATE TABLE inputs ("
+            "position INTEGER PRIMARY KEY, input_id TEXT NOT NULL, sequence TEXT NOT NULL)"
+        )
+        digest = hashlib.sha256()
+        count = 0
+        pending: list[tuple[int, str, str]] = []
+        try:
+            for position, item in enumerate(values):
+                record = _normalize_input_item(position, item)
+                for value in (record.id, record.sequence):
+                    encoded = value.encode("utf-8")
+                    digest.update(len(encoded).to_bytes(8, "big"))
+                    digest.update(encoded)
+                pending.append((position, record.id, record.sequence))
+                count += 1
+                if len(pending) == 1_024:
+                    self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+                    pending.clear()
+            if pending:
+                self._connection.executemany("INSERT INTO inputs VALUES (?, ?, ?)", pending)
+            if count == 0:
+                raise ValueError("inputs must contain at least one sequence.")
+            self._connection.commit()
+            self._connection.close()
+            self._connection = sqlite3.connect(
+                f"{self.path.resolve().as_uri()}?mode=ro",
+                uri=True,
+            )
+        except BaseException:
+            self.close()
+            raise
+        digest.update(count.to_bytes(8, "big"))
+        self.input_fingerprint = digest.hexdigest()
+        self._count = count
+
+    def _require_connection(self) -> sqlite3.Connection:
+        if self._connection is None:
+            raise RuntimeError("Input spool is closed.")
+        return self._connection
+
+    def __len__(self) -> int:
+        return self._count
+
+    def __iter__(self) -> Iterator[EmbeddingInput]:
+        cursor = self._require_connection().execute(
+            "SELECT input_id, sequence FROM inputs ORDER BY position"
+        )
+        while rows := cursor.fetchmany(1_024):
+            for input_id, sequence in rows:
+                yield EmbeddingInput(input_id, sequence)
+
+    @overload
+    def __getitem__(self, index: int, /) -> EmbeddingInput: ...
+
+    @overload
+    def __getitem__(self, index: slice, /) -> list[EmbeddingInput]: ...
+
+    def __getitem__(self, index: int | slice) -> EmbeddingInput | list[EmbeddingInput]:
+        connection = self._require_connection()
+
+        if isinstance(index, slice):
+            start, stop, step = index.indices(self._count)
+            if step != 1:
+                return [self[position] for position in range(start, stop, step)]
+            rows = connection.execute(
+                "SELECT input_id, sequence FROM inputs "
+                "WHERE position >= ? AND position < ? ORDER BY position",
+                (start, stop),
+            ).fetchall()
+            return [EmbeddingInput(input_id, sequence) for input_id, sequence in rows]
+        position = index + self._count if index < 0 else index
+        if position < 0 or position >= self._count:
+            raise IndexError(index)
+        row = connection.execute(
+            "SELECT input_id, sequence FROM inputs WHERE position = ?", (position,)
+        ).fetchone()
+        if row is None:
+            raise IndexError(index)
+        return EmbeddingInput(row[0], row[1])
+
+    def close(self) -> None:
+        connection = getattr(self, "_connection", None)
+        if connection is not None:
+            connection.close()
+            self._connection = None
+        temporary = getattr(self, "_temporary", None)
+        if temporary is not None:
+            temporary.cleanup()
+            self._temporary = None
+
+    def __del__(self) -> None:
+        self.close()
+
+
+def _normalize_inputs(
+    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
+    *,
+    disk_backed: bool,
+) -> Sequence[EmbeddingInput]:
+    is_fasta_path = isinstance(inputs, Path)
+    if isinstance(inputs, str):
+        try:
+            is_fasta_path = Path(inputs).is_file()
+        except OSError:
+            is_fasta_path = False
+    should_spool = disk_backed or is_fasta_path or not isinstance(inputs, (str, Sequence, Mapping))
+    values: Iterable[str | EmbeddingInput | tuple[str, str]]
+    if isinstance(inputs, Path):
+        values = iter_fasta(inputs)
+    elif isinstance(inputs, str):
+        values = iter_fasta(inputs) if is_fasta_path else [inputs]
+    elif isinstance(inputs, Mapping):
+        values = inputs.items()
+    else:
+        values = inputs
+    if should_spool:
+        return _InputSpool(values)
+    records: list[EmbeddingInput] = []
+    for position, item in enumerate(values):
+        records.append(_normalize_input_item(position, item))
+    if not records:
+        raise ValueError("inputs must contain at least one sequence.")
+    return records
+
+
+def _validate_untruncated_lengths(
+    records: Sequence[EmbeddingInput],
+    *,
+    max_length: int | None,
+    truncate: bool,
+) -> None:
+    """Fail before inference when a biological-residue limit would be exceeded."""
+
+    if max_length is None or truncate:
+        return
+    for position, record in enumerate(records):
+        residue_count = len(record.sequence)
+        if residue_count > max_length:
+            raise ValueError(
+                f"Input at position {position} with id {record.id!r} has "
+                f"{residue_count} biological residues, exceeding max_length={max_length} "
+                "while truncate=False."
+            )
+
+
+def _planned_batches(
+    records: Sequence[EmbeddingInput],
+    positions: range,
+    *,
+    batch_size: int,
+    max_tokens_per_batch: int | None,
+    max_length: int | None,
+    truncate: bool,
+) -> Iterator[list[int]]:
+    """Length-bucket one bounded window while retaining stable output positions."""
+
+    def effective_length(position: int) -> int:
+        length = len(records[position].sequence)
+        return min(length, max_length) if truncate and max_length is not None else length
+
+    ordered = sorted(positions, key=lambda position: (-effective_length(position), position))
+    batch: list[int] = []
+    longest = 0
+    for position in ordered:
+        length = effective_length(position)
+        if max_tokens_per_batch is not None and length > max_tokens_per_batch:
+            raise ValueError(
+                f"Input at position {position} has {length} residues, exceeding "
+                f"max_tokens_per_batch={max_tokens_per_batch}."
+            )
+        candidate_longest = max(longest, length)
+        exceeds_tokens = (
+            max_tokens_per_batch is not None
+            and candidate_longest * (len(batch) + 1) > max_tokens_per_batch
+        )
+        if batch and (len(batch) >= batch_size or exceeds_tokens):
+            yield batch
+            batch = []
+            longest = 0
+        batch.append(position)
+        longest = max(longest, length)
+    if batch:
+        yield batch

fastplms/embeddings/output.py

@@ -1,215 +1,215 @@
-"""Resume validation and transactional publication of ordered embedding windows."""
-
-from __future__ import annotations
-
-from collections.abc import Sequence
-from pathlib import Path
-from typing import Any
-
-from .identity import _RUN_FINGERPRINT_SCHEMA_VERSION
-from .pooling import Pooler
-from .storage import (
-    SafetensorsStreamWriter,
-    append_sqlite_records,
-    initialize_sqlite_run,
-    load_result,
-    load_sqlite_result,
-    safetensors_result_exists,
-    save_result,
-    tensor_sha256,
-    update_sqlite_run_metadata,
-)
-from .types import EmbeddingInput, EmbeddingRecord, EmbeddingResult, LazyTensorReference
-
-
-def _output_exists(path: str | Path, format: str) -> bool:
-    path = Path(path)
-    if format == "sqlite":
-        return path.is_file()
-    return safetensors_result_exists(path)
-
-
-def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]:
-    tensor = record.tensor
-    if isinstance(tensor, LazyTensorReference):
-        dtype = tensor.dtype
-        shape = tensor.shape
-        digest = tensor.sha256
-    else:
-        dtype = str(tensor.dtype).removeprefix("torch.")
-        shape = tuple(tensor.shape)
-        digest = tensor_sha256(tensor)
-    return {
-        "position": position,
-        "id": record.id,
-        "dtype": dtype,
-        "shape": shape,
-        "sha256": digest,
-    }
-
-
-class EmbeddingOutput:
-    """Own the resumable prefix and the commit state of one output destination."""
-
-    def __init__(
-        self,
-        records: Sequence[EmbeddingInput],
-        *,
-        output: str | Path | None,
-        format: str,
-        resume: bool,
-        shard_size: int,
-        run_fingerprint: str,
-        input_fingerprint: str,
-        model_state_fingerprint: str | None,
-        model_state_fingerprint_source: str,
-        pooler: Pooler | None,
-        pooling_names: Sequence[str],
-    ) -> None:
-        self.output = output
-        self.format = format
-        self.shard_size = shard_size
-        self.completed: EmbeddingResult | None = None
-        output_already_exists = output is not None and _output_exists(output, format)
-        existing: EmbeddingResult | None = None
-        self.start_position = 0
-        if output is not None and resume and output_already_exists:
-            if format == "sqlite":
-                try:
-                    existing = load_sqlite_result(output, run_id=run_fingerprint)
-                except KeyError:
-                    existing = load_result(output, format=format)
-            else:
-                existing = load_result(output, format=format)
-            if existing.metadata.get("fingerprint_schema_version") != (
-                _RUN_FINGERPRINT_SCHEMA_VERSION
-            ):
-                raise ValueError(
-                    "Existing embeddings use an incompatible run fingerprint schema; "
-                    "choose another output or set resume=False."
-                )
-            if existing.metadata.get("run_fingerprint") != run_fingerprint:
-                raise ValueError(
-                    "Existing embeddings were produced by a different run fingerprint; "
-                    "choose another output or set resume=False."
-                )
-            if len(existing) > len(records):
-                raise ValueError(
-                    "Existing embeddings are not an ordered prefix of the requested inputs."
-                )
-            prefix_matches = all(
-                (observed.id, observed.sequence) == (expected.id, expected.sequence)
-                for expected, observed in zip(records, existing, strict=False)
-            )
-            if not prefix_matches:
-                raise ValueError(
-                    "Existing embeddings are not an ordered prefix of the requested inputs."
-                )
-            if len(existing) == len(records) and existing.metadata.get("complete", True):
-                self.completed = existing
-                return
-            self.start_position = len(existing)
-
-        self.sqlite_run_id: str | None = None
-        self.sqlite_replace_on_first_commit = False
-        self.sqlite_initial_metadata: dict[str, Any] | None = None
-        if output is not None and format == "sqlite":
-            self.sqlite_initial_metadata = {
-                "format_version": 1,
-                "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-                "run_fingerprint": run_fingerprint,
-                "input_fingerprint": input_fingerprint,
-                "model_state_fingerprint": model_state_fingerprint,
-                "model_state_fingerprint_source": model_state_fingerprint_source,
-                "complete": False,
-            }
-            self.sqlite_run_id = run_fingerprint
-            if not resume and output_already_exists:
-                try:
-                    load_sqlite_result(output, run_id=run_fingerprint)
-                except KeyError:
-                    pass
-                else:
-                    # Keep an exact prior run readable until replacement inference
-                    # has produced the first complete commit window.
-                    self.sqlite_replace_on_first_commit = True
-            if not self.sqlite_replace_on_first_commit:
-                initialize_sqlite_run(
-                    output,
-                    self.sqlite_initial_metadata,
-                    resume=resume,
-                )
-
-        stream_safetensors = output is not None and format == "safetensors"
-        self.output_records: list[EmbeddingRecord] = (
-            [] if self.sqlite_run_id is not None or stream_safetensors else list(existing or ())
-        )
-        self.output_descriptors: list[dict[str, Any]] | None = [] if output is None else None
-        self.pool_slices: dict[str, tuple[int, int]] = {}
-        if existing and pooler is not None:
-            pooled_width = existing[0].load_tensor().shape[-1]
-            if pooled_width % len(pooling_names) != 0:
-                raise ValueError("Stored pooled width is inconsistent with pooling metadata.")
-            self.pool_slices = pooler.output_slices(pooled_width // len(pooling_names))
-
-        self.safetensors_writer: SafetensorsStreamWriter | None = None
-        if stream_safetensors:
-            if output is None:
-                raise RuntimeError(
-                    "Safetensors streaming was enabled without an output destination."
-                )
-            transactional_overwrite = output_already_exists and not resume
-            self.safetensors_writer = SafetensorsStreamWriter(
-                output,
-                {
-                    "format_version": 1,
-                    "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-                    "run_fingerprint": run_fingerprint,
-                    "input_fingerprint": input_fingerprint,
-                    "model_state_fingerprint": model_state_fingerprint,
-                    "model_state_fingerprint_source": model_state_fingerprint_source,
-                    "complete": False,
-                },
-                shard_size=shard_size,
-                existing=existing or (),
-                reuse_existing=bool(resume and existing is not None),
-                publish_initial=not transactional_overwrite,
-                publish_incremental=not transactional_overwrite,
-            )
-
-    def append(self, window_start: int, new_records: list[EmbeddingRecord]) -> None:
-        """Commit a complete ordered window at the storage format's granularity."""
-
-        if self.output_descriptors is not None:
-            self.output_descriptors.extend(
-                _output_descriptor(window_start + offset, record)
-                for offset, record in enumerate(new_records)
-            )
-        if self.output is not None and self.sqlite_run_id is not None:
-            append_sqlite_records(
-                self.output,
-                self.sqlite_run_id,
-                window_start,
-                new_records,
-                replace_metadata=(
-                    self.sqlite_initial_metadata if self.sqlite_replace_on_first_commit else None
-                ),
-            )
-            self.sqlite_replace_on_first_commit = False
-        elif self.safetensors_writer is not None:
-            self.safetensors_writer.append(new_records)
-        else:
-            self.output_records.extend(new_records)
-
-    def finish(self, metadata: dict[str, Any]) -> EmbeddingResult:
-        """Publish completion only after every window has committed."""
-
-        if self.output is not None and self.sqlite_run_id is not None:
-            update_sqlite_run_metadata(self.output, self.sqlite_run_id, metadata)
-            return load_sqlite_result(self.output, run_id=self.sqlite_run_id)
-        if self.safetensors_writer is not None:
-            return self.safetensors_writer.publish(complete=True, metadata=metadata)
-        result = EmbeddingResult(self.output_records, metadata)
-        if self.output is not None:
-            return save_result(result, self.output, format=self.format, shard_size=self.shard_size)
-        return result
+"""Resume validation and transactional publication of ordered embedding windows."""
+
+from __future__ import annotations
+
+from collections.abc import Sequence
+from pathlib import Path
+from typing import Any
+
+from .identity import _RUN_FINGERPRINT_SCHEMA_VERSION
+from .pooling import Pooler
+from .storage import (
+    SafetensorsStreamWriter,
+    append_sqlite_records,
+    initialize_sqlite_run,
+    load_result,
+    load_sqlite_result,
+    safetensors_result_exists,
+    save_result,
+    tensor_sha256,
+    update_sqlite_run_metadata,
+)
+from .types import EmbeddingInput, EmbeddingRecord, EmbeddingResult, LazyTensorReference
+
+
+def _output_exists(path: str | Path, format: str) -> bool:
+    path = Path(path)
+    if format == "sqlite":
+        return path.is_file()
+    return safetensors_result_exists(path)
+
+
+def _output_descriptor(position: int, record: EmbeddingRecord) -> dict[str, Any]:
+    tensor = record.tensor
+    if isinstance(tensor, LazyTensorReference):
+        dtype = tensor.dtype
+        shape = tensor.shape
+        digest = tensor.sha256
+    else:
+        dtype = str(tensor.dtype).removeprefix("torch.")
+        shape = tuple(tensor.shape)
+        digest = tensor_sha256(tensor)
+    return {
+        "position": position,
+        "id": record.id,
+        "dtype": dtype,
+        "shape": shape,
+        "sha256": digest,
+    }
+
+
+class EmbeddingOutput:
+    """Own the resumable prefix and the commit state of one output destination."""
+
+    def __init__(
+        self,
+        records: Sequence[EmbeddingInput],
+        *,
+        output: str | Path | None,
+        format: str,
+        resume: bool,
+        shard_size: int,
+        run_fingerprint: str,
+        input_fingerprint: str,
+        model_state_fingerprint: str | None,
+        model_state_fingerprint_source: str,
+        pooler: Pooler | None,
+        pooling_names: Sequence[str],
+    ) -> None:
+        self.output = output
+        self.format = format
+        self.shard_size = shard_size
+        self.completed: EmbeddingResult | None = None
+        output_already_exists = output is not None and _output_exists(output, format)
+        existing: EmbeddingResult | None = None
+        self.start_position = 0
+        if output is not None and resume and output_already_exists:
+            if format == "sqlite":
+                try:
+                    existing = load_sqlite_result(output, run_id=run_fingerprint)
+                except KeyError:
+                    existing = load_result(output, format=format)
+            else:
+                existing = load_result(output, format=format)
+            if existing.metadata.get("fingerprint_schema_version") != (
+                _RUN_FINGERPRINT_SCHEMA_VERSION
+            ):
+                raise ValueError(
+                    "Existing embeddings use an incompatible run fingerprint schema; "
+                    "choose another output or set resume=False."
+                )
+            if existing.metadata.get("run_fingerprint") != run_fingerprint:
+                raise ValueError(
+                    "Existing embeddings were produced by a different run fingerprint; "
+                    "choose another output or set resume=False."
+                )
+            if len(existing) > len(records):
+                raise ValueError(
+                    "Existing embeddings are not an ordered prefix of the requested inputs."
+                )
+            prefix_matches = all(
+                (observed.id, observed.sequence) == (expected.id, expected.sequence)
+                for expected, observed in zip(records, existing, strict=False)
+            )
+            if not prefix_matches:
+                raise ValueError(
+                    "Existing embeddings are not an ordered prefix of the requested inputs."
+                )
+            if len(existing) == len(records) and existing.metadata.get("complete", True):
+                self.completed = existing
+                return
+            self.start_position = len(existing)
+
+        self.sqlite_run_id: str | None = None
+        self.sqlite_replace_on_first_commit = False
+        self.sqlite_initial_metadata: dict[str, Any] | None = None
+        if output is not None and format == "sqlite":
+            self.sqlite_initial_metadata = {
+                "format_version": 1,
+                "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+                "run_fingerprint": run_fingerprint,
+                "input_fingerprint": input_fingerprint,
+                "model_state_fingerprint": model_state_fingerprint,
+                "model_state_fingerprint_source": model_state_fingerprint_source,
+                "complete": False,
+            }
+            self.sqlite_run_id = run_fingerprint
+            if not resume and output_already_exists:
+                try:
+                    load_sqlite_result(output, run_id=run_fingerprint)
+                except KeyError:
+                    pass
+                else:
+                    # Keep an exact prior run readable until replacement inference
+                    # has produced the first complete commit window.
+                    self.sqlite_replace_on_first_commit = True
+            if not self.sqlite_replace_on_first_commit:
+                initialize_sqlite_run(
+                    output,
+                    self.sqlite_initial_metadata,
+                    resume=resume,
+                )
+
+        stream_safetensors = output is not None and format == "safetensors"
+        self.output_records: list[EmbeddingRecord] = (
+            [] if self.sqlite_run_id is not None or stream_safetensors else list(existing or ())
+        )
+        self.output_descriptors: list[dict[str, Any]] | None = [] if output is None else None
+        self.pool_slices: dict[str, tuple[int, int]] = {}
+        if existing and pooler is not None:
+            pooled_width = existing[0].load_tensor().shape[-1]
+            if pooled_width % len(pooling_names) != 0:
+                raise ValueError("Stored pooled width is inconsistent with pooling metadata.")
+            self.pool_slices = pooler.output_slices(pooled_width // len(pooling_names))
+
+        self.safetensors_writer: SafetensorsStreamWriter | None = None
+        if stream_safetensors:
+            if output is None:
+                raise RuntimeError(
+                    "Safetensors streaming was enabled without an output destination."
+                )
+            transactional_overwrite = output_already_exists and not resume
+            self.safetensors_writer = SafetensorsStreamWriter(
+                output,
+                {
+                    "format_version": 1,
+                    "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+                    "run_fingerprint": run_fingerprint,
+                    "input_fingerprint": input_fingerprint,
+                    "model_state_fingerprint": model_state_fingerprint,
+                    "model_state_fingerprint_source": model_state_fingerprint_source,
+                    "complete": False,
+                },
+                shard_size=shard_size,
+                existing=existing or (),
+                reuse_existing=bool(resume and existing is not None),
+                publish_initial=not transactional_overwrite,
+                publish_incremental=not transactional_overwrite,
+            )
+
+    def append(self, window_start: int, new_records: list[EmbeddingRecord]) -> None:
+        """Commit a complete ordered window at the storage format's granularity."""
+
+        if self.output_descriptors is not None:
+            self.output_descriptors.extend(
+                _output_descriptor(window_start + offset, record)
+                for offset, record in enumerate(new_records)
+            )
+        if self.output is not None and self.sqlite_run_id is not None:
+            append_sqlite_records(
+                self.output,
+                self.sqlite_run_id,
+                window_start,
+                new_records,
+                replace_metadata=(
+                    self.sqlite_initial_metadata if self.sqlite_replace_on_first_commit else None
+                ),
+            )
+            self.sqlite_replace_on_first_commit = False
+        elif self.safetensors_writer is not None:
+            self.safetensors_writer.append(new_records)
+        else:
+            self.output_records.extend(new_records)
+
+    def finish(self, metadata: dict[str, Any]) -> EmbeddingResult:
+        """Publish completion only after every window has committed."""
+
+        if self.output is not None and self.sqlite_run_id is not None:
+            update_sqlite_run_metadata(self.output, self.sqlite_run_id, metadata)
+            return load_sqlite_result(self.output, run_id=self.sqlite_run_id)
+        if self.safetensors_writer is not None:
+            return self.safetensors_writer.publish(complete=True, metadata=metadata)
+        result = EmbeddingResult(self.output_records, metadata)
+        if self.output is not None:
+            return save_result(result, self.output, format=self.format, shard_size=self.shard_size)
+        return result

fastplms/embeddings/pooling.py

@@ -4,6 +4,7 @@ from __future__ import annotations
 
 import math
 import torch
+
 from collections.abc import Sequence
 from torch import Tensor
 

fastplms/embeddings/runner.py

@@ -1,425 +1,426 @@
-"""Coordinate input preparation, run identity, batch execution, and publication."""
-
-from __future__ import annotations
-
-import torch
-from collections.abc import Callable, Iterable, Mapping, Sequence
-from pathlib import Path
-from typing import Any
+"""Coordinate input preparation, run identity, batch execution, and publication."""
+
+from __future__ import annotations
+
+import torch
+
+from collections.abc import Callable, Iterable, Mapping, Sequence
+from pathlib import Path
+from typing import Any
 from torch import Tensor
 
 from . import identity
 from .batches import (
-    BatchExecutor,
-    _residue_embeddings as _residue_embeddings,
-    _temporary_eval,
-    select_hidden_state_embeddings as select_hidden_state_embeddings,
-)
-from .identity import (
-    _RUN_FINGERPRINT_SCHEMA_VERSION,
-    _adapter_identity_metadata,
-    _attention_backend,
-    _attention_kernel_metadata,
-    _embedding_context,
-    _execution_identity_metadata,
-    _fingerprint_jsonable,
-    _model_identity_metadata,
-    _run_fingerprint,
+    BatchExecutor,
+    _residue_embeddings as _residue_embeddings,
+    _temporary_eval,
+    select_hidden_state_embeddings as select_hidden_state_embeddings,
+)
+from .identity import (
+    _RUN_FINGERPRINT_SCHEMA_VERSION,
+    _adapter_identity_metadata,
+    _attention_backend,
+    _attention_kernel_metadata,
+    _embedding_context,
+    _execution_identity_metadata,
+    _fingerprint_jsonable,
+    _model_identity_metadata,
+    _run_fingerprint,
     _tokenizer_metadata,
-)
-from .inputs import (
-    _InputSpool,
-    _normalize_inputs,
-    _validate_untruncated_lengths,
-    iter_fasta as iter_fasta,
-    parse_fasta as parse_fasta,
-)
-from .output import EmbeddingOutput
-from .pooling import Pooler
-from .types import EmbeddingBatch, EmbeddingInput, EmbeddingResult
-
-
-_DEFAULT_BATCH_WINDOW_MULTIPLIER = 16
-_SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"})
-
-
-def embed_dataset(
-    model: Any,
-    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
-    *,
-    batch_size: int = 2,
-    pooling: str | Sequence[str] | None = None,
-    full_embeddings: bool = False,
-    output: str | Path | None = None,
-    format: str = "safetensors",
-    resume: bool = True,
-    tokenizer: Any | None = None,
-    max_length: int | None = None,
-    truncate: bool = True,
-    dtype: torch.dtype | None = torch.float32,
-    shard_size: int = 2 * 1024**3,
-    model_state_fingerprint: str | None = None,
-    batch_window_size: int | None = None,
-    max_tokens_per_batch: int | None = None,
-    hidden_state_source: str = "encoder",
-    decoder_inputs: Sequence[str] | None = None,
-    decoder_input_ids: Tensor | None = None,
-    decoder_attention_mask: Tensor | None = None,
-    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None = None,
-    _embedding_batch_identity: Mapping[str, Any] | None = None,
-    _allowed_unsupported_pooling: Sequence[str] = (),
-    **model_kwargs: Any,
-) -> EmbeddingResult:
-    """Embed protein sequences with stable ordering and residue-only pooling."""
-
-    for name, value in (
-        ("batch_size", batch_size),
-        ("shard_size", shard_size),
-    ):
-        if not isinstance(value, int) or isinstance(value, bool):
-            raise TypeError(f"{name} must be a positive integer.")
-        if value <= 0:
-            raise ValueError(f"{name} must be a positive integer.")
-    for optional_name, optional_value in (
-        ("max_length", max_length),
-        ("max_tokens_per_batch", max_tokens_per_batch),
-        ("batch_window_size", batch_window_size),
-    ):
-        if optional_value is not None and (
-            not isinstance(optional_value, int) or isinstance(optional_value, bool)
-        ):
-            raise TypeError(f"{optional_name} must be a positive integer when provided.")
-        if optional_value is not None and optional_value <= 0:
-            raise ValueError(f"{optional_name} must be a positive integer when provided.")
-    for name, value in (
-        ("full_embeddings", full_embeddings),
-        ("resume", resume),
-        ("truncate", truncate),
-    ):
-        if not isinstance(value, bool):
-            raise TypeError(f"{name} must be a boolean.")
-    if not isinstance(format, str):
-        raise TypeError("format must be a string.")
-    if output is not None and not isinstance(output, (str, Path)):
-        raise TypeError("output must be a path or None.")
-    if model_state_fingerprint is not None and (
-        not isinstance(model_state_fingerprint, str) or not model_state_fingerprint
-    ):
-        raise ValueError("model_state_fingerprint must be a non-empty string when provided.")
-    if hidden_state_source not in {"encoder", "decoder"}:
-        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
-    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
-    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
-        raise TypeError("hidden_state_index must be an integer.")
-    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
-    if not isinstance(store_all_hidden_states, bool):
-        raise TypeError("store_all_hidden_states must be a boolean.")
-    if decoder_input_ids is not None:
-        if not isinstance(decoder_input_ids, Tensor):
-            raise TypeError("decoder_input_ids must be a tensor.")
-        if decoder_input_ids.is_meta:
-            raise ValueError("decoder_input_ids cannot be a meta tensor.")
-        if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[1] == 0:
-            raise ValueError("decoder_input_ids must have non-empty shape (batch, sequence).")
-        if decoder_input_ids.dtype not in {torch.int32, torch.int64}:
-            raise TypeError("decoder_input_ids must use torch.int32 or torch.int64.")
-    if decoder_attention_mask is not None:
-        if not isinstance(decoder_attention_mask, Tensor):
-            raise TypeError("decoder_attention_mask must be a tensor.")
-        if decoder_attention_mask.is_meta:
-            raise ValueError("decoder_attention_mask cannot be a meta tensor.")
-        if decoder_attention_mask.is_complex() or not bool(
-            torch.isfinite(decoder_attention_mask).all()
-        ):
-            raise ValueError("decoder_attention_mask must contain finite binary values.")
-        if not bool(((decoder_attention_mask == 0) | (decoder_attention_mask == 1)).all()):
-            raise ValueError("decoder_attention_mask must contain finite binary values.")
-    pooling_names = (
-        (("mean",) if not full_embeddings else ())
-        if pooling is None
-        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
-    )
-    if full_embeddings and pooling is not None:
-        raise ValueError("full_embeddings=True cannot be combined with pooling.")
-    if not full_embeddings and not pooling_names:
-        raise ValueError("pooling is required unless full_embeddings=True.")
-    pooler = Pooler(pooling_names) if pooling_names else None
-
-    if batch_size <= 0:
-        raise ValueError("batch_size must be positive.")
-    if format == "pth" or (output is not None and Path(output).suffix.lower() == ".pth"):
-        raise ValueError("Writing pickle-based .pth embeddings is not supported.")
-    if format not in _SUPPORTED_STORAGE_FORMATS:
-        raise ValueError("format must be 'safetensors' or 'sqlite'.")
-    if max_length is not None and max_length <= 0:
-        raise ValueError("max_length must be positive when provided.")
-    if max_tokens_per_batch is not None and max_tokens_per_batch <= 0:
-        raise ValueError("max_tokens_per_batch must be positive when provided.")
-    if not isinstance(dtype, (torch.dtype, type(None))):
-        raise TypeError("dtype must be a torch.dtype or None.")
-    if batch_window_size is not None and batch_window_size <= 0:
-        raise ValueError("batch_window_size must be positive when provided.")
-    if _embedding_batch_fn is not None and not callable(_embedding_batch_fn):
-        raise TypeError("_embedding_batch_fn must be callable when provided.")
-    if _embedding_batch_fn is not None and _embedding_batch_identity is None:
-        raise ValueError(
-            "_embedding_batch_identity is required with _embedding_batch_fn so persisted "
-            "runs bind the family-specific embedding behavior."
-        )
-    if _embedding_batch_identity is not None and not isinstance(_embedding_batch_identity, Mapping):
-        raise TypeError("_embedding_batch_identity must be a mapping when provided.")
-    if isinstance(_allowed_unsupported_pooling, (str, bytes)) or not isinstance(
-        _allowed_unsupported_pooling, Sequence
-    ):
-        raise TypeError("_allowed_unsupported_pooling must be a sequence of pooler names.")
-    if not all(isinstance(name, str) for name in _allowed_unsupported_pooling):
-        raise TypeError("_allowed_unsupported_pooling must contain only strings.")
-    allowed_unsupported_pooling = frozenset(_allowed_unsupported_pooling)
-    if allowed_unsupported_pooling and _embedding_batch_fn is None:
-        raise ValueError(
-            "_allowed_unsupported_pooling is only valid with a family-specific _embedding_batch_fn."
-        )
-    resolved_batch_window_size = (
-        batch_size * _DEFAULT_BATCH_WINDOW_MULTIPLIER
-        if batch_window_size is None
-        else batch_window_size
-    )
-    if resolved_batch_window_size < batch_size:
-        raise ValueError("batch_window_size must be at least batch_size.")
-    records = _normalize_inputs(inputs, disk_backed=output is not None)
-    _validate_untruncated_lengths(
-        records,
-        max_length=max_length,
-        truncate=truncate,
-    )
-    pooling_names = (
-        (("mean",) if not full_embeddings else ())
-        if pooling is None
-        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
-    )
-    if full_embeddings:
-        if pooling is not None:
-            raise ValueError("full_embeddings=True cannot be combined with pooling.")
-    elif not pooling_names:
-        raise ValueError("pooling is required unless full_embeddings=True.")
-    store_all_hidden_states = bool(model_kwargs.get("store_all_hidden_states", False))
-    if store_all_hidden_states and not full_embeddings:
-        raise ValueError("store_all_hidden_states=True requires full_embeddings=True.")
-
-    unsupported = set(getattr(model, "embedding_unsupported_pooling", ()))
-    unknown_pooling_overrides = allowed_unsupported_pooling.difference(unsupported)
-    if unknown_pooling_overrides:
-        raise ValueError(
-            "_allowed_unsupported_pooling may only override poolers declared unsupported "
-            f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}."
-        )
-    unsupported.difference_update(allowed_unsupported_pooling)
-    requested_unsupported = unsupported.intersection(pooling_names)
-    if requested_unsupported:
-        raise ValueError(
-            f"{model.__class__.__name__} does not support pooling operations "
-            f"{sorted(requested_unsupported)}."
-        )
-
-    # Constructing the pooler validates names and duplicate operations before
-    # any checkpoint hashing, tokenization, or inference occurs.
-    pooler = Pooler(pooling_names) if pooling_names else None
-    embedding_context, normalized_decoder_inputs = _embedding_context(
-        model,
-        records,
-        hidden_state_source=hidden_state_source,
-        decoder_inputs=decoder_inputs,
-        decoder_input_ids=decoder_input_ids,
-        decoder_attention_mask=decoder_attention_mask,
-        model_kwargs=model_kwargs,
-    )
-    if _embedding_batch_identity is not None:
-        embedding_context["family_adapter"] = _fingerprint_jsonable(_embedding_batch_identity)
-        if allowed_unsupported_pooling:
-            embedding_context["family_adapter_pooling_override"] = sorted(
-                allowed_unsupported_pooling
-            )
-
-    # A pending automatic attention request settles here, inside the caller's
-    # autocast context, so the fingerprint records the backend that executes.
-    attention_resolution = getattr(model, "attention_resolution", None)
-    if attention_resolution is not None and attention_resolution.deferred:
-        model.resolve_attn_implementation()
-
-    tokenizer_metadata = _tokenizer_metadata(model, tokenizer)
-    (
-        input_fingerprint,
-        run_fingerprint,
-        resolved_model_state_fingerprint,
-        model_state_fingerprint_source,
-    ) = _run_fingerprint(
-        model,
-        records,
-        pooling=pooling_names,
-        full_embeddings=full_embeddings,
-        max_length=max_length,
-        truncate=truncate,
-        dtype=dtype,
-        model_kwargs=model_kwargs,
-        tokenizer_metadata=tokenizer_metadata,
-        model_state_fingerprint=model_state_fingerprint,
-        persist_output=output is not None,
-        embedding_context=embedding_context,
-        batch_size=batch_size,
-        batch_window_size=resolved_batch_window_size,
-        max_tokens_per_batch=max_tokens_per_batch,
-    )
-    destination = EmbeddingOutput(
-        records,
-        output=output,
-        format=format,
-        resume=resume,
-        shard_size=shard_size,
-        run_fingerprint=run_fingerprint,
-        input_fingerprint=input_fingerprint,
-        model_state_fingerprint=resolved_model_state_fingerprint,
-        model_state_fingerprint_source=model_state_fingerprint_source,
-        pooler=pooler,
-        pooling_names=pooling_names,
-    )
-    if destination.completed is not None:
-        return destination.completed
-
-    attention_backend = _attention_backend(model)
-    executor = BatchExecutor(
-        model=model,
-        batch_size=batch_size,
-        max_tokens_per_batch=max_tokens_per_batch,
-        max_length=max_length,
-        truncate=truncate,
-        model_kwargs=model_kwargs,
-        hidden_state_source=hidden_state_source,
-        normalized_decoder_inputs=normalized_decoder_inputs,
-        decoder_input_ids=decoder_input_ids,
-        decoder_attention_mask=decoder_attention_mask,
-        _embedding_batch_fn=_embedding_batch_fn,
-        tokenizer=tokenizer,
-        store_all_hidden_states=store_all_hidden_states,
-        full_embeddings=full_embeddings,
-        dtype=dtype,
-        pooler=pooler,
-        attention_backend=attention_backend,
-        need_attentions="parti" in pooling_names,
-    )
-    pool_slices = destination.pool_slices
-    with _temporary_eval(model), torch.inference_mode():
-        for window_start in range(
-            destination.start_position, len(records), resolved_batch_window_size
-        ):
-            window_stop = min(window_start + resolved_batch_window_size, len(records))
-            window_records = records[window_start:window_stop]
-            if not isinstance(window_records, Sequence):
-                raise RuntimeError("The immutable embedding spool returned a non-sequence window.")
-            new_records, pool_slices = executor.run_window(
-                window_records, window_start=window_start
-            )
-            destination.append(window_start, new_records)
-
+)
+from .inputs import (
+    _InputSpool,
+    _normalize_inputs,
+    _validate_untruncated_lengths,
+    iter_fasta as iter_fasta,
+    parse_fasta as parse_fasta,
+)
+from .output import EmbeddingOutput
+from .pooling import Pooler
+from .types import EmbeddingBatch, EmbeddingInput, EmbeddingResult
+
+
+_DEFAULT_BATCH_WINDOW_MULTIPLIER = 16
+_SUPPORTED_STORAGE_FORMATS = frozenset({"safetensors", "sqlite"})
+
+
+def embed_dataset(
+    model: Any,
+    inputs: (Iterable[str | EmbeddingInput | tuple[str, str]] | Mapping[str, str] | str | Path),
+    *,
+    batch_size: int = 2,
+    pooling: str | Sequence[str] | None = None,
+    full_embeddings: bool = False,
+    output: str | Path | None = None,
+    format: str = "safetensors",
+    resume: bool = True,
+    tokenizer: Any | None = None,
+    max_length: int | None = None,
+    truncate: bool = True,
+    dtype: torch.dtype | None = torch.float32,
+    shard_size: int = 2 * 1024**3,
+    model_state_fingerprint: str | None = None,
+    batch_window_size: int | None = None,
+    max_tokens_per_batch: int | None = None,
+    hidden_state_source: str = "encoder",
+    decoder_inputs: Sequence[str] | None = None,
+    decoder_input_ids: Tensor | None = None,
+    decoder_attention_mask: Tensor | None = None,
+    _embedding_batch_fn: Callable[..., EmbeddingBatch] | None = None,
+    _embedding_batch_identity: Mapping[str, Any] | None = None,
+    _allowed_unsupported_pooling: Sequence[str] = (),
+    **model_kwargs: Any,
+) -> EmbeddingResult:
+    """Embed protein sequences with stable ordering and residue-only pooling."""
+
+    for name, value in (
+        ("batch_size", batch_size),
+        ("shard_size", shard_size),
+    ):
+        if not isinstance(value, int) or isinstance(value, bool):
+            raise TypeError(f"{name} must be a positive integer.")
+        if value <= 0:
+            raise ValueError(f"{name} must be a positive integer.")
+    for optional_name, optional_value in (
+        ("max_length", max_length),
+        ("max_tokens_per_batch", max_tokens_per_batch),
+        ("batch_window_size", batch_window_size),
+    ):
+        if optional_value is not None and (
+            not isinstance(optional_value, int) or isinstance(optional_value, bool)
+        ):
+            raise TypeError(f"{optional_name} must be a positive integer when provided.")
+        if optional_value is not None and optional_value <= 0:
+            raise ValueError(f"{optional_name} must be a positive integer when provided.")
+    for name, value in (
+        ("full_embeddings", full_embeddings),
+        ("resume", resume),
+        ("truncate", truncate),
+    ):
+        if not isinstance(value, bool):
+            raise TypeError(f"{name} must be a boolean.")
+    if not isinstance(format, str):
+        raise TypeError("format must be a string.")
+    if output is not None and not isinstance(output, (str, Path)):
+        raise TypeError("output must be a path or None.")
+    if model_state_fingerprint is not None and (
+        not isinstance(model_state_fingerprint, str) or not model_state_fingerprint
+    ):
+        raise ValueError("model_state_fingerprint must be a non-empty string when provided.")
+    if hidden_state_source not in {"encoder", "decoder"}:
+        raise ValueError("hidden_state_source must be 'encoder' or 'decoder'.")
+    hidden_state_index = model_kwargs.get("hidden_state_index", -1)
+    if not isinstance(hidden_state_index, int) or isinstance(hidden_state_index, bool):
+        raise TypeError("hidden_state_index must be an integer.")
+    store_all_hidden_states = model_kwargs.get("store_all_hidden_states", False)
+    if not isinstance(store_all_hidden_states, bool):
+        raise TypeError("store_all_hidden_states must be a boolean.")
+    if decoder_input_ids is not None:
+        if not isinstance(decoder_input_ids, Tensor):
+            raise TypeError("decoder_input_ids must be a tensor.")
+        if decoder_input_ids.is_meta:
+            raise ValueError("decoder_input_ids cannot be a meta tensor.")
+        if decoder_input_ids.ndim != 2 or decoder_input_ids.shape[1] == 0:
+            raise ValueError("decoder_input_ids must have non-empty shape (batch, sequence).")
+        if decoder_input_ids.dtype not in {torch.int32, torch.int64}:
+            raise TypeError("decoder_input_ids must use torch.int32 or torch.int64.")
+    if decoder_attention_mask is not None:
+        if not isinstance(decoder_attention_mask, Tensor):
+            raise TypeError("decoder_attention_mask must be a tensor.")
+        if decoder_attention_mask.is_meta:
+            raise ValueError("decoder_attention_mask cannot be a meta tensor.")
+        if decoder_attention_mask.is_complex() or not bool(
+            torch.isfinite(decoder_attention_mask).all()
+        ):
+            raise ValueError("decoder_attention_mask must contain finite binary values.")
+        if not bool(((decoder_attention_mask == 0) | (decoder_attention_mask == 1)).all()):
+            raise ValueError("decoder_attention_mask must contain finite binary values.")
+    pooling_names = (
+        (("mean",) if not full_embeddings else ())
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if full_embeddings and pooling is not None:
+        raise ValueError("full_embeddings=True cannot be combined with pooling.")
+    if not full_embeddings and not pooling_names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    pooler = Pooler(pooling_names) if pooling_names else None
+
+    if batch_size <= 0:
+        raise ValueError("batch_size must be positive.")
+    if format == "pth" or (output is not None and Path(output).suffix.lower() == ".pth"):
+        raise ValueError("Writing pickle-based .pth embeddings is not supported.")
+    if format not in _SUPPORTED_STORAGE_FORMATS:
+        raise ValueError("format must be 'safetensors' or 'sqlite'.")
+    if max_length is not None and max_length <= 0:
+        raise ValueError("max_length must be positive when provided.")
+    if max_tokens_per_batch is not None and max_tokens_per_batch <= 0:
+        raise ValueError("max_tokens_per_batch must be positive when provided.")
+    if not isinstance(dtype, (torch.dtype, type(None))):
+        raise TypeError("dtype must be a torch.dtype or None.")
+    if batch_window_size is not None and batch_window_size <= 0:
+        raise ValueError("batch_window_size must be positive when provided.")
+    if _embedding_batch_fn is not None and not callable(_embedding_batch_fn):
+        raise TypeError("_embedding_batch_fn must be callable when provided.")
+    if _embedding_batch_fn is not None and _embedding_batch_identity is None:
+        raise ValueError(
+            "_embedding_batch_identity is required with _embedding_batch_fn so persisted "
+            "runs bind the family-specific embedding behavior."
+        )
+    if _embedding_batch_identity is not None and not isinstance(_embedding_batch_identity, Mapping):
+        raise TypeError("_embedding_batch_identity must be a mapping when provided.")
+    if isinstance(_allowed_unsupported_pooling, (str, bytes)) or not isinstance(
+        _allowed_unsupported_pooling, Sequence
+    ):
+        raise TypeError("_allowed_unsupported_pooling must be a sequence of pooler names.")
+    if not all(isinstance(name, str) for name in _allowed_unsupported_pooling):
+        raise TypeError("_allowed_unsupported_pooling must contain only strings.")
+    allowed_unsupported_pooling = frozenset(_allowed_unsupported_pooling)
+    if allowed_unsupported_pooling and _embedding_batch_fn is None:
+        raise ValueError(
+            "_allowed_unsupported_pooling is only valid with a family-specific _embedding_batch_fn."
+        )
+    resolved_batch_window_size = (
+        batch_size * _DEFAULT_BATCH_WINDOW_MULTIPLIER
+        if batch_window_size is None
+        else batch_window_size
+    )
+    if resolved_batch_window_size < batch_size:
+        raise ValueError("batch_window_size must be at least batch_size.")
+    records = _normalize_inputs(inputs, disk_backed=output is not None)
+    _validate_untruncated_lengths(
+        records,
+        max_length=max_length,
+        truncate=truncate,
+    )
+    pooling_names = (
+        (("mean",) if not full_embeddings else ())
+        if pooling is None
+        else ((pooling,) if isinstance(pooling, str) else tuple(pooling))
+    )
+    if full_embeddings:
+        if pooling is not None:
+            raise ValueError("full_embeddings=True cannot be combined with pooling.")
+    elif not pooling_names:
+        raise ValueError("pooling is required unless full_embeddings=True.")
+    store_all_hidden_states = bool(model_kwargs.get("store_all_hidden_states", False))
+    if store_all_hidden_states and not full_embeddings:
+        raise ValueError("store_all_hidden_states=True requires full_embeddings=True.")
+
+    unsupported = set(getattr(model, "embedding_unsupported_pooling", ()))
+    unknown_pooling_overrides = allowed_unsupported_pooling.difference(unsupported)
+    if unknown_pooling_overrides:
+        raise ValueError(
+            "_allowed_unsupported_pooling may only override poolers declared unsupported "
+            f"by the model; unknown overrides: {sorted(unknown_pooling_overrides)}."
+        )
+    unsupported.difference_update(allowed_unsupported_pooling)
+    requested_unsupported = unsupported.intersection(pooling_names)
+    if requested_unsupported:
+        raise ValueError(
+            f"{model.__class__.__name__} does not support pooling operations "
+            f"{sorted(requested_unsupported)}."
+        )
+
+    # Constructing the pooler validates names and duplicate operations before
+    # any checkpoint hashing, tokenization, or inference occurs.
+    pooler = Pooler(pooling_names) if pooling_names else None
+    embedding_context, normalized_decoder_inputs = _embedding_context(
+        model,
+        records,
+        hidden_state_source=hidden_state_source,
+        decoder_inputs=decoder_inputs,
+        decoder_input_ids=decoder_input_ids,
+        decoder_attention_mask=decoder_attention_mask,
+        model_kwargs=model_kwargs,
+    )
+    if _embedding_batch_identity is not None:
+        embedding_context["family_adapter"] = _fingerprint_jsonable(_embedding_batch_identity)
+        if allowed_unsupported_pooling:
+            embedding_context["family_adapter_pooling_override"] = sorted(
+                allowed_unsupported_pooling
+            )
+
+    # A pending automatic attention request settles here, inside the caller's
+    # autocast context, so the fingerprint records the backend that executes.
+    attention_resolution = getattr(model, "attention_resolution", None)
+    if attention_resolution is not None and attention_resolution.deferred:
+        model.resolve_attn_implementation()
+
+    tokenizer_metadata = _tokenizer_metadata(model, tokenizer)
+    (
+        input_fingerprint,
+        run_fingerprint,
+        resolved_model_state_fingerprint,
+        model_state_fingerprint_source,
+    ) = _run_fingerprint(
+        model,
+        records,
+        pooling=pooling_names,
+        full_embeddings=full_embeddings,
+        max_length=max_length,
+        truncate=truncate,
+        dtype=dtype,
+        model_kwargs=model_kwargs,
+        tokenizer_metadata=tokenizer_metadata,
+        model_state_fingerprint=model_state_fingerprint,
+        persist_output=output is not None,
+        embedding_context=embedding_context,
+        batch_size=batch_size,
+        batch_window_size=resolved_batch_window_size,
+        max_tokens_per_batch=max_tokens_per_batch,
+    )
+    destination = EmbeddingOutput(
+        records,
+        output=output,
+        format=format,
+        resume=resume,
+        shard_size=shard_size,
+        run_fingerprint=run_fingerprint,
+        input_fingerprint=input_fingerprint,
+        model_state_fingerprint=resolved_model_state_fingerprint,
+        model_state_fingerprint_source=model_state_fingerprint_source,
+        pooler=pooler,
+        pooling_names=pooling_names,
+    )
+    if destination.completed is not None:
+        return destination.completed
+
+    attention_backend = _attention_backend(model)
+    executor = BatchExecutor(
+        model=model,
+        batch_size=batch_size,
+        max_tokens_per_batch=max_tokens_per_batch,
+        max_length=max_length,
+        truncate=truncate,
+        model_kwargs=model_kwargs,
+        hidden_state_source=hidden_state_source,
+        normalized_decoder_inputs=normalized_decoder_inputs,
+        decoder_input_ids=decoder_input_ids,
+        decoder_attention_mask=decoder_attention_mask,
+        _embedding_batch_fn=_embedding_batch_fn,
+        tokenizer=tokenizer,
+        store_all_hidden_states=store_all_hidden_states,
+        full_embeddings=full_embeddings,
+        dtype=dtype,
+        pooler=pooler,
+        attention_backend=attention_backend,
+        need_attentions="parti" in pooling_names,
+    )
+    pool_slices = destination.pool_slices
+    with _temporary_eval(model), torch.inference_mode():
+        for window_start in range(
+            destination.start_position, len(records), resolved_batch_window_size
+        ):
+            window_stop = min(window_start + resolved_batch_window_size, len(records))
+            window_records = records[window_start:window_stop]
+            if not isinstance(window_records, Sequence):
+                raise RuntimeError("The immutable embedding spool returned a non-sequence window.")
+            new_records, pool_slices = executor.run_window(
+                window_records, window_start=window_start
+            )
+            destination.append(window_start, new_records)
+
     software_versions = identity._software_versions()
-    projection = getattr(model, "embedding_projection", None)
-    resolved_layer = getattr(
-        model,
-        "embedding_layer",
-        model_kwargs.get("hidden_state_index", -1),
-    )
-    token_policy = getattr(
-        model,
-        "embedding_token_policy",
-        {
-            "unit": "residue",
-            "include": ["biological residues"],
-            "exclude": [
-                "BOS",
-                "EOS",
-                "padding",
-                "chain delimiters",
-                "non-protein tokens",
-            ],
-        },
-    )
-    model_identity = _model_identity_metadata(model)
-    metadata: dict[str, Any] = {
-        "format_version": 1,
-        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
-        "run_fingerprint": run_fingerprint,
-        "input_fingerprint": input_fingerprint,
-        "model_state_fingerprint": resolved_model_state_fingerprint,
-        "model_state_fingerprint_source": model_state_fingerprint_source,
-        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
-        **model_identity,
-        "dtype": str(dtype).removeprefix("torch.") if dtype is not None else "model",
-        "attention_backend": attention_backend,
-        "attention_kernel": _attention_kernel_metadata(attention_backend),
-        "layer": resolved_layer,
-        "projection": projection,
-        "esmc_source": getattr(model, "_esmc_source", None),
-        "esmc_revision": getattr(model, "_esmc_source_revision", None),
-        "esmc_files": getattr(model, "_esmc_source_files", None),
-        "token_policy": token_policy,
-        "tokenizer": tokenizer_metadata,
-        **embedding_context,
-        "pooling": list(pooling_names),
-        "pool_slices": pool_slices,
-        "full_embeddings": full_embeddings,
-        "max_length": max_length,
-        "truncate": truncate,
-        "truncation": {"enabled": truncate, "max_length": max_length},
-        "batching": {
-            "batch_size": batch_size,
-            "batch_window_size": resolved_batch_window_size,
-            "max_tokens_per_batch": max_tokens_per_batch,
-            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
-            "ordering": "bounded-length-bucketed-stable-output",
-            "resume_commit_granularity": (
-                "not-applicable"
-                if output is None
-                else "batch-window"
-                if format == "sqlite"
-                else "shard-flush"
-            ),
-        },
-        "residue_mask_policy": "biological-residues-only",
-        "record_count": len(records),
-        "descriptor_index": (
-            "memory-metadata"
-            if output is None
-            else "sqlite-records"
-            if format == "sqlite"
-            else "safetensors-generation-index"
-        ),
-        "storage_format": format if output is not None else "memory",
-        "software": software_versions,
-        "execution": _execution_identity_metadata(model),
-        "adapter": _adapter_identity_metadata(model),
-        "torch_version": software_versions["torch"],
-        "transformers_version": software_versions["transformers"],
-        "complete": True,
-    }
-    if destination.output_descriptors is not None:
-        metadata["outputs"] = destination.output_descriptors
-        metadata["tensor_hashes"] = [item["sha256"] for item in destination.output_descriptors]
-    status = getattr(model, "esmc_precision_status", None)
-    if status is not None:
-        metadata["esmc_precision"] = status.as_dict() if hasattr(status, "as_dict") else status
-    return destination.finish(metadata)
-
-
-class EmbeddingMixin:
-    """Small delegation mixin shared by FastPLMs model classes."""
-
-    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
-        return embed_dataset(self, inputs, **kwargs)
-
-
-__all__ = [
-    "EmbeddingMixin",
-    "embed_dataset",
-    "iter_fasta",
-    "parse_fasta",
-    "select_hidden_state_embeddings",
-]
+    projection = getattr(model, "embedding_projection", None)
+    resolved_layer = getattr(
+        model,
+        "embedding_layer",
+        model_kwargs.get("hidden_state_index", -1),
+    )
+    token_policy = getattr(
+        model,
+        "embedding_token_policy",
+        {
+            "unit": "residue",
+            "include": ["biological residues"],
+            "exclude": [
+                "BOS",
+                "EOS",
+                "padding",
+                "chain delimiters",
+                "non-protein tokens",
+            ],
+        },
+    )
+    model_identity = _model_identity_metadata(model)
+    metadata: dict[str, Any] = {
+        "format_version": 1,
+        "fingerprint_schema_version": _RUN_FINGERPRINT_SCHEMA_VERSION,
+        "run_fingerprint": run_fingerprint,
+        "input_fingerprint": input_fingerprint,
+        "model_state_fingerprint": resolved_model_state_fingerprint,
+        "model_state_fingerprint_source": model_state_fingerprint_source,
+        "model_class": f"{model.__class__.__module__}.{model.__class__.__qualname__}",
+        **model_identity,
+        "dtype": str(dtype).removeprefix("torch.") if dtype is not None else "model",
+        "attention_backend": attention_backend,
+        "attention_kernel": _attention_kernel_metadata(attention_backend),
+        "layer": resolved_layer,
+        "projection": projection,
+        "esmc_source": getattr(model, "_esmc_source", None),
+        "esmc_revision": getattr(model, "_esmc_source_revision", None),
+        "esmc_files": getattr(model, "_esmc_source_files", None),
+        "token_policy": token_policy,
+        "tokenizer": tokenizer_metadata,
+        **embedding_context,
+        "pooling": list(pooling_names),
+        "pool_slices": pool_slices,
+        "full_embeddings": full_embeddings,
+        "max_length": max_length,
+        "truncate": truncate,
+        "truncation": {"enabled": truncate, "max_length": max_length},
+        "batching": {
+            "batch_size": batch_size,
+            "batch_window_size": resolved_batch_window_size,
+            "max_tokens_per_batch": max_tokens_per_batch,
+            "input_storage": ("disk-spool" if isinstance(records, _InputSpool) else "memory"),
+            "ordering": "bounded-length-bucketed-stable-output",
+            "resume_commit_granularity": (
+                "not-applicable"
+                if output is None
+                else "batch-window"
+                if format == "sqlite"
+                else "shard-flush"
+            ),
+        },
+        "residue_mask_policy": "biological-residues-only",
+        "record_count": len(records),
+        "descriptor_index": (
+            "memory-metadata"
+            if output is None
+            else "sqlite-records"
+            if format == "sqlite"
+            else "safetensors-generation-index"
+        ),
+        "storage_format": format if output is not None else "memory",
+        "software": software_versions,
+        "execution": _execution_identity_metadata(model),
+        "adapter": _adapter_identity_metadata(model),
+        "torch_version": software_versions["torch"],
+        "transformers_version": software_versions["transformers"],
+        "complete": True,
+    }
+    if destination.output_descriptors is not None:
+        metadata["outputs"] = destination.output_descriptors
+        metadata["tensor_hashes"] = [item["sha256"] for item in destination.output_descriptors]
+    status = getattr(model, "esmc_precision_status", None)
+    if status is not None:
+        metadata["esmc_precision"] = status.as_dict() if hasattr(status, "as_dict") else status
+    return destination.finish(metadata)
+
+
+class EmbeddingMixin:
+    """Small delegation mixin shared by FastPLMs model classes."""
+
+    def embed_dataset(self, inputs: Any, **kwargs: Any) -> EmbeddingResult:
+        return embed_dataset(self, inputs, **kwargs)
+
+
+__all__ = [
+    "EmbeddingMixin",
+    "embed_dataset",
+    "iter_fasta",
+    "parse_fasta",
+    "select_hidden_state_embeddings",
+]

fastplms/embeddings/storage.py

@@ -9,6 +9,7 @@ import sqlite3
 import struct
 import numpy as np
 import torch
+
 from bisect import bisect_right
 from collections.abc import Iterable, Iterator, Sequence
 from pathlib import Path

fastplms/models/esm3/modeling_esm3.py

@@ -16,16 +16,16 @@ import os
 import shutil
 import stat
 import tempfile
+import einops
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
 from contextlib import contextmanager
 from dataclasses import dataclass
 from pathlib import Path, PurePosixPath
 from typing import ClassVar
 from zipfile import ZIP_DEFLATED, ZipFile, ZipInfo
-
-import einops
-import torch
-import torch.nn as nn
-import torch.nn.functional as F
 from einops import rearrange
 from tokenizers import Tokenizer
 from tokenizers.models import BPE
@@ -968,7 +968,7 @@ def rbf(values: torch.Tensor, v_min: float, v_max: float, n_bins: int = 16) -> t
         n_bins,
         device=values.device,
         dtype=values.dtype,
-    )
+    )  # (n_bins,)
     centers = centers.view([1] * len(values.shape) + [-1])  # (..., n)
     std = (v_max - v_min) / n_bins
     z = (values.unsqueeze(-1) - centers) / std  # (..., n)
@@ -1078,12 +1078,12 @@ class RotaryEmbedding(nn.Module):
         return result
 
     def reset_parameters(self) -> None:
-        inv_freq = self._compute_inv_freq(self.device)
+        inv_freq = self._compute_inv_freq(self.device)  # (d / 2,)
         self.register_buffer("inv_freq", inv_freq, persistent=False)
-        arange = torch.arange(0, self.dim, 2, device=self.device, dtype=torch.float32)
+        arange = torch.arange(0, self.dim, 2, device=self.device, dtype=torch.float32)  # (d / 2,)
         scale = (
             (arange + 0.4 * self.dim) / (1.4 * self.dim) if self.scale_base is not None else None
-        )
+        )  # (d / 2,) or None
         self.register_buffer("scale", scale)
 
     def _compute_inv_freq(self, device: torch.device | None = None) -> torch.Tensor:
@@ -1109,22 +1109,22 @@ class RotaryEmbedding(nn.Module):
             # ``inv_freq`` is non-persistent and may have been materialized
             # without values after Transformers constructs this module on the
             # meta device. Recreate it deterministically on the first forward.
-            self.inv_freq = self._compute_inv_freq(device)
+            self.inv_freq = self._compute_inv_freq(device)  # (d / 2,)
             if self.pos_idx_in_fp32:
                 t = torch.arange(seqlen, device=device, dtype=torch.float32)  # (l,)
-                t /= self.scaling_factor
-                inv_freq = self.inv_freq
+                t /= self.scaling_factor  # (l,)
+                inv_freq = self.inv_freq  # (d / 2,)
             else:
                 t = torch.arange(
                     seqlen, device=device, dtype=self.inv_freq.dtype
                 )  # (l,)
-                t /= self.scaling_factor
-                inv_freq = self.inv_freq
+                t /= self.scaling_factor  # (l,)
+                inv_freq = self.inv_freq  # (d / 2,)
             freqs = torch.outer(t, inv_freq)  # (l, d / 2)
 
             if self.scale is None:
-                self._cos_cached = torch.cos(freqs).to(dtype)
-                self._sin_cached = torch.sin(freqs).to(dtype)
+                self._cos_cached = torch.cos(freqs).to(dtype)  # (l, d / 2)
+                self._sin_cached = torch.sin(freqs).to(dtype)  # (l, d / 2)
             else:
                 raise NotImplementedError("Scaled rotary embeddings are not used by ESM3.")
 
@@ -1169,19 +1169,19 @@ def fp32_autocast_context(device_type: str):
 class RotationMatrix:
     def __init__(self, rots: torch.Tensor) -> None:
         if rots.ndim >= 1 and rots.shape[-1] == 9:
-            rots = rots.unflatten(-1, (3, 3))
+            rots = rots.unflatten(-1, (3, 3))  # (..., 3, 3)
         if rots.ndim < 2 or tuple(rots.shape[-2:]) != (3, 3):
             raise ValueError(
                 "Rotation matrices must have trailing shape (3, 3) or flattened "
                 f"shape (9,); got {tuple(rots.shape)}."
             )
-        self._rots = rots.to(torch.float32)
+        self._rots = rots.to(torch.float32)  # (..., 3, 3)
 
     @classmethod
     def identity(cls, shape: tuple[int, ...], **tensor_kwargs) -> RotationMatrix:
-        rots = torch.eye(3, **tensor_kwargs)
-        rots = rots.view(*[1 for _ in range(len(shape))], 3, 3)
-        rots = rots.expand(*shape, -1, -1)
+        rots = torch.eye(3, **tensor_kwargs)  # (3, 3)
+        rots = rots.view(*[1 for _ in range(len(shape))], 3, 3)  # (1, ..., 1, 3, 3), with len(shape) leading singleton axes
+        rots = rots.expand(*shape, -1, -1)  # (*shape, 3, 3)
         return cls(rots)
 
     def __getitem__(self, idx) -> RotationMatrix:
@@ -1194,7 +1194,7 @@ class RotationMatrix:
 
     @property
     def tensor(self) -> torch.Tensor:
-        return self._rots.flatten(-2)
+        return self._rots.flatten(-2)  # (..., 9)
 
     @property
     def device(self) -> torch.device:
@@ -1220,15 +1220,15 @@ class RotationMatrix:
         eps: float = 1e-12,
     ) -> RotationMatrix:
         with fp32_autocast_context(x_axis.device.type):
-            e1 = xy_plane
-            denom = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps)
-            x_axis = x_axis / denom
-            dot = (x_axis * e1).sum(dim=-1, keepdim=True)
-            e1 = e1 - x_axis * dot
-            denom = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps)
-            e1 = e1 / denom
-            e2 = torch.cross(x_axis, e1, dim=-1)
-            return RotationMatrix(torch.stack([x_axis, e1, e2], dim=-1))
+            e1 = xy_plane  # (..., 3)
+            denom = torch.sqrt((x_axis**2).sum(dim=-1, keepdim=True) + eps)  # (..., 1)
+            x_axis = x_axis / denom  # (..., 3)
+            dot = (x_axis * e1).sum(dim=-1, keepdim=True)  # (..., 1)
+            e1 = e1 - x_axis * dot  # (..., 3)
+            denom = torch.sqrt((e1**2).sum(dim=-1, keepdim=True) + eps)  # (..., 1)
+            e1 = e1 / denom  # (..., 3)
+            e2 = torch.cross(x_axis, e1, dim=-1)  # (..., 3)
+            return RotationMatrix(torch.stack([x_axis, e1, e2], dim=-1))  # rotation tensor: (..., 3, 3)
 
 
 @dataclass(frozen=True)
@@ -1401,10 +1401,10 @@ class MultiHeadAttention(nn.Module):
     ) -> tuple[torch.Tensor, torch.Tensor | None]:
         # x: (b, l, d); attention_mask: (b, 1, l, l) or (b, 1, 1, l)
         qkv = self.layernorm_qkv(x)  # (b, l, 3 * d)
-        query, key, value = torch.chunk(qkv, 3, dim=-1)
-        query = self.q_ln(query).to(query.dtype)
-        key = self.k_ln(key).to(query.dtype)
-        query, key = self._apply_rotary(query, key)
+        query, key, value = torch.chunk(qkv, 3, dim=-1)  # each (b, l, d)
+        query = self.q_ln(query).to(query.dtype)  # (b, l, d)
+        key = self.k_ln(key).to(query.dtype)  # (b, l, d)
+        query, key = self._apply_rotary(query, key)  # each (b, l, d)
 
         reshaper = functools.partial(
             einops.rearrange,
@@ -1428,9 +1428,9 @@ class MultiHeadAttention(nn.Module):
                     ~mask,
                     torch.finfo(attn_scores.dtype).min,
                 )
-            attn_weights = torch.softmax(attn_scores, dim=-1)
+            attn_weights = torch.softmax(attn_scores, dim=-1)  # (b, h, l, l)
             if mask is not None:
-                attn_weights = attn_weights.masked_fill(~mask, 0.0)
+                attn_weights = attn_weights.masked_fill(~mask, 0.0)  # (b, h, l, l)
             context = torch.einsum(
                 "bhls,bhsd->bhld", attn_weights, value
             )  # (b, h, l, d_h)
@@ -1494,8 +1494,8 @@ class GeometricReasoningOriginalImpl(nn.Module):
         self.s_norm = nn.LayerNorm(c_s, bias=bias)
         self.proj = nn.Linear(c_s, projection_width, bias=bias)
         self.out_proj = nn.Linear(output_width, c_s, bias=bias)
-        self.distance_scale_per_head = nn.Parameter(torch.zeros(v_heads))
-        self.rotation_scale_per_head = nn.Parameter(torch.zeros(v_heads))
+        self.distance_scale_per_head = nn.Parameter(torch.zeros(v_heads))  # (h,), where h = v_heads
+        self.rotation_scale_per_head = nn.Parameter(torch.zeros(v_heads))  # (h,)
 
     def forward(
         self,
@@ -1505,28 +1505,30 @@ class GeometricReasoningOriginalImpl(nn.Module):
         sequence_id: torch.Tensor | None,
         chain_id: torch.Tensor,
     ) -> torch.Tensor:
+        # s: (b, l, c_s); masks and IDs: (b, l); affine frames: (b, l).
+        # h = v_heads; m = num_vector_messages; each vector has three xyz coordinates.
         if sequence_id is None:
-            sequence_id = torch.zeros_like(s[..., 0], dtype=torch.int64)
-        attn_bias = sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)
-        attn_bias = attn_bias.unsqueeze(1).float()
+            sequence_id = torch.zeros_like(s[..., 0], dtype=torch.int64)  # (b, l)
+        attn_bias = sequence_id.unsqueeze(-1) == sequence_id.unsqueeze(-2)  # (b, l, l)
+        attn_bias = attn_bias.unsqueeze(1).float()  # (b, 1, l, l)
         attn_bias = attn_bias.masked_fill(
             ~affine_mask[:, None, None, :],
             torch.finfo(attn_bias.dtype).min,
-        )
-        chain_id_mask = chain_id.unsqueeze(1) != chain_id.unsqueeze(2)
+        )  # (b, 1, l, l)
+        chain_id_mask = chain_id.unsqueeze(1) != chain_id.unsqueeze(2)  # (b, l, l)
         attn_bias = attn_bias.masked_fill(
             chain_id_mask.unsqueeze(1),
             torch.finfo(s.dtype).min,
-        )
+        )  # (b, 1, l, l)
 
-        ns = self.s_norm(s)
+        ns = self.s_norm(s)  # (b, l, c_s)
         vec_rot, vec_dist = self.proj(ns).split(
             [
                 self.v_heads * 2 * 3 + self.v_heads * 3 * self.num_vector_messages,
                 self.v_heads * 2 * 3,
             ],
             dim=-1,
-        )
+        )  # (b, l, 3 * h * (2 + m)), (b, l, 6 * h)
 
         query_rot, key_rot, value = (
             affine.rot[..., None]
@@ -1535,43 +1537,43 @@ class GeometricReasoningOriginalImpl(nn.Module):
                 [self.v_heads, self.v_heads, self.v_heads * self.num_vector_messages],
                 dim=-2,
             )
-        )
+        )  # (b, l, h, 3), (b, l, h, 3), (b, l, h * m, 3)
         query_dist, key_dist = (
             affine[..., None]
             .apply(rearrange(vec_dist, "... (h c) -> ... h c", c=3))
             .chunk(2, dim=-2)
-        )
+        )  # each (b, l, h, 3)
 
-        query_dist = rearrange(query_dist, "b s h d -> b h s 1 d")
-        key_dist = rearrange(key_dist, "b s h d -> b h 1 s d")
-        query_rot = rearrange(query_rot, "b s h d -> b h s d")
-        key_rot = rearrange(key_rot, "b s h d -> b h d s")
+        query_dist = rearrange(query_dist, "b s h d -> b h s 1 d")  # (b, h, l, 1, 3)
+        key_dist = rearrange(key_dist, "b s h d -> b h 1 s d")  # (b, h, 1, l, 3)
+        query_rot = rearrange(query_rot, "b s h d -> b h s d")  # (b, h, l, 3)
+        key_rot = rearrange(key_rot, "b s h d -> b h d s")  # (b, h, 3, l)
         value = rearrange(
             value,
             "b s (h m) d -> b h s (m d)",
             m=self.num_vector_messages,
-        )
+        )  # (b, h, l, m * 3)
 
-        distance_term = (query_dist - key_dist).norm(dim=-1) / math.sqrt(3)
-        rotation_term = query_rot.matmul(key_rot) / math.sqrt(3)
+        distance_term = (query_dist - key_dist).norm(dim=-1) / math.sqrt(3)  # (b, h, l, l)
+        rotation_term = query_rot.matmul(key_rot) / math.sqrt(3)  # (b, h, l, l)
         distance_term_weight = rearrange(
             F.softplus(self.distance_scale_per_head),
             "h -> h 1 1",
-        )
+        )  # (h, 1, 1)
         rotation_term_weight = rearrange(
             F.softplus(self.rotation_scale_per_head),
             "h -> h 1 1",
-        )
-        attn_weight = rotation_term * rotation_term_weight - distance_term * distance_term_weight
+        )  # (h, 1, 1)
+        attn_weight = rotation_term * rotation_term_weight - distance_term * distance_term_weight  # (b, h, l, l)
 
         s_q = attn_weight.size(2)
         s_k = attn_weight.size(3)
         offset_q = max(0, attn_bias.size(2) - s_q)
         offset_k = max(0, attn_bias.size(3) - s_k)
-        attn_bias = attn_bias[:, :, offset_q:, offset_k:]
-        attn_weight = torch.softmax(attn_weight + attn_bias, dim=-1)
+        attn_bias = attn_bias[:, :, offset_q:, offset_k:]  # (b, 1, s_q, s_k)
+        attn_weight = torch.softmax(attn_weight + attn_bias, dim=-1)  # (b, h, s_q, s_k)
 
-        attn_out = attn_weight.matmul(value)
+        attn_out = attn_weight.matmul(value)  # (b, h, s_q, m * 3)
         attn_out = (
             affine.rot[..., None]
             .invert()
@@ -1582,16 +1584,16 @@ class GeometricReasoningOriginalImpl(nn.Module):
                     m=self.num_vector_messages,
                 )
             )
-        )
+        )  # (b, s_q, h * m, 3)
         attn_out = rearrange(
             attn_out,
             "b s (h m) d -> b s (h m d)",
             m=self.num_vector_messages,
-        )
+        )  # (b, s_q, h * m * 3)
         if self.mask_and_zero_frameless:
-            attn_out = attn_out.masked_fill(~affine_mask[..., None], 0.0)
-        attn_out = attn_out.to(self.out_proj.weight.dtype)
-        return self.out_proj(attn_out)
+            attn_out = attn_out.masked_fill(~affine_mask[..., None], 0.0)  # (b, s_q, h * m * 3)
+        attn_out = attn_out.to(self.out_proj.weight.dtype)  # (b, s_q, h * m * 3)
+        return self.out_proj(attn_out)  # (b, s_q, c_s)
 
 
 def swiglu_correction_fn(expansion_ratio: float, d_model: int) -> int:
@@ -1770,7 +1772,7 @@ class TransformerStack(nn.Module):
     ]:
         *batch_dims, _ = x.shape
         if chain_id is None:
-            chain_id = torch.ones(size=batch_dims, dtype=torch.int64, device=x.device)
+            chain_id = torch.ones(size=batch_dims, dtype=torch.int64, device=x.device)  # (*batch_dims), matching x without its feature axis
         if affine is None or affine_mask is None:
             raise ValueError("affine and affine_mask are required for ESM3 transformer calls.")
         attention_mask, flex_block_mask, affine_mask, mask_semantics, effective_backend = (
@@ -1947,17 +1949,17 @@ class EncodeInputs(nn.Module):
         function_tokens: torch.Tensor,
         residue_annotation_tokens: torch.Tensor,
     ) -> torch.Tensor:
-        sequence_embed = self.sequence_embed(sequence_tokens)
+        sequence_embed = self.sequence_embed(sequence_tokens)  # (b, l, d)
         rbf_16_fn = functools.partial(rbf, v_min=0.0, v_max=1.0, n_bins=16)
         plddt_embed = self.plddt_projection(
             rbf_16_fn(average_plddt).to(self.plddt_projection.weight.dtype)
-        )
+        )  # (..., d), broadcast over residues when supplied once per sequence
         structure_per_res_plddt = self.structure_per_res_plddt_projection(
             rbf_16_fn(per_res_plddt).to(self.structure_per_res_plddt_projection.weight.dtype)
-        )
-        structure_embed = self.structure_tokens_embed(structure_tokens)
-        ss8_embed = self.ss8_embed(ss8_tokens)
-        sasa_embed = self.sasa_embed(sasa_tokens)
+        )  # (b, l, d)
+        structure_embed = self.structure_tokens_embed(structure_tokens)  # (b, l, d)
+        ss8_embed = self.ss8_embed(ss8_tokens)  # (b, l, d)
+        sasa_embed = self.sasa_embed(sasa_tokens)  # (b, l, d)
         function_embed = torch.cat(
             [
                 embed_fn(funcs)
@@ -1968,7 +1970,7 @@ class EncodeInputs(nn.Module):
                 )
             ],
             -1,
-        )
+        )  # (b, l, d)
 
         batch_size, seq_len, num_annotations = residue_annotation_tokens.shape
         residue_embed = self.residue_embed(
@@ -1979,13 +1981,13 @@ class EncodeInputs(nn.Module):
                 l=seq_len,
                 n=num_annotations,
             )
-        )
+        )  # (b * l, d)
         residue_embed = rearrange(
             residue_embed,
             "(b l) d -> b l d",
             b=batch_size,
             l=seq_len,
-        )
+        )  # (b, l, d)
 
         return (
             sequence_embed
@@ -2029,8 +2031,8 @@ class OutputHeads(nn.Module):
         hidden_states: tuple[torch.Tensor, ...] | None = None,
         attentions: tuple[torch.Tensor, ...] | None = None,
     ) -> ESM3CoreOutput:
-        function_logits = self.function_head(x)
-        function_logits = rearrange(function_logits, "... (k v) -> ... k v", k=8)
+        function_logits = self.function_head(x)  # (..., 8 * 260)
+        function_logits = rearrange(function_logits, "... (k v) -> ... k v", k=8)  # (..., 8, 260)
         return ESM3CoreOutput(
             sequence_logits=self.sequence_head(x),
             structure_logits=self.structure_head(x),
@@ -2123,23 +2125,23 @@ class ESM3Core(nn.Module):
                 RESIDUE_PAD_TOKEN,
                 dtype=torch.long,
                 device=device,
-            )
+            )  # (1, l, MAX_RESIDUE_ANNOTATIONS)
         if function_tokens is None:
             function_tokens = torch.full(
                 (1, seq_len, FUNCTION_TOKENS_DEPTH),
                 INTERPRO_PAD_TOKEN,
                 dtype=torch.long,
                 device=device,
-            )
+            )  # (1, l, FUNCTION_TOKENS_DEPTH)
         if structure_coords is None:
             structure_coords = torch.full(
                 (1, seq_len, 3, 3),
                 float("nan"),
                 dtype=torch.float,
                 device=device,
-            )
+            )  # (1, l, 3, 3): residue, backbone atom, xyz axes
 
-        structure_coords = structure_coords[..., :3, :]
+        structure_coords = structure_coords[..., :3, :]  # (..., 3, 3): first three backbone atoms
         affine, affine_mask = build_affine3d_from_coordinates(structure_coords)
 
         structure_tokens = defaults(structure_tokens, STRUCTURE_MASK_TOKEN)
@@ -2379,7 +2381,7 @@ class FastESM3Model(FastPLMTestTimeTrainingMixin, FastESM3PreTrainedModel, Embed
         token_ids = inputs["input_ids"] if isinstance(inputs, dict) else inputs
         single = token_ids.ndim == 1
         if single:
-            token_ids = token_ids.unsqueeze(0)
+            token_ids = token_ids.unsqueeze(0)  # (1, l)
         sequences = self.tokenizer.batch_decode(token_ids, skip_special_tokens=True)
         sequences = [sequence.replace(" ", "") for sequence in sequences]
         return sequences[0] if single else sequences
@@ -2449,7 +2451,7 @@ class FastESM3Model(FastPLMTestTimeTrainingMixin, FastESM3PreTrainedModel, Embed
         single_tensor = token_ids.ndim == 1
         if single_tensor:
             sequence_length = token_ids.shape[0]
-            token_ids = token_ids.unsqueeze(0)
+            token_ids = token_ids.unsqueeze(0)  # (1, l)
             conditioning = {
                 name: (
                     value.unsqueeze(0)
@@ -2462,7 +2464,7 @@ class FastESM3Model(FastPLMTestTimeTrainingMixin, FastESM3PreTrainedModel, Embed
         initial_mask = sampled_ids.eq(SEQUENCE_MASK_TOKEN)
         n_masked = int(initial_mask.sum().item())
         if n_masked == 0:
-            result = sampled_ids.squeeze(0) if single_tensor else sampled_ids
+            result = sampled_ids.squeeze(0) if single_tensor else sampled_ids  # (l,) for one input, otherwise (b, l)
             if return_strings:
                 decoded = self.decode(result)
                 return decoded[0] if single_string and isinstance(decoded, list) else decoded
@@ -2487,7 +2489,7 @@ class FastESM3Model(FastPLMTestTimeTrainingMixin, FastESM3PreTrainedModel, Embed
                     **conditioning,
                 )
             amino_acid_logits = output.sequence_logits[..., 4:29] / config.temperature
-            probabilities = amino_acid_logits.softmax(dim=-1)
+            probabilities = amino_acid_logits.softmax(dim=-1)  # (b, l, amino_acid_vocab_size)
             sampled = (
                 torch.multinomial(
                     probabilities.reshape(-1, probabilities.shape[-1]),
@@ -2495,21 +2497,21 @@ class FastESM3Model(FastPLMTestTimeTrainingMixin, FastESM3PreTrainedModel, Embed
                     generator=generator,
                 ).reshape_as(sampled_ids)
                 + 4
-            )
+            )  # (b, l)
 
             remaining_count = int(remaining.sum().item())
             steps_left = n_steps - step
             fill_count = max(1, (remaining_count + steps_left - 1) // steps_left)
-            confidence = probabilities.max(dim=-1).values.masked_fill(~remaining, -1.0)
-            selected = torch.zeros_like(remaining)
-            flat_selected = selected.reshape(-1)
-            chosen = confidence.reshape(-1).topk(min(fill_count, remaining_count)).indices
-            flat_selected[chosen] = True
-            sampled_ids[selected] = sampled[selected]
+            confidence = probabilities.max(dim=-1).values.masked_fill(~remaining, -1.0)  # (b, l)
+            selected = torch.zeros_like(remaining)  # (b, l)
+            flat_selected = selected.reshape(-1)  # (b * l,)
+            chosen = confidence.reshape(-1).topk(min(fill_count, remaining_count)).indices  # (min(fill_count, remaining_count),)
+            flat_selected[chosen] = True  # chosen positions in (b * l,) storage
+            sampled_ids[selected] = sampled[selected]  # selected positions in (b, l) storage
 
         if bool(sampled_ids.eq(SEQUENCE_MASK_TOKEN).any()):
             raise RuntimeError("generation ended before all sequence masks were filled")
-        result = sampled_ids.squeeze(0) if single_tensor else sampled_ids
+        result = sampled_ids.squeeze(0) if single_tensor else sampled_ids  # (l,) for one input, otherwise (b, l)
         if return_strings:
             decoded = self.decode(result)
             return decoded[0] if single_string and isinstance(decoded, list) else decoded
@@ -2650,19 +2652,19 @@ def _esm3_residue_mask(
 
     tokens = sequence_tokens if sequence_tokens is not None else input_ids
     if attention_mask is None:
-        mask = torch.ones(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)
+        mask = torch.ones(hidden_states.shape[:2], dtype=torch.bool, device=hidden_states.device)  # (b, l)
     else:
-        mask = attention_mask.to(device=hidden_states.device, dtype=torch.bool)
+        mask = attention_mask.to(device=hidden_states.device, dtype=torch.bool)  # (b, l)
     if tokens is not None:
-        tokens = tokens.to(hidden_states.device)
+        tokens = tokens.to(hidden_states.device)  # (b, l)
         special = (
             tokens.eq(SEQUENCE_BOS_TOKEN)
             | tokens.eq(SEQUENCE_PAD_TOKEN)
             | tokens.eq(SEQUENCE_EOS_TOKEN)
             | tokens.eq(SEQUENCE_CHAINBREAK_TOKEN)
-        )
-        mask = mask & ~special
-    return mask
+        )  # (b, l)
+        mask = mask & ~special  # (b, l)
+    return mask  # (b, l)
 
 
 def _esm3_problem_type(
@@ -2724,13 +2726,13 @@ def _esm3_token_classification_loss(
         )
 
     if problem_type == "regression" and num_labels == 1 and labels.ndim == 2:
-        labels = labels.unsqueeze(-1)
+        labels = labels.unsqueeze(-1)  # (b, l, 1)
     if labels.shape != logits.shape:
         raise ValueError(
             "Token regression and multilabel targets must match the logits shape; "
             f"received {tuple(labels.shape)} and {tuple(logits.shape)}."
         )
-    valid = residue_mask.unsqueeze(-1) & labels.ne(-100)
+    valid = residue_mask.unsqueeze(-1) & labels.ne(-100)  # (b, l, num_labels)
     if not bool(valid.any()):
         raise ValueError("Token labels do not contain a supervised biological residue.")
     if problem_type == "regression":
@@ -2798,7 +2800,7 @@ class FastESM3ForSequenceClassification(FastESM3Model):
             return_dict=True,
             **kwargs,
         )
-        hidden_states = output.last_hidden_state
+        hidden_states = output.last_hidden_state  # (b, l, d)
         if hidden_states is None:
             raise RuntimeError("ESM3 did not return final residue embeddings.")
         residue_mask = _esm3_residue_mask(
@@ -2806,13 +2808,13 @@ class FastESM3ForSequenceClassification(FastESM3Model):
             classifier_attention_mask,
             input_ids,
             sequence_tokens,
-        )
-        residue_counts = residue_mask.sum(dim=1, keepdim=True)
+        )  # (b, l)
+        residue_counts = residue_mask.sum(dim=1, keepdim=True)  # (b, 1)
         if not bool(residue_counts.all()):
             raise ValueError("Sequence classification requires one biological residue per row.")
-        pooled = (hidden_states * residue_mask.unsqueeze(-1)).sum(dim=1)
-        pooled = pooled / residue_counts.to(hidden_states.dtype)
-        logits = self.classifier(self.dropout(pooled))
+        pooled = (hidden_states * residue_mask.unsqueeze(-1)).sum(dim=1)  # (b, d)
+        pooled = pooled / residue_counts.to(hidden_states.dtype)  # (b, d)
+        logits = self.classifier(self.dropout(pooled))  # (b, num_labels)
 
         loss = None
         if labels is not None:
@@ -2892,7 +2894,7 @@ class FastESM3ForTokenClassification(FastESM3Model):
             return_dict=True,
             **kwargs,
         )
-        hidden_states = output.last_hidden_state
+        hidden_states = output.last_hidden_state  # (b, l, d)
         if hidden_states is None:
             raise RuntimeError("ESM3 did not return final residue embeddings.")
         residue_mask = _esm3_residue_mask(
@@ -2900,7 +2902,7 @@ class FastESM3ForTokenClassification(FastESM3Model):
             classifier_attention_mask,
             input_ids,
             sequence_tokens,
-        )
+        )  # (b, l)
         logits = self.classifier(self.dropout(hidden_states))
 
         loss = None

fastplms/models/ttt.py

@@ -6,6 +6,7 @@ import numbers
 import torch
 import torch.nn as nn
 import torch.nn.functional as F
+
 from collections.abc import Iterator, Mapping
 from dataclasses import asdict, dataclass, fields
 from typing import Any

fastplms/registry.py

@@ -9,6 +9,7 @@ from __future__ import annotations
 
 import re
 import tomllib
+
 from collections.abc import Iterator, Mapping
 from dataclasses import dataclass
 from functools import lru_cache

modeling_fastplms.py

@@ -13,7 +13,7 @@ from zipfile import ZIP_DEFLATED, ZipFile
 
 from .fastplms_bundle import RUNTIME_DATA, RUNTIME_HASH
 
-if RUNTIME_HASH != "63f353cc98282efb0d16e5c6f1e1a43f46ac01d151d91791d685e76c39d2f6b4":
+if RUNTIME_HASH != "fb292c86fa7e6e771d63e48e4f897c8a75dcfd01726fdbac030bb5355ddf3319":
     raise RuntimeError("FastPLMs runtime identity differs from the bridge.")
 
 _RUNTIME_TEMPORARIES = []