inference_v0.py

"""Standalone Bekko v0 typed-decision inference and remote-code interface.

Maintainer note: this source file is already the standalone runtime.
export_v0.py:runtime_source() copies it verbatim into exported model directories;
it does not convert imports or inline helpers. The same file can be distributed
with compatible v0 weights on Hugging Face. Do not add relative imports or
imports of bekko_system_one: consumers must not need the training package.
Keep embedded helpers aligned with their training counterparts and preserve the
isolated-process standalone loading test in tests/test_inference_v0.py.

Install the exported requirements.txt (and a CUDA-compatible PyTorch build for
GPU use). This source file includes its helpers and needs no Bekko training
package, datasets, PEFT or W&B. SDPA needs no external FlashAttention extension;
FA2 is an optional acceleration backend selected at model loading.

Load once, reuse for many requests
---------------------------------
From a downloaded export, import BekkoSentenceTransformer from inference_v0.
To load the class directly from a Hugging Face model repository::

    from transformers.dynamic_module_utils import get_class_from_dynamic_module

    repo = "YOUR_ORG/YOUR_MODEL"
    revision = "FULL_COMMIT_HASH"  # Pin the same revision for code and weights.
    Model = get_class_from_dynamic_module(
        "inference_v0.BekkoSentenceTransformer", repo,
        revision=revision, token=True,
    )
    model = Model(repo, revision=revision, token=True,
                  trust_remote_code=True, device="cuda")  # Or device="cpu".

Authenticate with Hugging Face before accessing a private model. token=True
uses your saved token. Plain SentenceTransformer(..., trust_remote_code=True)
also loads the export, but its typed entry point is model[0].predict(...).
BekkoSentenceTransformer exposes predict() directly. Raw-text encode() is not
a typed-decision API.

Input and output interface
--------------------------
Pass only a native input dict with exactly state_json and decisions; do not
pass a dataset row, targets, labels, or provenance. Fields ending in _json are
JSON-encoded strings, including when they contain a plain string. For example::

    import json
    request = {
        "state_json": json.dumps({"message": "Please refund a duplicate charge."}),
        "decisions": [{
            "id": "department", "kind": "judgment", "type": "choice",
            "instructions_json": json.dumps("Which department should respond?"),
            "system_prompt": "",
            "criteria": [
                {"id": "billing", "description_json": json.dumps("Payments and refunds"),
                 "value": None},
                {"id": "technical", "description_json": json.dumps("Technical failures"),
                 "value": None},
            ],
            "documents": [], "scoring": None,
        }],
    }
    result = model.predict(request, show_progress_bar=False)
    selected = result["department"]["selected_id"]

Decision IDs must be nonempty and unique within each request; they may repeat
across requests. Candidate IDs must be nonempty and distinct within a decision.
Judgments use kind="judgment", criteria, empty documents, and scoring=None:
* choice: returns selected_id and probabilities keyed by candidate ID.
* noul: use exactly true/false or yes/no IDs and authored descriptions of both
  meanings; returns probability_yes and probabilities.
* score: supply numeric value for each criterion (at least two distinct values);
  returns score (expected value on that scale), normalized_score in [0, 1],
  probabilities, and values. Values are never inferred from descriptions.
Relative ranking uses kind="ranking", type=None, scoring="relative", empty
criteria, and documents containing id and content_json. It returns probabilities
and order (document IDs sorted by descending probability).

Choose the batching interface
-----------------------------
Use predict(request) for one request: it returns a dict keyed by decision ID.
Use predict(requests) for a list of requests: it returns a list of result dicts
in input order. This is the usual throughput interface, including mixed tasks::

    results = model.predict(
        [request, request], batch_size=128, token_budget=64000,
        show_progress_bar=False,
    )

Prefer this list call to a Python loop of single-request calls. batch_size limits
both the requests rendered per window and the decisions tokenized per window;
it is not a fixed candidate count or GPU microbatch size. token_budget estimates
candidate_count * (max_query_tokens + max_document_tokens) for each microbatch.
An oversized decision runs alone, so this is not a strict memory limit. Candidate
groups are never split. Length bucketing reduces padding; original request and
candidate order is restored. Tokenization deduplicates text within each window,
and encoder prefixes are shared within each microbatch, not cached across calls.

Use predict_groups(groups) only if you already have rendered Group objects,
for example from input_groups(request). It returns one CPU FP32 probability
tensor per decision, in group/candidate order, without typed interpretation.
It is not required for ordinary list batching. Empty request lists return [];
a request with no decisions returns {}. Inputs are materialized, not streamed;
chunk very large datasets into lists in the caller. Progress defaults to stderr;
predict counts requests, predict_groups counts decisions.

Fast execution and memory tuning
--------------------------------
Reuse a loaded model, batch requests, and select device="cuda" when available.
BekkoSentenceTransformer(..., attn_implementation="auto") is the default:
CUDA capability 8.0+ with a compatible flash-attn library selects FlashAttention 2;
otherwise it selects PyTorch SDPA. Force a backend at model load time::

    model = Model(repo, revision=revision, device="cuda", trust_remote_code=True,
                  attn_implementation="flash_attention_2")
    # Or attn_implementation="sdpa" to require SDPA without importing flash-attn.
    # model_kwargs={"attn_implementation": "flash_attention_2"} is also accepted.
    print(model[0].attn_implementation)  # The backend actually selected.

Explicit FA2 requires CUDA capability 8.0+ and a compatible flash-attn wheel;
missing or binary-incompatible libraries raise during model loading, never silently
fall back. Install the optional fa2 extra in the training package, or a flash-attn
wheel matching your inference environment's Python, PyTorch and CUDA versions.
SDPA has no external attention dependency. Auto selection happens at loading;
load again with the desired device/backend when switching devices. Plain ST can
load the format, but use BekkoSentenceTransformer for backend selection.

For 17M models, SDPA and FA2 generally have little speed difference; SDPA is a
reasonable dependency-free choice. For 68M and larger models, prefer FA2 for
throughput, especially on long inputs. This is a sizing guideline, not a speed
guarantee for every larger checkpoint: measure representative inputs on your GPU,
excluding loading and warmup. BF16 backend rounding can change probabilities and
occasionally the selected candidate; the backends are not bitwise interchangeable.
The SDPA path reuses per-forward masks/rotary tensors and blocks local prefix
attention; FA2 keeps valid tokens packed through attention and feed-forward layers.
Both retain the same rendering, input budgets, candidate order and task heads.
CUDA encoder execution uses BF16 autocast; CPU uses FP32.
Heads and per-decision softmax use FP32. Small CPU/GPU differences are expected.
Tune batch_size for the tokenization/sorting window and token_budget for GPU
work per microbatch; reduce the latter when memory is tight. A single oversized
decision still runs alone and may require shorter inputs or fewer candidates.
For repeated workloads, optionally enable compilation before warmup::

    model.compile_inference()  # Lazy encoder-only torch.compile, default Inductor.
    model.predict([request, request], show_progress_bar=False)  # Warmup.
    results = model.predict([request, request], show_progress_bar=False)
    model.disable_compile()  # Return to eager encoder execution.

Compilation leaves rendering, tokenization, heads, packing, and output processing
eager. First calls include compilation cost; new shapes can recompile despite
dynamic=True. Benchmark representative warmed batches, synchronizing CUDA when
timing; no universal speedup is guaranteed and compile errors are not silently
converted to eager execution.

Context limits are separate from the microbatch token_budget. Fresh v0 exports
use adaptive-v1: reserve half the context for query and candidate (candidate gets
the odd token), then lend unused capacity subject to branch caps. The default
context is the backbone positional capacity (7,999 for v0 17M); candidates default
to min(3800, context - 3) tokens. Special tokens count. All candidates in a decision
share one truncated query. Queries follow saved balanced/right truncation and
candidates truncate on the right. Per-call context_length, query_length, and
document_length override limits without modifying the checkpoint. context_length
must not exceed positional capacity. prefix_layout overrides instruction_state
or state_instruction rendering; normally keep the exported default.

CLI: python inference_v0.py --model PATH_OR_HUB_ID --input requests.json
Add --device cuda, --attn-implementation flash_attention_2 (or sdpa/auto),
--compile, --batch-size 128, --token-budget 64000, or
--no-show-progress-bar as needed. Input JSON can be one request or an array;
results are JSON on stdout and progress is on stderr.
"""

