Bundle the qevi inference engine in-repo

qevi/core.py

@@ -0,0 +1,179 @@
+"""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)]