"""Inference for v-jev 1.4 (each size in its folder: 0.8b/, 4b/): a state and typed questions in, decisions and reads out, in one forward pass.

    from inference import VJev
    jev = VJev(size="4b")                         # the 4B model in 4b/ next to this file (default size: 0.8b)
    jev = VJev(repo="haohan84/v-jev1.4", size="4b")   # or from the Hub (downloads only that size)
    answers = jev.ask(state, questions)           # text or JSON state
    answers = jev.ask(state, questions, image="photo.jpg")   # an image in the state (decisions only)

`state` is text or any JSON value. `questions` and `answers` follow the Jev request shape (POST /v1/systemone):
  noul    {"type": "noul", "instructions": ...}                           -> {"type": "noul", "noul": p(yes)}
  choice  {"type": "choice", "instructions": ..., "criteria": {id: text}}  -> {"type": "choice", "probabilities": {id: p}}
  score   {"type": "score", "instructions": ..., "criteria": [levels]}     -> {"type": "score", "probabilities": {"0": p, ...}}
  read    {"type": "read", "instructions": ..., "kind": "text"}            -> {"type": "read", "value": text or None, "found": p}
A noul may also take "criteria": {"true": text, "false": text}. A choice takes 2 to 64 options; an option's text may
be empty (null or ""), then only its id is shown. A score takes 2 to 10 levels as a list, lowest first, or as an
object {id: text}, lowest first; its answer is then keyed by those ids.

A read takes a value out of the state, as written: "value" is the state's own characters from the chosen start token
to the chosen end token (trimmed of a trailing ":" ";" "," and unmatched brackets), or None when the model says the
value is not there. "found" is the model's probability that it is there; below the fitted threshold (vjev.json) the
value is None. "kind" is how the value is asked: "text" (default), "number" or "date", or one of the value types of
training (money, quantity, percent -> number; person, org, place, address, phone, email, url, id, time, duration,
category, attribute, geo -> text). For "number" and "date" the answer also has "parsed": the value as a float or as
YYYY-MM-DD, or None when the span cannot be parsed strictly (no arithmetic, no unit conversion).
`lock=True` asks every read three ways in the same batch: as served (packed with the request's other questions),
alone, and alone with the state's first line dropped. The answer then has "locked": the same span every time (or
nothing every time); when it is not locked, "value" is None. On real text a locked read is much more precise.

Images: `image` is a path or a PIL image; the state may also hold text beside it. Each decision question about an
image is asked in a prompt of its own, as in training. Reads are text only.

How it is run. Nothing is generated. The prompt is words: the state, then each question with its options numbered,
then one <ans> per decision and, after the decisions, each read's question line with a start and an end <ans>.
At each decision <ans>: a choice points at its options (the <ans> state against each option line's last token,
through the pointer head vjev.pointer.*); a noul or score reads the answer-token rows <no> <yes> / <lvl_k>. A read's
start and end slots point, through the span head vjev.span.* over [last layer ; the middle layer in vjev.json: 12 at 0.8B, 16 at 4B], at a token of the state or
at the "nothing" key (the ":" of "State:"). Decisions about one state share a prompt (at most 8 to a prompt), the
reads go after them in the same prompt (at most 8 reads to a prompt), and every prompt of a request runs in one batch.

The state and the question texts are data: any added token of the tokenizer written out in them is broken up ("<ans>"
-> "< ans>", "<|im_end|>" -> "< |im_end|>") before the prompt is built, so it stays text. A read's value is cut from
the state as the caller wrote it. Text in the state that imitates the prompt's own lines cannot add answer positions,
but it can sway the answers.

Needs torch, transformers (5.x, with Qwen3.5), safetensors, numpy and PIL.
"""
import json
import math
import os
import re
import warnings

import numpy as np
import torch
from safetensors import safe_open
from transformers import AutoModelForCausalLM, AutoTokenizer
from transformers.utils import logging as hf_logging

