Blame for jue_flash/encoding/pack.py

"""Single-pass multi-question bidirectional packing and MPS bucket collation for Jue-Flash.

Layout of one packed request in a bidirectional encoder:
    [CLS]
      [Q1_HEADER] <instruction_1>
      [OPT] <opt_1,1> [OPT] <opt_1,2> ... [OPT] <opt_1,K1>
      ...
      [QM_HEADER] <instruction_M>
      [OPT] <opt_M,1> ... [OPT] <opt_M,KM>
    [SEP]
      <state tokens (head + tail preserved if truncated)>
    [SEP]

Every question m records:
- ``q_span``: ``(start_idx, end_idx)`` of its instruction tokens inside ``input_ids``
- ``opt_spans``: list of ``(start_idx, end_idx)`` for each candidate option k in ``0..K_m-1``
- ``primitive_id``: ``0`` (choice), ``1`` (score), or ``2`` (noul)
And the request records:
- ``state_span``: ``(start_idx, end_idx)`` of the state tokens inside ``input_ids``.
"""

from __future__ import annotations

from collections.abc import Mapping, Sequence
from dataclasses import dataclass
import re
from typing import Any

import torch
from torch import Tensor

from jue_flash.api.schema import JSONValue, Question, option_keys
from jue_flash.encoding.render import PRIMITIVE_TO_ID, render, render_question_parts

MPS_BUCKETS: tuple[int, ...] = (256, 512, 1024, 2048, 4096)

_EMBEDDED_ABCD_RE = re.compile(
    r"\nA\.\s+(.*?)\n+B\.\s+(.*?)\n+C\.\s+(.*?)\n+D\.\s+(.*)\Z",
    flags=re.DOTALL,
)


def _extract_embedded_abcd(q_text: str) -> tuple[str, dict[str, str]]:
    """Extract trailing embedded A./B./C./D. option bodies from state['question'] if present."""
    m = _EMBEDDED_ABCD_RE.search(q_text)
    if not m:
        return q_text, {}
    stem = q_text[: m.start()].strip()
    opts = {
        "A": m.group(1).strip(),
        "B": m.group(2).strip(),
        "C": m.group(3).strip(),
        "D": m.group(4).strip(),
    }
    return stem, opts


@dataclass(frozen=True)
class EncodedQuestionSpan:
    """Token span metadata for one question inside a packed encoder sequence."""

    key: str
    qtype: str
    primitive_id: int
    option_keys: list[str]
    q_span: tuple[int, int]  # [start, end) in input_ids
    opt_spans: list[tuple[int, int]]  # len == K, each [start, end) in input_ids
    target: list[float] | None = None  # len == K if training/eval target is provided
    canon_prior: list[float] | None = None


@dataclass(frozen=True)
class EncodedRequest:
    """One packed request (one state + 1..M questions) ready for the bidirectional encoder."""

    group_id: str
    input_ids: list[int]
    state_span: tuple[int, int]
    questions: list[EncodedQuestionSpan]
    raw_total_tokens: int
    state_truncated: bool


@dataclass(frozen=True)
class PackedFlashBatch:
    """Padded batch of requests aligned to a fixed MPS sequence length bucket."""

    bucket_len: int
    input_ids: Tensor  # [B, L] long
    attention_mask: Tensor  # [B, L] long
    state_spans: Tensor  # [B, 2] long (start, end) per batch row
    # Flattened question indexing across the batch (N total questions across B requests):
    question_batch_idx: Tensor  # [N] long -> which row b in 0..B-1
    question_primitive_id: Tensor  # [N] long -> 0=choice, 1=score, 2=noul
    question_spans: Tensor  # [N, 2] long -> (start, end) of instruction
    opt_spans: Tensor  # [N, K_max, 2] long -> (start, end) of candidate k
    cand_mask: Tensor  # [N, K_max] bool -> True where candidate k is valid
    targets: Tensor | None  # [N, K_max] float32 -> normalized target distribution (if present)
    question_refs: list[tuple[str, str, str, list[str]]]  # [(group_id, q_key, qtype, option_keys), ...]
    canon_priors: Tensor | None = None


def select_mps_bucket(seq_len: int, buckets: Sequence[int] = MPS_BUCKETS, max_len: int = 4096) -> int:
    """Return the smallest fixed bucket >= seq_len, capped at max_len."""
    limit = min(max_len, max(buckets))
    for b in buckets:
        if b >= seq_len and b <= limit:
            return b
    return limit


