gpc1_server/compiler.py

"""Compile numeric requests and decode range predictions."""
from __future__ import annotations

from dataclasses import dataclass
import json
import math
from typing import Any, Sequence


@dataclass(frozen=True, slots=True)
class NumericField:
    key: str
    description: str
    minimum: float
    maximum: float
    unit: str
    aliases: tuple[str, ...] = ()
    reference: str = "declared_by_caller"

    def __post_init__(self) -> None:
        values=(self.key,self.description,self.unit,self.reference,*self.aliases)
        if any(not isinstance(v,str) or not v.strip() or v != v.strip() for v in values):
            raise ValueError("numeric names and descriptions must be nonempty and edge-whitespace-free")
        if self.key in self.aliases or len(self.aliases) != len(set(self.aliases)):
            raise ValueError("numeric aliases must be unique and exclude the canonical key")
        if not math.isfinite(self.minimum) or not math.isfinite(self.maximum) or self.minimum >= self.maximum:
            raise ValueError("numeric bounds must be finite and strictly increasing")


@dataclass(frozen=True, slots=True)
class CompiledNumeric:
    prompt: str
    input_ids: tuple[int, ...]
    positions: tuple[int, ...]
    candidate_token_ids: tuple[int, ...]


@dataclass(frozen=True, slots=True)
class CompiledVisionNumeric:
    numeric: CompiledNumeric
    processor_inputs: dict[str, Any]
    positions: tuple[int, ...]


def _query(field: NumericField) -> str:
    # Preserve formatting for tokenizer compatibility.
    return ("QUERY canonical_key=" + json.dumps(field.key, ensure_ascii=False)
            + " unit=" + json.dumps(field.unit, ensure_ascii=False)
            + f" bounds=[{field.minimum!r},{field.maximum!r}]\nAnswer code:\n")


