src/jolt/model.py

"""Parallel candidate scoring over shared text/image computation."""

import hashlib
import json
from pathlib import Path

import torch
import torch.nn.functional as F
from torch import Tensor, nn

from jolt.adapters import Qwen35Adapter
from jolt.recipe import IMAGE_PIXELS, MAX_LENGTH, normalize_architecture


def model_revision(directory):
    """Identify a downloaded Hub snapshot or a locally pinned training base."""
    directory = Path(directory)
    revision = directory / "revision.txt"
    if revision.is_file():
        return revision.read_text().strip()
    if directory.parent.name == "snapshots" and len(directory.name) == 40:
        return directory.name
    raise ValueError("The base model must be a pinned Hub snapshot or contain revision.txt")


class ResidualBottleneck(nn.Module):
    def __init__(self, hidden_size, width, activation):
        super().__init__()
        self.norm = nn.LayerNorm(hidden_size, eps=1e-5)
        self.down = nn.Linear(hidden_size, width, bias=False)
        self.up = nn.Linear(width, hidden_size, bias=False)
        self.activation = F.silu if activation == "silu" else F.gelu
        nn.init.xavier_uniform_(self.down.weight)
        nn.init.zeros_(self.up.weight)

    def forward(self, hidden):
        residual = hidden
        block_input = hidden.to(dtype=self.norm.weight.dtype)
        update = self.up(self.activation(self.down(self.norm(block_input))))
        return residual + update.to(dtype=residual.dtype)


