qevi/core.py

"""The validated Qevi readout, in one place.

Consolidates the routine that had been copy-pasted into eleven experiment scripts. Every design
decision here is one this project measured rather than assumed:

  packed layout + block-diagonal mask   plan/02 section 3; exact order invariance, research/19
  widened shared V-prefix               research/31: the image is wrapped once in
                                        <|im_start|>user\\n<|vision_start|> ... <|vision_end|> so
                                        every question's causal view reconstructs a real chat turn
  corrected M-RoPE text offset          research/18 section 6: text resumes at
                                        max(grid_h, grid_w) // spatial_merge_size, NOT at the raw
                                        image token count
  native LM-head readout                research/31/32: reading the trunk's own logits at each
                                        sentinel beats a trained typed head by 15-34 points on
                                        every task tested. Do NOT delete the LM head and train
                                        typed heads, which is what plan/02 section 4.4 originally
                                        specified.

One forward pass answers N questions. No text is generated. `forward_logits` returns, per question,
the raw logits over that question's own candidate answer tokens, so callers softmax over 2 logits
for noul, K for choice, L for score.

This module is import-safe without torch for the pure-python helpers at the bottom, which the
corpus builder and analysis scripts use.
"""

from dataclasses import dataclass
from typing import Sequence

PREFIX_PRE = "<|im_start|>user\n<|vision_start|>"
PREFIX_POST = "<|vision_end|>"

NOUL_TEMPLATE = "Statement: {s}\nIs this statement true of the image? Answer Yes or No."
LETTERS = ["A", "B", "C", "D", "E", "F", "G", "H"]


def wrap_turn(body: str) -> str:
    """Close the user turn and open the assistant turn, so the sentinel sits exactly where the
    model would begin generating its answer."""
    return f"{body}<|im_end|>\n<|im_start|>assistant\n"


def noul_question(statement: str) -> str:
    return wrap_turn(NOUL_TEMPLATE.format(s=statement))


def choice_question(options: Sequence[str]) -> str:
    opts = " ".join(f"{LETTERS[i]}) {o}" for i, o in enumerate(options))
    return wrap_turn(f"Which of these is in the image? {opts}\nAnswer with the letter.")


def score_question(levels: Sequence[str], subject: str = "people") -> str:
    lv = " / ".join(l.capitalize() for l in levels)
    return wrap_turn(f"How many {subject} are in this image? Answer with one word: {lv}.")


def choice_question_generic(text: str, options: Sequence[str]) -> str:
    """Like choice_question, but the caller's own question text is the prompt, not a hardcoded
    "which of these is in the image". Needed for the API (server.py), where a caller's choice
    question is arbitrary ("What product category does this belong to?"), not always about
    object presence the way every choice_question call site in the research scripts was."""
    opts = " ".join(f"{LETTERS[i]}) {o}" for i, o in enumerate(options))
    return wrap_turn(f"{text} {opts}\nAnswer with the letter.")


def score_question_generic(text: str, levels: Sequence[str]) -> str:
    """Like score_question, but the caller's own question text is the prompt, not the hardcoded
    "how many {subject}" counting template. Needed for the API, where a score question is
    arbitrary ("Rate the condition"), not always a count."""
    lv = " / ".join(l.capitalize() for l in levels)
    return wrap_turn(f"{text} Answer with one word: {lv}.")


@dataclass
class Readout:
    """Where to read, and which vocabulary entries constitute the answer alphabet."""
    question_text: str
    candidate_token_ids: list


def softmax(logits, T=1.0):
    import math
    m = max(l / T for l in logits)
    ex = [math.exp(l / T - m) for l in logits]
    s = sum(ex)
    return [e / s for e in ex]


def build_readouts(tokenizer, specs):
    """specs: list of (question_text, [answer strings]) -> list[Readout]."""
    out = []
    for qtext, answers in specs:
        ids = [tokenizer.encode(a, add_special_tokens=False)[0] for a in answers]
        out.append(Readout(question_text=qtext, candidate_token_ids=ids))
    return out