from __future__ import annotations

import argparse
import importlib
import json
import math
from dataclasses import asdict, dataclass
from pathlib import Path
from typing import Any, cast

import torch
from safetensors.torch import load_file, save_file
from sentence_transformers import SentenceTransformer
from sentence_transformers.base.modules import InputModule
from torch import nn
from torch.nn import functional as F
from tqdm.auto import tqdm
from transformers import AutoConfig, AutoModel, AutoTokenizer, PreTrainedTokenizerBase
from transformers.models.modernbert.modeling_modernbert import apply_rotary_pos_emb


@dataclass(frozen=True)
class QueryParts:
    """Explicit boundaries: never infer them from text inside the context."""

    instruction: str
    context: str
    system: str = ""
    layout: str = "instruction_state"

    def __post_init__(self):
        if not isinstance(self.instruction, str) or not self.instruction.strip():
            raise ValueError("instruction must be nonempty text")
        if not isinstance(self.context, str) or not isinstance(self.system, str):
            raise ValueError("context and system must be text")
        if self.layout not in {"instruction_state", "state_instruction"}:
            raise ValueError("Unsupported query layout")

    def render(self):
        prefix = f"{self.system}\n\n" if self.system.strip() else ""
        instruction = f"Instruction: {self.instruction}"
        context = f"State: {self.context}"
        body = (
            (instruction, context) if self.layout == "instruction_state" else (context, instruction)
        )
        return prefix + "\n".join(body)