_PLAIN_ENCODE_CACHE: dict[tuple[int, str], list[int]] = {}
_PLAIN_CACHE_MAX = 131072


def _encode_plain(tokenizer: Any, text: str) -> list[int]:
    tok_id = id(tokenizer)
    use_cache = len(text) <= 1024
    if use_cache:
        cached = _PLAIN_ENCODE_CACHE.get((tok_id, text))
        if cached is not None:
            return cached
    ids = tokenizer.encode(text, add_special_tokens=False)
    if not ids:
        # Ensure non-empty span (at least 1 token) so span pooling never divides by zero
        unk = getattr(tokenizer, "unk_token_id", None)
        res = [unk if unk is not None else 0]
    else:
        res = list(ids)
    if use_cache:
        if len(_PLAIN_ENCODE_CACHE) >= _PLAIN_CACHE_MAX:
            _PLAIN_ENCODE_CACHE.clear()
        _PLAIN_ENCODE_CACHE[(tok_id, text)] = res
    return res


def encode_group_request(
    tokenizer: Any,
    group_id: str,
    state: JSONValue,
    questions: Mapping[str, Question],
    targets: Mapping[str, Mapping[str, float]] | None = None,
    *,
    max_len: int = 2048,
    max_option_tokens: int = 96,
    max_inst_tokens: int = 160,
    max_questions_per_chunk: int = 8,
    head_ratio: float = 0.70,
    state_ids_override: list[int] | None = None,
) -> list[EncodedRequest]:
    """Encode a state and its questions into one or more :class:`EncodedRequest` objects within ``max_len``.

    If a group has many questions whose combined tokens would leave fewer than 64 tokens for ``state``
    (or exceed ``max_questions_per_chunk``), the questions are automatically chunked across sub-requests
    ``<group_id>#0``, ``<group_id>#1``, etc.
    """
    cls_id = tokenizer.cls_token_id if tokenizer.cls_token_id is not None else tokenizer.bos_token_id
    sep_id = tokenizer.sep_token_id if tokenizer.sep_token_id is not None else tokenizer.eos_token_id
    if cls_id is None:
        cls_id = 1
    if sep_id is None:
        sep_id = 2

    q_stem_from_state: str = ""
    embedded_abcd: dict[str, str] = {}
    precomputed_priors_by_q: dict[str, list[float]] = {}
    state_for_render: JSONValue = state
    if isinstance(state, Mapping) and isinstance(state.get("question"), str):
        raw_q_stem = state["question"].strip()
        q_stem_from_state, embedded_abcd = _extract_embedded_abcd(raw_q_stem)
        if embedded_abcd:
            state_for_render = {**state, "question": q_stem_from_state}

    if state_ids_override is not None:
        state_ids_full = state_ids_override
    else:
        state_str = render(state_for_render).strip() or "<empty_state>"
        state_ids_full = _encode_plain(tokenizer, state_str)

    # Pre-tokenize each question's instruction and options
    q_items: list[tuple[str, Question, list[int], list[tuple[str, list[int]]], int]] = []
    for q_key, q in questions.items():
        inst_text, opts = render_question_parts(q)
        if q_stem_from_state and (
            "multiple-choice question from a cybersecurity" in inst_text
            or inst_text.startswith("Evaluate ")
            or q_stem_from_state[:30] not in inst_text
        ):
            inst_text = f"Question: {q_stem_from_state}\n{inst_text}"
        header_inst = f"[{q.type.upper()}] {inst_text}"
        inst_ids = _encode_plain(tokenizer, header_inst)[:max_inst_tokens]
        num_opts = len(opts)
        if num_opts <= 2:
            eff_max_opt = max(max_option_tokens, 280)
        elif num_opts <= 4:
            eff_max_opt = max(max_option_tokens, 192)
        else:
            eff_max_opt = max_option_tokens
        opt_id_pairs: list[tuple[str, list[int]]] = []
        for opt_k, opt_str in opts:
            if embedded_abcd and opt_k in embedded_abcd and len(opt_str) < len(embedded_abcd[opt_k]):
                opt_str = f"{opt_k}: {embedded_abcd[opt_k]}"
            o_ids = _encode_plain(tokenizer, f"| {opt_str}")[:eff_max_opt]
            opt_id_pairs.append((opt_k, o_ids))
        q_tok_len = len(inst_ids) + sum(len(o_ids) for _, o_ids in opt_id_pairs) + 1  # +1 for sep after question
        q_items.append((q_key, q, inst_ids, opt_id_pairs, q_tok_len))

    # Reserve at least min(256, max_len // 3) tokens for state when chunking multiple questions
    min_state_budget = min(len(state_ids_full), max(64, max_len // 3))
    max_q_budget = max(64, max_len - 3 - min_state_budget)

    chunks: list[list[tuple[str, Question, list[int], list[tuple[str, list[int]]], int]]] = []
    cur_chunk: list[tuple[str, Question, list[int], list[tuple[str, list[int]]], int]] = []
    cur_q_tokens = 0
    for item in q_items:
        q_tok_len = item[4]
        if cur_chunk and (cur_q_tokens + q_tok_len > max_q_budget or len(cur_chunk) >= max_questions_per_chunk):
            chunks.append(cur_chunk)
            cur_chunk = []
            cur_q_tokens = 0
        cur_chunk.append(item)
        cur_q_tokens += q_tok_len
    if cur_chunk:
        chunks.append(cur_chunk)

    out_requests: list[EncodedRequest] = []
    for chunk_idx, chunk in enumerate(chunks):
        sub_gid = group_id if len(chunks) == 1 else f"{group_id}#{chunk_idx}"
        input_ids: list[int] = [cls_id]
        encoded_questions: list[EncodedQuestionSpan] = []

        # If a single question has very wide choices (e.g. K=104..166), reserve space for state_ids_full
        # and use a water-filling per-option cap so short options keep 100% of their tokens while only
        # overly verbose multi-alias options are trimmed to fit inside max_len.
        total_q_len = sum(x[4] for x in chunk)
        reserved_state = min(len(state_ids_full), max(128, max_len // 4))
        hard_q_cap = max(64, max_len - 8 - reserved_state)
        scale = min(1.0, hard_q_cap / max(1, total_q_len))

        for q_key, q, inst_ids, opt_id_pairs, q_tok_len in chunk:
            opt_cap: int | None = None
            if scale < 1.0:
                q_budget = max(16, int(round(hard_q_cap * (q_tok_len / max(1, total_q_len)))))
                inst_ids = inst_ids[: min(len(inst_ids), max(8, q_budget // 5))]
                opt_budget = max(len(opt_id_pairs) * 2, q_budget - len(inst_ids) - 1)
                c = max((len(o_ids) for _, o_ids in opt_id_pairs), default=2)
                while c > 2 and sum(min(len(o_ids), c) for _, o_ids in opt_id_pairs) > opt_budget:
                    c -= 1
                opt_cap = c
            q_start = len(input_ids)
            input_ids.extend(inst_ids)
            q_end = len(input_ids)

            opt_spans: list[tuple[int, int]] = []
            opt_keys_list: list[str] = []
            for opt_k, o_ids in opt_id_pairs:
                if opt_cap is not None and len(o_ids) > opt_cap:
                    o_ids = o_ids[:opt_cap]
                o_start = len(input_ids)
                input_ids.extend(o_ids)
                o_end = len(input_ids)
                opt_spans.append((o_start, o_end))
                opt_keys_list.append(opt_k)

            input_ids.append(sep_id)

            q_target: list[float] | None = None
            if targets is not None and q_key in targets:
                t_map = targets[q_key]
                keys = option_keys(q)
                q_target = [float(t_map[k]) for k in keys]

            q_prior = precomputed_priors_by_q.get(q_key)
            if q_prior is not None and len(q_prior) != len(opt_keys_list):
                q_prior = None

            encoded_questions.append(
                EncodedQuestionSpan(
                    key=q_key,
                    qtype=q.type,
                    primitive_id=PRIMITIVE_TO_ID[q.type],
                    option_keys=opt_keys_list,
                    q_span=(q_start, q_end),
                    opt_spans=opt_spans,
                    target=q_target,
                    canon_prior=q_prior,
                )
            )

        raw_total = len(input_ids) + len(state_ids_full) + 1
        rem_budget = max(1, max_len - len(input_ids) - 1)
        if len(state_ids_full) <= rem_budget:
            state_ids = state_ids_full
            truncated = False
        else:
            # Head + Tail truncation preserves both top-level metadata/headers and bottom error/verdict lines
            head_len = max(1, int(rem_budget * head_ratio))
            tail_len = rem_budget - head_len
            state_ids = state_ids_full[:head_len] + (state_ids_full[-tail_len:] if tail_len > 0 else [])
            truncated = True

        s_start = len(input_ids)
        input_ids.extend(state_ids)
        s_end = len(input_ids)
        input_ids.append(sep_id)

        out_requests.append(
            EncodedRequest(
                group_id=sub_gid,
                input_ids=input_ids,
                state_span=(s_start, s_end),
                questions=encoded_questions,
                raw_total_tokens=raw_total,
                state_truncated=truncated,
            )
        )

    return out_requests


def collate_encoded_requests(
    requests: Sequence[EncodedRequest],
    pad_token_id: int = 0,
    *,
    buckets: Sequence[int] = MPS_BUCKETS,
    max_len: int = 4096,
    fixed_bucket_len: int | None = None,
    device: str | torch.device = "cpu",
) -> PackedFlashBatch:
    """Collate a list of :class:`EncodedRequest` into a fixed-shape :class:`PackedFlashBatch`."""
    if not requests:
        raise ValueError("Cannot collate an empty sequence of EncodedRequest")

    max_seq_in_batch = max(len(r.input_ids) for r in requests)
    bucket_len = (
        fixed_bucket_len
        if fixed_bucket_len is not None
        else select_mps_bucket(max_seq_in_batch, buckets=buckets, max_len=max_len)
    )
    if bucket_len < max_seq_in_batch:
        bucket_len = max_seq_in_batch

    bsz = len(requests)
    input_ids = torch.full((bsz, bucket_len), pad_token_id, dtype=torch.long)
    attention_mask = torch.zeros((bsz, bucket_len), dtype=torch.long)
    state_spans = torch.zeros((bsz, 2), dtype=torch.long)

    all_q_batch_idx: list[int] = []
    all_q_prim_id: list[int] = []
    all_q_spans: list[tuple[int, int]] = []
    all_opt_spans: list[list[tuple[int, int]]] = []
    all_targets: list[list[float]] = []
    all_priors: list[list[float] | None] = []
    has_targets = True
    has_any_prior = False
    question_refs: list[tuple[str, str, str, list[str]]] = []

    for b, req in enumerate(requests):
        seq_len = len(req.input_ids)
        input_ids[b, :seq_len] = torch.tensor(req.input_ids, dtype=torch.long)
        attention_mask[b, :seq_len] = 1
        state_spans[b, 0] = req.state_span[0]
        state_spans[b, 1] = req.state_span[1]

        for q in req.questions:
            all_q_batch_idx.append(b)
            all_q_prim_id.append(q.primitive_id)
            all_q_spans.append(q.q_span)
            all_opt_spans.append(q.opt_spans)
            question_refs.append((req.group_id, q.key, q.qtype, q.option_keys))
            if q.target is None:
                has_targets = False
            else:
                all_targets.append(q.target)
            all_priors.append(q.canon_prior)
            if q.canon_prior is not None:
                has_any_prior = True

    n_questions = len(all_q_batch_idx)
    k_max = max(len(spans) for spans in all_opt_spans)

    opt_spans_tensor = torch.zeros((n_questions, k_max, 2), dtype=torch.long)
    cand_mask = torch.zeros((n_questions, k_max), dtype=torch.bool)
    targets_tensor = torch.zeros((n_questions, k_max), dtype=torch.float32) if has_targets else None
    priors_tensor = torch.zeros((n_questions, k_max), dtype=torch.float32) if has_any_prior else None

    for i, spans in enumerate(all_opt_spans):
        k = len(spans)
        cand_mask[i, :k] = True
        for j, (s, e) in enumerate(spans):
            opt_spans_tensor[i, j, 0] = s
            opt_spans_tensor[i, j, 1] = e
        if targets_tensor is not None:
            targets_tensor[i, :k] = torch.tensor(all_targets[i], dtype=torch.float32)
        if priors_tensor is not None and all_priors[i] is not None:
            priors_tensor[i, :k] = torch.tensor(all_priors[i], dtype=torch.float32)

    dev = torch.device(device)
    return PackedFlashBatch(
        bucket_len=bucket_len,
        input_ids=input_ids.to(dev),
        attention_mask=attention_mask.to(dev),
        state_spans=state_spans.to(dev),
        question_batch_idx=torch.tensor(all_q_batch_idx, dtype=torch.long, device=dev),
        question_primitive_id=torch.tensor(all_q_prim_id, dtype=torch.long, device=dev),
        question_spans=torch.tensor(all_q_spans, dtype=torch.long, device=dev),
        opt_spans=opt_spans_tensor.to(dev),
        cand_mask=cand_mask.to(dev),
        targets=targets_tensor.to(dev) if targets_tensor is not None else None,
        question_refs=question_refs,
        canon_priors=priors_tensor.to(dev) if priors_tensor is not None else None,
    )