OPT = [f"<opt_{i}>" for i in range(1, 17)]
LVL = [f"<lvl_{i}>" for i in range(10)]
NOUL = ["<no>", "<yes>"]
ANS = "<ans>"
MAX_OPTIONS = 64
TRAINED_OPTIONS = 16  # the decision anchor trained on at most 16 options in one question
PER_PROMPT = 8  # decisions per prompt; reads per prompt
TRAINED_TOKENS = 4608  # the longest prompt of training; longer ones are untested
KIND_LABEL = {"choice": "choice", "noul": "yes/no", "score": "score"}
STATE_PREFIX = "State:\n"
NOTHING_CHAR = len("State")  # the ":" of "State:": the read heads' "nothing" key
CAPS = {"text": 128, "number": 32, "date": 32}  # longest read span, in tokens
READ_KINDS = {"text": "text", "number": "number", "date": "date", "money": "number", "quantity": "number", "percent": "number"}
READ_KINDS.update({t: "text" for t in ("person", "org", "place", "address", "phone", "email", "url", "id", "time", "duration", "category", "attribute", "geo")})
VSTART, IMG, VEND = "<|vision_start|>", "<|image_pad|>", "<|vision_end|>"


class TooLong(ValueError):
    """A prompt is over max_tokens."""


# ---------- typed values of a read (strict: no arithmetic, no unit conversion) ----------

NUM = re.compile(r"[-+−]?\d[\d\s.,'’]*\d|[-+−]?\d")
SCALE = {"k": 1e3, "thousand": 1e3, "m": 1e6, "mn": 1e6, "million": 1e6, "b": 1e9, "bn": 1e9, "billion": 1e9}
MONTHS = {m: i for i, ms in enumerate([["jan", "january"], ["feb", "february"], ["mar", "march"], ["apr", "april"], ["may"], ["jun", "june"],
                                       ["jul", "july"], ["aug", "august"], ["sep", "sept", "september"], ["oct", "october"], ["nov", "november"],
                                       ["dec", "december"]], 1) for m in ms}


def _decimal_mark(context):
    dots = len(re.findall(r"\d,\d{3}\.\d", context)) + len(re.findall(r"\d\.\d{1,2}\b(?![.,]\d)", context))
    commas = len(re.findall(r"\d\.\d{3},\d", context)) + len(re.findall(r"\d,\d{1,2}\b(?![.,]\d)", context))
    return "." if dots > commas else "," if commas > dots else None


def parse_number(span, context=""):
    """The number written in `span` ("1 234,50 EUR" -> 1234.5, "2.5k" -> 2500.0), or None. Grouping and the decimal
    mark come from the span where it settles them, else from how the context writes numbers, else '.' is decimal."""
    m = NUM.search(span)
    if not m:
        return None
    t = re.sub(r"[\s'’]", "", m.group(0).replace("−", "-"))
    if "," in t and "." in t:
        mark = "." if t.rfind(".") > t.rfind(",") else ","
    elif "," in t or "." in t:
        sep = "," if "," in t else "."
        parts = t.lstrip("+-").split(sep)
        if len(parts) > 2:
            mark = "," if sep == "." else "."
        elif len(parts[1]) == 3 and parts[0] not in ("0", ""):
            mark = sep if _decimal_mark(context) == sep else ("." if sep == "," else ",")
        else:
            mark = sep
    else:
        mark = "."
    t = t.replace("," if mark == "." else ".", "").replace(mark, ".")
    try:
        v = float(t)
    except ValueError:
        return None
    w = re.match(r"([a-z]+)", span[m.end():].strip().lower())
    if w and w.group(1) in SCALE:
        v *= SCALE[w.group(1)]
    return v


def _day_first(context):
    df = mf = 0
    for a, b, _ in re.findall(r"\b(\d{1,2})[/.-](\d{1,2})[/.-](\d{2,4})\b", context):
        a, b = int(a), int(b)
        df += a > 12 >= b
        mf += b > 12 >= a
    return True if df > mf else False if mf > df else None


def _ymd(y, mo, d):
    return f"{y:04d}-{mo:02d}-{d:02d}" if 1 <= mo <= 12 and 1 <= d <= 31 and 1000 <= y <= 2999 else None


