"""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)