"""Strict CPU codec for the JSON-shaped 137-slot parallel target.""" from __future__ import annotations import json import math from collections.abc import Mapping, Sequence from dataclasses import dataclass from decimal import ROUND_HALF_UP, Decimal, InvalidOperation from .gui_owl_short_parallel_input import FIELD_WIDTHS, build_parallel_input VARIANT = "json137" ACTIONS = ( "key", "type", "mouse_move", "left_click", "left_click_drag", "right_click", "middle_click", "double_click", "triple_click", "scroll", "hscroll", "wait", "interact", "terminate", "answer", ) POINT_ACTIONS = frozenset({ "mouse_move", "left_click", "left_click_drag", "right_click", "middle_click", "double_click", "triple_click", }) TEXT_ACTIONS = frozenset({"type", "interact", "answer"}) @dataclass(frozen=True) class AdaptedTarget: variant: str canonical_action: Mapping[str, object] target_ids: tuple[int, ...] selected_spans: tuple[slice, ...] selected_field_kinds: tuple[str, ...] loss_mask: tuple[bool, ...] text_payload: str | None adaptation_metadata: Mapping[str, object] def _half_up(value: object, scale: int, maximum: int, error: str) -> int: if type(value) not in (int, float) or not math.isfinite(value): raise ValueError(error) try: raw = Decimal(str(value)) if not Decimal(0) <= raw <= Decimal(maximum) / scale: raise ValueError(error) scaled = (raw * scale).quantize(Decimal(1), rounding=ROUND_HALF_UP) except InvalidOperation as exc: raise ValueError(error) from exc result = int(scaled) if not 0 <= result <= maximum: raise ValueError(error) return result class Parallel137Codec: """Encode/decode the approved fixed JSON137 target without repair.""" variant = VARIANT output_slots = 137 def __init__(self, tokenizer, *, scroll_coordinates=True): self.scroll_coordinates = scroll_coordinates self.tokenizer = tokenizer built = build_parallel_input(tokenizer) self.input_ids = built.input_ids self.slot_positions = built.slot_positions self.mask_id = tokenizer.token_to_id("<|fim_pad|>") self.pad_id = tokenizer.token_to_id("<|endoftext|>") self.eos_id = tokenizer.token_to_id("<|im_end|>") self.vocab_size = tokenizer.get_vocab_size() if any(type(value) is not int for value in (self.pad_id, self.eos_id)): raise ValueError("native_padding_and_eos_tokens_are_required") added = tokenizer.get_added_tokens_decoder() self.special_ids = frozenset( token_id for token_id, token in added.items() if token.special ) | frozenset({self.pad_id, self.eos_id, self.mask_id}) offset = 0 self.field_slices = {} for name, width in FIELD_WIDTHS: self.field_slices[name] = slice(offset, offset + width) offset += width self.action_encodings = {} for action in ACTIONS: ids = tuple(tokenizer.encode(action, add_special_tokens=False).ids) if not 1 <= len(ids) <= 3 or self.pad_id in ids: raise ValueError("action_has_no_native_three_slot_encoding") encoded = ids + (self.pad_id,) * (3 - len(ids)) if encoded in self.action_encodings.values(): raise ValueError("action_encoding_collision") self.action_encodings[action] = encoded self.action_inverse = {value: key for key, value in self.action_encodings.items()} self.action_ids = self.action_encodings self.base_codec = self digit_tokens = {} for digit in "0123456789": ids = tuple(tokenizer.encode(digit, add_special_tokens=False).ids) if len(ids) != 1: raise ValueError("digits_must_be_native_single_tokens") digit_tokens[digit] = ids[0] self.digit_ids = tuple(digit_tokens[str(digit)] for digit in range(10)) self.inverse_digits = {value: str(digit) for digit, value in enumerate(self.digit_ids)} self.sign_ids = {sign: self._single(sign) for sign in ("+", "-")} self.inverse_signs = {value: key for key, value in self.sign_ids.items()} self.status_ids = {status: self._single(status) for status in ("success", "failure")} self.inverse_status = {value: key for key, value in self.status_ids.items()} self.legal_field_encodings = { "selector": frozenset(self.action_encodings.values()), "digit": frozenset((value,) for value in self.digit_ids), "sign": frozenset((value,) for value in self.sign_ids.values()), "status": frozenset((value,) for value in self.status_ids.values()), } def _single(self, text: str) -> int: ids = self.tokenizer.encode(text, add_special_tokens=False).ids if len(ids) != 1: raise ValueError(f"{text}_must_be_a_native_single_token") return ids[0] def _digits(self, value: int, width: int) -> list[int]: return [self.digit_ids[int(digit)] for digit in f"{value:0{width}d}"] def _decode_digits(self, ids: Sequence[int]) -> int: try: return int("".join(self.inverse_digits[value] for value in ids)) except KeyError as exc: raise ValueError("non_digit_in_numeric_field") from exc def _tokenize_exact(self, text: str, kind: str) -> list[int]: ids = self.tokenizer.encode(text, add_special_tokens=False).ids if any(value in self.special_ids for value in ids): raise ValueError(f"{kind}_contains_special_token") if self.tokenizer.decode(ids, skip_special_tokens=False) != text: raise ValueError(f"{kind}_does_not_roundtrip_exactly") return ids def adapt_target(self, target_action: Mapping[str, object], metadata: Mapping[str, object]): if not isinstance(target_action, Mapping) or not isinstance(metadata, Mapping): raise TypeError("target_action_and_metadata_must_be_mappings") action = target_action.get("action") if action not in self.action_encodings: raise ValueError("unsupported_source_action") target = [self.pad_id] * self.output_slots loss = [False] * self.output_slots spans = [self.field_slices["action"]] kinds = ["selector"] canonical: dict[str, object] = {"action": action} target[0:3] = self.action_encodings[action] loss[0:3] = [True] * 3 def select_digits(name: str, value: int, width: int): span = self.field_slices[name] target[span] = self._digits(value, width) loss[span] = [True] * width spans.append(span); kinds.append("coord4" if name.startswith("coordinate") else name) scroll = action in {"scroll", "hscroll"} if scroll and "coordinate" in target_action and not self.scroll_coordinates: raise ValueError("legacy_scroll_profile_cannot_preserve_coordinate") if action in POINT_ACTIONS or (scroll and "coordinate" in target_action): coordinate = target_action.get("coordinate") if not isinstance(coordinate, list) or len(coordinate) != 2: raise ValueError("point_action_requires_coordinate_pair") x = _half_up(coordinate[0], 1, 1000, "coordinate_must_be_0_to_1000") y = _half_up(coordinate[1], 1, 1000, "coordinate_must_be_0_to_1000") select_digits("coordinate.x", x, 4); select_digits("coordinate.y", y, 4) canonical["coordinate"] = [x, y] if scroll and self.scroll_coordinates and "coordinate" not in target_action: # All PAD means absent. Supervise absence, not arbitrary inactive logits. for name in ("coordinate.x", "coordinate.y"): span = self.field_slices[name] loss[span] = [True] * 4 spans.append(span); kinds.append("scroll_coordinate_absent") if action == "key": keys = target_action.get("keys") if not isinstance(keys, list) or not keys or any(not isinstance(key, str) for key in keys): raise ValueError("key_action_requires_nonempty_string_list") interior = json.dumps(keys, ensure_ascii=False, separators=(",", ":"))[1:-1] ids = self._tokenize_exact(interior, "keys") if len(ids) > 16: raise ValueError("keys_exceed_sixteen_native_tokens") span = self.field_slices["keys"] target[span.start:span.start + len(ids)] = ids # Strict decoding requires a contiguous PAD tail, so every key slot # is supervised, including all padding after the JSON interior. loss[span] = [True] * 16 spans.append(span); kinds.append("keys") canonical["keys"] = list(keys) elif action in {"scroll", "hscroll"}: pixels = target_action.get("pixels") if type(pixels) is not int or not -999 <= pixels <= 999: raise ValueError("pixels_must_be_integer_minus999_to_999") span = self.field_slices["pixels"] target[span.start] = self.sign_ids["-" if pixels < 0 else "+"] target[span.start + 1:span.stop] = self._digits(abs(pixels), 3) loss[span] = [True] * 4 spans.append(span); kinds.append("pixels") canonical["pixels"] = pixels elif action == "wait": hundredths = _half_up(target_action.get("time"), 100, 99999, "time_must_be_0_to_999_99") select_digits("time", hundredths, 5) canonical["time"] = hundredths / 100 elif action == "terminate": status = target_action.get("status") if status not in self.status_ids: raise ValueError("status_must_be_success_or_failure") span = self.field_slices["status"] target[span.start] = self.status_ids[status]; loss[span.start] = True spans.append(span); kinds.append("status"); canonical["status"] = status elif action in TEXT_ACTIONS: text = target_action.get("text") if not isinstance(text, str): raise TypeError("text_action_requires_exact_text") ids = self._tokenize_exact(text, "text") if len(ids) > 99: raise ValueError("text_exceeds_ninety_nine_native_tokens") span = self.field_slices["text"] values = ids + [self.eos_id] target[span.start:span.start + len(values)] = values loss[span.start:span.start + len(values)] = [True] * len(values) spans.append(slice(span.start, span.start + len(values))); kinds.append("text") canonical["text"] = text profile = { "scroll_coordinate_profile": "optional_xy_all_pad_absent_v2" if self.scroll_coordinates else "legacy_ignored_v1", "coordinate_profile": "native_integer_0_to_1000_half_up_no_clipping", "keys_profile": "compact_json_array_interior_native_tokens_pad_terminated", "pixels_profile": "signed_integer_minus999_to_999_no_quantization", "time_profile": "five_digits_implicit_hundredths_half_up", "text_profile": "exact_native_tokens_eos_inclusive_max99_content", } return AdaptedTarget(self.variant, canonical, tuple(target), tuple(spans), tuple(kinds), tuple(loss), text if action in TEXT_ACTIONS else None, profile) def adapt_record(self, record: Mapping[str, object]): if not isinstance(record, Mapping) or not isinstance(record.get("target"), Mapping): raise TypeError("record_requires_target_action") metadata = record.get("provenance", {}) if not isinstance(metadata, Mapping): raise TypeError("record_provenance_must_be_mapping") return self.adapt_target(record["target"], metadata) def decode_argmax(self, raw_ids: Sequence[int]): if isinstance(raw_ids, (str, bytes)) or len(raw_ids) != self.output_slots: raise ValueError("wrong_argmax_width") ids = list(raw_ids) if any(type(value) is not int or not 0 <= value < self.vocab_size for value in ids): raise ValueError("argmax_ids_must_be_in_vocabulary_integers") try: action = self.action_inverse[tuple(ids[0:3])] except KeyError as exc: raise ValueError("invalid_action_encoding") from exc result: dict[str, object] = {"action": action} scroll = action in {"scroll", "hscroll"} if action in POINT_ACTIONS or (scroll and self.scroll_coordinates and ids[3:11] != [self.pad_id] * 8): x = self._decode_digits(ids[3:7]); y = self._decode_digits(ids[7:11]) if x > 1000 or y > 1000: raise ValueError("coordinate_out_of_range") result["coordinate"] = [x, y] if action == "key": values = ids[11:27] try: end = values.index(self.pad_id) except ValueError: end = 16 if any(value != self.pad_id for value in values[end:]): raise ValueError("keys_padding_must_be_contiguous") text = self.tokenizer.decode(values[:end], skip_special_tokens=False) try: keys = json.loads(f"[{text}]") except json.JSONDecodeError as exc: raise ValueError("invalid_keys_json") from exc if not isinstance(keys, list) or not keys or any(not isinstance(key, str) for key in keys): raise ValueError("invalid_keys_json") if self._tokenize_exact(json.dumps(keys, ensure_ascii=False, separators=(",", ":"))[1:-1], "keys") != values[:end]: raise ValueError("noncanonical_keys_encoding") result["keys"] = keys elif action in {"scroll", "hscroll"}: try: sign = self.inverse_signs[ids[27]] except KeyError as exc: raise ValueError("invalid_pixel_sign") from exc magnitude = self._decode_digits(ids[28:31]) result["pixels"] = -magnitude if sign == "-" else magnitude elif action == "wait": result["time"] = self._decode_digits(ids[31:36]) / 100 elif action == "terminate": try: result["status"] = self.inverse_status[ids[36]] except KeyError as exc: raise ValueError("invalid_status") from exc elif action in TEXT_ACTIONS: values = ids[37:137] try: end = values.index(self.eos_id) except ValueError as exc: raise ValueError("text_requires_eos") from exc text = self.tokenizer.decode(values[:end], skip_special_tokens=False) if self._tokenize_exact(text, "text") != values[:end]: raise ValueError("text_does_not_roundtrip_exactly") result["text"] = text return result