def unwrap(model):
    """Reach the underlying Qwen3VL model whether or not `model` is PEFT/LoRA-wrapped.

    Needed because PEFT replaces `model.model` with a LoraModel and `model.get_rope_index` then
    does not exist. PEFT injects its adapter layers in place, so the base model's own forward and
    lm_head still route through them and gradients still reach the LoRA parameters; the wrapper is
    only needed for saving.
    """
    return model.get_base_model() if hasattr(model, "get_base_model") else model


def forward_logits(model, processor, image, readouts, need_grad=False):
    """One packed forward pass over `image` answering every readout independently.

    Returns list[list[float]] when need_grad is False, or a list of 1-D tensors when True (so a
    training loop can backprop through them). The image tokens are causally isolated from all
    question text, and each question is isolated from every other by the block-diagonal mask.

    Works with a plain Qwen3VLForConditionalGeneration or a PEFT-wrapped one.
    """
    import torch
    from .pack import Question, build_packed_batch

    base = unwrap(model)
    ctx = torch.enable_grad() if need_grad else torch.no_grad()
    with ctx:
        ii = processor.image_processor(images=[image], return_tensors="pt")
        pixel_values = ii["pixel_values"].to(model.device, model.dtype)
        grid = ii["image_grid_thw"].to(model.device)

        sms = model.config.vision_config.spatial_merge_size
        _, gh, gw = grid[0].tolist()
        gt = grid[0].tolist()[0]
        n_img = (gt * gh * gw) // (sms**2)
        img_tok = model.config.image_token_id

        tok = processor.tokenizer
        pre = tok(PREFIX_PRE, add_special_tokens=False).input_ids
        post = tok(PREFIX_POST, add_special_tokens=False).input_ids
        L1, L2 = len(pre), len(post)
        v_len = L1 + n_img + L2
        tso = max(gh, gw) // sms  # research/18 section 6

        toks = [tok(r.question_text, add_special_tokens=False).input_ids for r in readouts]
        qs = [Question(qtype="noul", context_len=len(t), readout_lens=[]) for t in toks]
        batch = build_packed_batch(image_len=v_len, questions=qs, position_scheme="reset",
                                   text_start_offset=L1 + tso + L2)

        input_ids = torch.full((batch.seq_len,), img_tok, dtype=torch.long)
        input_ids[:L1] = torch.tensor(pre)
        input_ids[L1 + n_img : v_len] = torch.tensor(post)
        cur, sentinels = v_len, []
        for t in toks:
            input_ids[cur : cur + len(t)] = torch.tensor(t)
            cur += len(t)
            sentinels.append(cur - 1)
        input_ids = input_ids.unsqueeze(0).to(model.device)

        ph = torch.full((1, n_img), img_tok, dtype=torch.long, device=model.device)
        tt = torch.ones((1, n_img), dtype=torch.int, device=model.device)
        ipos, _ = base.model.get_rope_index(input_ids=ph, mm_token_type_ids=tt,
                                            image_grid_thw=grid)

        pos = torch.zeros(3, 1, batch.seq_len, dtype=torch.long, device=model.device)
        pos[:, :, :L1] = torch.arange(L1, device=model.device)
        pos[:, :, L1 : L1 + n_img] = ipos + L1
        pos[:, :, L1 + n_img : v_len] = (L1 + tso) + torch.arange(L2, device=model.device)
        ts = batch.position_ids[0].to(model.device)
        pos[:, :, v_len:] = ts[v_len:]

        out = base.model(input_ids=input_ids,
                         attention_mask=batch.attn_bias.to(dtype=model.dtype,
                                                           device=model.device),
                         position_ids=pos, pixel_values=pixel_values, image_grid_thw=grid)
        hidden = out.last_hidden_state[0]
        stacked = torch.stack([hidden[p] for p in sentinels]).to(model.dtype)
        logits = base.lm_head(stacked).float()

        if need_grad:
            return [logits[i, r.candidate_token_ids] for i, r in enumerate(readouts)]
        return [[logits[i, c].item() for c in r.candidate_token_ids]
                for i, r in enumerate(readouts)]