def render_prefix(context: str, fields: Sequence[NumericField], codes: Sequence[str]) -> str:
    shared_map = [{"transport_code": codes[i], "percentage_index": i,
                   "normalized_point": i / 100.0} for i in range(101)]
    definitions = [{"canonical_key": f.key, "aliases": list(f.aliases),
                    "meaning": f.description, "range": [f.minimum, f.maximum],
                    "unit": f.unit, "reference": f.reference,
                    "affine_decode": "value=lower+normalized_point*(upper-lower)"}
                   for f in fields]
    return ("NATIVE_NUMERIC_SHARED_CODEBOOK_QUERY_SLOTS_V1\n"
            "Return exactly one listed opaque transport code for every point field. "
            "The shared map binds each transport code to one normalized point. "
            "Each field definition supplies its own affine bounds and unit. Transport strings "
            "are labels, not pretrained numeric semantics.\n"
            "All maps and field definitions precede all answer-query slots. "
            "No answer is present in this input.\nSHARED_TRANSPORT_MAP="
            + json.dumps(shared_map, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
            + "\nFIELD_DEFINITIONS="
            + json.dumps(definitions, ensure_ascii=False, sort_keys=True, separators=(",", ":"))
            + "\nCONTEXT_BEGIN\n" + context + "\nCONTEXT_END\nBEGIN_QUERY_SLOTS\n")


def compile_numeric(tokenizer: Any, context: str, fields: Sequence[NumericField],
                    codebook: Sequence[dict[str, Any]], max_input_tokens: int) -> CompiledNumeric:
    if not context.strip() or not fields:
        raise ValueError("context and fields must be nonempty")
    names=[name for field in fields for name in (field.key,*field.aliases)]
    if len(names) != len(set(names)):
        raise ValueError("numeric canonical keys and aliases must be globally unique")
    codes = [str(row["code"]) for row in codebook[:101]]
    prefix = render_prefix(context, fields, codes)
    cumulative = prefix
    previous = tuple(tokenizer.encode(cumulative, add_special_tokens=False))
    positions: list[int] = []
    for field in fields:
        cumulative += _query(field)
        current = tuple(tokenizer.encode(cumulative, add_special_tokens=False))
        if len(current) <= len(previous) or current[:len(previous)] != previous:
            raise ValueError("tokenizer resegmented an earlier query boundary")
        # Native LM logits at the last prompt token predict the opaque code.
        positions.append(len(current) - 1)
        previous = current
    if len(previous) > max_input_tokens:
        raise ValueError(f"request needs {len(previous)} tokens; limit is {max_input_tokens}; no truncation performed")
    return CompiledNumeric(cumulative, previous, tuple(positions),
                           tuple(int(row["token_id"]) for row in codebook[:101]))


def decode_numeric(probabilities: Sequence[float], field: NumericField) -> dict[str, Any]:
    if len(probabilities) != 101:
        raise ValueError("numeric probability vector must have 101 entries")
    total = math.fsum(float(p) for p in probabilities)
    if not math.isfinite(total) or total <= 0 or any(float(p) < 0 for p in probabilities):
        raise ValueError("invalid numeric probability vector")
    probs = [float(p) / total for p in probabilities]
    map_index = max(range(101), key=probs.__getitem__)
    expected_index = math.fsum(i * p for i, p in enumerate(probs))
    span = field.maximum - field.minimum
    return {
        "value": field.minimum + span * map_index / 100.0,
        "map_value": field.minimum + span * map_index / 100.0,
        "expected_value": field.minimum + span * expected_index / 100.0,
        "map_index": map_index,
        "normalized_grid": [i / 100.0 for i in range(101)],
        "probabilities": probs,
        "minimum": field.minimum, "maximum": field.maximum, "unit": field.unit,
        "calibrated": False,
    }


def _unique_subsequence_offset(values: Sequence[int], needle: Sequence[int]) -> int:
    """Find one token span in linear time, rejecting absent or repeated spans."""
    if not needle:
        raise ValueError("target-free numeric prompt must occur exactly once")
    failure = [0] * len(needle)
    matched = 0
    for index in range(1, len(needle)):
        while matched and needle[index] != needle[matched]:
            matched = failure[matched - 1]
        if needle[index] == needle[matched]:
            matched += 1
        failure[index] = matched
    matched = 0
    offset = -1
    for index, value in enumerate(values):
        while matched and value != needle[matched]:
            matched = failure[matched - 1]
        if value == needle[matched]:
            matched += 1
        if matched == len(needle):
            if offset >= 0:
                raise ValueError("target-free numeric prompt must occur exactly once")
            offset = index + 1 - len(needle)
            matched = failure[matched - 1]
    if offset < 0:
        raise ValueError("target-free numeric prompt must occur exactly once")
    return offset


def compile_vision_numeric(processor: Any, image: Any, context: str,
                           fields: Sequence[NumericField], codebook: Sequence[dict[str, Any]],
                           max_input_tokens: int, model_config: Any) -> CompiledVisionNumeric:
    """Compile one image and numeric fields using the native processor."""
    import torch
    width, height = getattr(image, "size", (0, 0))
    if width <= 0 or height <= 0:
        raise ValueError("source raster dimensions must be positive")
    numeric = compile_numeric(processor.tokenizer, context, fields, codebook, max_input_tokens)
    messages = [[
        {"role": "user", "content": [
            {"type": "image", "image": image},
            {"type": "text", "text": (
                "Inspect only this supplied full uncropped image. The following assistant "
                "schema contains target-free coordinate query slots for this image.")}]},
        {"role": "assistant", "content": [{"type": "text",
            "text": numeric.prompt + "END_TARGET_FREE_QUERY_BLOCK\n"}]},
    ]]
    batch = processor.apply_chat_template(messages, tokenize=True, truncation=False,
        add_generation_prompt=False, continue_final_message=False, enable_thinking=False,
        return_dict=True, return_tensors="pt", return_mm_token_type_ids=True)
    tensors = {k: v for k, v in batch.items() if isinstance(v, torch.Tensor)}
    required = {"input_ids", "attention_mask", "pixel_values", "image_grid_thw", "mm_token_type_ids"}
    if set(tensors) != required:
        raise ValueError(f"processor tensor keys must be exactly {sorted(required)}, got {sorted(tensors)}")
    ids, mask, mm = tensors["input_ids"], tensors["attention_mask"], tensors["mm_token_type_ids"]
    if ids.ndim != 2 or ids.shape[0] != 1 or mask.shape != ids.shape or mm.shape != ids.shape:
        raise ValueError("vision request requires aligned singleton text tensors")
    if not bool(mask.eq(1).all()) or ids.shape[1] + 1 > max_input_tokens:
        raise ValueError("vision request exceeds token admission or contains padding; no truncation performed")
    grid = tensors["image_grid_thw"]
    if grid.shape != (1, 3) or not bool((grid > 0).all()):
        raise ValueError("image_grid_thw must be one positive [t,h,w] row")
    patch_rows = int(grid.prod().item())
    if tensors["pixel_values"].ndim != 2 or tensors["pixel_values"].shape[0] != patch_rows:
        raise ValueError("pixel_values and processor grid disagree")
    config = model_config.to_dict() if hasattr(model_config, "to_dict") else dict(model_config)
    if config.get("model_type") != "qwen3_5_moe" or not isinstance(config.get("vision_config"), dict):
        raise ValueError("vision compiler requires the pinned multimodal Qwen3.5 MoE config")
    merge = int(processor.image_processor.merge_size)
    if patch_rows % (merge * merge): raise ValueError("image grid is not divisible by merge area")
    image_token = int(config["image_token_id"])
    active = [int(v) for v in ids[0].tolist()]
    expected_image_tokens = patch_rows // (merge * merge)
    if active.count(image_token) != expected_image_tokens:
        raise ValueError("image token count and processor grid disagree")
    vision_start = int(config["vision_start_token_id"]); vision_end = int(config["vision_end_token_id"])
    if active.count(vision_start) != 1 or active.count(vision_end) != 1:
        raise ValueError("vision boundary token count is not exactly one pair")
    mm_values = [int(v) for v in mm[0].tolist()]
    if any((kind == 1) != (token == image_token) for token, kind in zip(active, mm_values, strict=True)):
        raise ValueError("image tokens and mm_token_type_ids disagree")
    needle = numeric.input_ids
    offset = _unique_subsequence_offset(active, needle)
    positions = tuple(offset + p for p in numeric.positions)
    return CompiledVisionNumeric(numeric, tensors, positions)