def parse_date(span, context=""):
    """The date written in `span` as YYYY-MM-DD, or None. Numeric d/m/y or m/d/y takes its order from the context's
    other dates, else month first."""
    s = span.strip().lower().replace(",", " ")
    m = re.search(r"\b(\d{4})[-/.](\d{1,2})[-/.](\d{1,2})\b", s)
    if m:
        return _ymd(*map(int, m.groups()))
    m = re.search(r"\b(\d{1,2})(?:st|nd|rd|th)?\s+(?:of\s+)?([a-z]{3,9})\.?\s+(\d{4})\b", s)
    if m and m.group(2) in MONTHS:
        return _ymd(int(m.group(3)), MONTHS[m.group(2)], int(m.group(1)))
    m = re.search(r"\b([a-z]{3,9})\.?\s+(\d{1,2})(?:st|nd|rd|th)?\s+(\d{4})\b", s)
    if m and m.group(1) in MONTHS:
        return _ymd(int(m.group(3)), MONTHS[m.group(1)], int(m.group(2)))
    m = re.search(r"\b(\d{1,2})[/.-](\d{1,2})[/.-](\d{4}|\d{2})\b", s)
    if m:
        a, b, y = int(m.group(1)), int(m.group(2)), int(m.group(3))
        y += 2000 if y < 100 else 0
        if a > 12 >= b:
            return _ymd(y, b, a)
        if b > 12 >= a:
            return _ymd(y, a, b)
        return _ymd(y, b, a) if _day_first(context) else _ymd(y, a, b)
    return None


def trim(text):
    """A span's edges: a trailing ':' ';' ',' and an unmatched bracket at either end go; nothing inside changes."""
    t = re.sub(r"[:;,]+$", "", text.strip()).strip()
    for o, c in ("()", "[]", "{}", "<>"):
        if t.endswith(c) and t.count(o) < t.count(c):
            t = t[:-1].strip()
        if t.startswith(o) and t.count(c) < t.count(o):
            t = t[1:].strip()
    return t


# ---------- questions and the prompt ----------

def to_row(key, q, image=False):
    """A request question as the prompt's row: kind, question text, options, descriptions (and read kind)."""
    if not isinstance(q, dict) or not isinstance(q.get("instructions"), str):
        raise ValueError(f'question "{key}": needs "instructions" as text')
    kind, c = q.get("type"), q.get("criteria")
    row = {"key": key, "kind": kind, "question": q["instructions"], "descriptions": None}
    if kind == "noul":
        if c is not None:
            if not isinstance(c, dict) or not {"true", "false"} <= set(c):
                raise ValueError(f'question "{key}": a noul takes no "criteria", or an object with "true" and "false"')
            row["descriptions"] = [str(c["false"]), str(c["true"])]
        row["options"] = ["no", "yes"]
    elif kind == "choice":
        if not isinstance(c, dict) or not 2 <= len(c) <= MAX_OPTIONS:
            raise ValueError(f'question "{key}": a choice needs "criteria" as an object of 2 to {MAX_OPTIONS} options')
        row["options"] = [str(k) for k in c]
        row["descriptions"] = [None if v is None else str(v) for v in c.values()]
        row["names"] = list(c)
    elif kind == "score":
        levels = list(c.values()) if isinstance(c, dict) else c
        if not isinstance(levels, list) or not 2 <= len(levels) <= len(LVL):
            raise ValueError(f'question "{key}": a score needs "criteria" as a list (or object) of 2 to {len(LVL)} levels, lowest first')
        row["options"] = [str(x) for x in levels]
        row["names"] = list(c) if isinstance(c, dict) else [str(i) for i in range(len(levels))]
    elif kind == "read":
        if image:
            raise ValueError(f'question "{key}": reads are text only; ask it without the image (the state text alone)')
        rk = q.get("kind", "text")
        if rk not in READ_KINDS:
            raise ValueError(f'question "{key}": "kind" must be one of {sorted(READ_KINDS)}')
        row["read_kind"], row["options"] = READ_KINDS[rk], []
    else:
        raise ValueError(f'question "{key}": "type" must be noul, choice, score or read')
    return row


def question_lines(k, r):
    """The lines of decision k, and for a choice the index of each option line among them."""
    out = [f"q{k} ({KIND_LABEL[r['kind']]}): {r['question']}"]
    desc = r["descriptions"] or [None] * len(r["options"])
    opt_lines = []
    if r["kind"] == "choice":
        for i, (o, d) in enumerate(zip(r["options"], desc), 1):
            opt_lines.append(len(out))
            out.append(f"  {i}. {o}" + (f": {d}" if d else ""))
    elif r["kind"] == "score":
        out += [f"  {i}. {o}" for i, o in enumerate(r["options"])]
    elif r["descriptions"]:
        out += [f"  no: {desc[0]}", f"  yes: {desc[1]}"]
    return [line.rstrip() for line in out], opt_lines  # the chat template trims trailing space at the end


