code/src/jev_judge/model.py

"""JevJudge = qwen3_5 text backbone (bf16, frozen/LoRA) + 24-slot fp32 decision head (DESIGN §6).

Loading strategy (B200 / v0.7): we never instantiate the vision tower. The HF checkpoint of
`Qwen3_5ForConditionalGeneration` is streamed shard-by-shard into a `Qwen3_5ForCausalLM` built from
`config.text_config`, remapping `model.language_model.*` -> `model.*` and skipping `model.visual.*`
and `mtp.*`. The same loader reads our exported text-only checkpoints.
"""

from __future__ import annotations

import glob
import json
import logging
import os
import re
from dataclasses import dataclass, field
from typing import Iterable, Sequence

import torch
import torch.nn as nn
import torch.nn.functional as F
from safetensors import safe_open

from .head_init import VerbalizerTable, build_head_weight, resolve_verbalizers
from .template import KIND_TO_ID, KINDS, NUM_SLOTS, SLOT_RANGES, SlotLayout

log = logging.getLogger(__name__)

_SKIP_PREFIXES = ("model.visual.", "mtp.")
_REMAP = [(re.compile(r"^model\.language_model\."), "model.")]


@dataclass
class LoraSpec:
    r: int = 16
    alpha: int = 32
    dropout: float = 0.05
    # Linear leaves of the qwen3_5 decoder layer (DESIGN V3). Fused/unfused GDN projections are
    # both listed; peft only attaches to names that exist.
    target_modules: list[str] = field(
        default_factory=lambda: [
            "in_proj_qkvz", "in_proj_ba", "in_proj_qkv", "in_proj_z", "in_proj_a", "in_proj_b",
            "q_proj", "k_proj", "v_proj", "o_proj",
            "gate_proj", "up_proj", "down_proj", "out_proj",
        ]
    )


# --------------------------------------------------------------------------------------------
# Backbone loading
# --------------------------------------------------------------------------------------------
def _remap_key(k: str) -> str | None:
    if k.startswith(_SKIP_PREFIXES):
        return None
    for pat, rep in _REMAP:
        k = pat.sub(rep, k)
    return k


def load_text_causal_lm(model_path: str, device: str | torch.device = "cuda", dtype: torch.dtype = torch.bfloat16,
                        attn_implementation: str = "sdpa", require_lm_head: bool = True):
    """Build Qwen3_5ForCausalLM from text_config and stream text weights from safetensors shards.

    Works for the HF multimodal checkpoint (keys `model.language_model.*`, vision/mtp skipped) and for our
    exported text-only bundles (keys `model.*`, no lm_head -> pass require_lm_head=False)."""
    from accelerate import init_empty_weights
    from accelerate.utils import set_module_tensor_to_device
    from transformers import AutoConfig, Qwen3_5ForCausalLM

    cfg = AutoConfig.from_pretrained(model_path)
    text_cfg = getattr(cfg, "text_config", cfg)
    text_cfg._attn_implementation = attn_implementation
    with init_empty_weights(include_buffers=False):
        model = Qwen3_5ForCausalLM(text_cfg)
    expected = set(model.state_dict().keys())
    seen: set[str] = set()
    shards = sorted(glob.glob(os.path.join(model_path, "*.safetensors")))
    shards = [s for s in shards if os.path.basename(s) != "head.safetensors"]
    if not shards:
        raise FileNotFoundError(f"no safetensors shards under {model_path}")
    for shard in shards:
        with safe_open(shard, framework="pt", device="cpu") as f:
            for k in f.keys():
                nk = _remap_key(k)
                if nk is None or nk not in expected:
                    continue
                t = f.get_tensor(k)
                # keep fp32 tensors fp32 (GDN A_log / gated-norm weights are stored fp32 on purpose);
                # everything else (bf16 in the checkpoint) is cast to `dtype`
                tgt_dtype = torch.float32 if t.dtype == torch.float32 else dtype
                set_module_tensor_to_device(model, nk, device, value=t, dtype=tgt_dtype)
                seen.add(nk)
    missing = expected - seen
    if missing == {"lm_head.weight"} and not require_lm_head:
        model.lm_head = None
        missing = set()
    if missing:
        raise RuntimeError(f"{len(missing)} text weights missing from checkpoint, e.g. {sorted(missing)[:5]}")
    model.to(device)  # buffers (inv_freq etc.) were created on CPU
    model.eval()
    for p in model.parameters():
        p.requires_grad_(False)
    n_params = sum(p.numel() for p in model.parameters())
    log.info("loaded text backbone %s: %.2fB params, %d shards", model_path, n_params / 1e9, len(shards))
    return model, text_cfg