class CandidateInteraction(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        width = 128
        self.heads = 4
        self.head_width = width // self.heads
        self.norm_attention = nn.LayerNorm(hidden_size, eps=1e-5)
        self.query = nn.Linear(hidden_size, width)
        self.key = nn.Linear(hidden_size, width)
        self.value = nn.Linear(hidden_size, width)
        self.output = nn.Linear(width, hidden_size)
        self.norm_mlp = nn.LayerNorm(hidden_size, eps=1e-5)
        self.mlp_in = nn.Linear(hidden_size, 256)
        self.mlp_out = nn.Linear(256, hidden_size)
        for layer in (self.query, self.key, self.value, self.mlp_in):
            nn.init.xavier_uniform_(layer.weight)
            nn.init.zeros_(layer.bias)
        for layer in (self.output, self.mlp_out):
            nn.init.zeros_(layer.weight)
            nn.init.zeros_(layer.bias)

    def forward(self, hidden, valid):
        if valid is None:
            raise ValueError("Candidate attention requires a validity mask")
        output_dtype = hidden.dtype
        hidden = hidden.to(dtype=self.norm_attention.weight.dtype)
        normalized = self.norm_attention(hidden)
        batch, candidates, _ = hidden.shape

        def split_heads(tensor):
            return tensor.view(batch, candidates, self.heads, self.head_width).transpose(1, 2)

        attended = F.scaled_dot_product_attention(
            split_heads(self.query(normalized)),
            split_heads(self.key(normalized)),
            split_heads(self.value(normalized)),
            attn_mask=valid[:, None, None, :],
            dropout_p=0.0,
        )
        attended = attended.transpose(1, 2).reshape(batch, candidates, -1)
        hidden = hidden + self.output(attended)
        hidden = hidden + self.mlp_out(F.gelu(self.mlp_in(self.norm_mlp(hidden))))
        return hidden.masked_fill(~valid[..., None], 0).to(dtype=output_dtype)


class GeluScoringHead(nn.Module):
    def __init__(self, hidden_size):
        super().__init__()
        self.in_proj = nn.Linear(hidden_size, 256, bias=False)
        self.out_proj = nn.Linear(256, 1, bias=False)
        nn.init.xavier_uniform_(self.in_proj.weight)
        nn.init.normal_(self.out_proj.weight, std=0.01)

    def forward(self, hidden):
        return self.out_proj(F.gelu(self.in_proj(hidden.float()))).squeeze(-1)


class InterpolatedActivation(nn.Module):
    def __init__(self, initial_mix):
        super().__init__()
        self.mix = nn.Parameter(torch.tensor(float(initial_mix), device="cuda"))

    def forward(self, value):
        mix = self.mix.to(value.dtype)
        return (1 - mix) * F.silu(value) + mix * F.gelu(value)


class DecisionModel(nn.Module):
    def __init__(self, checkpoint: str | Path, *, adapter=None, architecture=None):
        super().__init__()
        self.adapter = adapter or Qwen35Adapter()
        self.base_path = Path(checkpoint)
        self.architecture = normalize_architecture(architecture)
        self.backbone = self.adapter.load(checkpoint, architecture=self.architecture)
        hidden_size = self.adapter.hidden_size(self.backbone)
        self.processor = self.adapter.processor(checkpoint)
        self.head = nn.Linear(hidden_size, 1, bias=False, device="cuda", dtype=torch.float32)
        nn.init.normal_(self.head.weight, std=0.01)
        self.lora_rank = 8
        self.lora_alpha = 16
        self.lora_dropout = 0.0
        self.marker_id = self.processor.tokenizer.convert_tokens_to_ids(self.adapter.marker)
        if self.processor.tokenizer.encode(self.adapter.marker, add_special_tokens=False) != [
            self.marker_id
        ]:
            raise ValueError("The candidate marker must be a single reserved token")

        # Added modules use a forked RNG so they do not shift the baseline LoRA
        # initialization that follows enable_lora().
        with torch.random.fork_rng(devices=[torch.cuda.current_device()]):
            if self.architecture["head"] == "gelu-mlp":
                self.head = GeluScoringHead(hidden_size).to("cuda")
            if self.architecture["candidate_attention"]:
                self.candidate_block = CandidateInteraction(hidden_size).to("cuda")
            else:
                self.candidate_block = None

            adapter_config = self.architecture["adapters"]
            self.layer_adapters = nn.ModuleDict(
                {
                    str(index): ResidualBottleneck(
                        hidden_size, adapter_config["width"], adapter_config["activation"]
                    ).to("cuda")
                    for index in adapter_config["layers"]
                }
            )

            self.layer_mixture_weights = (
                nn.Parameter(torch.zeros(3, device="cuda"))
                if self.architecture["readout"] == "layer-mixture"
                else None
            )

            activation_config = self.architecture["backbone_activation"]
            self.activation_mix_layers = tuple(activation_config["layers"])
            text_layers = self.adapter.language_layers(self.backbone)
            for index in self.activation_mix_layers:
                if not 0 <= index < len(text_layers):
                    raise ValueError(f"Backbone activation layer is out of range: {index}")
                text_layers[index].mlp.act_fn = InterpolatedActivation(
                    activation_config["initial_mix"]
                ).to("cuda")

        layer_count = len(self.adapter.language_layers(self.backbone))
        self.capture_layers = (20, 21, 22, 23) if self.layer_mixture_weights is not None else ()
        for index in self.capture_layers:
            if index >= layer_count:
                raise ValueError("Layer mixture requires decoder layers 21 through 24")
        for index in self.layer_adapters:
            if not 0 <= int(index) < layer_count:
                raise ValueError(f"Residual adapter layer is out of range: {index}")
        self._adapter_hooks = []

    def forward(
        self, inputs: dict[str, Tensor], positions: Tensor, *, candidate_mask: Tensor | None = None
    ) -> Tensor:
        hidden, offset, intermediate = self.adapter.forward(
            self.backbone, inputs, capture_layers=self.capture_layers
        )
        local_positions = (positions - offset).clamp_min(0)
        rows = torch.arange(hidden.shape[0], device=hidden.device)[:, None]
        if self.layer_mixture_weights is None:
            candidate_hidden = hidden[rows, local_positions]
        else:
            norm = self.adapter.final_norm(self.backbone)
            layer_features = torch.stack(
                [
                    norm(intermediate[index][rows, local_positions])
                    for index in self.capture_layers
                ],
                dim=0,
            )
            layer_features = layer_features.float()
            final = layer_features[-1]
            deltas = layer_features[:-1] - final
            candidate_hidden = final + (
                self.layer_mixture_weights.float()[:, None, None, None] * deltas
            ).sum(0)
        if self.candidate_block is not None:
            candidate_hidden = self.candidate_block(candidate_hidden, candidate_mask)
        if isinstance(self.head, GeluScoringHead):
            return self.head(candidate_hidden)
        return self.head(candidate_hidden.float()).squeeze(-1)

    def enable_lora(self, *, checkpointing=True, rank=8, alpha=16, dropout=0.0):
        self.lora_rank = rank
        self.lora_alpha = alpha
        self.lora_dropout = dropout
        self.adapter.adapt_language(
            self.backbone,
            training=checkpointing,
            rank=rank,
            alpha=alpha,
            dropout=dropout,
        )
        layers = self.adapter.language_layers(self.backbone)
        for index in self.activation_mix_layers:
            layers[index].mlp.act_fn.mix.requires_grad_(True)
        for handle in self._adapter_hooks:
            handle.remove()
        self._adapter_hooks.clear()
        for index, module in self.layer_adapters.items():
            layer = layers[int(index)]

            def apply_adapter(_layer, _inputs, output, adapter_module=module):
                if isinstance(output, tuple):
                    return (adapter_module(output[0]), *output[1:])
                return adapter_module(output)

            self._adapter_hooks.append(layer.register_forward_hook(apply_adapter))

    def load_adapter(self, directory):
        """Load the same verified unmerged weights for learning or deployment."""
        from safetensors.torch import load_file

        directory = Path(directory)
        metadata = json.loads((directory / "dohnuts.json").read_text())
        weights = directory / "adapter.safetensors"
        if (
            metadata.get("format_version") not in {1, 2, 3}
            or metadata.get("adapter", "qwen3.5") != self.adapter.name
            or metadata.get("base_revision") != model_revision(self.base_path)
            or (metadata["image_pixels"], metadata["max_length"], metadata["lora_rank"])
            != (IMAGE_PIXELS, MAX_LENGTH, self.lora_rank)
            or metadata.get("lora_alpha", 16) != self.lora_alpha
            or metadata.get("lora_dropout", 0.0) != self.lora_dropout
            or normalize_architecture(metadata.get("architecture")) != self.architecture
            or hashlib.sha256(weights.read_bytes()).hexdigest() != metadata["weights_sha256"]
        ):
            raise ValueError(
                "Checkpoint differs from the pinned model, training recipe, or checksum"
            )
        state = load_file(str(weights))
        parameters = {n: p for n, p in self.named_parameters() if p.requires_grad}
        if state.keys() != parameters.keys() or any(
            state[name].shape != parameter.shape for name, parameter in parameters.items()
        ):
            raise ValueError("Checkpoint trainable parameters differ")
        with torch.no_grad():
            for name, parameter in parameters.items():
                parameter.copy_(state[name])
        return metadata

    def merge(self):
        self.adapter.merge(self.backbone)
        self.requires_grad_(False).eval()

    def train(self, mode: bool = True):
        super().train(mode)
        self.adapter.freeze_vision(self.backbone)
        return self


def marker_positions(input_ids: Tensor, marker_id: int, candidates: int) -> Tensor:
    positions = []
    for row in input_ids:
        found = (row == marker_id).nonzero(as_tuple=True)[0]
        if found.numel() != candidates:
            raise ValueError(f"Expected {candidates} candidate markers, found {found.numel()}")
        positions.append(found)
    return torch.stack(positions)