def layout(state, rows, block=None):
    """Request text, answer text, per row its option-line ends (characters, in the request), the state's character
    range in the request, and the slot plan: (row, None) per decision, (row, "start"), (row, "end") per read.
    The decisions come first; each read follows them in the answer block. `block`: an image block before the state."""
    lines = ["State:"] + ([block] if block is not None else []) + [state, "", "Questions:"]
    ends, plan = [[] for _ in rows], []
    dec = [i for i, r in enumerate(rows) if r["kind"] != "read"]
    reads = [i for i, r in enumerate(rows) if r["kind"] == "read"]
    for k, i in enumerate(dec, 1):
        qs, opt = question_lines(k, rows[i])
        pos = len("\n".join(lines)) + 1
        for j, line in enumerate(qs):
            if j in opt:
                ends[i].append(pos + len(line))
            pos += len(line) + 1
        lines += qs
    answers = []
    for k, i in enumerate(dec, 1):
        answers.append(f"q{k}: {ANS}")
        plan.append((i, None))
    for k, i in enumerate(reads, len(dec) + 1):
        answers.append(f"q{k} (read {rows[i]['read_kind']}): {rows[i]['question']}".rstrip())
        answers += [f"q{k} start: {ANS}", f"q{k} end: {ANS}"]
        plan += [(i, "start"), (i, "end")]
    lo = len(STATE_PREFIX)
    return "\n".join(lines), "Answers:\n" + "\n".join(answers), ends, (lo, lo + len(state)), plan


def keys_of(offsets, state_range, nothing_char):
    """The read heads' keys: the nothing key, then every token overlapping the state block."""
    lo, hi = state_range
    nothing = next(i for i, (b, e) in enumerate(offsets) if b <= nothing_char < e)
    return [nothing] + [i for i, (b, e) in enumerate(offsets) if e > b and e > lo and b < hi]


def decode(logp_start, logp_end, key_offsets, lo, state, kind, threshold):
    """(found, start char, end char) of a read in `state`, or (found, None, None): the best start among the 20 most
    likely, each with its best end within the kind's cap and before the first line break."""
    found = 1.0 - math.exp(float(logp_start[0]))
    if found < threshold:
        return found, None, None
    cap = CAPS[kind]
    ls, le = np.asarray(logp_start[1:]), np.asarray(logp_end[1:])
    best, bs, be = -1e18, None, None
    for i in np.argsort(-ls)[:20]:
        span = cap
        for d in range(1, min(cap, len(le) - i)):
            b, e = key_offsets[i + d + 1]
            if "\n" in state[max(0, b - lo) : max(0, e - lo)]:
                span = d
                break
        j = i + int(np.argmax(le[i : i + span]))
        if ls[i] + le[j] > best:
            best, bs, be = ls[i] + le[j], i, j
    return found, max(0, key_offsets[bs + 1][0] - lo), min(len(state), key_offsets[be + 1][1] - lo)


def rope_positions(ids, mask, img_id, grids, merge):
    """3D rotary positions (3, B, L) as the checkpoint's get_rope_index gives them: text counts up on all three axes;
    an image of merged grid h x w gets (start, start + row, start + col) and advances the count by max(h, w)."""
    B, L = ids.shape
    pos = torch.zeros(3, B, L, dtype=torch.long, device=ids.device)
    it = iter(grids)
    for b in range(B):
        real = mask[b].bool()
        is_img = (ids[b][real] == img_id).tolist()
        parts, cur, i, n = [], 0, 0, len(is_img)
        while i < n:
            j = i
            while j < n and is_img[j] == is_img[i]:
                j += 1
            if not is_img[i]:
                parts.append(torch.arange(cur, cur + j - i, device=ids.device).expand(3, -1))
                cur += j - i
            else:
                t, H, W = (int(x) for x in next(it))
                h, w = H // merge, W // merge
                assert t == 1 and h * w == j - i, f"image block of {j - i} tokens for grid {(t, H, W)}"
                hh, ww = torch.meshgrid(torch.arange(h, device=ids.device), torch.arange(w, device=ids.device), indexing="ij")
                parts.append(torch.stack([torch.full((h * w,), cur, device=ids.device), cur + hh.reshape(-1), cur + ww.reshape(-1)]))
                cur += max(h, w)
            i = j
        pos[:, b, real] = torch.cat(parts, dim=1)
    return pos