def allocate_query_budget(instruction_length, context_length, budget):
    """Reserve half per field, transfer unused capacity, favor instruction on odd budgets."""
    if min(instruction_length, context_length, budget) < 0:
        raise ValueError("Lengths and budget must be nonnegative")
    instruction = min(instruction_length, (budget + 1) // 2)
    context = min(context_length, budget // 2)
    instruction += min(instruction_length - instruction, budget - instruction - context)
    context += min(context_length - context, budget - instruction - context)
    return instruction, context


def balanced_query_ids(tokenizer, parts, query_length):
    """Reserve system/markers first, then share the remaining budget between fields.

    Tokenize components independently so both layouts retain exactly the same
    content tokens. Truncate field tails, without decoding and re-tokenizing.
    """
    if tokenizer.truncation_side != "right":
        raise ValueError("balanced query truncation requires a right-truncating tokenizer")
    if any(not isinstance(p, QueryParts) for p in parts):
        raise ValueError("balanced query truncation requires QueryParts for every query")
    if not parts:
        return []
    limits = [query_length] * len(parts) if isinstance(query_length, int) else list(query_length)
    if len(limits) != len(parts) or any(limit < 3 for limit in limits):
        raise ValueError("One valid query budget is required per query")
    markers = tokenizer(["Instruction: ", "State: ", "\n"], add_special_tokens=False)["input_ids"]
    instruction_marker, context_marker, separator = markers
    texts = []
    for p in parts:
        texts.extend([f"{p.system}\n\n" if p.system.strip() else "", p.instruction, p.context])
    encoded = tokenizer(texts, add_special_tokens=False, truncation=True, max_length=max(limits))[
        "input_ids"
    ]
    result = []
    overhead = 2 + sum(map(len, markers))
    for i, p in enumerate(parts):
        system, instruction, context = encoded[3 * i : 3 * i + 3]
        budget = limits[i] - overhead - len(system)
        if budget < 2:
            raise ValueError("System prompt and query markers leave fewer than two content tokens")
        ni, nc = allocate_query_budget(len(instruction), len(context), budget)
        ins = instruction_marker + instruction[:ni]
        ctx = context_marker + context[:nc]
        body = ins + separator + ctx if p.layout == "instruction_state" else ctx + separator + ins
        result.append([tokenizer.cls_token_id, *system, *body, tokenizer.sep_token_id])
    return result


@dataclass(frozen=True)
class DecisionMetadata:
    """Non-tokenized candidate alignment and source identity; never model features."""

    candidate_ids: tuple[str, ...] | None = None
    candidate_values: tuple[float, ...] | None = None
    kind: str | None = None
    case_id: str | None = None
    group_id: str | None = None
    decision_id: str | None = None

    def __post_init__(self):
        if self.kind not in {None, "judgment", "ranking"}:
            raise ValueError("metadata kind must be judgment or ranking")
        if self.candidate_ids is not None:
            ids = tuple(self.candidate_ids)
            if not ids or any(not isinstance(i, str) or not i for i in ids):
                raise ValueError("candidate_ids must be nonempty strings")
            if len(set(ids)) != len(ids):
                raise ValueError("candidate_ids must be unique")
            object.__setattr__(self, "candidate_ids", ids)
        if self.candidate_values is not None:
            values = tuple(self.candidate_values)
            if not values or any(not math.isfinite(v) for v in values):
                raise ValueError("candidate_values must be finite")
            if len(set(values)) != len(values):
                raise ValueError("candidate_values must be unique")
            object.__setattr__(self, "candidate_values", values)
        for identity in (self.case_id, self.group_id, self.decision_id):
            if identity is not None and not isinstance(identity, str):
                raise ValueError("metadata identities must be strings")

    @property
    def yes_index(self):
        ids = self.candidate_ids
        if ids is not None and set(ids) in ({"true", "false"}, {"yes", "no"}):
            return ids.index("true" if "true" in ids else "yes")
        return None

    @property
    def score_values(self):
        if self.kind == "judgment" and self.candidate_values is not None:
            if len(self.candidate_values) >= 2:
                return self.candidate_values
        return None


@dataclass(frozen=True)
class Group:
    query: str
    candidates: list[str]
    task: str = "reranker"
    target: list[float] | None = None
    query_parts: QueryParts | None = None
    metadata: DecisionMetadata | None = None

    def __post_init__(self):
        if self.metadata is not None:
            for values in (self.metadata.candidate_ids, self.metadata.candidate_values):
                if values is not None and len(values) != len(self.candidates):
                    raise ValueError("candidate metadata must align with candidates")
        if self.query_parts is not None and self.query_parts.render() != self.query:
            raise ValueError("query must match query_parts.render()")
        if self.task not in {"reranker", "choice", "noul", "score"}:
            raise ValueError(f"Unknown task: {self.task}")
        if not isinstance(self.query, str) or not self.query.strip():
            raise ValueError("query must be nonempty text")
        if not self.candidates or any(not isinstance(c, str) for c in self.candidates):
            raise ValueError("candidates must contain text")
        if self.target is not None and (
            len(self.target) != len(self.candidates)
            or any(not math.isfinite(t) or t < 0 for t in self.target)
            or abs(sum(self.target) - 1) > 1e-5
        ):
            raise ValueError("target must be a normalized distribution aligned with candidates")

    @classmethod
    def from_dict(cls, row):
        # Metadata and targets are never concatenated into inference inputs.
        return cls(
            query=row["query"],
            candidates=row["candidates"],
            task=row.get("task", "reranker"),
            target=row.get("target"),
            query_parts=QueryParts(**row["query_parts"]) if row.get("query_parts") else None,
            metadata=DecisionMetadata(**row["metadata"]) if row.get("metadata") else None,
        )


@dataclass
class PreparedGroup:
    key: str | tuple[int, ...]
    task: str
    query: list[int]
    documents: list[list[int]]
    target: list[float] | None
    metadata: DecisionMetadata | None = None

    @property
    def cost(self):
        return len(self.query) * len(self.documents) + sum(map(len, self.documents))


def prepare_groups(groups, encoder):
    queries = list(dict.fromkeys((g.query, g.query_parts) for g in groups))
    documents = list(dict.fromkeys((g.task, d) for g in groups for d in g.candidates))
    qids, dids = encoder.tokenize_branches(
        [q for q, _ in queries],
        [d for _, d in documents],
        [t for t, _ in documents],
        query_parts=[p for _, p in queries],
    )
    qmap, dmap = dict(zip(queries, qids, strict=True)), dict(zip(documents, dids, strict=True))
    return [
        PreparedGroup(
            tuple(qmap[g.query, g.query_parts]),
            g.task,
            qmap[g.query, g.query_parts],
            [dmap[g.task, d] for d in g.candidates],
            g.target,
            g.metadata,
        )
        for g in groups
    ]


def collate_groups(groups, encoder):
    queries, docs, owners, seen, indices = [], [], [], {}, {}
    for group in groups:
        if group.key not in seen:
            seen[group.key] = len(queries)
            queries.append(group.query)
        indices.setdefault(group.task, []).extend(
            range(len(docs), len(docs) + len(group.documents))
        )
        docs.extend(group.documents)
        owners.extend([seen[group.key]] * len(group.documents))
    features = encoder.collate_tokens(queries, docs, owners)
    features["head_indices"] = {next(iter(indices)): None} if len(indices) == 1 else indices
    choice_rows, offset = [], 0
    for group in groups:
        count = len(group.documents)
        if group.task == "choice":
            choice_rows.append(list(range(offset, offset + count)))
        offset += count
    if choice_rows:
        width = max(map(len, choice_rows))
        features["choice_indices"] = torch.tensor(
            [row + [0] * (width - len(row)) for row in choice_rows], dtype=torch.long
        )
        features["choice_mask"] = torch.tensor(
            [[True] * len(row) + [False] * (width - len(row)) for row in choice_rows]
        )
    return features


def to_device(features, device):
    return {k: v.to(device) if isinstance(v, torch.Tensor) else v for k, v in features.items()}


def prepare_batch(model, groups):
    """Build a feature dictionary accepted by an ordinary SentenceTransformer."""
    return to_device(collate_groups(prepare_groups(groups, model[0]), model[0]), model.device)


def work_tokens(items, padded_documents=True):
    cost = sum(x.cost for x in items)
    if padded_documents and items:
        lengths = [len(d) for x in items for d in x.documents]
        cost += len(lengths) * max(lengths) - sum(lengths)
    return cost


def pack_groups(items, budget, padded_documents=True):
    """First-fit decreasing; never split the candidates of a training group."""
    if budget < 1:
        raise ValueError("Token budget must be positive")
    groups, statistics, ordered = [], [], []
    for item in items:
        lengths = [len(d) for d in item.documents]
        ordered.append((item, item.cost, len(lengths), sum(lengths), max(lengths)))
    for item, cost, count, total, maximum in sorted(
        ordered, key=lambda x: (x[4], x[1]), reverse=True
    ):
        for i, (old_cost, old_count, old_total, old_max) in enumerate(statistics):
            combined = old_cost + cost, old_count + count, old_total + total, max(old_max, maximum)
            estimate = combined[0]
            if padded_documents:
                estimate += combined[1] * combined[3] - combined[2]
            if estimate <= budget:
                groups[i].append(item)
                statistics[i] = combined
                break
        else:
            groups.append([item])
            statistics.append((cost, count, total, maximum))
    return groups


@dataclass(frozen=True)
class ChoicePrediction:
    selected_id: str
    probabilities: dict[str, float]


@dataclass(frozen=True)
class NoulPrediction:
    probability_yes: float
    probabilities: dict[str, float]


@dataclass(frozen=True)
class ScorePrediction:
    score: float
    normalized_score: float
    probabilities: dict[str, float]
    values: dict[str, float]


def validate_typed_group(group: Group):
    metadata = group.metadata
    if metadata is None or metadata.candidate_ids is None or metadata.kind == "ranking":
        raise ValueError(
            "Typed prediction requires judgment candidate IDs; use predict for ranking"
        )
    if group.task == "noul" and metadata.yes_index is None:
        raise ValueError("Noul requires true/false or yes/no candidate IDs")
    if group.task == "score" and metadata.score_values is None:
        raise ValueError(
            "Score requires judgment metadata and at least two distinct numeric values"
        )
    if group.task not in {"choice", "noul", "score"}:
        raise ValueError("Typed prediction supports choice, noul and ordinal score")


def interpret_prediction(
    group: Group, probabilities
) -> ChoicePrediction | NoulPrediction | ScorePrediction:
    """Map probabilities to IDs and values; never infer numbers from candidate text."""
    validate_typed_group(group)
    p = torch.as_tensor(probabilities, dtype=torch.float64).detach().cpu()
    if (
        p.ndim != 1
        or len(p) != len(group.candidates)
        or not torch.isfinite(p).all()
        or (p < 0).any()
        or abs(p.sum().item() - 1) > 1e-5
    ):
        raise ValueError("Expected normalized probabilities aligned with candidates")
    metadata = group.metadata
    assert metadata is not None and metadata.candidate_ids is not None
    mapping = dict(zip(metadata.candidate_ids, p.tolist(), strict=True))
    if group.task == "choice":
        return ChoicePrediction(metadata.candidate_ids[int(p.argmax())], mapping)
    if group.task == "noul":
        assert metadata.yes_index is not None
        return NoulPrediction(p[metadata.yes_index].item(), mapping)
    values = metadata.score_values
    assert values is not None
    score = sum(prob * value for prob, value in zip(p.tolist(), values, strict=True))
    return ScorePrediction(
        score,
        (score - min(values)) / (max(values) - min(values)),
        mapping,
        dict(zip(metadata.candidate_ids, values, strict=True)),
    )


class ChoiceInteraction(nn.Module):
    """One small attention block, isolated by decision rather than shared prefix.

    Zero-initializing only the final projection preserves the existing scorer.
    No candidate position embeddings or dropout are used.
    """

    def __init__(self, hidden_size, width=128, heads=4):
        super().__init__()
        if any(not isinstance(v, int) or isinstance(v, bool) or v < 1 for v in (width, heads)):
            raise ValueError("Choice width and heads must be positive integers")
        if width % heads:
            raise ValueError("Choice width must be divisible by heads")
        self.width, self.heads = width, heads
        self.project = nn.Linear(hidden_size, width)
        self.norm = nn.LayerNorm(width)
        self.qkv = nn.Linear(width, width * 3)
        self.attention_out = nn.Linear(width, width)
        self.ffn = nn.Sequential(
            nn.LayerNorm(width),
            nn.Linear(width, width * 2),
            nn.GELU(),
            nn.Linear(width * 2, width),
        )
        self.output = nn.Linear(width, 1)
        nn.init.zeros_(self.output.weight)
        nn.init.zeros_(self.output.bias)

    def forward(self, hidden, indices, mask):
        x = self.project(hidden)[indices]
        batch, count, _ = x.shape
        qkv = self.qkv(self.norm(x)).reshape(batch, count, 3, self.heads, -1)
        q, k, v = qkv.unbind(2)
        attended = (
            F.scaled_dot_product_attention(
                q.transpose(1, 2),
                k.transpose(1, 2),
                v.transpose(1, 2),
                attn_mask=mask[:, None, None, :],
            )
            .transpose(1, 2)
            .reshape(batch, count, self.width)
        )
        x = x + self.attention_out(attended)
        x = x + self.ffn(x)
        delta = self.output(x).masked_fill(~mask.unsqueeze(-1), 0)
        return hidden.new_zeros((hidden.shape[0], 1)).index_add(
            0, indices.flatten(), delta.flatten(0, 1)
        )


def _text(value):
    """Return strings unchanged; serialize other decoded JSON values deterministically."""
    return (
        value
        if isinstance(value, str)
        else json.dumps(value, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
    )


def input_groups(case_input, *, prefix_layout="instruction_state"):
    """Render one native request into ordered Group objects for predict_groups().

    Accept exactly state_json and decisions, with JSON strings for state,
    instructions and candidate descriptions/content. See the module input schema.
    This performs no tokenization or inference. Noul descriptions are embedded
    into the query; numeric Score values remain output metadata. Labels and row
    metadata are not accepted. prefix_layout controls instruction/state order.
    """
    if not isinstance(case_input, dict) or set(case_input) != {"state_json", "decisions"}:
        raise ValueError("Pass the input object only: state_json and decisions")
    state = case_input["state_json"]
    json.loads(state)
    groups, seen = [], set()
    for decision in case_input["decisions"]:
        did = decision["id"]
        if not isinstance(did, str) or not did or did in seen:
            raise ValueError("Decision IDs must be nonempty and unique")
        seen.add(did)
        kind, task = decision["kind"], decision["type"]
        if kind not in {"judgment", "ranking"}:
            raise ValueError("Unknown decision kind")
        ranking = kind == "ranking"
        if ranking and (task is not None or decision.get("scoring") != "relative"):
            raise ValueError("Ranking requires null type and scoring=relative")
        if ranking:
            task = "score"
        if not ranking and (decision.get("documents") or decision.get("scoring") is not None):
            raise ValueError("Judgments cannot contain ranking documents or scoring")
        candidates = decision["documents"] if ranking else decision["criteria"]
        if ranking and decision.get("criteria"):
            raise ValueError("Ranking cannot contain judgment criteria")
        ids = tuple(c["id"] for c in candidates)
        descriptions = [
            _text(json.loads(c["content_json" if ranking else "description_json"]))
            for c in candidates
        ]
        values = None if ranking or task != "score" else tuple(c["value"] for c in candidates)
        if values is not None and any(
            not isinstance(v, (int, float)) or isinstance(v, bool) for v in values
        ):
            raise ValueError("Score criteria require numeric values")
        context = state
        if task == "noul":
            if set(ids) not in ({"true", "false"}, {"yes", "no"}):
                raise ValueError("Noul requires authored true/false or yes/no criteria")
            mapping = dict(zip(ids, descriptions, strict=True))
            context = json.dumps(
                {
                    "noul": {
                        "yes": mapping["true" if "true" in ids else "yes"],
                        "no": mapping["false" if "false" in ids else "no"],
                    },
                    "state": json.loads(state),
                },
                ensure_ascii=False,
                separators=(",", ":"),
            )
        parts = QueryParts(
            _text(json.loads(decision["instructions_json"])),
            context,
            decision.get("system_prompt") or "",
            prefix_layout,
        )
        group = Group(
            parts.render(),
            [
                f"Document: {d}" if ranking else f"Candidate: {i}: {d}"
                for i, d in zip(ids, descriptions, strict=True)
            ],
            task,
            query_parts=parts,
            metadata=DecisionMetadata(
                candidate_ids=ids, candidate_values=values, kind=kind, decision_id=did
            ),
        )
        if not ranking:
            from_types = {"choice", "noul", "score"}
            if task not in from_types or (task == "score" and len(candidates) < 2):
                raise ValueError("Expected Choice, Noul or numeric Score judgment")
        groups.append(group)
    return groups


class _SDPAPrefix(nn.Module):
    """Shared-prefix SDPA with per-forward layouts and blocked local attention."""

    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        self.local_block_size = 256

    def qkv(self, layer, hidden, rotary):
        attn = layer.attn
        q, k, v = (
            attn.Wqkv(layer.attn_norm(hidden))
            .view(*hidden.shape[:2], 3, -1, attn.head_dim)
            .unbind(2)
        )
        q, k = apply_rotary_pos_emb(q.transpose(1, 2), k.transpose(1, 2), *rotary)
        return q.transpose(1, 2), k.transpose(1, 2), v

    @staticmethod
    def layout(qmask, kmask, window):
        allowed = kmask[:, None, :]
        if window is not None:
            qp = qmask.long().cumsum(1) - 1 + (kmask.sum(1) - qmask.sum(1))[:, None]
            kp = kmask.long().cumsum(1) - 1
            allowed = allowed & ((qp[:, :, None] - kp[:, None, :]).abs() <= window)
        return allowed[:, None]

    @staticmethod
    def attend(q, k, v, qmask, allowed):
        result = F.scaled_dot_product_attention(
            q.transpose(1, 2),
            k.transpose(1, 2),
            v.transpose(1, 2),
            attn_mask=allowed,
            dropout_p=0.0,
        ).transpose(1, 2)
        return result * qmask[:, :, None, None]

    def blocked_layout(self, mask, window):
        """Bound local attention work using exact overlapping key windows."""
        length = mask.shape[1]
        block = self.local_block_size
        count = (length + block - 1) // block
        starts = torch.arange(count, device=mask.device)[:, None] * block
        queries = starts + torch.arange(block, device=mask.device)[None]
        keys = starts + torch.arange(-window, block + window, device=mask.device)[None]
        indices = keys.clamp(0, length - 1)
        valid = (keys >= 0) & (keys < length)
        allowed = mask[:, indices][:, :, None, :] & valid[None, :, None, :]
        allowed = allowed & ((queries[:, :, None] - keys[:, None, :]).abs() <= window)[None]
        return indices, allowed[:, :, None], count * block - length

    def blocked_attend(self, q, k, v, mask, layout):
        indices, allowed, padding = layout
        batch, length, heads, dim = q.shape
        count, width = indices.shape
        q = (
            F.pad(q, (0, 0, 0, 0, 0, padding))
            .reshape(batch, count, self.local_block_size, heads, dim)
            .permute(0, 1, 3, 2, 4)
        )
        k = k[:, indices].permute(0, 1, 3, 2, 4)
        v = v[:, indices].permute(0, 1, 3, 2, 4)
        result = F.scaled_dot_product_attention(
            q.flatten(0, 1),
            k.flatten(0, 1),
            v.flatten(0, 1),
            attn_mask=allowed.flatten(0, 1),
            dropout_p=0.0,
        )
        result = (
            result.reshape(batch, count, heads, self.local_block_size, dim)
            .permute(0, 1, 3, 2, 4)
            .reshape(batch, -1, heads, dim)[:, :length]
        )
        return result * mask[:, :, None, None]

    @staticmethod
    def update(layer, hidden, attended):
        hidden = hidden + layer.attn.out_drop(layer.attn.Wo(attended.flatten(2)))
        return hidden + layer.mlp(layer.mlp_norm(hidden))

    def forward(self, prefix_ids, prefix_mask, doc_ids, doc_mask, owners):
        prefix = self.backbone.embeddings(prefix_ids)
        hidden = self.backbone.embeddings(doc_ids)
        pp = (prefix_mask.long().cumsum(1) - 1).clamp_min(0)
        pm = prefix_mask[owners]
        dp = (doc_mask.long().cumsum(1) - 1).clamp_min(0) + pm.sum(1)[:, None]
        combined = torch.cat((pm, doc_mask), dim=1)
        layouts, rotations = {}, {}
        for layer in self.backbone.layers:
            kind = layer.attention_type
            window = None if layer.attn.sliding_window is None else layer.attn.sliding_window - 1
            blocked = window is not None and prefix.shape[1] > self.local_block_size + 2 * window
            if window not in layouts:
                pl = (
                    self.blocked_layout(prefix_mask, window)
                    if blocked
                    else self.layout(prefix_mask, prefix_mask, window)
                )
                layouts[window] = pl, self.layout(doc_mask, combined, window)
            if kind not in rotations:
                rotations[kind] = (
                    self.backbone.rotary_emb(prefix, pp, kind),
                    self.backbone.rotary_emb(hidden, dp, kind),
                )
            pq, pk, pv = self.qkv(layer, prefix, rotations[kind][0])
            pl, dl = layouts[window]
            attended = (
                self.blocked_attend(pq, pk, pv, prefix_mask, pl)
                if blocked
                else self.attend(pq, pk, pv, prefix_mask, pl)
            )
            prefix = self.update(layer, prefix, attended)
            q, k, v = self.qkv(layer, hidden, rotations[kind][1])
            k, v = torch.cat((pk[owners], k), 1), torch.cat((pv[owners], v), 1)
            hidden = self.update(layer, hidden, self.attend(q, k, v, doc_mask, dl))
        return self.backbone.final_norm(hidden)


def _load_fa2():
    """Load the optional native wheel, including its binary compatibility check."""
    try:
        return importlib.import_module("flash_attn").flash_attn_varlen_func
    except (ImportError, OSError, RuntimeError, AttributeError) as error:
        raise RuntimeError(
            "flash_attention_2 requires a compatible flash-attn wheel for this PyTorch/CUDA "
            "installation; install the fa2 extra or select attn_implementation='sdpa'"
        ) from error


class _TokenLayout:
    def __init__(self, mask):
        self.indices = mask.flatten().nonzero().flatten()
        self.lengths = mask.sum(1)
        self.cumulative = F.pad(self.lengths.cumsum(0).to(torch.int32), (1, 0))
        self.positions = (mask.long().cumsum(1) - 1).flatten()[self.indices]
        self.maximum = int(self.lengths.max())


class _FA2Prefix(nn.Module):
    """Keep real tokens packed through attention, projections and feed-forward layers."""

    def __init__(self, backbone):
        super().__init__()
        self.backbone = backbone
        self.attention = _load_fa2()

    def rotary(self, hidden, positions):
        return {
            kind: tuple(
                x.squeeze(0) for x in self.backbone.rotary_emb(hidden, positions[None], kind)
            )
            for kind in set(self.backbone.config.layer_types)
        }

    @staticmethod
    def qkv(layer, hidden, rotary):
        q, k, v = (
            layer.attn.Wqkv(layer.attn_norm(hidden))
            .view(hidden.shape[0], 3, -1, layer.attn.head_dim)
            .unbind(1)
        )
        q, k = apply_rotary_pos_emb(q, k, *rotary, unsqueeze_dim=1)
        return q, k, v

    @staticmethod
    def update(layer, hidden, attended):
        hidden = hidden + layer.attn.out_drop(layer.attn.Wo(attended.flatten(1)))
        return hidden + layer.mlp(layer.mlp_norm(hidden))

    def forward(self, prefix_ids, prefix_mask, doc_ids, doc_mask, owners):
        if not prefix_ids.is_cuda:
            raise RuntimeError("FA2 encoder requires CUDA")
        pl, dl = _TokenLayout(prefix_mask), _TokenLayout(doc_mask)
        prefix_lengths = pl.lengths[owners]
        kv_lengths = prefix_lengths + dl.lengths
        cuk = F.pad(kv_lengths.cumsum(0).to(torch.int32), (1, 0))
        total, maximum = int(cuk[-1]), int(kv_lengths.max())
        branches = torch.repeat_interleave(
            torch.arange(len(owners), device=owners.device), kv_lengths, output_size=total
        )
        offsets = torch.arange(total, device=owners.device) - cuk[branches]
        gather = torch.where(
            offsets < prefix_lengths[branches],
            pl.cumulative[owners[branches]] + offsets,
            pl.indices.numel() + dl.cumulative[branches] + offsets - prefix_lengths[branches],
        ).long()
        documents = torch.repeat_interleave(
            torch.arange(len(owners), device=owners.device),
            dl.lengths,
            output_size=dl.indices.numel(),
        )
        prefix = self.backbone.embeddings(prefix_ids.flatten()[pl.indices])
        hidden = self.backbone.embeddings(doc_ids.flatten()[dl.indices])
        pr, dr = (
            self.rotary(prefix, pl.positions),
            self.rotary(hidden, dl.positions + prefix_lengths[documents]),
        )
        for layer in self.backbone.layers:
            kind = layer.attention_type
            window = (
                (-1, -1)
                if layer.attn.sliding_window is None
                else (layer.attn.sliding_window - 1,) * 2
            )
            pq, pk, pv = self.qkv(layer, prefix, pr[kind])
            attended = self.attention(
                pq,
                pk,
                pv,
                pl.cumulative,
                pl.cumulative,
                pl.maximum,
                pl.maximum,
                dropout_p=0.0,
                causal=False,
                window_size=window,
            )
            prefix = self.update(layer, prefix, attended)
            q, k, v = self.qkv(layer, hidden, dr[kind])
            k = torch.cat((pk, k)).index_select(0, gather)
            v = torch.cat((pv, v)).index_select(0, gather)
            attended = self.attention(
                q,
                k,
                v,
                dl.cumulative,
                cuk,
                dl.maximum,
                maximum,
                dropout_p=0.0,
                causal=False,
                window_size=window,
            )
            hidden = self.update(layer, hidden, attended)
        hidden = self.backbone.final_norm(hidden)
        output = hidden.new_zeros((doc_ids.numel(), hidden.shape[-1]))
        return output.index_copy(0, dl.indices, hidden).view(*doc_ids.shape, hidden.shape[-1])


class BekkoInference(InputModule):
    """Self-contained ST module with typed predict() and optional torch.compile()."""

    tokenizer: PreTrainedTokenizerBase
    config_file_name = "inference_config.json"
    save_in_root = False
    budget_policy = "adaptive-v1"

    def __init__(
        self,
        backbone,
        tokenizer,
        *,
        tasks,
        query_length=None,
        document_length=None,
        context_length=None,
        query_truncation="balanced",
        task_tokens=None,
        choice_interaction=None,
        prefix_layout="instruction_state",
    ):
        """Build a runtime from a ModernBERT backbone, tokenizer, heads and saved limits."""
        super().__init__()
        if backbone.config.model_type != "modernbert":
            raise ValueError("Expected ModernBERT-compatible weights")
        maximum = backbone.config.max_position_embeddings
        capacity = maximum if context_length is None else context_length
        if not 5 <= capacity <= maximum:
            raise ValueError("Context length exceeds the backbone positional capacity")
        self.context_length = capacity
        query_length = capacity - 2 if query_length is None else query_length
        document_length = min(3800, capacity - 3) if document_length is None else document_length
        self._validate_limits(query_length, document_length, capacity)
        if query_truncation not in {"right", "balanced"} or prefix_layout not in {
            "instruction_state",
            "state_instruction",
        }:
            raise ValueError("Invalid truncation or prefix layout")
        if (
            not tasks
            or len(set(tasks)) != len(tasks)
            or set(tasks) - {"choice", "noul", "score", "reranker"}
        ):
            raise ValueError("Expected distinct supported task heads")
        self.encoder = _SDPAPrefix(backbone)
        self.attn_implementation = "sdpa"
        self.tokenizer = tokenizer
        self.query_length, self.document_length = query_length, document_length
        self.query_truncation, self.tasks = query_truncation, list(tasks)
        self.prefix_layout = prefix_layout
        self.task_tokens = dict(task_tokens or {})
        self.task_token_ids = {
            task: tokenizer.convert_tokens_to_ids(marker)
            for task, marker in self.task_tokens.items()
        }
        if any(marker not in tokenizer.get_vocab() for marker in self.task_tokens.values()):
            raise ValueError("Checkpoint tokenizer is missing a task token")
        self.hidden_size = backbone.config.hidden_size
        self.heads = nn.ModuleDict({task: nn.Linear(self.hidden_size, 1) for task in tasks})
        self.choice_interaction_config = choice_interaction
        self.choice_interaction = (
            ChoiceInteraction(self.hidden_size, **choice_interaction)
            if choice_interaction
            else None
        )
        self._compiled_forward = None
        self.settings = dict(
            tasks=list(tasks),
            query_length=query_length,
            document_length=document_length,
            context_length=capacity,
            query_truncation=query_truncation,
            task_tokens=self.task_tokens,
            choice_interaction=choice_interaction,
            prefix_layout=prefix_layout,
        )

    def set_attention_implementation(self, implementation="auto", *, device=None):
        """Select an encoder without changing weights, budgets or the saved configuration.

        Auto prefers compatible FA2 on CUDA capability 8+ and otherwise uses SDPA.
        Explicit FA2 never falls back. Selection clears any compiled encoder.
        """
        if implementation not in {"auto", "sdpa", "flash_attention_2"}:
            raise ValueError("attn_implementation must be auto, sdpa or flash_attention_2")
        device = torch.device(device) if device is not None else next(self.parameters()).device
        if implementation != "sdpa":
            supported = device.type == "cuda" and torch.cuda.get_device_capability(device)[0] >= 8
            if not supported:
                if implementation == "flash_attention_2":
                    raise ValueError("flash_attention_2 requires a CUDA GPU with capability 8.0+")
                implementation = "sdpa"
            else:
                try:
                    _load_fa2()
                except RuntimeError:
                    if implementation == "flash_attention_2":
                        raise
                    implementation = "sdpa"
                else:
                    implementation = "flash_attention_2"
        if implementation != self.attn_implementation:
            encoder_class = _FA2Prefix if implementation == "flash_attention_2" else _SDPAPrefix
            self.encoder = encoder_class(self.encoder.backbone).train(self.training)
            self._compiled_forward = None
            self.attn_implementation = implementation
        return self

    def compile_inference(self, *, mode="default", dynamic=True, backend="inductor"):
        """Enable lazy compilation of encoder.forward and return this runtime.

        mode, dynamic and backend are forwarded to torch.compile; defaults are
        "default", True and "inductor". Heads, tokenization, rendering and output
        interpretation stay eager. First calls pay compilation cost and changing
        shapes may recompile. Warm representative batches before timing. Compile
        failures propagate; call disable_compile() to explicitly use eager mode.
        """
        self.eval().requires_grad_(False)
        self._compiled_forward = torch.compile(
            self.encoder.forward, mode=mode, dynamic=dynamic, backend=backend
        )
        return self

    def disable_compile(self):
        """Clear the compiled encoder callable and restore eager execution; return None."""
        self._compiled_forward = None

    def preprocess(self, inputs, prompt=None, **kwargs):
        """Reject raw-text encode(); typed decisions require predict() or predict_groups()."""
        raise ValueError(
            "Typed decisions require candidate groups; use model.predict(input_object)"
        )

    @staticmethod
    def _validate_limits(query_length, document_length, capacity):
        """Validate branch caps including special tokens against shared positional capacity."""
        if not 3 <= query_length <= capacity - 2 or not 2 <= document_length <= capacity - 3:
            raise ValueError("Invalid query/document limits for this backbone")

    def _limits(self, query_length, document_length, context_length=None):
        """Resolve per-call caps, clamping saved defaults to a smaller requested context."""
        capacity = self.context_length if context_length is None else context_length
        if not 5 <= capacity <= self.encoder.backbone.config.max_position_embeddings:
            raise ValueError("Context length exceeds the backbone positional capacity")
        q = min(self.query_length, capacity - 2) if query_length is None else query_length
        d = min(self.document_length, capacity - 3) if document_length is None else document_length
        self._validate_limits(q, d, capacity)
        return q, d, capacity

    def _documents(self, documents, tasks, limit):
        """Tokenize candidates with task markers and final SEP, right-truncated to limit."""
        if len(tasks) != len(documents):
            raise ValueError("Tasks and candidates must align")
        if not documents:
            return []
        encoded = self.tokenizer(
            documents,
            add_special_tokens=False,
            truncation=True,
            max_length=limit - 1,
        )["input_ids"]
        return [
            [self.task_token_ids[t], *d[: limit - 2], self.tokenizer.sep_token_id]
            if t in self.task_token_ids
            else [*d, self.tokenizer.sep_token_id]
            for t, d in zip(tasks, encoded, strict=True)
        ]

    def _queries(self, queries, parts, limits):
        """Tokenize queries under per-query caps with balanced or right truncation."""
        if not queries:
            return []
        if self.query_truncation == "balanced":
            if (
                parts is None
                or len(parts) != len(queries)
                or any(p is None or p.render() != q for p, q in zip(parts, queries, strict=True))
            ):
                raise ValueError("Balanced truncation requires matching QueryParts")
            # Reconstruct at the boundary: callers may use training-side dataclasses.
            parts = [QueryParts(p.instruction, p.context, p.system, p.layout) for p in parts]
            return balanced_query_ids(self.tokenizer, parts, limits)
        encoded = self.tokenizer(
            queries,
            add_special_tokens=False,
            truncation=True,
            max_length=max(limits) - 2,
        )["input_ids"]
        return [
            [self.tokenizer.cls_token_id, *q[: limit - 2], self.tokenizer.sep_token_id]
            for q, limit in zip(encoded, limits, strict=True)
        ]

    def tokenize_branches(
        self,
        queries,
        documents,
        document_tasks=None,
        *,
        query_parts=None,
        query_length=None,
        document_length=None,
        context_length=None,
    ):
        """Tokenize separate branch lists and return (query_ids, document_ids).

        This low-level helper has no decision ownership, so allocation uses the
        longest branches conservatively. Use prepare_groups() for per-decision
        adaptive allocation or predict() for the full native-input interface.
        Balanced truncation requires query_parts matching each rendered query.
        """
        qlimit, dlimit, capacity = self._limits(query_length, document_length, context_length)
        dids = self._documents(documents, document_tasks or ["reranker"] * len(documents), dlimit)
        qids = self._queries(queries, query_parts, [qlimit] * len(queries))
        dsize, qsize = allocate_query_budget(
            max(map(len, dids), default=2),
            max(map(len, qids), default=3),
            capacity,
        )
        if any(len(q) > qsize for q in qids):
            qids = self._queries(queries, query_parts, [qsize] * len(queries))
        return qids, [d[: dsize - 1] + [d[-1]] if len(d) > dsize else d for d in dids]

    def prepare_groups(
        self, groups, *, query_length=None, document_length=None, context_length=None
    ):
        """Share positions equally, then lend unused capacity; candidates win odd tokens.

        Candidate caps apply before sharing. All candidates in a decision see the
        same query. Allocation does not depend on neighboring decisions or batches.
        Special tokens count toward limits. Training settings are not modified.
        """
        qlimit, dlimit, capacity = self._limits(query_length, document_length, context_length)
        documents = list(dict.fromkeys((g.task, d) for g in groups for d in g.candidates))
        dids = self._documents([d for _, d in documents], [t for t, _ in documents], dlimit)
        dmap = dict(zip(documents, dids, strict=True))
        queries = list(dict.fromkeys((g.query, g.query_parts) for g in groups))
        qids = self._queries(
            [q for q, _ in queries], [p for _, p in queries], [qlimit] * len(queries)
        )
        original = dict(zip(queries, qids, strict=True))
        keys, document_limits = [], []
        for g in groups:
            if not g.candidates:
                raise ValueError("A decision requires at least one candidate")
            ds, qs = allocate_query_budget(
                max(len(dmap[g.task, d]) for d in g.candidates),
                len(original[g.query, g.query_parts]),
                capacity,
            )
            keys.append((g.query, g.query_parts, qs))
            document_limits.append(ds)
        unique = list(dict.fromkeys(keys))
        qmap = {key: original[key[:2]] for key in unique if len(original[key[:2]]) <= key[2]}
        reduced = [key for key in unique if key not in qmap]
        encoded = self._queries(
            [q for q, _, _ in reduced], [p for _, p, _ in reduced], [n for _, _, n in reduced]
        )
        qmap.update(zip(reduced, encoded, strict=True))
        return [
            PreparedGroup(
                tuple(qmap[key]),
                g.task,
                qmap[key],
                [
                    dmap[g.task, d][: limit - 1] + [dmap[g.task, d][-1]]
                    if len(dmap[g.task, d]) > limit
                    else dmap[g.task, d]
                    for d in g.candidates
                ],
                g.target,
                g.metadata,
            )
            for g, key, limit in zip(groups, keys, document_limits, strict=True)
        ]

    def collate_tokens(self, queries, documents, owners):
        """Pad token lists on CPU and map each document to a unique prefix via owners."""

        def pad(rows):
            ids = torch.nn.utils.rnn.pad_sequence(
                [torch.tensor(r, dtype=torch.long) for r in rows],
                batch_first=True,
                padding_value=self.tokenizer.pad_token_id,
            )
            return ids, torch.arange(ids.shape[1])[None] < torch.tensor([len(r) for r in rows])[
                :, None
            ]

        pi, pm = pad(queries)
        di, dm = pad(documents)
        return dict(
            prefix_ids=pi,
            prefix_mask=pm,
            doc_ids=di,
            doc_mask=dm,
            owners=torch.tensor(owners, dtype=torch.long),
        )

    def forward(self, features, **kwargs):
        """Score collated tensor features and return the same dict with raw logits.

        Requires prefix_ids, prefix_mask, doc_ids, doc_mask and owners; multiple
        heads require head_indices. Optional Choice interaction also uses its
        group indices/mask. Adds scores and sentence_embedding, both (N, 1).
        These are logits, not typed predictions or normalized probabilities.
        Prefer predict()/predict_groups(), which also manage batching and
        inference mode. CUDA encoder autocast is BF16; task heads use FP32.
        """
        device = next(self.parameters()).device
        with torch.autocast(device.type, dtype=torch.bfloat16, enabled=device.type == "cuda"):
            execute = self._compiled_forward or self.encoder
            hidden = execute(
                **{
                    k: features[k]
                    for k in ("prefix_ids", "prefix_mask", "doc_ids", "doc_mask", "owners")
                }
            )
            mask = features["doc_mask"]
            hidden = (hidden * mask.unsqueeze(-1)).sum(1) / mask.sum(1)[:, None]
        with torch.autocast(device.type, enabled=False):
            indices = features.get("head_indices")
            if indices is None:
                if len(self.tasks) != 1:
                    raise ValueError("Multi-task forward requires head_indices")
                indices = {self.tasks[0]: None}
            scores = hidden.new_zeros((hidden.shape[0], 1), dtype=torch.float32)
            for task, index in indices.items():
                if task not in self.heads:
                    raise ValueError(f"Checkpoint has no {task} head")
                if index is None:
                    scores = self.heads[task](hidden.float())
                else:
                    ix = torch.as_tensor(index, device=device, dtype=torch.long)
                    scores = scores.index_copy(
                        0, ix, self.heads[task](hidden.index_select(0, ix).float())
                    )
            if self.choice_interaction is not None and "choice" in indices:
                scores = scores + self.choice_interaction(
                    hidden.float(), features["choice_indices"], features["choice_mask"]
                )
        features["scores"] = features["sentence_embedding"] = scores
        return features

    @torch.inference_mode()
    def predict_groups(
        self,
        groups,
        *,
        batch_size=128,
        token_budget=64000,
        query_length=None,
        document_length=None,
        context_length=None,
        show_progress_bar=True,
    ):
        """Predict probability tensors for already-rendered decision groups.

        Use predict(list_of_requests) for ordinary native-input batching. This
        lower-level API accepts an iterable of Group objects (materialized),
        such as input_groups(request), and skips typed result interpretation.
        With balanced query truncation, each group needs matching QueryParts.

        Returns a list of CPU FP32 tensors, one per group, each shaped
        (number_of_candidates,) and normalized within that group. Group order
        and candidate order match the input, regardless of internal bucketing.
        Supported tasks must have heads in this checkpoint; empty input gives [].

        batch_size (128) bounds decisions tokenized/sorted per window.
        token_budget (64,000) bounds estimated padded microbatch work; a decision
        exceeding it runs alone, with all candidates together. Length overrides
        are per call; see predict() and the module's adaptive context description.
        show_progress_bar counts completed decisions on stderr, not requests.
        """
        if batch_size < 1 or token_budget < 1:
            raise ValueError("batch_size and token_budget must be positive")
        groups = list(groups)
        self.eval()
        # Validate overrides even for empty requests.
        _, _, capacity = self._limits(query_length, document_length, context_length)
        if any(g.task not in self.tasks for g in groups):
            raise ValueError("Requested task is absent from this checkpoint")
        device = next(self.parameters()).device
        result = [None] * len(groups)
        with tqdm(total=len(groups), desc="Batches", disable=not show_progress_bar) as progress:
            for start in range(0, len(groups), batch_size):
                prepared = self.prepare_groups(
                    groups[start : start + batch_size],
                    query_length=query_length,
                    document_length=document_length,
                    context_length=context_length,
                )
                ordered = sorted(
                    enumerate(prepared),
                    key=lambda pair: (
                        len(pair[1].query),
                        max(map(len, pair[1].documents)),
                    ),
                    reverse=True,
                )
                batch, positions = [], []
                count = max_query = max_document = 0

                def flush():
                    scores = self(to_device(collate_groups(batch, self), device))[
                        "scores"
                    ].flatten()
                    lengths = [len(g.documents) for g in batch]
                    # One host transfer/synchronization per microbatch, including
                    # the finite check; preserve FP32 per-decision softmax.
                    probabilities = torch.cat(
                        [p.float().softmax(0) for p in scores.split(lengths)]
                    ).cpu()
                    if not torch.isfinite(probabilities).all():
                        raise FloatingPointError("Nonfinite inference probabilities")
                    for index, values in zip(positions, probabilities.split(lengths), strict=True):
                        result[start + index] = values
                    progress.update(len(batch))

                for index, group in ordered:
                    nq = max(max_query, len(group.query))
                    nd = max(max_document, max(map(len, group.documents)))
                    nc = count + len(group.documents)
                    if batch and (nc * (nq + nd) > token_budget or nq + nd > capacity):
                        flush()
                        batch, positions = [], []
                        nq, nd, nc = (
                            len(group.query),
                            max(map(len, group.documents)),
                            len(group.documents),
                        )
                    batch.append(group)
                    positions.append(index)
                    max_query, max_document, count = nq, nd, nc
                if batch:
                    flush()
        return result

    @staticmethod
    def _interpret(group, probabilities):
        """Convert one group distribution into typed fields or stable relative ranking."""
        assert group.metadata is not None and group.metadata.candidate_ids is not None
        if group.metadata.kind == "ranking":
            return {
                "probabilities": dict(
                    zip(
                        group.metadata.candidate_ids,
                        probabilities.tolist(),
                        strict=True,
                    )
                ),
                "order": [
                    group.metadata.candidate_ids[i]
                    for i in probabilities.argsort(descending=True, stable=True).tolist()
                ],
            }
        return asdict(interpret_prediction(group, probabilities))

    def predict(
        self,
        inputs,
        *,
        batch_size=128,
        token_budget=64000,
        query_length=None,
        document_length=None,
        context_length=None,
        prefix_layout=None,
        show_progress_bar=True,
    ):
        """Predict typed decisions for one request or a batch of native requests.

        Args:
            inputs: Dict with exactly state_json and decisions, or an iterable
                of those dicts (materialized as a list). See the module example.
                Use a list for ordinary batched inference, including mixed tasks.
            batch_size: Positive limit on requests rendered per window and on
                decisions tokenized per window (default 128), not candidate count.
            token_budget: Positive padded-work estimate per microbatch (64,000).
                Complete decisions stay together; oversized ones run alone.
                Lower this to reduce microbatch work; it is not a memory cap.
            query_length: Optional query token cap, including special tokens.
            document_length: Optional candidate token cap, including special tokens.
            context_length: Optional shared query/candidate positional limit.
                All length overrides apply only to this call; see adaptive-v1
                allocation in the module docstring.
            prefix_layout: Optional instruction_state or state_instruction;
                None uses the exported rendering order.
            show_progress_bar: Show completed requests on stderr (default True).

        Returns:
            One dict keyed by decision ID for a dict input, otherwise a list of
            those dicts in input order. Each value contains the typed fields
            documented above (Choice, Noul, Score, or relative ranking).
            No decisions yields {}; an empty request list yields [].

        For already-rendered Group objects and raw probability tensors, use
        predict_groups(). For repeated GPU workloads, compile_inference() is
        optional; warm representative batches before measuring throughput.
        """
        single = isinstance(inputs, dict)
        cases = [inputs] if single else list(inputs)
        options = dict(
            batch_size=batch_size,
            token_budget=token_budget,
            query_length=query_length,
            document_length=document_length,
            context_length=context_length,
            show_progress_bar=False,
        )
        self.predict_groups([], **options)
        results = []
        with tqdm(total=len(cases), desc="Predict", disable=not show_progress_bar) as progress:
            for start in range(0, len(cases), batch_size):
                groups, owners = [], []
                chunk = cases[start : start + batch_size]
                outputs = [{} for _ in chunk]
                for index, case in enumerate(chunk):
                    rendered = input_groups(case, prefix_layout=prefix_layout or self.prefix_layout)
                    groups.extend(rendered)
                    owners.extend([index] * len(rendered))
                probabilities = self.predict_groups(groups, **options)
                for owner, group, p in zip(owners, groups, probabilities, strict=True):
                    outputs[owner][group.metadata.decision_id] = self._interpret(group, p)
                results.extend(outputs)
                progress.update(len(chunk))
        return results[0] if single else results

    def get_sentence_embedding_dimension(self):
        """Return the ST-compatible scalar logit dimension; this is not a text embedding."""
        return 1

    def get_config_dict(self):
        """Return serializable runtime settings used by the ST module configuration."""
        return self.settings

    def save(self, output_path, *args, **kwargs):
        """Save module config, backbone config, safetensors and tokenizer to a directory."""
        path = Path(output_path)
        path.mkdir(parents=True, exist_ok=True)
        self.save_config(str(path))
        config = self.encoder.backbone.config.to_dict()
        config.pop("_name_or_path", None)
        (path / "backbone_config.json").write_text(json.dumps(config, indent=2))
        save_file(
            {k: v.detach().cpu().contiguous() for k, v in self.state_dict().items()},
            str(path / "model.safetensors"),
        )
        self.save_tokenizer(str(path / "tokenizer"))

    @classmethod
    def load(
        cls,
        model_name_or_path,
        subfolder="",
        token=None,
        cache_folder=None,
        revision=None,
        local_files_only=False,
        init_defaults=None,
        **kwargs,
    ):
        """Load this ST module from a local directory or Hugging Face repository.

        Called by SentenceTransformer with subfolder, authentication, revision,
        cache and offline settings. Restores FP32 weights with strict matching,
        SDPA attention, eval mode and gradients disabled. Prefer constructing
        BekkoSentenceTransformer for the public typed interface and device setup.
        """
        hub: dict[str, Any] = dict(
            subfolder=subfolder,
            token=token,
            cache_folder=cache_folder,
            revision=revision,
            local_files_only=local_files_only,
        )
        settings = cls.load_config(model_name_or_path, **hub)
        config_path = cls.load_file_path(model_name_or_path, "backbone_config.json", **hub)
        weights = cls.load_file_path(model_name_or_path, "model.safetensors", **hub)
        tokenizer_hub: dict[str, Any] = {**hub, "subfolder": f"{subfolder}/tokenizer".lstrip("/")}
        tok = cls.load_dir_path(model_name_or_path, **tokenizer_hub)
        if config_path is None or weights is None or tok is None:
            raise FileNotFoundError("Incomplete portable checkpoint")
        config = json.loads(Path(config_path).read_text())
        backbone = AutoModel.from_config(
            AutoConfig.for_model(config.pop("model_type"), **config),
            attn_implementation="sdpa",
            dtype=torch.float32,
        )
        model = cls(backbone, AutoTokenizer.from_pretrained(tok, local_files_only=True), **settings)
        model.load_state_dict(load_file(weights), strict=True)
        return model.eval().requires_grad_(False)


class BekkoSentenceTransformer(SentenceTransformer):
    """SentenceTransformer with a public typed-decision inference interface.

    Load an exported v0 checkpoint using the usual SentenceTransformer constructor
    arguments, including device, revision, local_files_only and trust_remote_code.
    The same class is included in the standalone exported inference_v0.py.
    """

    def __init__(self, *args, attn_implementation="auto", **kwargs):
        """Load native weights, then select SDPA or optional FA2 on the final device."""
        model_kwargs = dict(kwargs.get("model_kwargs") or {})
        nested = model_kwargs.pop("attn_implementation", None)
        if nested is not None:
            if attn_implementation != "auto" and attn_implementation != nested:
                raise ValueError("Conflicting attn_implementation arguments")
            attn_implementation = nested
        if attn_implementation not in {"auto", "sdpa", "flash_attention_2"}:
            raise ValueError("attn_implementation must be auto, sdpa or flash_attention_2")
        kwargs["model_kwargs"] = model_kwargs
        super().__init__(*args, **kwargs)
        # Remote-code loading creates a distinct Python class identity, so check
        # the portable module contract instead of using isinstance.
        if len(self) != 1 or getattr(self[0], "config_file_name", None) != "inference_config.json":
            raise ValueError("BekkoSentenceTransformer requires an exported v0 checkpoint")
        cast(Any, self[0]).set_attention_implementation(attn_implementation, device=self.device)

    def predict(
        self,
        inputs,
        *,
        batch_size=128,
        token_budget=64000,
        query_length=None,
        document_length=None,
        context_length=None,
        prefix_layout=None,
        show_progress_bar=True,
    ):
        """Predict typed decisions for one request or a batch of native requests.

        Args:
            inputs: Dict with exactly state_json and decisions, or an iterable
                of those dicts (materialized as a list). See the module example.
                Use a list for ordinary batched inference, including mixed tasks.
            batch_size: Positive limit on requests rendered per window and on
                decisions tokenized per window (default 128), not candidate count.
            token_budget: Positive padded-work estimate per microbatch (64,000).
                Complete decisions stay together; oversized ones run alone.
                Lower this to reduce microbatch work; it is not a memory cap.
            query_length: Optional query token cap, including special tokens.
            document_length: Optional candidate token cap, including special tokens.
            context_length: Optional shared query/candidate positional limit.
                All length overrides apply only to this call; see adaptive-v1
                allocation in the module docstring.
            prefix_layout: Optional instruction_state or state_instruction;
                None uses the exported rendering order.
            show_progress_bar: Show completed requests on stderr (default True).

        Returns:
            One dict keyed by decision ID for a dict input, otherwise a list of
            those dicts in input order. Each value contains the typed fields
            documented above (Choice, Noul, Score, or relative ranking).
            No decisions yields {}; an empty request list yields [].

        For already-rendered Group objects and raw probability tensors, use
        predict_groups(). For repeated GPU workloads, compile_inference() is
        optional; warm representative batches before measuring throughput.
        """
        return cast(Any, self[0]).predict(
            inputs,
            batch_size=batch_size,
            token_budget=token_budget,
            query_length=query_length,
            document_length=document_length,
            context_length=context_length,
            prefix_layout=prefix_layout,
            show_progress_bar=show_progress_bar,
        )

    def predict_groups(
        self,
        groups,
        *,
        batch_size=128,
        token_budget=64000,
        query_length=None,
        document_length=None,
        context_length=None,
        show_progress_bar=True,
    ):
        """Predict probability tensors for already-rendered decision groups.

        Use predict(list_of_requests) for ordinary native-input batching. This
        lower-level API accepts an iterable of Group objects (materialized),
        such as input_groups(request), and skips typed result interpretation.
        With balanced query truncation, each group needs matching QueryParts.

        Returns a list of CPU FP32 tensors, one per group, each shaped
        (number_of_candidates,) and normalized within that group. Group order
        and candidate order match the input, regardless of internal bucketing.
        Supported tasks must have heads in this checkpoint; empty input gives [].

        batch_size (128) bounds decisions tokenized/sorted per window.
        token_budget (64,000) bounds estimated padded microbatch work; a decision
        exceeding it runs alone, with all candidates together. Length overrides
        are per call; see predict() and the module's adaptive context description.
        show_progress_bar counts completed decisions on stderr, not requests.
        """
        return cast(Any, self[0]).predict_groups(
            groups,
            batch_size=batch_size,
            token_budget=token_budget,
            query_length=query_length,
            document_length=document_length,
            context_length=context_length,
            show_progress_bar=show_progress_bar,
        )

    def compile_inference(self, *, mode="default", dynamic=True, backend="inductor"):
        """Enable lazy encoder compilation and return self for optional chaining.

        For repeated inference, call once before representative warmup batches.
        mode="default", dynamic=True and backend="inductor" pass to torch.compile.
        Shapes may recompile; first-call latency includes compilation. Rendering,
        tokenization, heads and typed outputs stay eager. See the module guide.
        """
        cast(Any, self[0]).compile_inference(mode=mode, dynamic=dynamic, backend=backend)
        return self

    def disable_compile(self):
        """Restore eager tensor execution and return this model."""
        cast(Any, self[0]).disable_compile()
        return self


def main():
    """Run standalone inference from a JSON request or request array.

    --model accepts a local export or Hub ID; private Hub access uses saved
    authentication. --compile enables lazy encoder compilation. Writes results
    to stdout and optional request progress to stderr. See --help for controls.
    """
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--model", required=True)
    parser.add_argument(
        "--input", required=True, type=Path, help="JSON file containing only native inference input"
    )
    parser.add_argument("--device", default="cpu")
    parser.add_argument("--compile", action="store_true")
    parser.add_argument(
        "--attn-implementation", choices=["auto", "sdpa", "flash_attention_2"], default="auto"
    )
    parser.add_argument("--token-budget", type=int, default=64000)
    parser.add_argument("--batch-size", type=int, default=128)
    parser.add_argument("--query-length", type=int)
    parser.add_argument("--document-length", type=int)
    parser.add_argument("--context-length", type=int)
    parser.add_argument("--show-progress-bar", action=argparse.BooleanOptionalAction, default=True)
    parser.add_argument("--local-files-only", action="store_true")
    args = parser.parse_args()
    model = BekkoSentenceTransformer(
        args.model,
        device=args.device,
        attn_implementation=args.attn_implementation,
        trust_remote_code=True,
        local_files_only=args.local_files_only,
    )
    if args.compile:
        model.compile_inference()
    print(
        json.dumps(
            model.predict(
                json.loads(args.input.read_text()),
                token_budget=args.token_budget,
                batch_size=args.batch_size,
                query_length=args.query_length,
                document_length=args.document_length,
                context_length=args.context_length,
                show_progress_bar=args.show_progress_bar,
            ),
            ensure_ascii=False,
            indent=2,
        )
    )


if __name__ == "__main__":
    main()