runtime/kev/visual_adapter.py

"""Strict state-dict remapper: released text Kev -> the Qwen3.5 multimodal language tower.

A released Kev adapter carries its LoRA tensors under the text backbone's PEFT path
`base_model.model.layers.<i>.<module>.lora_{A,B}.weight`. Injecting the same LoRA into a Qwen3.5 image-text-to-text
model moves the language tower to `base_model.model.model.language_model.layers.<i>.<module>.lora_{A,B}.weight` — the
text path with `model.language_model.` spliced in after `base_model.model.`. The safetensors audit for `kev-0.8b`
found a strict 372-key bijection between the two (WORKLOG, 2026-09-23), so warming up the multimodal language tower
is a pure key rewrite: same tensors, same shapes, only the prefix moves. This module does that rewrite and nothing
else.

It refuses — never silently repairs — anything that would let a half-loaded adapter still load and still train, which
is the exact hazard `Checkpoint.warm_start` guards against in kev/checkpoint.py. It imports no PEFT, loads no weights
and mutates nothing: it is a `Mapping[str, Tensor] -> dict[str, Tensor]` function that runs on CPU.

    state = remap_state_dict(text_adapter, expected_language_tower)
    summarize_state_dict(text_adapter, expected_language_tower)      # -> RemapSummary(mapped=372, rewritten=372, ...)
"""
from collections.abc import Mapping
from dataclasses import dataclass

import torch

TEXT_PREFIX = "base_model.model.layers."
LANGUAGE_PREFIX = "base_model.model.model.language_model.layers."


class VisualAdapterError(ValueError):
    """A source state dict that will not cleanly warm-start the multimodal language tower.

    `reason` names the gate that failed — `outside_prefix`, `not_a_tensor`, `nonfinite`, `collision`, `missing`,
    `unexpected` or `shape` — and `keys` lists the offending state-dict keys, sorted, so the message is deterministic."""

    def __init__(self, reason, message, keys=()):
        super().__init__(message)
        self.reason = reason
        self.keys = tuple(sorted(keys))


@dataclass(frozen=True)
class RemapSummary:
    """Deterministic counts over a successful remap. `passed_through` must be zero:
    already-remapped keys are not valid inputs to this released-text adapter helper."""

    mapped: int
    rewritten: int
    passed_through: int


def _target_key(key: str) -> str | None:
    """The language-tower key for a released text key, or None otherwise."""
    if key.startswith(TEXT_PREFIX):
        return LANGUAGE_PREFIX + key[len(TEXT_PREFIX):]
    return None


def remap_state_dict(source: Mapping[str, torch.Tensor], expected: Mapping[str, torch.Tensor]) -> dict[str, torch.Tensor]:
    """-> a new dict mapping every language-tower key to its source tensor. Raises VisualAdapterError unless the
    rewrite is an exact, finite, collision-free bijection onto `expected`. Neither argument is mutated."""
    return _remap(source, expected)[0]


def summarize_state_dict(source: Mapping[str, torch.Tensor], expected: Mapping[str, torch.Tensor]) -> RemapSummary:
    """The same strict pass as remap_state_dict, returning the RemapSummary instead of the tensors (it re-checks every
    gate, so it raises on anything remap_state_dict would)."""
    return _remap(source, expected)[1]


def _remap(source, expected):
    state, origin, rewritten = {}, {}, 0
    for key in sorted(source):
        value = source[key]
        target = _target_key(key)
        if target is None:
            raise VisualAdapterError("outside_prefix",
                                     f"{key!r} is outside {TEXT_PREFIX!r}: only text-tower keys move to the "
                                     f"multimodal language tower", keys=(key,))
        if not isinstance(value, torch.Tensor):
            raise VisualAdapterError("not_a_tensor", f"{key!r} is {type(value).__name__}, not a torch.Tensor", keys=(key,))
        if value.is_floating_point() and not torch.isfinite(value).all():
            raise VisualAdapterError("nonfinite", f"{key!r} carries a non-finite value", keys=(key,))
        if target in state:
            raise VisualAdapterError("collision",
                                     f"{key!r} and {origin[target]!r} both map to {target!r}", keys=(key, origin[target]))
        state[target] = value
        origin[target] = key
        rewritten += target != key
    _match(state, expected)
    return state, RemapSummary(len(state), rewritten, len(state) - rewritten)


def _match(state, expected):
    if missing := sorted(set(expected) - set(state)):
        raise VisualAdapterError("missing",
                                 f"the language tower needs {len(missing)} tensor(s) the source does not cover "
                                 f"(e.g. {missing[:2]}); check --lora_targets", keys=missing)
    if unexpected := sorted(set(state) - set(expected)):
        raise VisualAdapterError("unexpected",
                                 f"the source carries {len(unexpected)} tensor(s) the language tower does not have "
                                 f"(e.g. {unexpected[:2]}); check the adapter_config.json", keys=unexpected)
    if shapes := [k for k in sorted(state) if state[k].shape != expected[k].shape]:
        k = shapes[0]
        raise VisualAdapterError("shape",
                                 f"{len(shapes)} tensor(s) disagree on shape, first {k}: source "
                                 f"{tuple(state[k].shape)} vs target {tuple(expected[k].shape)} (e.g. {shapes[:2]})",
                                 keys=shapes)