# ---------- the model ----------

SIZES = {"0.8b": "0.8b", "4b": "4b"}  # the folder of each size in this repo

class VJev:
    def __init__(self, folder=None, repo=None, size="0.8b", device=None, batch_size=8, max_tokens=None):
        """`folder`: a model folder (default: the `size` folder next to this file, e.g. 4b/). `repo`: a
        Hugging Face repo ID instead, downloaded to the local cache (only the files of `size`). `size`: "0.8b" or
        "4b", used when `folder` is not given. `device`: the GPU if there is one, else the CPU. `batch_size`: prompts per
        forward pass. `max_tokens`: longer prompts are refused; default the backbone's context length (training used
        at most 4,608 tokens)."""
        if size not in SIZES:
            raise ValueError(f"size must be one of {sorted(SIZES)}")
        if repo:
            from huggingface_hub import snapshot_download
            folder = os.path.join(snapshot_download(repo, allow_patterns=[SIZES[size] + "/*"]), SIZES[size])
        self.folder = folder or os.path.join(os.path.dirname(os.path.abspath(__file__)), SIZES[size])
        self.device = torch.device(device or ("cuda" if torch.cuda.is_available() else "cpu"))
        self.batch_size = batch_size
        self.cfg = json.load(open(os.path.join(self.folder, "vjev.json")))
        self.tok = AutoTokenizer.from_pretrained(self.folder)
        self.tok.padding_side = "right"
        self.special = re.compile("|".join(sorted(map(re.escape, self.tok.get_added_vocab()), key=len, reverse=True)))
        # The text model only: the vision tower and the heads in the same file are loaded on their own, so the
        # loader's report of those keys as unexpected is silenced.
        level = hf_logging.get_verbosity()
        hf_logging.set_verbosity_error()
        try:
            self.lm = AutoModelForCausalLM.from_pretrained(self.folder, dtype=torch.bfloat16, device_map=self.device).eval()
        finally:
            hf_logging.set_verbosity(level)
        self.lm.requires_grad_(False)
        self.max_tokens = max_tokens or self.lm.config.max_position_embeddings
        first = self.cfg["first_new_token"]
        assert self.tok.convert_tokens_to_ids(OPT[0]) == first and self.tok.convert_tokens_to_ids("<opt_64>") == first + self.cfg["answer_tokens"] - 1
        self.ans_id = self.tok.convert_tokens_to_ids(ANS)
        self.img_id = self.tok.convert_tokens_to_ids(IMG)
        # Output rows of the answer tokens (tied embedding), in float32 as in training.
        emb = self.lm.get_input_embeddings().weight
        self.rows = {t: emb[self.tok.convert_tokens_to_ids(t)].float() for t in NOUL + LVL}
        with safe_open(os.path.join(self.folder, self.cfg["weights"]), "pt") as f:
            h = {k: f.get_tensor(k).float().to(self.device) for k in ("vjev.pointer.wq", "vjev.pointer.wk", "vjev.span.wq", "vjev.span.wk")}
        self.wq, self.wk, self.wq_span, self.wk_span = h["vjev.pointer.wq"], h["vjev.pointer.wk"], h["vjev.span.wq"], h["vjev.span.wk"]
        self.threshold = self.cfg["found_threshold"]
        self._mid = None
        layer = self.lm.model.layers[self.cfg["mid_layer"] - 1]  # the span head also reads this layer's output
        layer.register_forward_hook(lambda m, i, o: setattr(self, "_mid", o[0] if isinstance(o, tuple) else o))
        self.visual = None

    # ---------- text as data ----------

    def plain(self, text):
        """Text with every added token in it broken up, and for each of its characters the caller's character."""
        out, where, last = [], [], 0
        for m in self.special.finditer(text):
            out.append(text[last : m.start() + 1] + " ")
            where += list(range(last, m.start() + 1)) + [m.start()]
            last = m.start() + 1
        out.append(text[last:])
        where += list(range(last, len(text)))
        return "".join(out), where

    def clean(self, text):
        return self.plain(text)[0]

    def clean_row(self, r):
        r = {**r, "question": self.clean(r["question"]), "options": [self.clean(o) for o in r["options"]]}
        if r["descriptions"]:
            r["descriptions"] = [None if d is None else self.clean(d) for d in r["descriptions"]]
        return r

    # ---------- images ----------

    def load_vision(self):
        """The backbone's own vision encoder (frozen, unchanged), from the same weight file."""
        from transformers import AutoConfig, AutoImageProcessor
        from transformers.models.qwen3_5.modeling_qwen3_5 import Qwen3_5VisionModel
        vcfg = AutoConfig.from_pretrained(self.folder).vision_config
        vcfg._attn_implementation = "sdpa"
        visual = Qwen3_5VisionModel(vcfg).to(self.device)
        for p in visual.parameters():  # weights only: the rotary buffers stay float32
            p.data = p.data.to(torch.bfloat16)
        prefix = "model.visual."
        with safe_open(os.path.join(self.folder, self.cfg["weights"]), "pt") as f:
            state = {k[len(prefix):]: f.get_tensor(k) for k in f.keys() if k.startswith(prefix)}
        missing, unexpected = visual.load_state_dict(state, strict=False)
        assert not unexpected and not [m for m in missing if "rotary" not in m], (missing, unexpected)
        self.visual = visual.requires_grad_(False).eval()
        self.improc = AutoImageProcessor.from_pretrained(self.folder, backend="pil")
        self.merge = vcfg.spatial_merge_size

    @torch.no_grad()
    def encode(self, image):
        """An image's token embeddings (N, d) and its patch grid (t, H, W)."""
        from PIL import Image
        if self.visual is None:
            self.load_vision()
        img = Image.open(image) if isinstance(image, (str, os.PathLike)) else image
        o = self.improc(images=[img.convert("RGB")], return_tensors="pt")
        grid = o["image_grid_thw"].to(self.device)
        out = self.visual(o["pixel_values"].to(self.device, torch.bfloat16), grid_thw=grid)
        return out.pooler_output, grid[0]

    # ---------- one batch of prompts ----------

    def build(self, state, rows, image_feats=None):
        """One prompt: (text, option-line ends, state range, nothing char, plan), all in the prompt's characters."""
        block = None
        if image_feats is not None:
            t, H, W = (int(x) for x in image_feats[1])
            block = VSTART + IMG * (t * H * W // self.merge**2) + VEND
        request, answers, ends, (lo, hi), plan = layout(state, rows, block)
        head = self.tok.apply_chat_template([{"role": "user", "content": request}], tokenize=False, add_generation_prompt=True, enable_thinking=False)
        shift = head.index(request)  # the request sits inside the template
        return head + answers, [[shift + e for e in es] for es in ends], (shift + lo, shift + hi), shift + NOTHING_CHAR, plan

    @torch.no_grad()
    def run(self, jobs):
        """jobs: [(state, rows, image features or None)]; per job, per row: a decision's probabilities (numpy) or a
        read's (found, start char, end char) in that job's state."""
        built = [self.build(*j) for j in jobs]
        lengths = [len(x) for x in self.tok([b[0] for b in built], add_special_tokens=False)["input_ids"]]
        for n in lengths:
            if n > self.max_tokens:
                raise TooLong(f"a prompt is {n} tokens; the limit is {self.max_tokens} (max_tokens). Send fewer questions or a shorter state.")
        if max(lengths) > TRAINED_TOKENS:
            warnings.warn(f"a prompt is {max(lengths)} tokens; training used at most {TRAINED_TOKENS}, so the answers are untested at this length", stacklevel=3)
        order = sorted(range(len(jobs)), key=lambda k: lengths[k])
        out = [None] * len(jobs)
        for s in range(0, len(order), self.batch_size):
            part = order[s : s + self.batch_size]
            for k, res in zip(part, self.forward([jobs[k] for k in part], [built[k] for k in part])):
                out[k] = res
        return out

    def forward(self, jobs, built):
        enc = self.tok([b[0] for b in built], return_tensors="pt", padding=True, add_special_tokens=False, return_offsets_mapping=True)
        offsets = enc.pop("offset_mapping")
        ids, mask = enc["input_ids"].to(self.device), enc["attention_mask"].to(self.device)
        emb = self.lm.get_input_embeddings()(ids)
        pos = None
        feats = [j[2] for j in jobs if j[2] is not None]
        if feats:
            emb = emb.masked_scatter((ids == self.img_id)[..., None], torch.cat([f[0] for f in feats]).to(emb.dtype))
            pos = rope_positions(ids, mask, self.img_id, [f[1] for f in feats], self.merge)
        with torch.autocast(self.device.type, dtype=torch.bfloat16):
            hidden = self.lm.model(inputs_embeds=emb, attention_mask=mask, position_ids=pos).last_hidden_state
        mid = self._mid
        out = []
        with torch.autocast(self.device.type, enabled=False):  # the heads in float32, as in training
            for b, ((state, rows, _), (_, ends, rng, nothing, plan)) in enumerate(zip(jobs, built)):
                n = int(mask[b].sum())
                h, idb, off = hidden[b, :n].float(), ids[b, :n], offsets[b, :n].tolist()
                slots = (idb == self.ans_id).nonzero().squeeze(-1).tolist()
                if len(slots) != len(plan):  # answers would be read from the wrong positions
                    raise RuntimeError(f"the prompt has {len(slots)} answer positions for a plan of {len(plan)}")
                tok_end = np.maximum.accumulate(np.array([e for _, e in off]))
                res = [None] * len(rows)
                kk = None
                for p, (i, part) in zip(slots, plan):
                    r = rows[i]
                    if part is None:
                        if r["kind"] == "choice":
                            keys = [int(np.searchsorted(tok_end, e, side="left")) for e in ends[i]]  # each option line's last token
                            z = (h[p] @ self.wq) @ (h[keys] @ self.wk).T / math.sqrt(self.wq.shape[1])
                        else:
                            labels = NOUL if r["kind"] == "noul" else LVL[: len(r["options"])]
                            z = h[p] @ torch.stack([self.rows[t] for t in labels]).T
                        res[i] = torch.softmax(z, -1).cpu().numpy()
                        continue
                    if kk is None:
                        hs = torch.cat([h, mid[b, :n].float()], -1)
                        key_idx = keys_of(off, rng, nothing)
                        kk, key_off = hs[key_idx] @ self.wk_span, [tuple(off[x]) for x in key_idx]
                    z = (hs[p] @ self.wq_span) @ kk.T / math.sqrt(self.wq_span.shape[1])
                    res[i] = {**(res[i] or {}), part: torch.log_softmax(z, -1).cpu().numpy()}
                for i, r in enumerate(rows):
                    if r["kind"] == "read":
                        x = res[i]
                        res[i] = decode(x["start"], x["end"], key_off, rng[0], state, r["read_kind"], self.threshold)
                out.append(res)
        return out

    # ---------- requests ----------

    def ask(self, state, questions, image=None, lock=False):
        """Answers to `questions` ({id: question}) about `state` (text or a JSON value; may be None with an image)."""
        rows = [to_row(k, q, image is not None) for k, q in questions.items()]
        if not rows:
            return {}
        if any(r["kind"] == "choice" and len(r["options"]) > TRAINED_OPTIONS for r in rows):
            warnings.warn(f"a choice has more than {TRAINED_OPTIONS} options; training used at most {TRAINED_OPTIONS} in one question, so it is untested", stacklevel=2)
        raw = "" if state is None else state if isinstance(state, str) else json.dumps(state, ensure_ascii=False)
        text, where = self.plain(raw)
        clean = [self.clean_row(r) for r in rows]
        dec = [r for r in clean if r["kind"] != "read"]
        reads = [r for r in clean if r["kind"] == "read"]
        jobs, meta = [], []  # meta: per job, the variant (0 = as served, 1 = alone, 2 = alone without the first line) and its source text
        if image is not None:
            feats = self.encode(image)
            for r in dec:
                jobs.append((text, [r], feats))
                meta.append((0, raw, where))
        else:
            dg = [dec[i : i + PER_PROMPT] for i in range(0, len(dec), PER_PROMPT)]
            rg = [reads[i : i + PER_PROMPT] for i in range(0, len(reads), PER_PROMPT)]
            for g in range(max(len(dg), len(rg))):
                jobs.append((text, (dg[g] if g < len(dg) else []) + (rg[g] if g < len(rg) else []), None))
                meta.append((0, raw, where))
            if lock and reads:
                rest = raw.split("\n", 1)[1] if "\n" in raw else raw
                body, body_where = self.plain(rest)
                for r in reads:
                    jobs += [(text, [r], None), (body, [r], None)]
                    meta += [(1, raw, where), (2, rest, body_where)]
        got, cut = {}, {}
        for (_, rs, _), (v, src, idx), res in zip(jobs, meta, self.run(jobs)):
            for r, x in zip(rs, res):
                if r["kind"] == "read":
                    found, a, z = x
                    span = trim(src[idx[a] : idx[z - 1] + 1]) if a is not None and z > a else None
                    cut.setdefault(r["key"], {})[v] = (found, span or None)
                else:
                    got[r["key"]] = x
        answers = {}
        for r in rows:
            key = r["key"]
            if r["kind"] == "read":
                found, value = cut[key][0]
                a = {"type": "read", "value": value, "found": found}
                if r["read_kind"] != "text":
                    a["parsed"] = None if value is None else (parse_number if r["read_kind"] == "number" else parse_date)(value, raw)
                if lock:
                    texts = {(cut[key][v][1] or "").strip().lower() for v in (0, 1, 2)}
                    a["locked"] = len(texts) == 1
                    if not a["locked"]:
                        a["value"] = None
                        if "parsed" in a:
                            a["parsed"] = None
                answers[key] = a
            elif r["kind"] == "noul":
                answers[key] = {"type": "noul", "noul": float(got[key][1])}
            else:
                answers[key] = {"type": r["kind"], "probabilities": {n: float(p) for n, p in zip(r["names"], got[key])}}
        return answers


if __name__ == "__main__":
    import argparse

    ap = argparse.ArgumentParser(description="An example request; with an image path, also three questions about the image.")
    ap.add_argument("image", nargs="?")
    ap.add_argument("--size", choices=sorted(SIZES), default="0.8b")
    args = ap.parse_args()
    jev = VJev(size=args.size)
    state = ("Support ticket #48213\n"
             "Opened 3 Nov 2026 by Dana Whitfield (dana.whitfield@northwind.example).\n"
             "Our nightly export has failed three nights in a row and month-end close is on Friday. The last good export "
             "was on 31 October. We pay 1,250.00 EUR a month for the Pro plan. If this is not fixed by Thursday we will "
             "have to look at other vendors.")
    questions = {
        "export_failing": {"type": "noul", "instructions": "Is a scheduled data export currently failing?"},
        "refund_asked": {"type": "noul", "instructions": "Does the customer ask for a refund?"},
        "team": {"type": "choice", "instructions": "Which team should own the issue?",
                 "criteria": {"engineering": "Fixes broken product functionality", "billing": "Handles invoices, charges and refunds", "sales": "Handles plans, upgrades and renewals"}},
        "churn_risk": {"type": "score", "instructions": "How likely is the customer to leave?",
                       "criteria": ["No sign of leaving", "Mild dissatisfaction", "Mentions alternatives", "States an intent to leave"]},
        "email": {"type": "read", "instructions": "What is the customer's email address?", "kind": "email"},
        "fee": {"type": "read", "instructions": "How much does the customer pay per month?", "kind": "money"},
        "last_good": {"type": "read", "instructions": "When was the last successful export?", "kind": "date"},
        "phone": {"type": "read", "instructions": "What is the customer's phone number?", "kind": "phone"},
    }
    print(json.dumps(jev.ask(state, questions, lock=True), indent=2))
    if args.image:
        image_questions = {
            "photo": {"type": "noul", "instructions": "Is this a photograph (not a drawing or a diagram)?"},
            "people": {"type": "score", "instructions": "How many people are visible?", "criteria": ["None", "One", "Two to five", "More than five"]},
            "setting": {"type": "choice", "instructions": "Where was this taken?", "criteria": {"indoors": None, "outdoors": None}},
        }
        print(json.dumps(jev.ask(None, image_questions, image=args.image), indent=2))