# --------------------------------------------------------------------------------------------
# Decision head
# --------------------------------------------------------------------------------------------
class DecisionHead(nn.Module):
    """Pure linear fp32 head H -> 24 slots. `softcap` is only for non-qwen backbones (kept None here)."""

    def __init__(self, hidden_size: int, num_slots: int = NUM_SLOTS, softcap: float | None = None):
        super().__init__()
        self.proj = nn.Linear(hidden_size, num_slots, bias=True, dtype=torch.float32)
        self.softcap = softcap

    def forward(self, h: torch.Tensor) -> torch.Tensor:
        z = self.proj(h.to(torch.float32))
        if self.softcap:
            z = self.softcap * torch.tanh(z / self.softcap)
        return z

    @torch.no_grad()
    def init_from_lm_head(self, lm_head_weight: torch.Tensor, verbalizer_ids: Sequence[int]) -> None:
        W, b = build_head_weight(lm_head_weight, list(verbalizer_ids))
        self.proj.weight.copy_(W.to(self.proj.weight.device))
        self.proj.bias.copy_(b.to(self.proj.bias.device))


def slot_mask(kind_ids: torch.Tensor, n_options: torch.Tensor, num_slots: int = NUM_SLOTS) -> torch.Tensor:
    """bool [B, 24]: active slots per sample (noul 0-1, score 2-7, choice 8..8+n-1)."""
    B = kind_ids.shape[0]
    ar = torch.arange(num_slots, device=kind_ids.device).unsqueeze(0).expand(B, -1)
    starts = torch.empty(B, dtype=torch.long, device=kind_ids.device)
    ends = torch.empty_like(starts)
    for k, kid in KIND_TO_ID.items():
        s, e = SLOT_RANGES[k]
        sel = kind_ids == kid
        starts[sel] = s
        ends[sel] = (s + n_options[sel]) if k == "choice" else e
    return (ar >= starts.unsqueeze(1)) & (ar < ends.unsqueeze(1))


def masked_log_softmax(z: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    """log-probs over active slots; inactive slots stay -inf (so .exp() gives exact 0 there)."""
    z = z.to(torch.float32).masked_fill(~mask, float("-inf"))
    return F.log_softmax(z, dim=-1)


def masked_probs(z: torch.Tensor, mask: torch.Tensor) -> torch.Tensor:
    return masked_log_softmax(z, mask).exp()


# --------------------------------------------------------------------------------------------
# Judge
# --------------------------------------------------------------------------------------------
class JevJudge(nn.Module):
    def __init__(self, causal_lm, head: DecisionHead, tokenizer, verbalizers: VerbalizerTable, layout: SlotLayout,
                 base_model_path: str):
        super().__init__()
        self.lm = causal_lm  # Qwen3_5ForCausalLM (possibly peft-wrapped); .model is the text backbone
        self.head = head
        self.tokenizer = tokenizer
        self.verbalizers = verbalizers
        self.layout = layout
        self.base_model_path = base_model_path

    # -- construction --------------------------------------------------------------------------
    @classmethod
    def from_base(cls, model_path: str, device: str = "cuda", dtype: torch.dtype = torch.bfloat16,
                  keep_lm_head: bool = False, attn_implementation: str = "sdpa") -> "JevJudge":
        from transformers import AutoTokenizer

        tok = AutoTokenizer.from_pretrained(model_path)
        if tok.pad_token_id is None:
            tok.pad_token = tok.eos_token
        lm, text_cfg = load_text_causal_lm(model_path, device=device, dtype=dtype, attn_implementation=attn_implementation)
        vt = resolve_verbalizers(tok, strict=True)
        head = DecisionHead(text_cfg.hidden_size).to(device)
        head.init_from_lm_head(lm.lm_head.weight, vt.ids)
        judge = cls(lm, head, tok, vt, SlotLayout(), base_model_path=model_path)
        if not keep_lm_head:
            judge.drop_lm_head()
        return judge

    @classmethod
    def from_export(cls, export_dir: str, device: str = "cuda", dtype: torch.dtype = torch.bfloat16,
                    attn_implementation: str = "sdpa", merge_adapter: bool = True) -> tuple["JevJudge", dict]:
        """Load a bundle written by scripts/export.py (merged text-only weights + head + tokenizer)."""
        from safetensors.torch import load_file
        from transformers import AutoTokenizer

        with open(os.path.join(export_dir, "judge_config.json")) as f:
            jc = json.load(f)
        tok = AutoTokenizer.from_pretrained(export_dir)
        if tok.pad_token_id is None:
            tok.pad_token = tok.eos_token
        lm, text_cfg = load_text_causal_lm(export_dir, device=device, dtype=dtype, attn_implementation=attn_implementation,
                                           require_lm_head=False)
        adapter = jc.get("adapter_subfolder")
        if adapter and os.path.isdir(os.path.join(export_dir, adapter)):
            # dual-mode bundle: pristine base weights + LoRA adapter. For decision serving we fold the
            # adapter into the in-memory weights (same latency as a merged bundle); the files on disk
            # stay pristine so `AutoModelForCausalLM.from_pretrained(bundle)` is exactly the base model.
            from peft import PeftModel

            lm = PeftModel.from_pretrained(lm, os.path.join(export_dir, adapter), is_trainable=False)
            if merge_adapter:
                lm = lm.merge_and_unload()
        vt = resolve_verbalizers(tok, strict=True)
        if jc.get("verbalizer_ids") != list(vt.ids):
            raise RuntimeError("verbalizer ids in judge_config.json do not match the bundled tokenizer")
        head = DecisionHead(text_cfg.hidden_size, softcap=jc.get("softcap")).to(device)
        head.load_state_dict({k: v.to(device) for k, v in load_file(os.path.join(export_dir, "head.safetensors")).items()})
        judge = cls(lm, head, tok, vt, SlotLayout(), base_model_path=jc.get("base_model_path", export_dir))
        judge.drop_lm_head()  # decisions never touch the vocab projection; free it if the bundle ships one
        judge.eval()
        return judge, jc

    def drop_lm_head(self) -> None:
        """Free the vocab projection (≈2-2.5 GB); the head already holds the 24 rows we need."""
        if getattr(self.lm, "lm_head", None) is not None:
            self.lm.lm_head = None
            torch.cuda.empty_cache()

    @property
    def backbone(self):
        m = self.lm
        if hasattr(m, "get_base_model"):
            m = m.get_base_model()
        return m.model

    @property
    def hidden_size(self) -> int:
        return self.head.proj.in_features

    @property
    def device(self) -> torch.device:
        return self.head.proj.weight.device

    # -- LoRA -----------------------------------------------------------------------------------
    def attach_lora(self, spec: LoraSpec) -> list[str]:
        from peft import LoraConfig, get_peft_model

        present = {n.rsplit(".", 1)[-1] for n, m in self.lm.named_modules() if isinstance(m, nn.Linear)}
        targets = [t for t in spec.target_modules if t in present]
        if not targets:
            raise RuntimeError(f"none of the LoRA targets exist; present Linear leaves: {sorted(present)}")
        cfg = LoraConfig(r=spec.r, lora_alpha=spec.alpha, lora_dropout=spec.dropout, bias="none",
                         target_modules=targets, task_type=None)
        self.lm = get_peft_model(self.lm, cfg)
        # peft casts nothing here; keep adapters fp32 for stable small-lr updates.
        for n, p in self.lm.named_parameters():
            if "lora_" in n:
                p.data = p.data.to(torch.float32)
                p.requires_grad_(True)
        return targets

    def trainable_parameters(self, stage: str) -> dict[str, list[nn.Parameter]]:
        groups = {"head": list(self.head.parameters())}
        if stage == "s2":
            groups["lora"] = [p for n, p in self.lm.named_parameters() if "lora_" in n]
        elif stage == "s3":
            groups["backbone"] = [p for n, p in self.lm.named_parameters() if "lora_" not in n]
        return groups

    def set_stage_grads(self, stage: str) -> None:
        for p in self.lm.parameters():
            p.requires_grad_(False)
        for name, ps in self.trainable_parameters(stage).items():
            for p in ps:
                p.requires_grad_(True)

    # -- forward --------------------------------------------------------------------------------
    def hidden_last(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, lengths: torch.Tensor) -> torch.Tensor:
        """Final-norm hidden state at the last real token, fp32 [B, H]."""
        out = self.backbone(input_ids=input_ids, attention_mask=attention_mask, use_cache=False)
        hs = out.last_hidden_state  # [B, T, H] bf16
        idx = (lengths - 1).clamp_min(0)
        h = hs[torch.arange(hs.shape[0], device=hs.device), idx]
        return h.to(torch.float32)

    def forward(self, input_ids: torch.Tensor, attention_mask: torch.Tensor, lengths: torch.Tensor,
                kind_ids: torch.Tensor, n_options: torch.Tensor) -> tuple[torch.Tensor, torch.Tensor]:
        """Returns (logits fp32 [B,24], mask bool [B,24]). Masked positions are meaningless."""
        h = self.hidden_last(input_ids, attention_mask, lengths)
        z = self.head(h)
        return z, slot_mask(kind_ids, n_options)

    # -- reference for the equivalence gate -----------------------------------------------------
    @torch.no_grad()
    def restricted_reference(self, input_ids, attention_mask, lengths) -> torch.Tensor:
        """Zero-shot restricted next-token logits over the 24 verbalizers via lm_head (fp32 recompute).

        Requires keep_lm_head=True at construction."""
        lm_head = getattr(self.lm, "lm_head", None)
        if lm_head is None:
            raise RuntimeError("lm_head was dropped; construct with keep_lm_head=True for the gate")
        h = self.hidden_last(input_ids, attention_mask, lengths)
        rows = lm_head.weight[torch.as_tensor(self.verbalizers.ids, device=h.device)].to(torch.float32)
        return h @ rows.T

    # -- housekeeping ---------------------------------------------------------------------------
    def judge_config(self, extra: dict | None = None) -> dict:
        d = {
            "base_model_path": self.base_model_path,
            "hidden_size": self.hidden_size,
            "slots": self.layout.to_dict(),
            "verbalizer_ids": list(self.verbalizers.ids),
            "kinds": list(KINDS),
        }
        if extra:
            d.update(extra)
        return d

    def linear_leaf_table(self) -> dict[str, int]:
        """Unique Linear leaf names inside decoder layers with counts (DESIGN V3 module table)."""
        counts: dict[str, int] = {}
        for n, m in self.backbone.named_modules():
            if isinstance(m, nn.Linear) and ("." + n).count(".layers.") > 0:
                leaf = n.rsplit(".", 1)[-1]
                counts[leaf] = counts.get(leaf, 0) + 1
        return dict(sorted(counts.items()))