From 4ebf89e1e5ee29e76dd52d29db2cd92aaa4eb0c5 Mon Sep 17 00:00:00 2001 From: Tianyao Wu Date: Thu, 8 Oct 2026 11:34:50 +0800 Subject: [PATCH 1/4] omnijev: add the request contract, preparation, heads and export First step of OmniJev-4B v1.1 (#114), on the CPU: the request contract for one image and Choice, Noul or Score questions, prompt rendering and tokenization, the reference's branch layout and readout positions, Qwen3.5's rotary positions for the image, the trained heads in FP32, the six LM features and the reference's response finishing. A CPU-only exporter merges the language LoRA in BF16, restores the option tokens' embeddings and writes the vision tower, heads and calibration. Fixtures come from the pinned reference's own code without its backbone; the tokenizer and heads checks are opt-in. --- Cargo.lock | 21 + Cargo.toml | 2 +- mkdocs.yml | 1 + recipe/README.md | 2 + recipe/omnijev/README.md | 70 +++ recipe/omnijev/export.py | 120 +++++ src/models/omnijev/README.md | 141 ++++++ src/models/omnijev/native/Cargo.toml | 36 ++ src/models/omnijev/native/LICENSE.omnijev | 208 ++++++++ src/models/omnijev/native/src/contract.rs | 496 ++++++++++++++++++++ src/models/omnijev/native/src/heads.rs | 269 +++++++++++ src/models/omnijev/native/src/lib.rs | 5 + src/models/omnijev/native/src/processing.rs | 307 ++++++++++++ tests/omnijev/common.rs | 56 +++ tests/omnijev/contract.rs | 138 ++++++ tests/omnijev/data/features.json | 1 + tests/omnijev/data/finish.json | 1 + tests/omnijev/data/heads.json | 1 + tests/omnijev/data/processing.json | 1 + tests/omnijev/finish.rs | 105 +++++ tests/omnijev/generate_fixtures.py | 333 +++++++++++++ tests/omnijev/heads.rs | 116 +++++ tests/omnijev/processing.rs | 117 +++++ 23 files changed, 2546 insertions(+), 1 deletion(-) create mode 100644 recipe/omnijev/README.md create mode 100644 recipe/omnijev/export.py create mode 100644 src/models/omnijev/README.md create mode 100644 src/models/omnijev/native/Cargo.toml create mode 100644 src/models/omnijev/native/LICENSE.omnijev create mode 100644 src/models/omnijev/native/src/contract.rs create mode 100644 src/models/omnijev/native/src/heads.rs create mode 100644 src/models/omnijev/native/src/lib.rs create mode 100644 src/models/omnijev/native/src/processing.rs create mode 100644 tests/omnijev/common.rs create mode 100644 tests/omnijev/contract.rs create mode 100644 tests/omnijev/data/features.json create mode 100644 tests/omnijev/data/finish.json create mode 100644 tests/omnijev/data/heads.json create mode 100644 tests/omnijev/data/processing.json create mode 100644 tests/omnijev/finish.rs create mode 100644 tests/omnijev/generate_fixtures.py create mode 100644 tests/omnijev/heads.rs create mode 100644 tests/omnijev/processing.rs diff --git a/Cargo.lock b/Cargo.lock index c43ff226..8f85ea60 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -971,6 +971,12 @@ dependencies = [ "windows-link", ] +[[package]] +name = "libm" +version = "0.2.16" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6d2cec3eae94f9f509c767b45932f1ada8350c4bdb85af2fcab4a3c14807981" + [[package]] name = "linux-raw-sys" version = "0.12.1" @@ -1213,6 +1219,21 @@ dependencies = [ "tokio", ] +[[package]] +name = "omni-omnijev-native" +version = "0.1.0" +dependencies = [ + "anyhow", + "base64 0.22.1", + "image", + "libm", + "omni-qwen3-5-native", + "safetensors 0.8.0", + "serde_json", + "sha2 0.10.9", + "tokenizers 0.22.2", +] + [[package]] name = "omni-open-jev-native" version = "0.1.0" diff --git a/Cargo.toml b/Cargo.toml index dc0bef34..ea6e1654 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend", "src/runtime", "src/models/clm", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/jev_vl/native", "src/models/laya", "src/backends/cuda"] +members = ["src/frontend", "src/runtime", "src/models/clm", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/native", "src/models/jev_vl/native", "src/models/omnijev/native", "src/models/laya", "src/backends/cuda"] resolver = "3" diff --git a/mkdocs.yml b/mkdocs.yml index ef244fd2..e364a634 100644 --- a/mkdocs.yml +++ b/mkdocs.yml @@ -70,5 +70,6 @@ nav: - Open-Jev-9B validation: recipe/open_jev/validation-9b.md - JEV-27B-VL experimental worker: recipe/jev_vl/README.md - JEV-27B-VL validation status: recipe/jev_vl/validation.md + - OmniJev preparation: recipe/omnijev/README.md - CLM stub-encoder contract: recipe/clm/README.md - Contributing: CONTRIBUTING.md diff --git a/recipe/README.md b/recipe/README.md index bffec66d..e3ca0819 100644 --- a/recipe/README.md +++ b/recipe/README.md @@ -19,6 +19,8 @@ For a first real decision, follow the [complete CPU walkthrough](../docs/getting compares the 9B worker with the reference. - [JEV-27B-VL experimental worker](jev_vl/README.md): native text and preencoded image decisions; bounded H800 validation with explicit deployment limits. +- [OmniJev-4B v1.1 preparation](omnijev/README.md): export the checkpoint and + regenerate the CPU fixtures; the native worker is not available yet. - [CLM behind the frontend](clm/README.md): run CLM's own server behind the frontend on CPU with a stub encoder, and what the response comparison has to allow for. diff --git a/recipe/omnijev/README.md b/recipe/omnijev/README.md new file mode 100644 index 00000000..39b19400 --- /dev/null +++ b/recipe/omnijev/README.md @@ -0,0 +1,70 @@ +# OmniJev-4B v1.1 preparation + +This recipe prepares OmniJev-4B v1.1 ([#114](https://github.com/ThinkFlowLab/system1-omni/issues/114)) +for the native worker: it exports the checkpoint and regenerates the CPU fixtures +that the contract and preparation tests use. The native worker itself is not +available yet; [the model README](../../src/models/omnijev/README.md) describes what +is implemented. Run the commands from the repository root. + +## Reference environment + +The export and the fixtures use the reference's own code and Python packages on the +CPU; no GPU is needed. + +```bash +python3.10 -m venv .venv-omnijev +.venv-omnijev/bin/pip install torch==2.14.0 torchvision==0.29.0 transformers==5.17.0 peft==0.19.1 accelerate pillow +git clone https://github.com/tinnel123666888/OmniJev +git -C OmniJev checkout 14dbec4f71e194852c8d7b88ab36ef639493f400 +``` + +Leave `fla-core` uninstalled, so transformers runs its PyTorch Gated DeltaNet code. +Keep transformers at 5.7.0 or later: older releases do not continue the +linear-attention state across the reference's prefix. The fixtures were generated +with Python 3.10.12; Python 3.12 changed how `sum` adds floats, which the reference's +calibration uses. + +## Download the pinned files + +```bash +hf download tinnel123/OmniJev --revision ffe5f436eaf22e20e2f041f8e74e121fd057a6cb --local-dir omnijev-v1.1 +hf download Qwen/Qwen3.5-4B --revision 851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a --local-dir qwen3.5-4b +``` + +## Export + +```bash +.venv-omnijev/bin/python recipe/omnijev/export.py --base qwen3.5-4b --checkpoint omnijev-v1.1 --out omnijev-export +``` + +The export checks the checkpoint files it reads against the release's +`release_manifest.json` and the base shards against their pinned SHA-256, then +writes: + +- the language model in BF16 with the adapter merged, its tied embedding resized to + the tokenizer's 248,079 entries and the two option tokens' saved rows restored, and + `config.json` with the vision configuration; +- `vision.safetensors`, the base vision tower, which the adapter does not change; +- `heads.safetensors`, the decision and ordinal heads in FP32; +- the tokenizer and chat template; +- `omnijev_export.json`, written last, with the revisions, calibration, image budget, + and the SHA-256 of the checked inputs and of every output file. + +It needs about 18 GB of host RAM and writes 9 GB; on a 48-core workstation it takes +under a minute. + +## Fixtures and tests + +```bash +.venv-omnijev/bin/python tests/omnijev/generate_fixtures.py \ + --reference OmniJev --checkpoint omnijev-v1.1 --base qwen3.5-4b --out tests/omnijev/data +cargo test --locked -p omni-omnijev-native +OMNIJEV_CHECKPOINT=$PWD/omnijev-v1.1 OMNIJEV_EXPORT=$PWD/omnijev-export \ + cargo test --locked -p omni-omnijev-native -- --include-ignored +``` + +The generator runs the reference's preparation, layout, rotary positions, heads and +finishing without loading the backbone. The ordinary tests need only the committed +fixtures; the opt-in ones compare token ids with the checkpoint's tokenizer and the +heads with the export. Cargo runs the tests from the crate's directory, so the two +paths must be absolute. diff --git a/recipe/omnijev/export.py b/recipe/omnijev/export.py new file mode 100644 index 00000000..890e1eb6 --- /dev/null +++ b/recipe/omnijev/export.py @@ -0,0 +1,120 @@ +"""Export OmniJev-4B v1.1 for the native worker on CPU. + +Merges the language LoRA into the pinned Qwen3.5-4B in BF16, restores the option +tokens' embeddings, and writes the base vision tensors, the FP32 heads, the tokenizer +and the calibration next to it. Use the reference environment in README.md; no GPU +is needed. Needs about 18 GB of host RAM and writes 9 GB. +""" + +import argparse +import hashlib +import json +from pathlib import Path + +import torch +from peft import PeftModel +from safetensors.torch import save_file +from transformers import AutoModelForImageTextToText, AutoTokenizer + +MODEL_ID = "tinnel123/OmniJev" +CHECKPOINT_REVISION = "ffe5f436eaf22e20e2f041f8e74e121fd057a6cb" +BASE_MODEL_ID = "Qwen/Qwen3.5-4B" +BASE_REVISION = "851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a" +BASE_SHARDS = { + "model.safetensors-00001-of-00002.safetensors": "26a93f066e1916adb13453dae5a0c707c0fbc71299ed98779571a907b8e74c61", + "model.safetensors-00002-of-00002.safetensors": "cb544bd9bfae93dc59b0f22b292f5933573854a7f9b97835c67060d7d910e188", +} +OPTION_TOKENS = {"<|opt|>": 248077, "<|/opt|>": 248078} +VOCABULARY = 248079 # the tokenizer's length; the reference resizes the tied embedding to it +MAX_PIXELS = 768 * 28 * 28 # MSO1's default, which overrides the processor's 401,408 +NEEDED = ["adapter_config.json", "adapter_model.safetensors", "chat_template.jinja", "head.pt", "head_meta.json", + "new_tok_emb.pt", "ord.pt", "processor_config.json", "tokenizer.json", "tokenizer_config.json"] + + +def sha256(path): + h = hashlib.sha256() + with open(path, "rb") as f: + for block in iter(lambda: f.read(1 << 24), b""): + h.update(block) + return h.hexdigest() + + +def main(): + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("--base", required=True, type=Path, help=f"{BASE_MODEL_ID} @ {BASE_REVISION}") + parser.add_argument("--checkpoint", required=True, type=Path, help=f"{MODEL_ID} @ {CHECKPOINT_REVISION}") + parser.add_argument("--out", required=True, type=Path) + args = parser.parse_args() + if args.out.exists(): + raise ValueError("output already exists; choose a new export directory") + + release = json.loads((args.checkpoint / "release_manifest.json").read_text()) + if release.get("release") != "v1.1" or release.get("base_model") != BASE_MODEL_ID: + raise ValueError("expected the OmniJev v1.1 release manifest") + inputs = {} + for name in NEEDED: + digest = sha256(args.checkpoint / name) + if digest != release["files"][name]["sha256"]: + raise ValueError(f"{name} does not match the release manifest") + inputs[f"checkpoint/{name}"] = digest + for name, expected in BASE_SHARDS.items(): + if sha256(args.base / name) != expected: + raise ValueError(f"{name} does not match {BASE_MODEL_ID} @ {BASE_REVISION}") + inputs[f"base/{name}"] = expected + meta = json.loads((args.checkpoint / "head_meta.json").read_text()) + if not (meta["norm"] == "softmax" and meta["lm_feats"] and meta["ordinal"] and meta["branch"]): + raise ValueError("expected softmax heads with LM features, the ordinal head and branching") + extra = set(meta.get("biases", {})) - {"noul"} + if extra: + raise ValueError(f"unsupported calibration biases: {sorted(extra)}") + + tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, local_files_only=True) + if len(tokenizer) != VOCABULARY or any(tokenizer.convert_tokens_to_ids(t) != i for t, i in OPTION_TOKENS.items()): + raise ValueError("unexpected tokenizer") + model = AutoModelForImageTextToText.from_pretrained( + args.base, dtype=torch.bfloat16, device_map={"": "cpu"}, local_files_only=True) + # mso.records.add_option_tokens: resize the tied embedding to the tokenizer and + # restore the saved option-token rows. + model.resize_token_embeddings(VOCABULARY) + saved = torch.load(args.checkpoint / "new_tok_emb.pt", map_location="cpu", weights_only=True) + if saved.shape != (2, model.config.text_config.hidden_size): + raise ValueError("unexpected option-token embeddings") + with torch.no_grad(): + embedding = model.get_input_embeddings().weight + for row, token in zip(saved, OPTION_TOKENS.values()): + embedding[token] = row.to(embedding.dtype) + if model.get_output_embeddings().weight.data_ptr() != embedding.data_ptr(): + raise ValueError("expected tied input and output embeddings") + + model = PeftModel.from_pretrained(model, args.checkpoint).merge_and_unload(safe_merge=True) + model.model.language_model.save_pretrained(args.out, max_shard_size="5GB") + model.config.to_json_file(args.out / "config.json") + vision = {f"model.visual.{k}": v.contiguous() for k, v in model.model.visual.state_dict().items()} + save_file(vision, args.out / "vision.safetensors") + + heads = {} + for prefix, name in (("head", "head.pt"), ("ord", "ord.pt")): + state = torch.load(args.checkpoint / name, map_location="cpu", weights_only=True) + for key, value in state.items(): + if not torch.isfinite(value).all(): + raise ValueError(f"non-finite {prefix}.{key}") + heads[f"{prefix}.{key}"] = value.float().contiguous() + save_file(heads, args.out / "heads.safetensors") + tokenizer.save_pretrained(args.out) + + outputs = {p.name: sha256(p) for p in sorted(args.out.iterdir()) if p.is_file()} + # Written last: the worker refuses incomplete exports. + (args.out / "omnijev_export.json").write_text(json.dumps({ + "format": "omnijev-merged/1", + "model_id": MODEL_ID, "checkpoint_revision": CHECKPOINT_REVISION, + "base_model_id": BASE_MODEL_ID, "base_revision": BASE_REVISION, + "lora": {"rank": 32, "alpha": 64, "merged": "bfloat16"}, + "vocabulary": VOCABULARY, "option_tokens": OPTION_TOKENS, "max_pixels": MAX_PIXELS, + "heads": {"norm": meta["norm"], "lm_features": True, "ordinal": True}, + "calibration": {"temperatures": meta["temperatures"], "noul_bias": meta["biases"]["noul"]}, + "inputs": inputs, "outputs": outputs, + }, indent=1, allow_nan=False) + "\n") + + +if __name__ == "__main__": + main() diff --git a/src/models/omnijev/README.md b/src/models/omnijev/README.md new file mode 100644 index 00000000..05cd376b --- /dev/null +++ b/src/models/omnijev/README.md @@ -0,0 +1,141 @@ +# OmniJev + +OmniJev-4B v1.1 ([#114](https://github.com/ThinkFlowLab/system1-omni/issues/114)) +answers Choice, Noul and Score questions about one image with a Qwen3.5-4B backbone, +a merged LoRA adapter, two option-marker tokens and trained FP32 heads. This directory +holds the native implementation's CPU side so far: the request contract, preparation +and the branch layout, the heads, and response finishing, each checked against the +pinned reference. Vision and language execution, the worker and frontend integration +come next. The [recipe](../../../recipe/omnijev/README.md) exports the checkpoint and +regenerates the fixtures. + +## Pinned artifacts + +| Artifact | Revision | +| --- | --- | +| Reference code: [`tinnel123666888/OmniJev`](https://github.com/tinnel123666888/OmniJev/tree/14dbec4f71e194852c8d7b88ab36ef639493f400) | `14dbec4f71e194852c8d7b88ab36ef639493f400` | +| Adapter, heads, tokenizer and processor: [`tinnel123/OmniJev`](https://huggingface.co/tinnel123/OmniJev/tree/ffe5f436eaf22e20e2f041f8e74e121fd057a6cb), tag `v1.1` | `ffe5f436eaf22e20e2f041f8e74e121fd057a6cb` | +| Base model: [`Qwen/Qwen3.5-4B`](https://huggingface.co/Qwen/Qwen3.5-4B/tree/851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a) | `851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a` | + +The adapter names its base without a revision; `851bf6e` is the one Cua-S1 pins. The +reference environment is transformers 5.17.0, torch 2.14.0, torchvision 0.29.0 and +PEFT 0.19.1, without `fla-core`, so transformers runs its PyTorch Gated DeltaNet code. +The reference's branch path needs transformers 5.7.0 or later: 5.5 and 5.6 start the +linear-attention state of every row from zero, and 5.2 to 5.4 fail in it. + +## Request and response + +```json +{"model": "tinnel123/OmniJev", + "state": {"images": ["data:image/png;base64,..."]}, + "questions": { + "visible": {"type": "noul", "instructions": "A person is visible."}, + "next": {"type": "choice", "instructions": "Which operation comes next?", + "criteria": {"click": null, "type": "Enter text into the focused field"}}, + "risk": {"type": "score", "instructions": "How risky is acting on this screen?", + "levels": ["safe", "check first", "dangerous"]}}} +``` + +- `model` may be omitted, `tinnel123/OmniJev`, `omnijev` or `omnijev-4b-v1.1`. +- `state.images` holds exactly one inline PNG or JPEG. +- Noul takes `instructions` and an optional `region` (`{"box": [x1, y1, x2, y2]}`). +- Choice takes Jev's `criteria` map (a key alone, or `key: rubric`), or the reference's + `options` list, whose entries carry `key`, `text`, `region` or `abstain`. Abstain + entries are skipped, since the head has its own abstain probability; a missing key + falls back to the text, then to `option_`. +- Score takes `levels` in order, or the keys of `criteria`. + +Option text is cut to 200 characters, and region coordinates are rounded half to even, +both as the reference does. Answers keep the reference's fields: `noul`; `choice`, +`probabilities`, `abstain`, `valid` and `confidence`; or `score`, `probabilities` and +`confidence`; each with `latency_s`, and with `latency_total_s` once the worker adds +it, both from the native timing. Probabilities are rounded to four decimals. The +layout also gives the reference's input-token count, the shared prefix once plus every +row, for the worker's `usage`. + +Deliberate differences from the reference: + +- One image per request. The reference tiles several still images into a numbered + panel and reads a video as a frame mosaic; both come later. +- Text containing `<|opt|>`, `<|/opt|>`, `<|image_pad|>`, `<|video_pad|>`, + `<|vision_start|>` or `<|vision_end|>` is refused. The reference would read it as + option or image markers and misplace its readouts. +- Images arrive as data URLs instead of file paths. +- Explicit limits: a 12 MiB body; an 8 MiB image of at most 4096 pixels per side and + an aspect ratio of at most 200, which the processor also requires; 64 questions with + ids of at most 256 characters; 1 to 255 options per question and 1,024 in all; + 8,192 characters per text field; 8,192 tokens per question; and 65,536 processed + tokens per request. +- Fields the reference ignores (Noul `criteria`, `criteria` next to Choice `options` or + Score `levels`, unknown fields), and values it would turn into text (numbers as + levels, for example), are refused. + +## Preparation and layout + +[`processing.rs`](native/src/processing.rs) renders each question as the reference's +chat template does: + +```text +<|im_start|>user +<|vision_start|><|image_pad|><|vision_end|>{instructions}<|im_end|> +<|im_start|>assistant + +<|opt|>{option 1}<|/opt|><|opt|>{option 2}<|/opt|>... +``` + +The template trims the user content, so trailing whitespace in the instructions goes, +by Python's definition of whitespace. The image placeholder expands to one token per +2×2 merged patch of the processor's resize, with `MSO1`'s 602,112-pixel budget (it +overrides the processor config's 401,408) and the 65,536-pixel minimum. + +The layout follows `MSO1.ask_branch`. The shared prefix is the questions' common token +prefix, cut at least one token before any question's first `<|opt|>`. Each row is the +rest of its question plus one option block. A row's readouts are `u` at its closing +marker, `zq` at the token before its option marker, and, for the LM features, each +option-text token with the position that predicts it. Rotary positions are Qwen3.5's +T/H/W positions of the single-question sequence; after the prefix, rows continue from +its largest position plus one. Since every row equals a plain forward over prefix plus +row up to rounding, execution may choose its own split points. + +## Heads and finishing + +[`heads.rs`](native/src/heads.rs) runs the trained heads in FP32 on the CPU. Choice +and Noul use `OptionScorer`: gated option and question MLPs, the LM features, a +per-type temperature, and for Choice a softmax that includes the abstain logit; Noul +is a sigmoid. Score uses only the cumulative-link ordinal head. The six LM +features of an option come from its text tokens' log-probabilities over the full +248,079-entry vocabulary, which execution computes. A plain Noul option has no text, +so its features are zero. + +`contract::answer` follows `MSO1._finish`: Noul's log-odds bias before its +temperature, the saved temperatures for each type, Choice's abstain and validity, +Score's renormalization, the float32 values the reference stores in between, Jev's +confidence, and four-decimal rounding. The float32 sum of more than four head outputs +can add in a different order than PyTorch's, which moves the fourth decimal by one at +most; the reference's own sum on a GPU has the same freedom. + +[`processing::image_grid`](native/src/processing.rs) and `processing::image_positions` +follow Transformers 5.17.0's `Qwen2VLImageProcessor` resize and Qwen3.5 rotary +positions (Apache-2.0, Copyright 2025 The Qwen team, Alibaba Group and the HuggingFace +Inc. team), as Cua-S1's copies do. Both move to the shared Qwen crate once its image +preprocessing is shared. + +## Validation + +[`generate_fixtures.py`](../../../tests/omnijev/generate_fixtures.py) runs the +reference's own preparation, layout, rotary positions, heads and finishing on the CPU +without loading the backbone. The tests in [`tests/omnijev/`](../../../tests/omnijev/) +compare against those fixtures: prompts, answer keys and option text, marker +positions, image grids, rotary positions, the prefix, rows, readouts and token counts +for seven requests, exactly; the LM features within 1e-6; and finished answers for +eighteen head outputs, exactly up to four options and within one unit of the fourth +decimal beyond. Two checks are opt-in: token ids against the checkpoint's tokenizer, +exactly, and the heads against an export, within 1e-5 (relative for values above one). Contract tests cover +refused requests, with the reason, and the limits. + +No GPU work runs in this crate, and nothing here establishes parity of the model's +outputs; that comes with execution. + +The Rust code adapts OmniJev's Apache-2.0 code; its copyright and license are retained +in [`native/LICENSE.omnijev`](native/LICENSE.omnijev). No model weights are +distributed here. diff --git a/src/models/omnijev/native/Cargo.toml b/src/models/omnijev/native/Cargo.toml new file mode 100644 index 00000000..2d19a8d2 --- /dev/null +++ b/src/models/omnijev/native/Cargo.toml @@ -0,0 +1,36 @@ +[package] +name = "omni-omnijev-native" +version = "0.1.0" +edition = "2024" +publish = false +description = "OmniJev-4B v1.1 request contract, preparation, heads and finishing" + +[dependencies] +anyhow = "1.0.100" +base64 = "0.22.1" +image = { version = "0.25", default-features = false, features = ["png", "jpeg"] } +libm = "0.2" +omni-qwen3-5-native = { path = "../../qwen3_5/native" } +safetensors = "0.8.0" +serde_json = { version = "1.0.149", features = ["float_roundtrip", "preserve_order"] } +# the onig regex backend, as in the Python tokenizers wheel +tokenizers = { version = "=0.22.2", default-features = false, features = ["onig"] } + +[[test]] +name = "contract" +path = "../../../../tests/omnijev/contract.rs" + +[[test]] +name = "processing" +path = "../../../../tests/omnijev/processing.rs" + +[[test]] +name = "finish" +path = "../../../../tests/omnijev/finish.rs" + +[[test]] +name = "heads" +path = "../../../../tests/omnijev/heads.rs" + +[dev-dependencies] +sha2 = "0.10" diff --git a/src/models/omnijev/native/LICENSE.omnijev b/src/models/omnijev/native/LICENSE.omnijev new file mode 100644 index 00000000..f5bc30ca --- /dev/null +++ b/src/models/omnijev/native/LICENSE.omnijev @@ -0,0 +1,208 @@ +OmniJev - inference code and model weights +Copyright 2026 Beijing Zhongguancun Academy, Institute of Automation of the Chinese Academy of Sciences, and Zevo +(北京中关村学院、中国科学院自动化研究所、智进化) + +Licensed under the Apache License, Version 2.0 (full text below). The backbone model (Qwen3-VL-4B-Instruct) keeps its own license. + + + Apache License + Version 2.0, January 2004 + http://www.apache.org/licenses/ + + TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION + + 1. Definitions. + + "License" shall mean the terms and conditions for use, reproduction, + and distribution as defined by Sections 1 through 9 of this document. + + "Licensor" shall mean the copyright owner or entity authorized by + the copyright owner that is granting the License. + + "Legal Entity" shall mean the union of the acting entity and all + other entities that control, are controlled by, or are under common + control with that entity. For the purposes of this definition, + "control" means (i) the power, direct or indirect, to cause the + direction or management of such entity, whether by contract or + otherwise, or (ii) ownership of fifty percent (50%) or more of the + outstanding shares, or (iii) beneficial ownership of such entity. + + "You" (or "Your") shall mean an individual or Legal Entity + exercising permissions granted by this License. + + "Source" form shall mean the preferred form for making modifications, + including but not limited to software source code, documentation + source, and configuration files. + + "Object" form shall mean any form resulting from mechanical + transformation or translation of a Source form, including but + not limited to compiled object code, generated documentation, + and conversions to other media types. + + "Work" shall mean the work of authorship, whether in Source or + Object form, made available under the License, as indicated by a + copyright notice that is included in or attached to the work + (an example is provided in the Appendix below). + + "Derivative Works" shall mean any work, whether in Source or Object + form, that is based on (or derived from) the Work and for which the + editorial revisions, annotations, elaborations, or other modifications + represent, as a whole, an original work of authorship. For the purposes + of this License, Derivative Works shall not include works that remain + separable from, or merely link (or bind by name) to the interfaces of, + the Work and Derivative Works thereof. + + "Contribution" shall mean any work of authorship, including + the original version of the Work and any modifications or additions + to that Work or Derivative Works thereof, that is intentionally + submitted to Licensor for inclusion in the Work by the copyright owner + or by an individual or Legal Entity authorized to submit on behalf of + the copyright owner. For the purposes of this definition, "submitted" + means any form of electronic, verbal, or written communication sent + to the Licensor or its representatives, including but not limited to + communication on electronic mailing lists, source code control systems, + and issue tracking systems that are managed by, or on behalf of, the + Licensor for the purpose of discussing and improving the Work, but + excluding communication that is conspicuously marked or otherwise + designated in writing by the copyright owner as "Not a Contribution." + + "Contributor" shall mean Licensor and any individual or Legal Entity + on behalf of whom a Contribution has been received by Licensor and + subsequently incorporated within the Work. + + 2. Grant of Copyright License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + copyright license to reproduce, prepare Derivative Works of, + publicly display, publicly perform, sublicense, and distribute the + Work and such Derivative Works in Source or Object form. + + 3. Grant of Patent License. Subject to the terms and conditions of + this License, each Contributor hereby grants to You a perpetual, + worldwide, non-exclusive, no-charge, royalty-free, irrevocable + (except as stated in this section) patent license to make, have made, + use, offer to sell, sell, import, and otherwise transfer the Work, + where such license applies only to those patent claims licensable + by such Contributor that are necessarily infringed by their + Contribution(s) alone or by combination of their Contribution(s) + with the Work to which such Contribution(s) was submitted. If You + institute patent litigation against any entity (including a + cross-claim or counterclaim in a lawsuit) alleging that the Work + or a Contribution incorporated within the Work constitutes direct + or contributory patent infringement, then any patent licenses + granted to You under this License for that Work shall terminate + as of the date such litigation is filed. + + 4. Redistribution. You may reproduce and distribute copies of the + Work or Derivative Works thereof in any medium, with or without + modifications, and in Source or Object form, provided that You + meet the following conditions: + + (a) You must give any other recipients of the Work or + Derivative Works a copy of this License; and + + (b) You must cause any modified files to carry prominent notices + stating that You changed the files; and + + (c) You must retain, in the Source form of any Derivative Works + that You distribute, all copyright, patent, trademark, and + attribution notices from the Source form of the Work, + excluding those notices that do not pertain to any part of + the Derivative Works; and + + (d) If the Work includes a "NOTICE" text file as part of its + distribution, then any Derivative Works that You distribute must + include a readable copy of the attribution notices contained + within such NOTICE file, excluding those notices that do not + pertain to any part of the Derivative Works, in at least one + of the following places: within a NOTICE text file distributed + as part of the Derivative Works; within the Source form or + documentation, if provided along with the Derivative Works; or, + within a display generated by the Derivative Works, if and + wherever such third-party notices normally appear. The contents + of the NOTICE file are for informational purposes only and + do not modify the License. You may add Your own attribution + notices within Derivative Works that You distribute, alongside + or as an addendum to the NOTICE text from the Work, provided + that such additional attribution notices cannot be construed + as modifying the License. + + You may add Your own copyright statement to Your modifications and + may provide additional or different license terms and conditions + for use, reproduction, or distribution of Your modifications, or + for any such Derivative Works as a whole, provided Your use, + reproduction, and distribution of the Work otherwise complies with + the conditions stated in this License. + + 5. Submission of Contributions. Unless You explicitly state otherwise, + any Contribution intentionally submitted for inclusion in the Work + by You to the Licensor shall be under the terms and conditions of + this License, without any additional terms or conditions. + Notwithstanding the above, nothing herein shall supersede or modify + the terms of any separate license agreement you may have executed + with Licensor regarding such Contributions. + + 6. Trademarks. This License does not grant permission to use the trade + names, trademarks, service marks, or product names of the Licensor, + except as required for reasonable and customary use in describing the + origin of the Work and reproducing the content of the NOTICE file. + + 7. Disclaimer of Warranty. Unless required by applicable law or + agreed to in writing, Licensor provides the Work (and each + Contributor provides its Contributions) on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or + implied, including, without limitation, any warranties or conditions + of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A + PARTICULAR PURPOSE. You are solely responsible for determining the + appropriateness of using or redistributing the Work and assume any + risks associated with Your exercise of permissions under this License. + + 8. Limitation of Liability. In no event and under no legal theory, + whether in tort (including negligence), contract, or otherwise, + unless required by applicable law (such as deliberate and grossly + negligent acts) or agreed to in writing, shall any Contributor be + liable to You for damages, including any direct, indirect, special, + incidental, or consequential damages of any character arising as a + result of this License or out of the use or inability to use the + Work (including but not limited to damages for loss of goodwill, + work stoppage, computer failure or malfunction, or any and all + other commercial damages or losses), even if such Contributor + has been advised of the possibility of such damages. + + 9. Accepting Warranty or Additional Liability. While redistributing + the Work or Derivative Works thereof, You may choose to offer, + and charge a fee for, acceptance of support, warranty, indemnity, + or other liability obligations and/or rights consistent with this + License. However, in accepting such obligations, You may act only + on Your own behalf and on Your sole responsibility, not on behalf + of any other Contributor, and only if You agree to indemnify, + defend, and hold each Contributor harmless for any liability + incurred by, or claims asserted against, such Contributor by reason + of your accepting any such warranty or additional liability. + + END OF TERMS AND CONDITIONS + + APPENDIX: How to apply the Apache License to your work. + + To apply the Apache License to your work, attach the following + boilerplate notice, with the fields enclosed by brackets "[]" + replaced with your own identifying information. (Don't include + the brackets!) The text should be enclosed in the appropriate + comment syntax for the file format. We also recommend that a + file or class name and description of purpose be included on the + same "printed page" as the copyright notice for easier + identification within third-party archives. + + Copyright [yyyy] [name of copyright owner] + + Licensed under the Apache License, Version 2.0 (the "License"); + you may not use this file except in compliance with the License. + You may obtain a copy of the License at + + http://www.apache.org/licenses/LICENSE-2.0 + + Unless required by applicable law or agreed to in writing, software + distributed under the License is distributed on an "AS IS" BASIS, + WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied. + See the License for the specific language governing permissions and + limitations under the License. diff --git a/src/models/omnijev/native/src/contract.rs b/src/models/omnijev/native/src/contract.rs new file mode 100644 index 00000000..3e66fecf --- /dev/null +++ b/src/models/omnijev/native/src/contract.rs @@ -0,0 +1,496 @@ +//! OmniJev's typed request contract: one image and Choice, Noul or Score questions. +//! Question shapes and option rendering follow OmniJev @ 14dbec4 (`mso/infer.py` +//! `MSO1._options`, `mso/records.py` `render_option`); see ../LICENSE.omnijev. + +use std::io::Cursor; + +use anyhow::{Context, Result, bail, ensure}; +use base64::Engine; +use omni_qwen3_5_native::json; +use serde_json::{Map, Value}; + +pub const MODEL_ID: &str = "tinnel123/OmniJev"; +const ALIASES: [&str; 2] = ["omnijev", "omnijev-4b-v1.1"]; + +pub const MAX_BODY: usize = 12 * 1024 * 1024; +pub const MAX_IMAGE_BYTES: usize = 8 * 1024 * 1024; +pub const MAX_IMAGE_SIDE: usize = 4096; +pub const MAX_QUESTIONS: usize = 64; +pub const MAX_OPTIONS: usize = 255; +pub const MAX_ROWS: usize = 1024; +pub const MAX_TEXT_CHARS: usize = 8192; +pub const MAX_ID_CHARS: usize = 256; +/// Rendered option text is cut to this many characters, as the reference does. +pub const OPTION_CHARS: usize = 200; + +/// Text the reference would read as image or option markers, which would move +/// the readout positions; requests containing it are refused. +pub const RESERVED_TOKENS: [&str; 6] = [ + "<|opt|>", + "<|/opt|>", + "<|image_pad|>", + "<|video_pad|>", + "<|vision_start|>", + "<|vision_end|>", +]; + +#[derive(Clone, Copy, Debug, PartialEq, Eq)] +pub enum Kind { + Noul, + Choice, + Score, +} + +impl Kind { + /// The decision head's question-type index. + pub fn type_id(self) -> usize { + match self { + Kind::Noul => 0, + Kind::Choice => 1, + Kind::Score => 2, + } + } +} + +pub struct Question { + pub id: String, + pub kind: Kind, + pub instructions: String, + /// Answer keys, one per option, in caller order. + pub keys: Vec, + /// Each option's text as written between the option markers. + pub options: Vec, +} + +/// The request's image: its encoded bytes and dimensions. Pixel decoding comes +/// with the vision path. +pub struct Image { + pub format: image::ImageFormat, + pub width: usize, + pub height: usize, + pub bytes: Vec, +} + +pub struct Request { + pub image: Image, + pub questions: Vec, +} + +pub fn compile(raw: &[u8]) -> Result { + ensure!(raw.len() <= MAX_BODY, "request body exceeds 12 MiB"); + let request = json::parse(raw).map_err(anyhow::Error::msg)?; + ensure!( + request + .keys() + .all(|k| ["model", "state", "questions"].contains(&k.as_str())), + "request fields must be model, state and questions" + ); + ensure!( + match request.get("model") { + None | Some(Value::Null) => true, + Some(Value::String(model)) => model == MODEL_ID || ALIASES.contains(&model.as_str()), + Some(_) => false, + }, + "requested model is not loaded" + ); + let image = image(request.get("state").context("request requires state")?)?; + let questions = request + .get("questions") + .and_then(Value::as_object) + .context("questions must be an object")?; + ensure!( + (1..=MAX_QUESTIONS).contains(&questions.len()), + "questions must contain 1 to {MAX_QUESTIONS} questions" + ); + let questions = questions + .iter() + .map(|(id, q)| { + ensure!( + id.chars().count() <= MAX_ID_CHARS, + "question ids must be at most {MAX_ID_CHARS} characters" + ); + question(id, q).with_context(|| format!("question {id}")) + }) + .collect::>>()?; + let rows: usize = questions.iter().map(|q| q.options.len()).sum(); + ensure!( + rows <= MAX_ROWS, + "request exceeds {MAX_ROWS} options in all" + ); + Ok(Request { image, questions }) +} + +fn image(state: &Value) -> Result { + let state = state + .as_object() + .context("state must be an object with images")?; + ensure!( + state.keys().all(|k| k == "images"), + "state supports images only" + ); + let images = state + .get("images") + .and_then(Value::as_array) + .context("state.images must be an array")?; + ensure!( + images.len() == 1, + "state.images must contain exactly one image" + ); + let url = images[0].as_str().context("an image must be a data URL")?; + let (prefix, encoded) = url.split_once(',').context("invalid image data URL")?; + let format = match prefix { + "data:image/png;base64" => image::ImageFormat::Png, + "data:image/jpeg;base64" => image::ImageFormat::Jpeg, + _ => bail!("only inline PNG/JPEG images are supported"), + }; + ensure!( + encoded.len() <= 4 * MAX_IMAGE_BYTES.div_ceil(3), + "image exceeds 8 MiB" + ); + let bytes = base64::engine::general_purpose::STANDARD + .decode(encoded) + .context("invalid base64 image")?; + ensure!(bytes.len() <= MAX_IMAGE_BYTES, "image exceeds 8 MiB"); + ensure!( + image::guess_format(&bytes).ok() == Some(format), + "image format does not match its MIME type" + ); + let mut limits = image::Limits::default(); + limits.max_image_width = Some(MAX_IMAGE_SIDE as u32); + limits.max_image_height = Some(MAX_IMAGE_SIDE as u32); + let mut reader = image::ImageReader::with_format(Cursor::new(&bytes), format); + reader.limits(limits); + let (width, height) = reader.into_dimensions().context("invalid image")?; + let (width, height) = (width as usize, height as usize); + ensure!( + width > 0 && height > 0 && width <= MAX_IMAGE_SIDE && height <= MAX_IMAGE_SIDE, + "image must be at most {MAX_IMAGE_SIDE} pixels per side" + ); + // The processor refuses aspect ratios above 200. + ensure!( + width.max(height) <= 200 * width.min(height), + "image aspect ratio must be at most 200" + ); + Ok(Image { + format, + width, + height, + bytes, + }) +} + +fn text<'a>(value: Option<&'a Value>, what: &str) -> Result<&'a str> { + let text = value + .and_then(Value::as_str) + .with_context(|| format!("{what} must be a string"))?; + ensure!( + text.chars().count() <= MAX_TEXT_CHARS, + "{what} exceeds {MAX_TEXT_CHARS} characters" + ); + Ok(text) +} + +/// Python's truthiness for the JSON values a question may carry. +fn truthy(value: Option<&Value>) -> bool { + match value { + None | Some(Value::Null) => false, + Some(Value::Bool(b)) => *b, + Some(Value::Number(n)) => n.as_f64() != Some(0.0), + Some(Value::String(s)) => !s.is_empty(), + Some(Value::Array(a)) => !a.is_empty(), + Some(Value::Object(o)) => !o.is_empty(), + } +} + +/// `region [x1,y1,x2,y2]`, each coordinate rounded as Python's `round` does. +fn region(value: &Value) -> Result { + let corners = value + .as_object() + .filter(|r| r.len() == 1) + .and_then(|r| r.get("box")) + .and_then(Value::as_array) + .filter(|b| b.len() == 4) + .context("a region must be {\"box\": [x1, y1, x2, y2]}")?; + let mut rounded = Vec::with_capacity(4); + for v in corners { + let v = v.as_f64().context("region coordinates must be numbers")?; + ensure!(v.abs() <= 1e9, "region coordinates must be within ±1e9"); + rounded.push(v.round_ties_even() as i64); + } + Ok(format!( + "region [{},{},{},{}]", + rounded[0], rounded[1], rounded[2], rounded[3] + )) +} + +fn cut(text: &str) -> String { + text.chars().take(OPTION_CHARS).collect() +} + +fn question(id: &str, q: &Value) -> Result { + let q = q.as_object().context("a question must be an object")?; + let kind = match q.get("type").and_then(Value::as_str) { + Some("noul") => Kind::Noul, + Some("choice") => Kind::Choice, + Some("score") => Kind::Score, + _ => bail!("question type must be choice, score, or noul"), + }; + let instructions = text(q.get("instructions"), "instructions")?.to_owned(); + let (keys, options) = match kind { + Kind::Noul => { + ensure!( + q.keys() + .all(|k| ["type", "instructions", "region"].contains(&k.as_str())), + "Noul supports type, instructions and region" + ); + let option = if truthy(q.get("region")) { + region(&q["region"])? + } else { + String::new() + }; + (vec!["yes".to_owned()], vec![option]) + } + Kind::Score => { + ensure!( + q.keys() + .all(|k| ["type", "instructions", "levels", "criteria"].contains(&k.as_str())), + "Score supports type, instructions, levels and criteria" + ); + ensure!( + !(q.contains_key("levels") && q.contains_key("criteria")), + "Score takes levels or criteria, not both" + ); + let levels: Vec = if q.contains_key("levels") { + q["levels"] + .as_array() + .context("levels must be an array of strings")? + .iter() + .map(|l| text(Some(l), "a level").map(str::to_owned)) + .collect::>()? + } else { + q.get("criteria") + .and_then(Value::as_object) + .context("Score requires levels, or criteria as an object keyed by level")? + .keys() + .map(|k| { + ensure!( + k.chars().count() <= MAX_TEXT_CHARS, + "a level exceeds {MAX_TEXT_CHARS} characters" + ); + Ok(k.clone()) + }) + .collect::>()? + }; + let options = levels.iter().map(|l| cut(l)).collect(); + (levels, options) + } + Kind::Choice => { + ensure!( + q.keys() + .all(|k| ["type", "instructions", "options", "criteria"].contains(&k.as_str())), + "Choice supports type, instructions, options and criteria" + ); + ensure!( + !(q.contains_key("options") && q.contains_key("criteria")), + "Choice takes options or criteria, not both" + ); + if q.contains_key("options") { + choice_options(&q["options"])? + } else { + choice_criteria(q.get("criteria"))? + } + } + }; + ensure!( + (1..=MAX_OPTIONS).contains(&options.len()), + "a question needs 1 to {MAX_OPTIONS} options" + ); + for (i, key) in keys.iter().enumerate() { + ensure!(!keys[..i].contains(key), "answer key {key:?} is repeated"); + } + for text in std::iter::once(&instructions).chain(&options) { + if let Some(token) = RESERVED_TOKENS.iter().find(|t| text.contains(*t)) { + bail!("text must not contain {token}"); + } + } + Ok(Question { + id: id.to_owned(), + kind, + instructions, + keys, + options, + }) +} + +/// The reference's option list: `key`, `text`, `region` and `abstain` entries. +fn choice_options(list: &Value) -> Result<(Vec, Vec)> { + let list = list.as_array().context("options must be an array")?; + ensure!(list.len() <= 4 * MAX_OPTIONS, "too many options"); + let (mut keys, mut options) = (Vec::new(), Vec::new()); + for (i, o) in list.iter().enumerate() { + let o: &Map = o.as_object().context("an option must be an object")?; + ensure!( + o.keys() + .all(|k| ["key", "text", "region", "abstain"].contains(&k.as_str())), + "an option supports key, text, region and abstain" + ); + ensure!( + matches!(o.get("abstain"), None | Some(Value::Bool(_))), + "abstain must be a boolean" + ); + if truthy(o.get("abstain")) { + continue; // the reference's abstain probability is the head's own + } + for field in ["key", "text"] { + if o.contains_key(field) { + text(o.get(field), field)?; + } + } + let key = [o.get("key"), o.get("text")] + .into_iter() + .find(|v| truthy(*v)) + .and_then(|v| v.and_then(Value::as_str)) + .map_or_else(|| format!("option_{i}"), str::to_owned); + let option = if o.contains_key("region") { + region(&o["region"])? + } else { + cut(o.get("text").and_then(Value::as_str).unwrap_or("")) + }; + keys.push(key); + options.push(option); + } + Ok((keys, options)) +} + +/// Jev's criteria map: a key alone, or `key: rubric`. +fn choice_criteria(criteria: Option<&Value>) -> Result<(Vec, Vec)> { + let criteria = criteria + .and_then(Value::as_object) + .context("Choice requires options or criteria")?; + let mut options = Vec::with_capacity(criteria.len()); + for (key, rubric) in criteria { + let rubric = match rubric { + Value::Null => "", + Value::String(s) => s.as_str(), + _ => bail!("criteria values must be strings or null"), + }; + ensure!( + key.chars().count() + rubric.chars().count() <= MAX_TEXT_CHARS, + "criteria exceed {MAX_TEXT_CHARS} characters" + ); + options.push(if rubric.is_empty() { + cut(key) + } else { + cut(&format!("{key}: {rubric}")) + }); + } + Ok((criteria.keys().cloned().collect(), options)) +} + +/// Post-hoc calibration saved with the heads (`head_meta.json`). +#[derive(Clone, Copy, Debug)] +pub struct Calibration { + /// Noul, Choice and Score temperatures. + pub temperatures: [f64; 3], + /// Noul's log-odds bias, applied before its temperature. + pub noul_bias: f64, +} + +impl Calibration { + /// `MSO1._scale`: `p^(1/T)`, renormalized. + fn scale(&self, kind: Kind, values: &[f64]) -> Vec { + let t = self.temperatures[kind.type_id()]; + if (t - 1.0).abs() < 1e-6 { + return values.to_vec(); + } + let w: Vec = values.iter().map(|&x| x.max(1e-9).powf(1.0 / t)).collect(); + let z: f64 = w.iter().sum(); + w.iter().map(|x| x / z).collect() + } +} + +/// Python's `round(x, 4)`: the nearest four-decimal value, ties to even. +pub fn round4(x: f64) -> f64 { + format!("{x:.4}").parse().expect("formatted float") +} + +/// Jev's confidence `(K p_max - 1) / (K - 1)`, in float32 as the reference computes it. +fn confidence(probs: &[f32]) -> f64 { + let k = probs.len(); + if k < 2 { + return 1.0; + } + let max = probs.iter().copied().fold(f32::NEG_INFINITY, f32::max); + (((k as f32) * max - 1.0) / (k as f32 - 1.0)).clamp(0.0, 1.0) as f64 +} + +/// The first maximum, as Python's `max(range(n), key=...)` picks it. +fn first_max(values: &[f64]) -> usize { + (1..values.len()).fold(0, |m, i| if values[i] > values[m] { i } else { m }) +} + +/// `MSO1._finish`: one question's answer from its head output `mu`. +pub fn answer( + question: &Question, + mu: &[f32], + calibration: &Calibration, + latency: f64, +) -> Result { + ensure!( + !mu.is_empty() && mu.len() == question.keys.len() && mu.iter().all(|v| v.is_finite()), + "invalid head output" + ); + let latency = round4(latency); + if question.kind == Kind::Noul { + let mut p0 = mu[0] as f64; + if calibration.noul_bias != 0.0 { + p0 = p0.clamp(1e-9, 1.0 - 1e-9); + let z = ((p0 / (1.0 - p0)).ln() + calibration.noul_bias).clamp(-40.0, 40.0); + p0 = 1.0 / (1.0 + (-z).exp()); + } + let p = calibration.scale(Kind::Noul, &[p0, 1.0 - p0])[0]; + return Ok(serde_json::json!({"noul": round4(p), "latency_s": latency})); + } + // choice_outputs, in float32 + let sum = mu.iter().fold(0.0f32, |a, &x| a + x); + let valid = sum <= 1.0 + 1e-5; + let (probs, abstain): (Vec, f32) = if question.kind == Kind::Score { + (mu.iter().map(|&x| x / sum.max(1e-6)).collect(), 0.0) + } else if valid { + (mu.to_vec(), (1.0 - sum).max(0.0)) + } else { + (mu.iter().map(|&x| x / sum).collect(), 0.0) + }; + let mut full: Vec = probs.iter().map(|&p| p as f64).collect(); + full.push(abstain as f64); + let full = calibration.scale(question.kind, &full); + // The reference stores the scaled values back into float32 tensors. + let probs: Vec = full[..full.len() - 1].iter().map(|&p| p as f32).collect(); + let abstain = full[full.len() - 1] as f32; + let mut probs64: Vec = probs.iter().map(|&p| p as f64).collect(); + let probabilities = |values: &[f64]| -> Map { + question + .keys + .iter() + .zip(values) + .map(|(k, &p)| (k.clone(), serde_json::json!(round4(p)))) + .collect() + }; + if question.kind == Kind::Score { + let total: f64 = probs64.iter().sum(); + let total = if total == 0.0 { 1.0 } else { total }; + probs64.iter_mut().for_each(|p| *p /= total); + let mode = first_max(&probs64); + let as_f32: Vec = probs64.iter().map(|&p| p as f32).collect(); + return Ok(serde_json::json!({ + "score": question.keys[mode], "probabilities": probabilities(&probs64), + "confidence": round4(confidence(&as_f32)), "latency_s": latency})); + } + let mode = first_max(&probs64); + let mut with_abstain = probs.clone(); + with_abstain.push(abstain); + Ok(serde_json::json!({ + "choice": question.keys[mode], "probabilities": probabilities(&probs64), + "abstain": round4(abstain as f64), "valid": valid, + "confidence": round4(confidence(&with_abstain)), "latency_s": latency})) +} diff --git a/src/models/omnijev/native/src/heads.rs b/src/models/omnijev/native/src/heads.rs new file mode 100644 index 00000000..5d5cbb02 --- /dev/null +++ b/src/models/omnijev/native/src/heads.rs @@ -0,0 +1,269 @@ +//! OmniJev's trained FP32 heads on the CPU and the backbone log-probability features +//! they take, following OmniJev @ 14dbec4 (`mso/head.py`, `mso/branch.py`). + +use std::fs; +use std::path::Path; + +use anyhow::{Context, Result, ensure}; +use safetensors::{Dtype, SafeTensors}; + +pub const HIDDEN: usize = 2560; +const SCORER_WIDTH: usize = 1024; +const ORDINAL_WIDTH: usize = 512; +pub const FEATURES: usize = 6; + +struct Linear { + weight: Vec, + bias: Vec, + inputs: usize, +} + +impl Linear { + fn load(tensors: &SafeTensors, name: &str, outputs: usize, inputs: usize) -> Result { + Ok(Self { + weight: tensor(tensors, &format!("{name}.weight"), &[outputs, inputs])?, + bias: tensor(tensors, &format!("{name}.bias"), &[outputs])?, + inputs, + }) + } + + fn apply(&self, x: &[f32]) -> Vec { + debug_assert_eq!(x.len(), self.inputs); + self.weight + .chunks_exact(self.inputs) + .zip(&self.bias) + .map(|(row, &b)| dot(row, x) + b) + .collect() + } +} + +/// Linear, exact GELU, Linear. +struct Mlp(Linear, Linear); + +impl Mlp { + fn load( + tensors: &SafeTensors, + name: &str, + inputs: usize, + width: usize, + outputs: usize, + ) -> Result { + Ok(Self( + Linear::load(tensors, &format!("{name}.0"), width, inputs)?, + Linear::load(tensors, &format!("{name}.2"), outputs, width)?, + )) + } + + fn apply(&self, x: &[f32]) -> Vec { + let hidden: Vec = self.0.apply(x).into_iter().map(gelu).collect(); + self.1.apply(&hidden) + } +} + +fn dot(a: &[f32], b: &[f32]) -> f32 { + // Eight partial sums, which the compiler vectorizes. + let mut acc = [0.0f32; 8]; + let (a8, a_rest) = a.split_at(a.len() / 8 * 8); + let (b8, b_rest) = b.split_at(a8.len()); + for (x, y) in a8.chunks_exact(8).zip(b8.chunks_exact(8)) { + for i in 0..8 { + acc[i] += x[i] * y[i]; + } + } + let tail: f32 = a_rest.iter().zip(b_rest).map(|(x, y)| x * y).sum(); + acc.iter().sum::() + tail +} + +fn gelu(x: f32) -> f32 { + 0.5 * x * (1.0 + libm::erff(x * std::f32::consts::FRAC_1_SQRT_2)) +} + +fn sigmoid(x: f32) -> f32 { + 1.0 / (1.0 + (-x).exp()) +} + +/// PyTorch's softplus with beta 1 and threshold 20. +fn softplus(x: f32) -> f32 { + if x > 20.0 { x } else { x.exp().ln_1p() } +} + +fn tensor(tensors: &SafeTensors, name: &str, shape: &[usize]) -> Result> { + let view = tensors + .tensor(name) + .with_context(|| format!("missing {name}"))?; + ensure!( + view.dtype() == Dtype::F32 && view.shape() == shape, + "{name} must be F32 {shape:?}" + ); + Ok(view + .data() + .chunks_exact(4) + .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .collect()) +} + +/// `OptionScorer` with softmax normalization, and `OrdinalScoreHead`. +pub struct Heads { + option: Mlp, + question: Mlp, + score: Linear, + log_tau: Vec, + abstain: Linear, + feature: Linear, + feature_logit: Linear, + ordinal_z: Mlp, + ordinal_cut: Mlp, +} + +impl Heads { + /// Load `head.*` and `ord.*` from the export's FP32 `heads.safetensors`. + pub fn load(path: &Path) -> Result { + let bytes = fs::read(path).with_context(|| format!("reading {}", path.display()))?; + let t = SafeTensors::deserialize(&bytes)?; + Ok(Self { + option: Mlp::load(&t, "head.opt", HIDDEN, SCORER_WIDTH, SCORER_WIDTH)?, + question: Mlp::load(&t, "head.que", HIDDEN, SCORER_WIDTH, SCORER_WIDTH)?, + score: Linear::load(&t, "head.score", 1, SCORER_WIDTH)?, + log_tau: tensor(&t, "head.log_tau", &[3])?, + abstain: Linear::load(&t, "head.abstain", 1, SCORER_WIDTH)?, + feature: Linear::load(&t, "head.feat", SCORER_WIDTH, FEATURES)?, + feature_logit: Linear::load(&t, "head.feat_lin", 1, FEATURES)?, + ordinal_z: Mlp::load(&t, "ord.z", HIDDEN, ORDINAL_WIDTH, 1)?, + ordinal_cut: Mlp::load(&t, "ord.cut", HIDDEN, ORDINAL_WIDTH, 1)?, + }) + } + + /// `OptionScorer.logits`: one logit per option, then the abstain logit for Choice + /// and Score or zero for Noul. `options` holds one hidden row per option. + pub fn option_logits( + &self, + options: &[Vec], + question: &[f32], + type_id: usize, + features: &[[f32; FEATURES]], + ) -> Result> { + ensure!( + !options.is_empty() + && options.len() == features.len() + && options.iter().all(|u| u.len() == HIDDEN) + && question.len() == HIDDEN + && type_id < 3, + "head inputs must be one {HIDDEN}-wide row and feature row per option" + ); + let gate: Vec = self + .question + .apply(question) + .into_iter() + .map(f32::tanh) + .collect(); + let tau = self.log_tau[type_id].exp(); + let mut logits: Vec = options + .iter() + .zip(features) + .map(|(u, f)| { + let mut h: Vec = self + .option + .apply(u) + .iter() + .zip(&gate) + .map(|(a, g)| a * g) + .collect(); + for (h, f) in h.iter_mut().zip(self.feature.apply(f)) { + *h += f; + } + (self.score.apply(&h)[0] + self.feature_logit.apply(f)[0]) / tau + }) + .collect(); + logits.push(if type_id == 0 { + 0.0 + } else { + self.abstain.apply(&gate)[0] / tau + }); + Ok(logits) + } + + /// `OptionScorer.forward`: the options' probabilities; their shortfall from one is + /// the abstain probability. Noul is a sigmoid of its one logit. + pub fn option_probabilities( + &self, + options: &[Vec], + question: &[f32], + type_id: usize, + features: &[[f32; FEATURES]], + ) -> Result> { + let logits = self.option_logits(options, question, type_id, features)?; + let k = logits.len() - 1; + if type_id == 0 { + return Ok(logits[..k].iter().map(|&x| sigmoid(x)).collect()); + } + let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let exp: Vec = logits.iter().map(|&x| (x - max).exp()).collect(); + let total: f32 = exp.iter().sum(); + Ok(exp[..k].iter().map(|e| e / total).collect()) + } + + /// `OrdinalScoreHead`: cumulative-link probabilities of the levels in caller order. + pub fn ordinal(&self, question: &[f32], levels: &[Vec]) -> Result> { + ensure!( + !levels.is_empty() + && levels.iter().all(|u| u.len() == HIDDEN) + && question.len() == HIDDEN, + "ordinal inputs must be {HIDDEN}-wide rows" + ); + let z = self.ordinal_z.apply(question)[0]; + let raw: Vec = levels + .iter() + .map(|u| self.ordinal_cut.apply(u)[0]) + .collect(); + let mut theta = Vec::with_capacity(raw.len()); + let mut sum = 0.0f32; + for (i, &r) in raw.iter().enumerate() { + if i > 0 { + sum += softplus(r); + } + theta.push(raw[0] + sum); + } + let mut cdf: Vec = theta.iter().map(|&t| sigmoid(t - z)).collect(); + *cdf.last_mut().expect("at least one level") = 1.0; + Ok((0..cdf.len()) + .map(|i| if i == 0 { cdf[0] } else { cdf[i] - cdf[i - 1] }.max(1e-8)) + .collect()) + } +} + +/// The six backbone features of every option of one question: the summed, mean and +/// counted log-probabilities of its text tokens, their log-softmax across the +/// question's options, the first token's log-probability from `zq`, and validity. +/// `picked[i]` holds option `i`'s text-token log-probabilities, `first[i]` the first +/// one's from `zq`. +pub fn lm_features(picked: &[Vec], first: &[Option]) -> Vec<[f32; FEATURES]> { + let sums: Vec = picked + .iter() + .map(|p| p.iter().fold(0.0f32, |a, &x| a + x)) + .collect(); + let masked: Vec = picked + .iter() + .zip(&sums) + .map(|(p, &s)| if p.is_empty() { -1e4 } else { s }) + .collect(); + let max = masked.iter().copied().fold(f32::NEG_INFINITY, f32::max); + let log_total = masked.iter().map(|&m| (m - max).exp()).sum::().ln(); + picked + .iter() + .zip(&sums) + .zip(&masked) + .zip(first) + .map(|(((p, &sum), &m), &first)| { + let count = p.len() as f32; + let valid = !p.is_empty(); + [ + sum / 10.0, + if valid { sum / count } else { 0.0 }, + count / 10.0, + m - max - log_total, + first.unwrap_or(0.0), + if valid { 1.0 } else { 0.0 }, + ] + }) + .collect() +} diff --git a/src/models/omnijev/native/src/lib.rs b/src/models/omnijev/native/src/lib.rs new file mode 100644 index 00000000..1be2f939 --- /dev/null +++ b/src/models/omnijev/native/src/lib.rs @@ -0,0 +1,5 @@ +//! OmniJev-4B v1.1 decisions on one image: request contract, preparation and the +//! branch layout, trained heads and response finishing. +pub mod contract; +pub mod heads; +pub mod processing; diff --git a/src/models/omnijev/native/src/processing.rs b/src/models/omnijev/native/src/processing.rs new file mode 100644 index 00000000..b08e0ce0 --- /dev/null +++ b/src/models/omnijev/native/src/processing.rs @@ -0,0 +1,307 @@ +//! Prompt rendering, tokenization and the branch layout, following OmniJev @ 14dbec4 +//! (`MSO1._prompt`, `MSO1._encode_many`, `MSO1.ask_branch`, `mso/branch.py`). + +use std::path::Path; + +use anyhow::{Context, Result, ensure}; +use tokenizers::Tokenizer; + +use crate::contract::{self, Image, Question, Request}; + +pub const IMAGE_TOKEN: u32 = 248_056; +pub const OPTION_OPEN: u32 = 248_077; +pub const OPTION_CLOSE: u32 = 248_078; +/// `MSO1`'s default image budget, which overrides the processor config's 401,408. +pub const MAX_PIXELS: usize = 768 * 28 * 28; +pub const MIN_PIXELS: usize = 65_536; +const FACTOR: usize = 32; // patch 16 × spatial merge 2 +/// A single question's tokens: prompt, image and every option. +pub const MAX_SEQUENCE: usize = 8192; +/// Tokens the reference processes for a request: its prefix once, then every row. +pub const MAX_PROCESSED: usize = 65_536; + +const PAD: &str = "<|image_pad|>"; + +/// `[1, height / 16, width / 16]` after the processor's resize: Transformers' +/// `Qwen2VLImageProcessor` smart resize (Apache-2.0; see ../../README.md), as Cua-S1's. +pub fn image_grid(width: usize, height: usize) -> [usize; 3] { + // Python's round uses ties-to-even. + let mut w = (width as f64 / FACTOR as f64).round_ties_even() as usize * FACTOR; + let mut h = (height as f64 / FACTOR as f64).round_ties_even() as usize * FACTOR; + if w * h > MAX_PIXELS { + let beta = ((width * height) as f64 / MAX_PIXELS as f64).sqrt(); + w = ((width as f64 / beta / FACTOR as f64).floor() as usize * FACTOR).max(FACTOR); + h = ((height as f64 / beta / FACTOR as f64).floor() as usize * FACTOR).max(FACTOR); + } else if w * h < MIN_PIXELS { + let beta = (MIN_PIXELS as f64 / (width * height) as f64).sqrt(); + w = (width as f64 * beta / FACTOR as f64).ceil() as usize * FACTOR; + h = (height as f64 * beta / FACTOR as f64).ceil() as usize * FACTOR; + } + [1, h / 16, w / 16] +} + +/// Image tokens after the 2×2 spatial merge. +pub fn image_tokens(grid: [usize; 3]) -> usize { + grid[1] / 2 * (grid[2] / 2) +} + +/// Python's `str.isspace`, which Jinja's `trim` uses. +fn python_space(c: char) -> bool { + c.is_whitespace() || ('\u{1c}'..='\u{1f}').contains(&c) +} + +/// The chat-template text of one question with its option blocks appended, with one +/// image placeholder that tokenization expands. +pub fn prompt(question: &Question) -> String { + // The template trims the user content, which starts with the image, so only the + // instructions' trailing whitespace goes. + let content = format!( + "<|vision_start|>{PAD}<|vision_end|>{}", + question.instructions.trim_end_matches(python_space) + ); + let mut text = + format!("<|im_start|>user\n{content}<|im_end|>\n<|im_start|>assistant\n\n"); + for option in &question.options { + text += "<|opt|>"; + text += option; + text += "<|/opt|>"; + } + text +} + +/// One (question, option) row: the rest of the question after the shared prefix, +/// then one option block. +#[derive(Debug, PartialEq)] +pub struct Row { + pub tokens: Vec, + /// The token before the option marker, where the question state `zq` is read. + pub zq: usize, + /// The closing marker, where the option state `u` is read; the row's last token. + pub u: usize, + /// For each option-text token: the position that predicts it and the token. + pub targets: Vec<(usize, u32)>, + /// The first option-text token, predicted from `zq`. + pub first: Option, +} + +pub struct PreparedQuestion { + /// The single-question sequence: prompt, image and every option block. + pub token_ids: Vec, + /// `[3, sequence]` T/H/W rotary positions of `token_ids`. + pub positions: [Vec; 3], + pub rows: Vec, +} + +/// The reference's split of a request into a shared prefix and rows. +pub struct Layout { + /// The prefix length: the questions' common token prefix, cut at least one token + /// before any question's first option marker. + pub prefix: usize, + /// The rotary position of the first token after the prefix. + pub next_position: i64, + /// The reference's input-token count: the prefix once, then every row. + pub processed_tokens: usize, +} + +/// Option-marker positions in a single-question sequence. +pub fn spans(ids: &[u32]) -> (Vec, Vec) { + let at = |token| (0..ids.len()).filter(|&i| ids[i] == token).collect(); + (at(OPTION_OPEN), at(OPTION_CLOSE)) +} + +/// Rows and readouts of every question, as `ask_branch` and `branch_forward` build them. +pub fn layout(sequences: &[Vec], grid: [usize; 3]) -> Result<(Layout, Vec>)> { + ensure!(!sequences.is_empty(), "no questions"); + let shortest = sequences.iter().map(Vec::len).min().unwrap_or(0); + let mut prefix = 0; + while prefix < shortest && sequences.iter().all(|s| s[prefix] == sequences[0][prefix]) { + prefix += 1; + } + let mut spans_of = Vec::with_capacity(sequences.len()); + for ids in sequences { + let (opens, closes) = spans(ids); + ensure!( + !opens.is_empty() + && opens.len() == closes.len() + && opens.iter().zip(&closes).all(|(o, c)| o < c) + && closes.iter().zip(&opens[1..]).all(|(c, o)| c < o), + "every option needs one opening and one closing marker, in order" + ); + // Every row keeps at least the token before its option marker. + ensure!(opens[0] >= 2, "options must follow the prompt"); + prefix = prefix.min(opens[0] - 1); + spans_of.push((opens, closes)); + } + let prefix = prefix.max(1); + let mut rows = Vec::with_capacity(sequences.len()); + let mut processed = prefix; + for (ids, (opens, closes)) in sequences.iter().zip(&spans_of) { + let head = &ids[prefix..opens[0]]; + let question_rows: Vec = opens + .iter() + .zip(closes) + .map(|(&o, &c)| { + let tokens: Vec = head.iter().chain(&ids[o..=c]).copied().collect(); + let open = head.len(); + let last = tokens.len() - 1; + let inner = open + 1..last; + Row { + zq: open - 1, + u: last, + targets: inner.clone().map(|t| (t - 1, tokens[t])).collect(), + first: inner.clone().next().map(|t| tokens[t]), + tokens, + } + }) + .collect(); + processed += question_rows.iter().map(|r| r.tokens.len()).sum::(); + rows.push(question_rows); + } + let positions = image_positions(&sequences[0], grid)?; + let next_position = positions + .iter() + .map(|axis| axis[prefix - 1]) + .max() + .unwrap_or(0) + + 1; + Ok(( + Layout { + prefix, + next_position, + processed_tokens: processed, + }, + rows, + )) +} + +/// Qwen3.5's T/H/W rotary positions (Transformers' `get_rope_index`, Apache-2.0; see +/// ../../README.md) for a sequence with one contiguous image span: +/// text before it counts up, the image's patches take the position of its first token +/// plus their row (H) and column (W), and text after it continues from the image's +/// largest position plus one. +pub fn image_positions(ids: &[u32], grid: [usize; 3]) -> Result<[Vec; 3]> { + let start = ids + .iter() + .position(|&id| id == IMAGE_TOKEN) + .context("missing image placeholders")?; + let count = ids[start..] + .iter() + .take_while(|&&id| id == IMAGE_TOKEN) + .count(); + let end = start + count; + ensure!( + !ids[end..].contains(&IMAGE_TOKEN), + "image placeholders must be one contiguous span" + ); + let [t, h, w] = grid; + ensure!( + t == 1 && h > 0 && w > 0 && h % 2 == 0 && w % 2 == 0 && h / 2 * (w / 2) == count, + "image span does not match its grid" + ); + let mut positions: [Vec; 3] = std::array::from_fn(|_| (0..start as i64).collect()); + for y in 0..h / 2 { + for x in 0..w / 2 { + positions[0].push(start as i64); + positions[1].push((start + y) as i64); + positions[2].push((start + x) as i64); + } + } + let next = start + h.max(w) / 2; + for axis in &mut positions { + axis.extend((next..next + ids.len() - end).map(|p| p as i64)); + } + Ok(positions) +} + +pub struct Processor { + tokenizer: Tokenizer, +} + +/// A request ready for execution; `questions` keeps identity and order for finishing. +pub struct PreparedRequest { + pub image: Image, + pub grid: [usize; 3], + pub questions: Vec, + pub inputs: Vec, + pub layout: Layout, +} + +impl Processor { + pub fn load(dir: &Path) -> Result { + let tokenizer = + Tokenizer::from_file(dir.join("tokenizer.json")).map_err(anyhow::Error::msg)?; + for (text, id) in [ + ("<|image_pad|>", IMAGE_TOKEN), + ("<|opt|>", OPTION_OPEN), + ("<|/opt|>", OPTION_CLOSE), + ] { + ensure!( + tokenizer.token_to_id(text) == Some(id), + "tokenizer does not map {text} to {id}" + ); + } + Ok(Self { tokenizer }) + } + + /// The single-question sequence: the prompt's text tokens around the image's + /// expanded placeholder. + pub fn tokenize(&self, question: &Question, grid: [usize; 3]) -> Result> { + let text = prompt(question); + let (before, after) = text.split_once(PAD).context("prompt has no image")?; + let encode = |part: &str| -> Result> { + Ok(self + .tokenizer + .encode(part, false) + .map_err(anyhow::Error::msg)? + .get_ids() + .to_vec()) + }; + let mut ids = encode(before)?; + ids.extend(std::iter::repeat_n(IMAGE_TOKEN, image_tokens(grid))); + ids.extend(encode(after)?); + Ok(ids) + } + + /// Validate and lay out the whole request before any device work. + pub fn prepare(&self, raw: &[u8]) -> Result { + let Request { image, questions } = contract::compile(raw)?; + let grid = image_grid(image.width, image.height); + let sequences = questions + .iter() + .map(|q| { + let ids = self.tokenize(q, grid)?; + ensure!( + ids.len() <= MAX_SEQUENCE, + "question {}: {} tokens exceeds {MAX_SEQUENCE}", + q.id, + ids.len() + ); + Ok(ids) + }) + .collect::>>()?; + let (layout, rows) = layout(&sequences, grid)?; + ensure!( + layout.processed_tokens <= MAX_PROCESSED, + "request needs {} tokens, above {MAX_PROCESSED}", + layout.processed_tokens + ); + let inputs = sequences + .into_iter() + .zip(rows) + .map(|(token_ids, rows)| { + Ok(PreparedQuestion { + positions: image_positions(&token_ids, grid)?, + token_ids, + rows, + }) + }) + .collect::>>()?; + Ok(PreparedRequest { + image, + grid, + questions, + inputs, + layout, + }) + } +} diff --git a/tests/omnijev/common.rs b/tests/omnijev/common.rs new file mode 100644 index 00000000..ad36aa14 --- /dev/null +++ b/tests/omnijev/common.rs @@ -0,0 +1,56 @@ +//! Fixture loading shared by the OmniJev tests. +#![allow(dead_code)] + +use std::io::Cursor; + +use base64::Engine; +use serde_json::Value; + +pub fn fixture(name: &str) -> Value { + let path = format!( + "{}/../../../../tests/omnijev/data/{name}", + env!("CARGO_MANIFEST_DIR") + ); + serde_json::from_str(&std::fs::read_to_string(path).unwrap()).unwrap() +} + +/// A solid-colour PNG as a data URL; every encoder gives these pixels. +pub fn png_url(width: u32, height: u32, colour: [u8; 3]) -> String { + let image = image::RgbImage::from_pixel(width, height, image::Rgb(colour)); + let mut bytes = Vec::new(); + image + .write_to(&mut Cursor::new(&mut bytes), image::ImageFormat::Png) + .unwrap(); + format!( + "data:image/png;base64,{}", + base64::engine::general_purpose::STANDARD.encode(bytes) + ) +} + +/// A processing fixture's request, with its image generated. +pub fn request(case: &Value) -> Vec { + let mut request = case["request"].clone(); + let size = &case["image_size"]; + let colour = &case["image_colour"]; + let rgb = [0, 1, 2].map(|i| colour[i].as_u64().unwrap() as u8); + request["state"]["images"][0] = Value::String(png_url( + size[0].as_u64().unwrap() as u32, + size[1].as_u64().unwrap() as u32, + rgb, + )); + serde_json::to_vec(&request).unwrap() +} + +/// Fixture token ids, with the image run (stored as its negated length) expanded. +pub fn token_ids(collapsed: &Value) -> Vec { + let mut ids = Vec::new(); + for v in collapsed.as_array().unwrap() { + let v = v.as_i64().unwrap(); + if v < 0 { + ids.extend(std::iter::repeat_n(248_056, (-v) as usize)); + } else { + ids.push(v as u32); + } + } + ids +} diff --git a/tests/omnijev/contract.rs b/tests/omnijev/contract.rs new file mode 100644 index 00000000..1976ecb8 --- /dev/null +++ b/tests/omnijev/contract.rs @@ -0,0 +1,138 @@ +//! Request validation: what the native contract accepts and refuses. +mod common; + +use omni_omnijev_native::contract::{self, Kind}; +use serde_json::{Value, json}; + +fn body(state: Value, questions: Value) -> Vec { + serde_json::to_vec( + &json!({"model": "tinnel123/OmniJev", "state": state, "questions": questions}), + ) + .unwrap() +} + +fn image(width: u32, height: u32) -> Value { + json!({"images": [common::png_url(width, height, [1, 2, 3])]}) +} + +fn noul() -> Value { + json!({"q": {"type": "noul", "instructions": "x"}}) +} + +/// The full error chain, so a case fails for the reason it names. +fn refused(raw: &[u8]) -> String { + format!("{:#}", contract::compile(raw).err().expect("accepted")) +} + +#[test] +fn accepts_models_and_reads_the_image() { + for model in [ + json!("tinnel123/OmniJev"), + json!("omnijev"), + json!("omnijev-4b-v1.1"), + Value::Null, + ] { + let raw = serde_json::to_vec( + &json!({"model": model, "state": image(64, 48), "questions": noul()}), + ) + .unwrap(); + let request = contract::compile(&raw).unwrap(); + assert_eq!((request.image.width, request.image.height), (64, 48)); + assert_eq!(request.questions[0].kind, Kind::Noul); + assert_eq!(request.questions[0].options, [""]); + } + let raw = serde_json::to_vec(&json!({"state": image(64, 48), "questions": noul()})).unwrap(); + assert!(contract::compile(&raw).is_ok()); + // The processor accepts aspect ratios up to 200. + assert!(contract::compile(&body(image(400, 2), noul())).is_ok()); + assert!(contract::compile(&body(image(4096, 64), noul())).is_ok()); +} + +/// Each case is ` => `; `{png}`, `{wide}` and `{thin}` stand +/// for generated images, `{long}` for 8,193 characters and `{id}` for a 257-character id. +fn check_refusals(cases: &[&str], build: impl Fn(Value) -> Vec) { + let png = common::png_url(64, 48, [1, 2, 3]); + for case in cases { + let (json, reason) = case.split_once(" => ").unwrap(); + let json = json + .replace("{png}", &png) + .replace("{jpeg}", &png.replace("data:image/png", "data:image/jpeg")) + .replace("{wide}", &common::png_url(4097, 64, [0, 0, 0])) + .replace("{thin}", &common::png_url(401, 2, [0, 0, 0])) + .replace("{long}", &"x".repeat(8193)) + .replace("{id}", &"q".repeat(257)); + let message = refused(&build(serde_json::from_str(&json).unwrap())); + assert!(message.contains(reason), "{case}: {message}"); + } +} + +#[test] +fn refuses_unsupported_images() { + let cases = [ + r#"{"images": []} => exactly one image"#, + r#"{"images": ["{png}"], "video": {}} => images only"#, + r#"{"images": ["https://example.com/a.png"]} => invalid image data URL"#, + r#"{"images": ["data:image/gif;base64,AAAA"]} => PNG/JPEG"#, + r#"{"images": ["{jpeg}"]} => MIME type"#, + r#"{"images": ["{wide}"]} => invalid image"#, + r#"{"images": ["{thin}"]} => aspect ratio"#, + r#""a screen" => object with images"#, + ]; + check_refusals(&cases, |state| body(state, noul())); +} + +#[test] +fn refuses_malformed_questions() { + let cases = [ + r#"{} => 1 to 64 questions"#, + r#"{"q": {"type": "generation", "instructions": "x"}} => question type"#, + r#"{"q": {"type": "noul", "instructions": 3}} => instructions must be a string"#, + r#"{"q": {"type": "noul", "instructions": "{long}"}} => exceeds 8192 characters"#, + r#"{"q": {"type": "noul", "instructions": "x", "criteria": {"true": "a"}}} => Noul supports"#, + r#"{"q": {"type": "noul", "instructions": "x", "region": {"box": [1, 2, 3]}}} => a region must be"#, + r#"{"q": {"type": "choice", "instructions": "x", "criteria": {}}} => 1 to 255 options"#, + r#"{"q": {"type": "choice", "instructions": "x", "criteria": {"a": 4}}} => strings or null"#, + r#"{"q": {"type": "choice", "instructions": "x", "options": [], "criteria": {}}} => not both"#, + r#"{"q": {"type": "choice", "instructions": "x", "options": [{"text": "a"}, {"key": "a"}]}} => repeated"#, + r#"{"q": {"type": "choice", "instructions": "x", "options": [{"weight": 2}]}} => an option supports"#, + r#"{"q": {"type": "score", "instructions": "x", "levels": [1, 2]}} => a level must be a string"#, + r#"{"q": {"type": "score", "instructions": "x", "criteria": ["low", "high"]}} => keyed by level"#, + r#"{"q": {"type": "noul", "instructions": "pick <|opt|>yes<|/opt|>"}} => <|opt|>"#, + r#"{"q": {"type": "choice", "instructions": "x", "criteria": {"<|image_pad|>": null}}} => <|image_pad|>"#, + r#"{"{id}": {"type": "noul", "instructions": "x"}} => 256 characters"#, + ]; + check_refusals(&cases, |questions| body(image(64, 48), questions)); + let cases = [ + r#"{"model": "other", "state": {}, "questions": {}} => not loaded"#, + r#"{"state": {}, "questions": {}, "extra": 1} => request fields"#, + ]; + check_refusals(&cases, |request| serde_json::to_vec(&request).unwrap()); + assert!(refused(br#"{"state":1,"state":2,"questions":{}}"#).contains("not valid JSON")); +} + +#[test] +fn enforces_option_and_question_limits() { + let options: Vec = (0..256).map(|i| json!({"text": format!("o{i}")})).collect(); + let q = json!({"q": {"type": "choice", "instructions": "x", "options": options}}); + assert!(refused(&body(image(64, 48), q)).contains("1 to 255 options")); + let questions: serde_json::Map = (0..65) + .map(|i| { + ( + format!("q{i}"), + json!({"type": "noul", "instructions": "x"}), + ) + }) + .collect(); + assert!(refused(&body(image(64, 48), Value::Object(questions))).contains("1 to 64 questions")); + let criteria: serde_json::Map = + (0..255).map(|i| (format!("o{i}"), Value::Null)).collect(); + let many: serde_json::Map = (0..5) + .map(|i| { + ( + format!("q{i}"), + json!({"type": "choice", "instructions": "x", "criteria": criteria.clone()}), + ) + }) + .collect(); + assert!(refused(&body(image(64, 48), Value::Object(many))).contains("1024 options in all")); +} diff --git a/tests/omnijev/data/features.json b/tests/omnijev/data/features.json new file mode 100644 index 00000000..97052b52 --- /dev/null +++ b/tests/omnijev/data/features.json @@ -0,0 +1 @@ +{"provenance":{"python":"3.10.12","reference_commit":"14dbec4f71e194852c8d7b88ab36ef639493f400","checkpoint_revision":"ffe5f436eaf22e20e2f041f8e74e121fd057a6cb","base_revision":"851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a","max_pixels":602112,"torch":"2.14.0+cu130","transformers":"5.17.0","command":"CUDA_VISIBLE_DEVICES= python tests/omnijev/generate_fixtures.py"},"groups":[[0,1,2],[3],[4]],"rows":[{"row":0,"picked":[-43.88062286376953,-43.723384857177734,-32.77771759033203],"first":-21.340444564819336,"feats":[-12.038171768188477,-40.12723922729492,0.30000001192092896,-73.82730102539062,-21.340444564819336,1.0]},{"row":1,"picked":[-46.554420471191406],"first":-50.642723083496094,"feats":[-4.655442237854004,-46.554420471191406,0.10000000149011612,0.0,-50.642723083496094,1.0]},{"row":2,"picked":[],"first":null,"feats":[0.0,0.0,0.0,-9953.4453125,0.0,0.0]},{"row":3,"picked":[-33.383216857910156,-33.89513397216797,-50.14861297607422,-58.03985595703125,-14.797294616699219,-7.693144798278809],"first":-41.983184814453125,"feats":[-19.795726776123047,-32.99287796020508,0.6000000238418579,0.0,-41.98318099975586,1.0]},{"row":4,"picked":[],"first":null,"feats":[0.0,0.0,0.0,0.0,0.0,0.0]}]} \ No newline at end of file diff --git a/tests/omnijev/data/finish.json b/tests/omnijev/data/finish.json new file mode 100644 index 00000000..4356c046 --- /dev/null +++ b/tests/omnijev/data/finish.json @@ -0,0 +1 @@ +{"provenance":{"python":"3.10.12","reference_commit":"14dbec4f71e194852c8d7b88ab36ef639493f400","checkpoint_revision":"ffe5f436eaf22e20e2f041f8e74e121fd057a6cb","base_revision":"851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a","max_pixels":602112,"torch":"2.14.0+cu130","transformers":"5.17.0","command":"CUDA_VISIBLE_DEVICES= python tests/omnijev/generate_fixtures.py"},"temperatures":{"noul":1.1991,"choice":1.1199,"score":1.1909},"biases":{"noul":0.05310856608843584},"cases":[{"type":"noul","keys":["yes"],"mu":[0.5],"answer":{"noul":0.5111,"latency_s":0.0}},{"type":"noul","keys":["yes"],"mu":[0.48669999837875366],"answer":{"noul":0.5,"latency_s":0.0}},{"type":"noul","keys":["yes"],"mu":[9.999999960041972e-13],"answer":{"noul":0.0,"latency_s":0.0}},{"type":"noul","keys":["yes"],"mu":[0.9999989867210388],"answer":{"noul":1.0,"latency_s":0.0}},{"type":"choice","keys":["a","b","c"],"mu":[0.20000000298023224,0.30000001192092896,0.10000000149011612],"answer":{"choice":"b","probabilities":{"a":0.207,"b":0.2973,"c":0.1115},"abstain":0.3843,"valid":true,"confidence":0.1791,"latency_s":0.0}},{"type":"choice","keys":["a","b"],"mu":[0.44999998807907104,0.44999998807907104],"answer":{"choice":"a","probabilities":{"a":0.4423,"b":0.4423},"abstain":0.1155,"valid":true,"confidence":0.1634,"latency_s":0.0}},{"type":"choice","keys":["a","b","c","d"],"mu":[9.999999747378752e-05,0.8999999761581421,0.05000000074505806,0.0],"answer":{"choice":"b","probabilities":{"a":0.0003,"b":0.8684,"c":0.0657,"d":0.0},"abstain":0.0656,"valid":true,"confidence":0.8355,"latency_s":0.0}},{"type":"choice","keys":["only"],"mu":[0.699999988079071],"answer":{"choice":"only","probabilities":{"only":0.6806},"abstain":0.3194,"valid":true,"confidence":0.3612,"latency_s":0.0}},{"type":"score","keys":["low","mid","high"],"mu":[0.20000000298023224,0.5,0.30000001192092896],"answer":{"score":"mid","probabilities":{"low":0.2191,"mid":0.4729,"high":0.308},"confidence":0.2094,"latency_s":0.0}},{"type":"score","keys":["0","1"],"mu":[9.99999993922529e-09,1.0],"answer":{"score":"1","probabilities":{"0":0.0,"1":1.0},"confidence":1.0,"latency_s":0.0}},{"type":"score","keys":["a","b","c","d"],"mu":[0.25,0.25,0.25,0.25],"answer":{"score":"a","probabilities":{"a":0.25,"b":0.25,"c":0.25,"d":0.25},"confidence":0.0,"latency_s":0.0}},{"type":"choice","keys":["o0","o1","o2","o3","o4"],"mu":[0.01543367002159357,0.09875854104757309,0.12689365446567535,0.05173622444272041,0.2497350424528122],"answer":{"choice":"o4","probabilities":{"o0":0.0207,"o1":0.1084,"o2":0.1356,"o3":0.0609,"o4":0.2482},"abstain":0.4262,"valid":true,"confidence":0.3114,"latency_s":0.0}},{"type":"score","keys":["l0","l1","l2","l3","l4"],"mu":[0.2119274139404297,0.26993995904922485,0.06472436338663101,0.03269525617361069,0.42071300745010376],"answer":{"score":"l4","probabilities":{"l0":0.2183,"l1":0.2675,"l2":0.0806,"l3":0.0454,"l4":0.3882},"confidence":0.2353,"latency_s":0.0}},{"type":"choice","keys":["o0","o1","o2","o3","o4","o5","o6","o7","o8","o9","o10","o11"],"mu":[0.013112403452396393,0.07969974726438522,0.3021763265132904,0.07818466424942017,0.047772157937288284,0.0009574098512530327,0.004682920873165131,0.020482001826167107,0.007647403981536627,0.034267622977495193,0.37759122252464294,0.01658891886472702],"answer":{"choice":"o10","probabilities":{"o0":0.0172,"o1":0.0862,"o2":0.2834,"o3":0.0847,"o4":0.0546,"o5":0.0017,"o6":0.0069,"o7":0.0256,"o8":0.0106,"o9":0.0406,"o10":0.3458,"o11":0.0212},"abstain":0.0215,"valid":true,"confidence":0.2913,"latency_s":0.0}},{"type":"score","keys":["l0","l1","l2","l3","l4","l5","l6","l7","l8","l9","l10","l11"],"mu":[0.16312673687934875,0.010625014081597328,0.1264624297618866,0.07897239178419113,0.11969983577728271,0.1022866740822792,0.012198138982057571,0.02941376529633999,0.04900258034467697,0.04420036822557449,0.14411836862564087,0.11989373713731766],"answer":{"score":"l0","probabilities":{"l0":0.151,"l1":0.0152,"l2":0.122,"l3":0.0821,"l4":0.1165,"l5":0.1021,"l6":0.0171,"l7":0.0358,"l8":0.055,"l9":0.0504,"l10":0.1361,"l11":0.1166},"confidence":0.0738,"latency_s":0.0}},{"type":"choice","keys":["o0","o1","o2","o3","o4","o5","o6","o7","o8","o9","o10","o11","o12","o13","o14","o15","o16","o17","o18","o19","o20","o21","o22","o23","o24","o25","o26","o27","o28","o29","o30","o31"],"mu":[0.002471247222274542,0.013411770574748516,0.012321481481194496,0.04999684542417526,0.0033170029055327177,0.0001840989280026406,0.041995126754045486,0.004115985706448555,0.16607055068016052,0.11407449096441269,0.0008663962944410741,0.060479845851659775,0.011369680054485798,0.06621436029672623,0.02663538232445717,0.0036796426866203547,0.001582379569299519,0.0002584033354651183,0.02693624421954155,0.0004841393092647195,0.0006544597563333809,0.0048047746531665325,0.004691817332059145,0.11469565331935883,0.0054892138577997684,0.010964952409267426,0.004664125386625528,0.00020721186592709273,0.006677792873233557,0.0009046547929756343,0.1823570728302002,0.03962162137031555],"answer":{"choice":"o30","probabilities":{"o0":0.0035,"o1":0.016,"o2":0.0149,"o3":0.0519,"o4":0.0046,"o5":0.0003,"o6":0.0444,"o7":0.0056,"o8":0.1515,"o9":0.1084,"o10":0.0014,"o11":0.0615,"o12":0.0138,"o13":0.0667,"o14":0.0296,"o15":0.005,"o16":0.0024,"o17":0.0005,"o18":0.0299,"o19":0.0008,"o20":0.0011,"o21":0.0064,"o22":0.0063,"o23":0.1089,"o24":0.0072,"o25":0.0134,"o26":0.0062,"o27":0.0004,"o28":0.0086,"o29":0.0014,"o30":0.1647,"o31":0.0421},"abstain":0.0206,"valid":true,"confidence":0.1386,"latency_s":0.0}},{"type":"score","keys":["l0","l1","l2","l3","l4","l5","l6","l7","l8","l9","l10","l11","l12","l13","l14","l15","l16","l17","l18","l19","l20","l21","l22","l23","l24","l25","l26","l27","l28","l29","l30","l31"],"mu":[0.03723496198654175,0.016620468348264694,0.0044656009413301945,0.007466810755431652,0.03569923713803291,0.03728441521525383,0.025642553344368935,0.008165847510099411,0.00760251097381115,0.045065972954034805,0.05315669625997543,0.015480493195354939,0.05399483069777489,0.012857469730079174,0.024664053693413734,0.027251185849308968,0.010871858336031437,0.03479168564081192,0.04937032237648964,0.04610868915915489,0.03389234095811844,0.04787628725171089,0.042349617928266525,0.039838120341300964,0.03504049777984619,0.039074283093214035,0.028010699898004532,0.024581212550401688,0.04388950392603874,0.03585708513855934,0.023984218016266823,0.051810480654239655],"answer":{"score":"l12","probabilities":{"l0":0.0368,"l1":0.0187,"l2":0.0062,"l3":0.0096,"l4":0.0356,"l5":0.0369,"l6":0.0269,"l7":0.0103,"l8":0.0097,"l9":0.0433,"l10":0.0497,"l11":0.0176,"l12":0.0503,"l13":0.0151,"l14":0.0261,"l15":0.0284,"l16":0.0131,"l17":0.0348,"l18":0.0467,"l19":0.0441,"l20":0.0341,"l21":0.0455,"l22":0.0411,"l23":0.039,"l24":0.035,"l25":0.0384,"l26":0.029,"l27":0.026,"l28":0.0423,"l29":0.0357,"l30":0.0255,"l31":0.0486},"confidence":0.0197,"latency_s":0.0}},{"type":"choice","keys":["a","b","c","d","e"],"mu":[0.30000001192092896,0.30000001192092896,0.20000000298023224,0.15000000596046448,0.10000000149011612],"answer":{"choice":"a","probabilities":{"a":0.277,"b":0.277,"c":0.1929,"d":0.1492,"e":0.1039},"abstain":0.0,"valid":false,"confidence":0.1324,"latency_s":0.0}}]} \ No newline at end of file diff --git a/tests/omnijev/data/heads.json b/tests/omnijev/data/heads.json new file mode 100644 index 00000000..12a3d3d0 --- /dev/null +++ b/tests/omnijev/data/heads.json @@ -0,0 +1 @@ +{"provenance":{"python":"3.10.12","reference_commit":"14dbec4f71e194852c8d7b88ab36ef639493f400","checkpoint_revision":"ffe5f436eaf22e20e2f041f8e74e121fd057a6cb","base_revision":"851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a","max_pixels":602112,"torch":"2.14.0+cu130","transformers":"5.17.0","command":"CUDA_VISIBLE_DEVICES= python tests/omnijev/generate_fixtures.py"},"cases":[{"name":"noul","seeds":[0,1,2],"type_id":0,"options":1,"logits":[0.30570778250694275,0.0],"mu":[0.575837254524231],"ordinal":[1.0]},{"name":"choice2","seeds":[3,4,5],"type_id":1,"options":2,"logits":[-0.03552265837788582,1.2622759342193604,-4.937353134155273],"mu":[0.2141941785812378,0.7842137813568115],"ordinal":[0.7478449940681458,0.25215500593185425]},{"name":"choice5","seeds":[6,7,8],"type_id":1,"options":5,"logits":[-0.03548160940408707,1.2623884677886963,-2.62677001953125,-3.0899574756622314,-2.590874195098877,-4.9373931884765625],"mu":[0.20539557933807373,0.7520537376403809,0.015388943254947662,0.009683871641755104,0.015951376408338547],"ordinal":[0.7479049563407898,0.21396982669830322,0.013586699962615967,0.006382167339324951,0.018156349658966064]},{"name":"score4","seeds":[9,10,11],"type_id":2,"options":4,"logits":[-0.0537799596786499,1.915781021118164,-3.949354648590088,-4.688468933105469,-7.492297649383545],"mu":[0.12197862565517426,0.8742853403091431,0.0024800312239676714,0.0011843058746308088],"ordinal":[0.7479646801948547,0.21393007040023804,0.013128340244293213,0.024976909160614014]}]} \ No newline at end of file diff --git a/tests/omnijev/data/processing.json b/tests/omnijev/data/processing.json new file mode 100644 index 00000000..85258d18 --- /dev/null +++ b/tests/omnijev/data/processing.json @@ -0,0 +1 @@ +{"provenance":{"python":"3.10.12","reference_commit":"14dbec4f71e194852c8d7b88ab36ef639493f400","checkpoint_revision":"ffe5f436eaf22e20e2f041f8e74e121fd057a6cb","base_revision":"851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a","max_pixels":602112,"torch":"2.14.0+cu130","transformers":"5.17.0","command":"CUDA_VISIBLE_DEVICES= python tests/omnijev/generate_fixtures.py"},"cases":[{"name":"single_noul","request":{"model":"tinnel123/OmniJev","state":{"images":[""]},"questions":{"visible":{"type":"noul","instructions":"A person is visible."}}},"image_size":[1600,900],"image_colour":[40,90,160],"image_grid_thw":[1,36,64],"prefix_length":592,"input_tokens":595,"pixel_values_sha256":"439a2c9ec89d25c0cf7669301bfabd1ebf56e1657790afd2545cad85fac583e3","pixel_values_shape":[2304,1536],"questions":[{"id":"visible","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>A person is visible.<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-576,248054,32,1637,369,9155,13,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[593],"closes":[594],"positions_sha256":"7e6dc6043b9dd86dd01828448e906bffaf4008274e146e85a429d359daf87f89","rows":[{"tokens":[198,248077,248078],"zq":0,"u":2,"targets":[],"first":null}]}]},{"name":"mixed_speed_bench","request":{"model":"tinnel123/OmniJev","state":{"images":[""]},"questions":{"q0":{"type":"choice","instructions":"Which operation comes next?","options":[{"key":"click","text":"click"},{"key":"type text","text":"type text"},{"key":"scroll","text":"scroll"},{"key":"wait","text":"wait"},{"abstain":true}]},"q1":{"type":"noul","instructions":"This screen shows an error dialog."},"q2":{"type":"score","instructions":"How irreversible is the next action?","levels":["harmless","needs care","irreversible"]},"q3":{"type":"choice","instructions":"What is the main object?","options":[{"key":"a board","text":"a board"},{"key":"a screen","text":"a screen"},{"key":"a person","text":"a person"},{"key":"a robot","text":"a robot"},{"key":"a street","text":"a street"},{"abstain":true}]},"q4":{"type":"noul","instructions":"The scene is indoors."},"q5":{"type":"score","instructions":"How cluttered is the scene?","levels":["empty","some","busy","packed"]},"q6":{"type":"choice","instructions":"Which region is most relevant?","options":[{"key":"r0","region":{"box":[0,100,180,300]}},{"key":"r1","region":{"box":[200,100,380,300]}},{"key":"r2","region":{"box":[400,100,580,300]}},{"key":"r3","region":{"box":[600,100,780,300]}},{"abstain":true}]},"q7":{"type":"noul","instructions":"Something in this image is moving."},"q8":{"type":"score","instructions":"How risky is acting on this screen?","levels":["safe","check first","dangerous"]},"q9":{"type":"choice","instructions":"What should happen next?","options":[{"key":"continue","text":"continue"},{"key":"stop","text":"stop"},{"key":"ask a human","text":"ask a human"},{"abstain":true}]},"q10":{"type":"noul","instructions":"The task looks finished."},"q11":{"type":"choice","instructions":"Which colour dominates?","options":[{"key":"white","text":"white"},{"key":"black","text":"black"},{"key":"blue","text":"blue"},{"key":"green","text":"green"},{"key":"red","text":"red"},{"abstain":true}]}}},"image_size":[1600,900],"image_colour":[40,90,160],"image_grid_thw":[1,36,64],"prefix_length":581,"input_tokens":1214,"pixel_values_sha256":"439a2c9ec89d25c0cf7669301bfabd1ebf56e1657790afd2545cad85fac583e3","pixel_values_shape":[2304,1536],"questions":[{"id":"q0","type":"choice","keys":["click","type text","scroll","wait"],"option_text":["click","type text","scroll","wait"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Which operation comes next?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>click<|/opt|><|opt|>type text<|/opt|><|opt|>scroll<|/opt|><|opt|>wait<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,3552,248078,248077,1267,1414,248078,248077,12242,248078,248077,11158,248078],"opens":[593,596,600,603],"closes":[595,599,602,605],"positions_sha256":"045266f4d5bb5d70b5f53edfa728763004a3bb0b59e35b934d4ceb3c816b27ec","rows":[{"tokens":[22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,3552,248078],"zq":11,"u":14,"targets":[[12,3552]],"first":3552},{"tokens":[22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,1267,1414,248078],"zq":11,"u":15,"targets":[[12,1267],[13,1414]],"first":1267},{"tokens":[22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,12242,248078],"zq":11,"u":14,"targets":[[12,12242]],"first":12242},{"tokens":[22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,11158,248078],"zq":11,"u":14,"targets":[[12,11158]],"first":11158}]},{"id":"q1","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>This screen shows an error dialog.<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-576,248054,1919,4034,4774,449,1412,7027,13,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[595],"closes":[596],"positions_sha256":"5dedc312e03d4751587ad21e3b2377a3345fa0f0e88c9b25e5a64ad95b8272f3","rows":[{"tokens":[1919,4034,4774,449,1412,7027,13,248046,198,248045,74455,198,248068,198,248077,248078],"zq":13,"u":15,"targets":[],"first":null}]},{"id":"q2","type":"score","keys":["harmless","needs care","irreversible"],"option_text":["harmless","needs care","irreversible"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>How irreversible is the next action?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>harmless<|/opt|><|opt|>needs care<|/opt|><|opt|>irreversible<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,4199,89013,369,279,1727,1852,30,248046,198,248045,74455,198,248068,198,248077,71,2109,1657,248078,248077,53390,2373,248078,248077,396,265,62127,248078],"opens":[595,600,604],"closes":[599,603,608],"positions_sha256":"d060c61fa5ed41dd13fc711bfb6aabc8dc821ab3043b9ba0232fad9872e7a912","rows":[{"tokens":[4199,89013,369,279,1727,1852,30,248046,198,248045,74455,198,248068,198,248077,71,2109,1657,248078],"zq":13,"u":18,"targets":[[14,71],[15,2109],[16,1657]],"first":71},{"tokens":[4199,89013,369,279,1727,1852,30,248046,198,248045,74455,198,248068,198,248077,53390,2373,248078],"zq":13,"u":17,"targets":[[14,53390],[15,2373]],"first":53390},{"tokens":[4199,89013,369,279,1727,1852,30,248046,198,248045,74455,198,248068,198,248077,396,265,62127,248078],"zq":13,"u":18,"targets":[[14,396],[15,265],[16,62127]],"first":396}]},{"id":"q3","type":"choice","keys":["a board","a screen","a person","a robot","a street"],"option_text":["a board","a screen","a person","a robot","a street"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What is the main object?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>a board<|/opt|><|opt|>a screen<|/opt|><|opt|>a person<|/opt|><|opt|>a robot<|/opt|><|opt|>a street<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,3710,369,279,1822,1576,30,248046,198,248045,74455,198,248068,198,248077,64,4331,248078,248077,64,4034,248078,248077,64,1637,248078,248077,64,11952,248078,248077,64,8329,248078],"opens":[594,598,602,606,610],"closes":[597,601,605,609,613],"positions_sha256":"298130328c97a5589416a450a4b09249363619c05e9d68f96011f6a7b0e9227b","rows":[{"tokens":[3710,369,279,1822,1576,30,248046,198,248045,74455,198,248068,198,248077,64,4331,248078],"zq":12,"u":16,"targets":[[13,64],[14,4331]],"first":64},{"tokens":[3710,369,279,1822,1576,30,248046,198,248045,74455,198,248068,198,248077,64,4034,248078],"zq":12,"u":16,"targets":[[13,64],[14,4034]],"first":64},{"tokens":[3710,369,279,1822,1576,30,248046,198,248045,74455,198,248068,198,248077,64,1637,248078],"zq":12,"u":16,"targets":[[13,64],[14,1637]],"first":64},{"tokens":[3710,369,279,1822,1576,30,248046,198,248045,74455,198,248068,198,248077,64,11952,248078],"zq":12,"u":16,"targets":[[13,64],[14,11952]],"first":64},{"tokens":[3710,369,279,1822,1576,30,248046,198,248045,74455,198,248068,198,248077,64,8329,248078],"zq":12,"u":16,"targets":[[13,64],[14,8329]],"first":64}]},{"id":"q4","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>The scene is indoors.<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-576,248054,760,5918,369,52239,13,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[593],"closes":[594],"positions_sha256":"7e6dc6043b9dd86dd01828448e906bffaf4008274e146e85a429d359daf87f89","rows":[{"tokens":[760,5918,369,52239,13,248046,198,248045,74455,198,248068,198,248077,248078],"zq":11,"u":13,"targets":[],"first":null}]},{"id":"q5","type":"score","keys":["empty","some","busy","packed"],"option_text":["empty","some","busy","packed"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>How cluttered is the scene?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>empty<|/opt|><|opt|>some<|/opt|><|opt|>busy<|/opt|><|opt|>packed<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,3092,248078,248077,14265,248078,248077,75778,248078,248077,48632,248078],"opens":[595,598,601,604],"closes":[597,600,603,606],"positions_sha256":"06ed9d7b99cc7d6a5738864434c9e628df5de389fb052b6aae03bced201ac6b8","rows":[{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,3092,248078],"zq":13,"u":16,"targets":[[14,3092]],"first":3092},{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,14265,248078],"zq":13,"u":16,"targets":[[14,14265]],"first":14265},{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,75778,248078],"zq":13,"u":16,"targets":[[14,75778]],"first":75778},{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,48632,248078],"zq":13,"u":16,"targets":[[14,48632]],"first":48632}]},{"id":"q6","type":"choice","keys":["r0","r1","r2","r3"],"option_text":["region [0,100,180,300]","region [200,100,380,300]","region [400,100,580,300]","region [600,100,780,300]"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Which region is most relevant?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>region [0,100,180,300]<|/opt|><|opt|>region [200,100,380,300]<|/opt|><|opt|>region [400,100,580,300]<|/opt|><|opt|>region [600,100,780,300]<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,22365,5364,369,1379,9468,30,248046,198,248045,74455,198,248068,198,248077,3810,498,15,11,16,15,15,11,16,23,15,11,18,15,15,60,248078,248077,3810,498,17,15,15,11,16,15,15,11,18,23,15,11,18,15,15,60,248078,248077,3810,498,19,15,15,11,16,15,15,11,20,23,15,11,18,15,15,60,248078,248077,3810,498,21,15,15,11,16,15,15,11,22,23,15,11,18,15,15,60,248078],"opens":[594,612,632,652],"closes":[611,631,651,671],"positions_sha256":"0fc96057bbc2470c8a7d05de2ecf3ddf125c3e733099ad5637fd1cdf0845a396","rows":[{"tokens":[22365,5364,369,1379,9468,30,248046,198,248045,74455,198,248068,198,248077,3810,498,15,11,16,15,15,11,16,23,15,11,18,15,15,60,248078],"zq":12,"u":30,"targets":[[13,3810],[14,498],[15,15],[16,11],[17,16],[18,15],[19,15],[20,11],[21,16],[22,23],[23,15],[24,11],[25,18],[26,15],[27,15],[28,60]],"first":3810},{"tokens":[22365,5364,369,1379,9468,30,248046,198,248045,74455,198,248068,198,248077,3810,498,17,15,15,11,16,15,15,11,18,23,15,11,18,15,15,60,248078],"zq":12,"u":32,"targets":[[13,3810],[14,498],[15,17],[16,15],[17,15],[18,11],[19,16],[20,15],[21,15],[22,11],[23,18],[24,23],[25,15],[26,11],[27,18],[28,15],[29,15],[30,60]],"first":3810},{"tokens":[22365,5364,369,1379,9468,30,248046,198,248045,74455,198,248068,198,248077,3810,498,19,15,15,11,16,15,15,11,20,23,15,11,18,15,15,60,248078],"zq":12,"u":32,"targets":[[13,3810],[14,498],[15,19],[16,15],[17,15],[18,11],[19,16],[20,15],[21,15],[22,11],[23,20],[24,23],[25,15],[26,11],[27,18],[28,15],[29,15],[30,60]],"first":3810},{"tokens":[22365,5364,369,1379,9468,30,248046,198,248045,74455,198,248068,198,248077,3810,498,21,15,15,11,16,15,15,11,22,23,15,11,18,15,15,60,248078],"zq":12,"u":32,"targets":[[13,3810],[14,498],[15,21],[16,15],[17,15],[18,11],[19,16],[20,15],[21,15],[22,11],[23,22],[24,23],[25,15],[26,11],[27,18],[28,15],[29,15],[30,60]],"first":3810}]},{"id":"q7","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Something in this image is moving.<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-576,248054,22367,303,411,2099,369,6992,13,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[595],"closes":[596],"positions_sha256":"5dedc312e03d4751587ad21e3b2377a3345fa0f0e88c9b25e5a64ad95b8272f3","rows":[{"tokens":[22367,303,411,2099,369,6992,13,248046,198,248045,74455,198,248068,198,248077,248078],"zq":13,"u":15,"targets":[],"first":null}]},{"id":"q8","type":"score","keys":["safe","check first","dangerous"],"option_text":["safe","check first","dangerous"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>How risky is acting on this screen?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>safe<|/opt|><|opt|>check first<|/opt|><|opt|>dangerous<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,4199,38841,369,14912,383,411,4034,30,248046,198,248045,74455,198,248068,198,248077,18112,248078,248077,1960,1118,248078,248077,30416,757,248078],"opens":[596,599,603],"closes":[598,602,606],"positions_sha256":"06ed9d7b99cc7d6a5738864434c9e628df5de389fb052b6aae03bced201ac6b8","rows":[{"tokens":[4199,38841,369,14912,383,411,4034,30,248046,198,248045,74455,198,248068,198,248077,18112,248078],"zq":14,"u":17,"targets":[[15,18112]],"first":18112},{"tokens":[4199,38841,369,14912,383,411,4034,30,248046,198,248045,74455,198,248068,198,248077,1960,1118,248078],"zq":14,"u":18,"targets":[[15,1960],[16,1118]],"first":1960},{"tokens":[4199,38841,369,14912,383,411,4034,30,248046,198,248045,74455,198,248068,198,248077,30416,757,248078],"zq":14,"u":18,"targets":[[15,30416],[16,757]],"first":30416}]},{"id":"q9","type":"choice","keys":["continue","stop","ask a human"],"option_text":["continue","stop","ask a human"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>What should happen next?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>continue<|/opt|><|opt|>stop<|/opt|><|opt|>ask a human<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,3710,1220,3420,1727,30,248046,198,248045,74455,198,248068,198,248077,9252,248078,248077,9215,248078,248077,1038,264,3611,248078],"opens":[593,596,599],"closes":[595,598,603],"positions_sha256":"4299b60e618ec8452eb3baa88aa2f170cbbeba907b88291888df7602065ef597","rows":[{"tokens":[3710,1220,3420,1727,30,248046,198,248045,74455,198,248068,198,248077,9252,248078],"zq":11,"u":14,"targets":[[12,9252]],"first":9252},{"tokens":[3710,1220,3420,1727,30,248046,198,248045,74455,198,248068,198,248077,9215,248078],"zq":11,"u":14,"targets":[[12,9215]],"first":9215},{"tokens":[3710,1220,3420,1727,30,248046,198,248045,74455,198,248068,198,248077,1038,264,3611,248078],"zq":11,"u":16,"targets":[[12,1038],[13,264],[14,3611]],"first":1038}]},{"id":"q10","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>The task looks finished.<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-576,248054,760,3274,5686,7807,13,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[593],"closes":[594],"positions_sha256":"7e6dc6043b9dd86dd01828448e906bffaf4008274e146e85a429d359daf87f89","rows":[{"tokens":[760,3274,5686,7807,13,248046,198,248045,74455,198,248068,198,248077,248078],"zq":11,"u":13,"targets":[],"first":null}]},{"id":"q11","type":"choice","keys":["white","black","blue","green","red"],"option_text":["white","black","blue","green","red"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Which colour dominates?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>white<|/opt|><|opt|>black<|/opt|><|opt|>blue<|/opt|><|opt|>green<|/opt|><|opt|>red<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,22365,12106,80018,30,248046,198,248045,74455,198,248068,198,248077,5603,248078,248077,11124,248078,248077,11855,248078,248077,12865,248078,248077,1114,248078],"opens":[592,595,598,601,604],"closes":[594,597,600,603,606],"positions_sha256":"06ed9d7b99cc7d6a5738864434c9e628df5de389fb052b6aae03bced201ac6b8","rows":[{"tokens":[22365,12106,80018,30,248046,198,248045,74455,198,248068,198,248077,5603,248078],"zq":10,"u":13,"targets":[[11,5603]],"first":5603},{"tokens":[22365,12106,80018,30,248046,198,248045,74455,198,248068,198,248077,11124,248078],"zq":10,"u":13,"targets":[[11,11124]],"first":11124},{"tokens":[22365,12106,80018,30,248046,198,248045,74455,198,248068,198,248077,11855,248078],"zq":10,"u":13,"targets":[[11,11855]],"first":11855},{"tokens":[22365,12106,80018,30,248046,198,248045,74455,198,248068,198,248077,12865,248078],"zq":10,"u":13,"targets":[[11,12865]],"first":12865},{"tokens":[22365,12106,80018,30,248046,198,248045,74455,198,248068,198,248077,1114,248078],"zq":10,"u":13,"targets":[[11,1114]],"first":1114}]}]},{"name":"jev_forms","request":{"model":"tinnel123/OmniJev","state":{"images":[""]},"questions":{"next":{"type":"choice","instructions":"Which operation comes next?","criteria":{"click":null,"type":"Enter text into the focused field","scroll":""}},"risk":{"type":"score","instructions":"How risky is acting on this screen?","criteria":{"safe":"nothing can go wrong","check first":null,"dangerous":null}},"done":{"type":"noul","instructions":"The task looks finished."}}},"image_size":[640,480],"image_colour":[40,90,160],"image_grid_thw":[1,30,40],"prefix_length":305,"input_tokens":427,"pixel_values_sha256":"1f9530f57100c378360118c45ab535c690d2e12e3019ec71b4da55606c73f23d","pixel_values_shape":[1200,1536],"questions":[{"id":"next","type":"choice","keys":["click","type","scroll"],"option_text":["click","type: Enter text into the focused field","scroll"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Which operation comes next?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>click<|/opt|><|opt|>type: Enter text into the focused field<|/opt|><|opt|>scroll<|/opt|>","token_ids":[248045,846,198,248053,-300,248054,22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,3552,248078,248077,1267,25,10925,1414,1083,279,10419,2002,248078,248077,12242,248078],"opens":[317,320,330],"closes":[319,329,332],"positions_sha256":"42ef974e3f377e8437b9eea997f8cf1cab4e03b88f9d496c97117e45bd53b76f","rows":[{"tokens":[22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,3552,248078],"zq":11,"u":14,"targets":[[12,3552]],"first":3552},{"tokens":[22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,1267,25,10925,1414,1083,279,10419,2002,248078],"zq":11,"u":21,"targets":[[12,1267],[13,25],[14,10925],[15,1414],[16,1083],[17,279],[18,10419],[19,2002]],"first":1267},{"tokens":[22365,5488,3905,1727,30,248046,198,248045,74455,198,248068,198,248077,12242,248078],"zq":11,"u":14,"targets":[[12,12242]],"first":12242}]},{"id":"risk","type":"score","keys":["safe","check first","dangerous"],"option_text":["safe","check first","dangerous"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>How risky is acting on this screen?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>safe<|/opt|><|opt|>check first<|/opt|><|opt|>dangerous<|/opt|>","token_ids":[248045,846,198,248053,-300,248054,4199,38841,369,14912,383,411,4034,30,248046,198,248045,74455,198,248068,198,248077,18112,248078,248077,1960,1118,248078,248077,30416,757,248078],"opens":[320,323,327],"closes":[322,326,330],"positions_sha256":"bf9e318d0eaa940a67be5deba7047908fe8be1ac7c6eab4d4726065b564a8e5a","rows":[{"tokens":[4199,38841,369,14912,383,411,4034,30,248046,198,248045,74455,198,248068,198,248077,18112,248078],"zq":14,"u":17,"targets":[[15,18112]],"first":18112},{"tokens":[4199,38841,369,14912,383,411,4034,30,248046,198,248045,74455,198,248068,198,248077,1960,1118,248078],"zq":14,"u":18,"targets":[[15,1960],[16,1118]],"first":1960},{"tokens":[4199,38841,369,14912,383,411,4034,30,248046,198,248045,74455,198,248068,198,248077,30416,757,248078],"zq":14,"u":18,"targets":[[15,30416],[16,757]],"first":30416}]},{"id":"done","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>The task looks finished.<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-300,248054,760,3274,5686,7807,13,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[317],"closes":[318],"positions_sha256":"3a7564eeef6a28871d009b504ae7d99d0ab70307913ccb3aa42afea57e5818ea","rows":[{"tokens":[760,3274,5686,7807,13,248046,198,248045,74455,198,248068,198,248077,248078],"zq":11,"u":13,"targets":[],"first":null}]}]},{"name":"shared_instruction_start","request":{"model":"tinnel123/OmniJev","state":{"images":[""]},"questions":{"a":{"type":"noul","instructions":"Is the button on the left side blue?"},"b":{"type":"noul","instructions":"Is the button on the left side red?"},"c":{"type":"choice","instructions":"Is the button on the left side small or large?","options":[{"key":"small","text":"small"},{"key":"large","text":"large"}]}}},"image_size":[64,64],"image_colour":[40,90,160],"image_grid_thw":[1,16,16],"prefix_length":76,"input_tokens":126,"pixel_values_sha256":"0d912295f0695ed5021ffdae57dfca9234c38d57a25d811543d88c97b860445b","pixel_values_shape":[256,1536],"questions":[{"id":"a","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Is the button on the left side blue?<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-64,248054,3742,279,3037,383,279,2047,3008,6105,30,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[85],"closes":[86],"positions_sha256":"aca225aa566af5935167d09ad12e38c987bee20dc45ef98b887f36d549252042","rows":[{"tokens":[6105,30,248046,198,248045,74455,198,248068,198,248077,248078],"zq":8,"u":10,"targets":[],"first":null}]},{"id":"b","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Is the button on the left side red?<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-64,248054,3742,279,3037,383,279,2047,3008,2438,30,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[85],"closes":[86],"positions_sha256":"aca225aa566af5935167d09ad12e38c987bee20dc45ef98b887f36d549252042","rows":[{"tokens":[2438,30,248046,198,248045,74455,198,248068,198,248077,248078],"zq":8,"u":10,"targets":[],"first":null}]},{"id":"c","type":"choice","keys":["small","large"],"option_text":["small","large"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Is the button on the left side small or large?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>small<|/opt|><|opt|>large<|/opt|>","token_ids":[248045,846,198,248053,-64,248054,3742,279,3037,383,279,2047,3008,2526,466,3349,30,248046,198,248045,74455,198,248068,198,248077,8732,248078,248077,16275,248078],"opens":[87,90],"closes":[89,92],"positions_sha256":"1532e904676d184289d107ece0eb8e03fe90d7aa1d86188c877fb6036fd54e1a","rows":[{"tokens":[2526,466,3349,30,248046,198,248045,74455,198,248068,198,248077,8732,248078],"zq":10,"u":13,"targets":[[11,8732]],"first":8732},{"tokens":[2526,466,3349,30,248046,198,248045,74455,198,248068,198,248077,16275,248078],"zq":10,"u":13,"targets":[[11,16275]],"first":16275}]}]},{"name":"options_list_forms","request":{"model":"tinnel123/OmniJev","state":{"images":[""]},"questions":{"pick":{"type":"choice","instructions":"Which region shows the error?","options":[{"key":"top","region":{"box":[0,0,1000,120.4]}},{"text":"the dialog in the middle"},{},{"abstain":true},{"key":"long","text":"选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项option text that runs past the two-hundred character cut option text that runs past the two-hundred character cut option text that runs past the two-hundred character cut 🙂end"},{"key":"astral","text":"xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx🙂yyyyyyyyyy"}]},"where":{"type":"noul","instructions":"The cursor is inside this region.","region":{"box":[10.5,20.49,300,400.5]}},"level":{"type":"score","instructions":"How cluttered is the scene?","levels":["empty","some","busy","packed","非常拥挤","🙂"]}}},"image_size":[1920,1080],"image_colour":[40,90,160],"image_grid_thw":[1,36,64],"prefix_length":581,"input_tokens":914,"pixel_values_sha256":"439a2c9ec89d25c0cf7669301bfabd1ebf56e1657790afd2545cad85fac583e3","pixel_values_shape":[2304,1536],"questions":[{"id":"pick","type":"choice","keys":["top","the dialog in the middle","option_2","long","astral"],"option_text":["region [0,0,1000,120]","the dialog in the middle","","选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项option text that runs past the two-hundred character cut option text that runs p","xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx🙂"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Which region shows the error?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>region [0,0,1000,120]<|/opt|><|opt|>the dialog in the middle<|/opt|><|opt|><|/opt|><|opt|>选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项选项option text that runs past the two-hundred character cut option text that runs p<|/opt|><|opt|>xxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxxx🙂<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,22365,5364,4774,279,1412,30,248046,198,248045,74455,198,248068,198,248077,3810,498,15,11,15,11,16,15,15,15,11,16,17,15,60,248078,248077,1719,7027,303,279,5956,248078,248077,248078,248077,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,1979,1414,421,8213,3162,279,1330,2741,6893,3542,3799,2904,1414,421,8213,280,248078,248077,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,17962,23295,169171,248078],"opens":[594,611,618,620,698],"closes":[610,617,619,697,726],"positions_sha256":"23ad2eb05b2e4aa910a57b5431677fc241c58d855fde56b4169b382abac36a5a","rows":[{"tokens":[22365,5364,4774,279,1412,30,248046,198,248045,74455,198,248068,198,248077,3810,498,15,11,15,11,16,15,15,15,11,16,17,15,60,248078],"zq":12,"u":29,"targets":[[13,3810],[14,498],[15,15],[16,11],[17,15],[18,11],[19,16],[20,15],[21,15],[22,15],[23,11],[24,16],[25,17],[26,15],[27,60]],"first":3810},{"tokens":[22365,5364,4774,279,1412,30,248046,198,248045,74455,198,248068,198,248077,1719,7027,303,279,5956,248078],"zq":12,"u":19,"targets":[[13,1719],[14,7027],[15,303],[16,279],[17,5956]],"first":1719},{"tokens":[22365,5364,4774,279,1412,30,248046,198,248045,74455,198,248068,198,248077,248078],"zq":12,"u":14,"targets":[],"first":null},{"tokens":[22365,5364,4774,279,1412,30,248046,198,248045,74455,198,248068,198,248077,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,103982,1979,1414,421,8213,3162,279,1330,2741,6893,3542,3799,2904,1414,421,8213,280,248078],"zq":12,"u":90,"targets":[[13,103982],[14,103982],[15,103982],[16,103982],[17,103982],[18,103982],[19,103982],[20,103982],[21,103982],[22,103982],[23,103982],[24,103982],[25,103982],[26,103982],[27,103982],[28,103982],[29,103982],[30,103982],[31,103982],[32,103982],[33,103982],[34,103982],[35,103982],[36,103982],[37,103982],[38,103982],[39,103982],[40,103982],[41,103982],[42,103982],[43,103982],[44,103982],[45,103982],[46,103982],[47,103982],[48,103982],[49,103982],[50,103982],[51,103982],[52,103982],[53,103982],[54,103982],[55,103982],[56,103982],[57,103982],[58,103982],[59,103982],[60,103982],[61,103982],[62,103982],[63,103982],[64,103982],[65,103982],[66,103982],[67,103982],[68,103982],[69,103982],[70,103982],[71,103982],[72,103982],[73,1979],[74,1414],[75,421],[76,8213],[77,3162],[78,279],[79,1330],[80,2741],[81,6893],[82,3542],[83,3799],[84,2904],[85,1414],[86,421],[87,8213],[88,280]],"first":103982},{"tokens":[22365,5364,4774,279,1412,30,248046,198,248045,74455,198,248068,198,248077,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,42638,17962,23295,169171,248078],"zq":12,"u":41,"targets":[[13,42638],[14,42638],[15,42638],[16,42638],[17,42638],[18,42638],[19,42638],[20,42638],[21,42638],[22,42638],[23,42638],[24,42638],[25,42638],[26,42638],[27,42638],[28,42638],[29,42638],[30,42638],[31,42638],[32,42638],[33,42638],[34,42638],[35,42638],[36,42638],[37,17962],[38,23295],[39,169171]],"first":42638}]},{"id":"where","type":"noul","keys":["yes"],"option_text":["region [10,20,300,400]"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>The cursor is inside this region.<|im_end|>\n<|im_start|>assistant\n\n<|opt|>region [10,20,300,400]<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,760,7872,369,4613,411,5364,13,248046,198,248045,74455,198,248068,198,248077,3810,498,16,15,11,17,15,11,18,15,15,11,19,15,15,60,248078],"opens":[595],"closes":[612],"positions_sha256":"d8939649b808d686f3b687b135293a170519cf6e768f9f92602c98c8da2c7e1f","rows":[{"tokens":[760,7872,369,4613,411,5364,13,248046,198,248045,74455,198,248068,198,248077,3810,498,16,15,11,17,15,11,18,15,15,11,19,15,15,60,248078],"zq":13,"u":31,"targets":[[14,3810],[15,498],[16,16],[17,15],[18,11],[19,17],[20,15],[21,11],[22,18],[23,15],[24,15],[25,11],[26,19],[27,15],[28,15],[29,60]],"first":3810}]},{"id":"level","type":"score","keys":["empty","some","busy","packed","非常拥挤","🙂"],"option_text":["empty","some","busy","packed","非常拥挤","🙂"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>How cluttered is the scene?<|im_end|>\n<|im_start|>assistant\n\n<|opt|>empty<|/opt|><|opt|>some<|/opt|><|opt|>busy<|/opt|><|opt|>packed<|/opt|><|opt|>非常拥挤<|/opt|><|opt|>🙂<|/opt|>","token_ids":[248045,846,198,248053,-576,248054,4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,3092,248078,248077,14265,248078,248077,75778,248078,248077,48632,248078,248077,96762,117114,248078,248077,169171,248078],"opens":[595,598,601,604,607,611],"closes":[597,600,603,606,610,613],"positions_sha256":"298130328c97a5589416a450a4b09249363619c05e9d68f96011f6a7b0e9227b","rows":[{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,3092,248078],"zq":13,"u":16,"targets":[[14,3092]],"first":3092},{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,14265,248078],"zq":13,"u":16,"targets":[[14,14265]],"first":14265},{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,75778,248078],"zq":13,"u":16,"targets":[[14,75778]],"first":75778},{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,48632,248078],"zq":13,"u":16,"targets":[[14,48632]],"first":48632},{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,96762,117114,248078],"zq":13,"u":17,"targets":[[14,96762],[15,117114]],"first":96762},{"tokens":[4199,52010,290,369,279,5918,30,248046,198,248045,74455,198,248068,198,248077,169171,248078],"zq":13,"u":16,"targets":[[14,169171]],"first":169171}]}]},{"name":"unicode_and_special_text","request":{"model":"tinnel123/OmniJev","state":{"images":[""]},"questions":{"zh":{"type":"choice","instructions":"屏幕上显示的是什么？<|im_end|> literal","criteria":{"登录页面":null,"设置菜单":"带有齿轮图标","错误提示":null}},"emoji":{"type":"noul","instructions":"There is a 🚀 icon in the toolbar."}}},"image_size":[800,1200],"image_colour":[40,90,160],"image_grid_thw":[1,58,38],"prefix_length":556,"input_tokens":632,"pixel_values_sha256":"47565c2d667ddc598410289cbd084547b7240c9f863998032e55b94ad4b19988","pixel_values_shape":[2204,1536],"questions":[{"id":"zh","type":"choice","keys":["登录页面","设置菜单","错误提示"],"option_text":["登录页面","设置菜单: 带有齿轮图标","错误提示"],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>屏幕上显示的是什么？<|im_end|> literal<|im_end|>\n<|im_start|>assistant\n\n<|opt|>登录页面<|/opt|><|opt|>设置菜单: 带有齿轮图标<|/opt|><|opt|>错误提示<|/opt|>","token_ids":[248045,846,198,248053,-551,248054,122830,97628,117015,10992,248046,22419,248046,198,248045,74455,198,248068,198,248077,99132,108586,248078,248077,98621,104643,25,220,109163,110751,111782,248078,248077,99147,99162,248078],"opens":[569,573,582],"closes":[572,581,585],"positions_sha256":"efc991675c8b8461db1aa57cffd20b161f623947aab4d315e72b5072b9d39d4c","rows":[{"tokens":[122830,97628,117015,10992,248046,22419,248046,198,248045,74455,198,248068,198,248077,99132,108586,248078],"zq":12,"u":16,"targets":[[13,99132],[14,108586]],"first":99132},{"tokens":[122830,97628,117015,10992,248046,22419,248046,198,248045,74455,198,248068,198,248077,98621,104643,25,220,109163,110751,111782,248078],"zq":12,"u":21,"targets":[[13,98621],[14,104643],[15,25],[16,220],[17,109163],[18,110751],[19,111782]],"first":98621},{"tokens":[122830,97628,117015,10992,248046,22419,248046,198,248045,74455,198,248068,198,248077,99147,99162,248078],"zq":12,"u":16,"targets":[[13,99147],[14,99162]],"first":99147}]},{"id":"emoji","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>There is a 🚀 icon in the toolbar.<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-551,248054,3733,369,264,10838,248,222,4454,303,279,25112,13,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[574],"closes":[575],"positions_sha256":"97c6b8b8b0e9b793334be4e386aa7c4d192f36a79f955da3ce227862060da01e","rows":[{"tokens":[3733,369,264,10838,248,222,4454,303,279,25112,13,248046,198,248045,74455,198,248068,198,248077,248078],"zq":17,"u":19,"targets":[],"first":null}]}]},{"name":"whitespace","request":{"model":"tinnel123/OmniJev","state":{"images":[""]},"questions":{"dark":{"type":"noul","instructions":" Is it dark? \n\u001f "},"pick":{"type":"choice","instructions":"Pick one.　","criteria":{"a ":null," b":"rubric "}}}},"image_size":[300,300],"image_colour":[40,90,160],"image_grid_thw":[1,18,18],"prefix_length":86,"input_tokens":131,"pixel_values_sha256":"eb6b06987ef0b50bb69086e27a6ea293ab13856e41eb581b7fa174d53141a23f","pixel_values_shape":[324,1536],"questions":[{"id":"dark","type":"noul","keys":["yes"],"option_text":[""],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|> Is it dark?<|im_end|>\n<|im_start|>assistant\n\n<|opt|><|/opt|>","token_ids":[248045,846,198,248053,-81,248054,220,2091,424,6121,30,248046,198,248045,74455,198,248068,198,248077,248078],"opens":[98],"closes":[99],"positions_sha256":"97d7607c59adc8017b8a9238531d941f6463869e7a28f6aa9cccbc2416c59748","rows":[{"tokens":[220,2091,424,6121,30,248046,198,248045,74455,198,248068,198,248077,248078],"zq":11,"u":13,"targets":[],"first":null}]},{"id":"pick","type":"choice","keys":["a "," b"],"option_text":["a "," b: rubric "],"prompt":"<|im_start|>user\n<|vision_start|><|image_pad|><|vision_end|>Pick one.<|im_end|>\n<|im_start|>assistant\n\n<|opt|>a <|/opt|><|opt|> b: rubric <|/opt|>","token_ids":[248045,846,198,248053,-81,248054,35728,799,13,248046,198,248045,74455,198,248068,198,248077,64,220,248078,248077,292,25,9967,2147,220,248078],"opens":[96,100],"closes":[99,106],"positions_sha256":"b354dd55ac13da11f75c573824a779640d5b459f2cc96eb498ca4fa1ced1e991","rows":[{"tokens":[35728,799,13,248046,198,248045,74455,198,248068,198,248077,64,220,248078],"zq":9,"u":13,"targets":[[10,64],[11,220]],"first":64},{"tokens":[35728,799,13,248046,198,248045,74455,198,248068,198,248077,292,25,9967,2147,220,248078],"zq":9,"u":16,"targets":[[10,292],[11,25],[12,9967],[13,2147],[14,220]],"first":292}]}]}]} \ No newline at end of file diff --git a/tests/omnijev/finish.rs b/tests/omnijev/finish.rs new file mode 100644 index 00000000..b8f03785 --- /dev/null +++ b/tests/omnijev/finish.rs @@ -0,0 +1,105 @@ +//! Response finishing against `MSO1._finish` (tests/omnijev/data/finish.json). +mod common; + +use omni_omnijev_native::contract::{self, Calibration, Kind, Question}; +use serde_json::Value; + +fn calibration(fixture: &Value) -> Calibration { + let t = &fixture["temperatures"]; + Calibration { + temperatures: ["noul", "choice", "score"].map(|k| t[k].as_f64().unwrap()), + noul_bias: fixture["biases"]["noul"].as_f64().unwrap(), + } +} + +fn question(kind: Kind, keys: &[String]) -> Question { + Question { + id: "q".into(), + kind, + instructions: "x".into(), + options: keys.to_vec(), + keys: keys.to_vec(), + } +} + +/// Values within one unit of the fourth decimal, and the same answer key unless the +/// reference's top two are that close. +fn near(actual: &Value, expected: &Value, case: &Value) { + let probs = &expected["probabilities"]; + let top: Vec = { + let mut p: Vec = probs + .as_object() + .map(|m| m.values().map(|v| v.as_f64().unwrap()).collect()) + .unwrap_or_default(); + p.sort_by(|a, b| b.partial_cmp(a).unwrap()); + p + }; + for (key, e) in expected.as_object().unwrap() { + let a = &actual[key]; + match e { + Value::Number(_) => { + assert!( + (a.as_f64().unwrap() - e.as_f64().unwrap()).abs() <= 1.0001e-4, + "{key}: {case}" + ) + } + Value::Object(m) => { + for (k, v) in m { + assert!( + (a[k].as_f64().unwrap() - v.as_f64().unwrap()).abs() <= 1.0001e-4, + "{k}: {case}" + ); + } + } + _ if top.len() > 1 && top[0] - top[1] <= 2e-4 => {} + _ => assert_eq!(a, e, "{key}: {case}"), + } + } +} + +#[test] +fn answers_match_reference() { + let fixture = common::fixture("finish.json"); + let calibration = calibration(&fixture); + for case in fixture["cases"].as_array().unwrap() { + let kind = match case["type"].as_str().unwrap() { + "noul" => Kind::Noul, + "choice" => Kind::Choice, + _ => Kind::Score, + }; + let keys: Vec = case["keys"] + .as_array() + .unwrap() + .iter() + .map(|k| k.as_str().unwrap().into()) + .collect(); + let mu: Vec = case["mu"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap() as f32) + .collect(); + let answer = contract::answer(&question(kind, &keys), &mu, &calibration, 0.0).unwrap(); + // Up to four options the float32 sums add in the same order as PyTorch's; with + // more, PyTorch's vectorized order can move the fourth decimal by one. + if keys.len() <= 4 { + assert_eq!(answer, case["answer"], "{case}"); + } else { + near(&answer, &case["answer"], case); + } + } +} + +#[test] +fn round4_matches_python() { + // round(x, 4) on exactly representable ties goes to the even digit. + for (x, expected) in [ + (0.03125, 0.0312), + (0.09375, 0.0938), + (0.5, 0.5), + (0.99995, 1.0), + (1e-9, 0.0), + ] { + assert_eq!(contract::round4(x), expected, "{x}"); + } +} diff --git a/tests/omnijev/generate_fixtures.py b/tests/omnijev/generate_fixtures.py new file mode 100644 index 00000000..66431901 --- /dev/null +++ b/tests/omnijev/generate_fixtures.py @@ -0,0 +1,333 @@ +#!/usr/bin/env python3 +"""Generate the OmniJev CPU fixtures from the pinned reference. + +Runs the reference's own preparation, branch split, rotary positions, heads and +response finishing without loading the backbone, so it needs no GPU: + + CUDA_VISIBLE_DEVICES= python tests/omnijev/generate_fixtures.py \ + --reference \ + --checkpoint \ + --base \ + --out tests/omnijev/data + +Environment: the reference environment in recipe/omnijev/README.md (generated with +Python 3.10.12, torch 2.14.0, torchvision 0.29.0, transformers 5.17.0, peft 0.19.1). +""" +import argparse +import copy +import hashlib +import io +import json +import os +import sys +import tempfile +import types + +import torch + +REFERENCE_COMMIT = "14dbec4f71e194852c8d7b88ab36ef639493f400" +CHECKPOINT_REVISION = "ffe5f436eaf22e20e2f041f8e74e121fd057a6cb" +BASE_REVISION = "851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a" +MAX_PIXELS = 768 * 28 * 28 # MSO1's default, which overrides the processor's 401,408 + + +def png(width, height, colour): + from PIL import Image + buf = io.BytesIO() + Image.new("RGB", (width, height), colour).save(buf, format="PNG") + return buf.getvalue() + + +IMAGE_PAD = 248056 +COLOUR = (40, 90, 160) # every test image is one solid colour, so any PNG encoder gives the same pixels + + +def collapse(ids): + """token ids with the run of image placeholders replaced by its negated length""" + out, n = [], 0 + for t in ids: + if t == IMAGE_PAD: + n += 1 + continue + if n: + out.append(-n) + n = 0 + out.append(t) + return out + + +def sha256_i64(rows): + """SHA-256 of a [3, T] position array as little-endian int64, row-major""" + import struct + return hashlib.sha256(b"".join(struct.pack(" literal", + "criteria": {"登录页面": None, "设置菜单": "带有齿轮图标", "错误提示": None}}, + "emoji": {"type": "noul", "instructions": "There is a 🚀 icon in the toolbar."}, + }), + ("whitespace", (300, 300), { + "dark": {"type": "noul", "instructions": " Is it dark? \n\x1f "}, + "pick": {"type": "choice", "instructions": "Pick one.\u3000", + "criteria": {"a ": None, " b": "rubric "}}, + }), +] + + +def speed_bench_questions(reference): + sys.path.insert(0, reference) + from bench.speed_bench import Q + return {k: copy.deepcopy(v) for k, v in Q} + + +def make_reference(reference, checkpoint): + """An MSO1 with the processor and collator only: no backbone, no device.""" + sys.path.insert(0, reference) + from transformers import AutoProcessor + import mso.records as T + from mso.infer import MSO1 + obj = MSO1.__new__(MSO1) + obj.proc = AutoProcessor.from_pretrained(checkpoint, max_pixels=MAX_PIXELS) + T.add_option_tokens(types.SimpleNamespace( + get_input_embeddings=lambda: types.SimpleNamespace(weight=torch.zeros(len(obj.proc.tokenizer), 1)), + resize_token_embeddings=lambda n: None), obj.proc) + obj.coll = T.Collator(obj.proc, MAX_PIXELS) + obj.dev = torch.device("cpu") + obj.pad_tok = "<|image_pad|>" + obj._cache, obj._cache_n = {}, 4 + obj.open_ids = obj.coll.tok(T.OPT_OPEN, add_special_tokens=False)["input_ids"] + obj.close_ids = obj.coll.tok(T.OPT_CLOSE, add_special_tokens=False)["input_ids"] + assert obj.open_ids == [248077] and obj.close_ids == [248078], (obj.open_ids, obj.close_ids) + assert len(obj.proc.tokenizer) == 248079 + return obj + + +def rope_owner(base): + """Qwen3.5's get_rope_index on a model built on the meta device (no weights).""" + from transformers import AutoConfig, AutoModelForImageTextToText + import mso.records as T + cfg = AutoConfig.from_pretrained(base) + with torch.device("meta"): + model = AutoModelForImageTextToText.from_config(cfg) + return types.SimpleNamespace(_rope=T._find_rope_owner(model)) + + +def processing_case(obj, rope, name, size, questions): + import mso.branch as BR + import mso.records as T + image = png(size[0], size[1], COLOUR) + with tempfile.TemporaryDirectory() as tmp: + path = os.path.join(tmp, "state.png") + with open(path, "wb") as f: + f.write(image) + obj._cache, obj._vc_key = {}, None + qids = list(questions) + keys_l, opts_l, ids_l, spans_l, pix = obj._encode_many(path, None, questions, qids) + prompts = [obj._prompt(questions[q], opts)[0] for q, opts in zip(qids, opts_l)] + # ask_branch's split + L, m = 0, min(len(x) for x in ids_l) + while L < m and all(x[L] == ids_l[0][L] for x in ids_l): + L += 1 + L = max(1, min([L] + [sp[0][0] - 1 for sp in spans_l if sp[0]])) + qrows = [BR.rows_from_ids(row, L, sp[0], sp[1], k=len(opts)) for row, sp, opts in zip(ids_l, spans_l, opts_l)] + usage = int(L + sum(len(r) for qr in qrows for r in qr)) + + def positions(ids): + enc = {"input_ids": torch.tensor(ids)[None], "image_grid_thw": pix["image_grid_thw"], + "attention_mask": torch.ones(1, len(ids), dtype=torch.long)} + return T.MSO.rope_positions(rope, enc, None, False)[:, 0, :] + + out_q = [] + for qid, keys, opts, ids, (opens, closes), prompt, rows in zip(qids, keys_l, opts_l, ids_l, spans_l, prompts, qrows): + pos = positions(ids) + # branch semantics: prefix positions from the prefix alone, then text positions from max + 1 + ppos = positions(ids[:L]) + assert torch.equal(ppos, pos[:, :L]), name + p0 = int(ppos.max()) + 1 + row_readouts = [] + for r in rows: + o, c, inner = BR._marks(r, obj.open_ids, obj.close_ids) + row_readouts.append({"tokens": r, "zq": max(0, o - 1), "u": c, + "targets": [[t - 1, r[t]] for t in inner], + "first": r[inner[0]] if inner else None}) + # the row's positions in a full sequence equal p0 + offset + assert torch.equal(positions(list(ids[:L]) + r)[:, L:], torch.arange(p0, p0 + len(r))[None].expand(3, -1)) + assert prompt.count("<|image_pad|>") == 1 # the processor expands it to the grid's token count + out_q.append({"id": qid, "type": questions[qid]["type"], "keys": keys, + "option_text": [T.render_option(o) for o in opts], + "prompt": prompt, "token_ids": collapse(ids), + "opens": opens, "closes": closes, "positions_sha256": sha256_i64(pos.tolist()), + "rows": row_readouts}) + request = {"model": "tinnel123/OmniJev", "state": {"images": [""]}, "questions": questions} + pv = pix["pixel_values"].float().numpy() # for the vision step's pixel check + return {"name": name, "request": request, "image_size": list(size), "image_colour": list(COLOUR), + "image_grid_thw": pix["image_grid_thw"][0].tolist(), "prefix_length": L, "input_tokens": usage, + "pixel_values_sha256": hashlib.sha256(pv.tobytes()).hexdigest(), "pixel_values_shape": list(pv.shape), + "questions": out_q} + + +def head_cases(checkpoint): + from mso.head import OptionScorer, OrdinalScoreHead + head = OptionScorer(2560, d_hidden=1024, norm="softmax") + head.load_state_dict(torch.load(os.path.join(checkpoint, "head.pt"), map_location="cpu"), strict=True) + ordh = OrdinalScoreHead(2560, d_hidden=512) + ordh.load_state_dict(torch.load(os.path.join(checkpoint, "ord.pt"), map_location="cpu"), strict=True) + head, ordh = head.float().eval(), ordh.float().eval() + cases = [] + for seed, (name, type_id, k) in enumerate((("noul", 0, 1), ("choice2", 1, 2), ("choice5", 1, 5), ("score4", 2, 4))): + u = sequence(3 * seed, k * 2560).reshape(k, 2560) + zq = sequence(3 * seed + 1, 2560) + feats = torch.zeros(k, 6) if type_id == 0 else sequence(3 * seed + 2, k * 6).reshape(k, 6) + with torch.no_grad(): + logits = head.logits(u, zq, type_id, feats) + mu = head(u, zq, type_id, feats) + ordinal = ordh(zq, u) + cases.append({"name": name, "seeds": [3 * seed, 3 * seed + 1, 3 * seed + 2], "type_id": type_id, "options": k, + "logits": logits.tolist(), "mu": mu.tolist(), "ordinal": ordinal.tolist()}) + return cases + + +def feature_cases(): + """branch_forward's six LM features, from a stub backbone and a small lm_head.""" + import mso.branch as BR + g = torch.Generator().manual_seed(85) + V, D = 300, 16 + lm_head = torch.nn.Linear(D, V, bias=False) + with torch.no_grad(): + lm_head.weight.copy_(torch.randn(V, D, generator=g)) + open_ids, close_ids = [298], [299] + head = [5, 6, 7] + rows = [head + [298, 11, 12, 13, 299], head + [298, 14, 299], head + [298, 299], + [8, 9] + [298, 20, 21, 22, 23, 24, 25, 299], [8, 9] + [298, 299]] + groups = [[0, 1, 2], [3], [4]] + hidden = {i: torch.randn(len(r), D, generator=g) * 3 for i, r in enumerate(rows)} + + class Cache: + layers = [] + + class Backbone: + def __call__(self, input_ids, **kw): + if "past_key_values" not in kw: + return types.SimpleNamespace(past_key_values=Cache()) + n = input_ids.shape[1] + h = torch.zeros(input_ids.shape[0], n, D) + for b in range(input_ids.shape[0]): + key = next(i for i, r in enumerate(rows) if input_ids[b, :len(r)].tolist() == r) + h[b, :len(rows[key])] = hidden[key] + return types.SimpleNamespace(hidden_states=[h]) + + penc = {"input_ids": torch.tensor([[1, 2, 3]])} + with torch.no_grad(): + out = BR.branch_forward(Backbone(), penc, rows, torch.device("cpu"), open_ids, close_ids, lm_head=lm_head, + groups=groups, chunk=2, pad_id=0, mrope_axes=3) + cases = [] + for i, r in enumerate(rows): + o, c, inner = BR._marks(r, open_ids, close_ids) + lp = torch.log_softmax(lm_head(hidden[i]).float(), dim=-1) + picked = [float(lp[t - 1, r[t]]) for t in inner] + first = float(lp[o - 1, r[inner[0]]]) if inner else None + cases.append({"row": i, "picked": picked, "first": first, "feats": out["feats"][i].tolist()}) + return {"groups": groups, "rows": cases} + + +def finish_cases(reference, checkpoint): + from mso.infer import MSO1 + meta = json.load(open(os.path.join(checkpoint, "head_meta.json"))) + stub = MSO1.__new__(MSO1) + stub.temp = float(meta.get("temperature", 1.0)) + stub.temps = {k: float(meta["temperatures"][k]) for k in ("noul", "choice", "score")} + stub.biases = {"noul": float(meta["biases"]["noul"])} + cases = [] + inputs = [ + ("noul", ["yes"], [0.5]), ("noul", ["yes"], [0.4867]), ("noul", ["yes"], [1e-12]), ("noul", ["yes"], [0.999999]), + ("choice", ["a", "b", "c"], [0.2, 0.3, 0.1]), ("choice", ["a", "b"], [0.45, 0.45]), + ("choice", ["a", "b", "c", "d"], [0.0001, 0.9, 0.05, 0.0]), ("choice", ["only"], [0.7]), + ("score", ["low", "mid", "high"], [0.2, 0.5, 0.3]), ("score", ["0", "1"], [1e-8, 1.0]), + ("score", ["a", "b", "c", "d"], [0.25, 0.25, 0.25, 0.25]), + ] + g = torch.Generator().manual_seed(61) + for k in (5, 12, 32): + logits = torch.randn(k + 1, generator=g) * 2 + inputs.append(("choice", [f"o{i}" for i in range(k)], torch.softmax(logits, 0)[:k].tolist())) + levels = torch.rand(k, generator=g) + 1e-3 + inputs.append(("score", [f"l{i}" for i in range(k)], (levels / levels.sum()).tolist())) + inputs.append(("choice", ["a", "b", "c", "d", "e"], [0.3, 0.3, 0.2, 0.15, 0.1])) # sums above one: invalid + for qtype, keys, mu in inputs: + q = {"type": qtype, "instructions": "x"} + out = MSO1._finish(stub, q, keys, [{"text": k} for k in keys], torch.tensor(mu), 0.0) + cases.append({"type": qtype, "keys": keys, "mu": torch.tensor(mu).tolist(), "answer": out}) + return {"temperatures": stub.temps, "biases": stub.biases, "cases": cases} + + +def main(): + ap = argparse.ArgumentParser() + ap.add_argument("--reference", required=True) + ap.add_argument("--checkpoint", required=True) + ap.add_argument("--base", required=True) + ap.add_argument("--out", required=True) + a = ap.parse_args() + import transformers + obj = make_reference(a.reference, a.checkpoint) + rope = rope_owner(a.base) + cases = [] + for name, size, questions in CASES: + qs = speed_bench_questions(a.reference) if questions is None else copy.deepcopy(questions) + cases.append(processing_case(obj, rope, name, size, qs)) + provenance = {"python": sys.version.split()[0], "reference_commit": REFERENCE_COMMIT, "checkpoint_revision": CHECKPOINT_REVISION, + "base_revision": BASE_REVISION, "max_pixels": MAX_PIXELS, "torch": torch.__version__, + "transformers": transformers.__version__, + "command": "CUDA_VISIBLE_DEVICES= python tests/omnijev/generate_fixtures.py"} + os.makedirs(a.out, exist_ok=True) + for fname, payload in (("processing.json", {"provenance": provenance, "cases": cases}), + ("heads.json", {"provenance": provenance, "cases": head_cases(a.checkpoint)}), + ("features.json", {"provenance": provenance, **feature_cases()}), + ("finish.json", {"provenance": provenance, **finish_cases(a.reference, a.checkpoint)})): + with open(os.path.join(a.out, fname), "w") as f: + json.dump(payload, f, ensure_ascii=False, separators=(",", ":")) + print("wrote", fname, os.path.getsize(os.path.join(a.out, fname)), "bytes") + + +if __name__ == "__main__": + main() diff --git a/tests/omnijev/heads.rs b/tests/omnijev/heads.rs new file mode 100644 index 00000000..1ef169d9 --- /dev/null +++ b/tests/omnijev/heads.rs @@ -0,0 +1,116 @@ +//! Heads and LM features against the pinned reference (tests/omnijev/data). +mod common; + +use omni_omnijev_native::heads::{self, Heads}; + +/// The fixture generator's inputs in [-2, 2): integer arithmetic, then one rounding. +fn sequence(seed: u64, n: usize) -> Vec { + (0..n as u64) + .map(|i| { + let x = (i * 2_654_435_761 + seed * 40_503) % (1 << 32); + (x as f64 / 4_294_967_296.0 * 4.0 - 2.0) as f32 + }) + .collect() +} + +fn close(actual: &[f32], expected: &serde_json::Value, tolerance: f32, what: &str) { + let expected: Vec = expected + .as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap() as f32) + .collect(); + assert_eq!(actual.len(), expected.len(), "{what}"); + for (a, e) in actual.iter().zip(&expected) { + assert!( + (a - e).abs() <= tolerance * e.abs().max(1.0), + "{what}: {actual:?} vs {expected:?}" + ); + } +} + +#[test] +fn lm_features_match_reference() { + let fixture = common::fixture("features.json"); + let rows = fixture["rows"].as_array().unwrap(); + for group in fixture["groups"].as_array().unwrap() { + let members: Vec = group + .as_array() + .unwrap() + .iter() + .map(|v| v.as_u64().unwrap() as usize) + .collect(); + let picked: Vec> = members + .iter() + .map(|&i| { + rows[i]["picked"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_f64().unwrap() as f32) + .collect() + }) + .collect(); + let first: Vec> = members + .iter() + .map(|&i| rows[i]["first"].as_f64().map(|v| v as f32)) + .collect(); + let features = heads::lm_features(&picked, &first); + for (f, &i) in features.iter().zip(&members) { + close(f, &rows[i]["feats"], 1e-6, &format!("row {i}")); + } + } +} + +/// Needs an export: `OMNIJEV_EXPORT=`. +#[test] +#[ignore] +fn heads_match_reference() { + let dir = std::env::var("OMNIJEV_EXPORT").expect("set OMNIJEV_EXPORT"); + let heads = Heads::load(&std::path::Path::new(&dir).join("heads.safetensors")).unwrap(); + let fixture = common::fixture("heads.json"); + for case in fixture["cases"].as_array().unwrap() { + let name = case["name"].as_str().unwrap(); + let seeds: Vec = case["seeds"] + .as_array() + .unwrap() + .iter() + .map(|v| v.as_u64().unwrap()) + .collect(); + let k = case["options"].as_u64().unwrap() as usize; + let type_id = case["type_id"].as_u64().unwrap() as usize; + let u: Vec> = sequence(seeds[0], k * heads::HIDDEN) + .chunks(heads::HIDDEN) + .map(<[f32]>::to_vec) + .collect(); + let zq = sequence(seeds[1], heads::HIDDEN); + let features: Vec<[f32; 6]> = if type_id == 0 { + vec![[0.0; 6]; k] + } else { + sequence(seeds[2], k * 6) + .chunks(6) + .map(|c| c.try_into().unwrap()) + .collect() + }; + close( + &heads.option_logits(&u, &zq, type_id, &features).unwrap(), + &case["logits"], + 1e-5, + name, + ); + close( + &heads + .option_probabilities(&u, &zq, type_id, &features) + .unwrap(), + &case["mu"], + 1e-5, + name, + ); + close( + &heads.ordinal(&zq, &u).unwrap(), + &case["ordinal"], + 1e-5, + name, + ); + } +} diff --git a/tests/omnijev/processing.rs b/tests/omnijev/processing.rs new file mode 100644 index 00000000..46d37e97 --- /dev/null +++ b/tests/omnijev/processing.rs @@ -0,0 +1,117 @@ +//! Preparation against fixtures from the pinned reference (tests/omnijev/generate_fixtures.py). +mod common; + +use omni_omnijev_native::contract; +use omni_omnijev_native::processing::{self, Processor}; +use sha2::{Digest, Sha256}; + +fn positions_sha256(positions: &[Vec; 3]) -> String { + let mut hash = Sha256::new(); + for axis in positions { + for v in axis { + hash.update(v.to_le_bytes()); + } + } + format!("{:x}", hash.finalize()) +} + +#[test] +fn questions_prompts_and_layout_match_reference() { + let fixture = common::fixture("processing.json"); + for case in fixture["cases"].as_array().unwrap() { + let name = case["name"].as_str().unwrap(); + let request = contract::compile(&common::request(case)).unwrap(); + let grid = processing::image_grid(request.image.width, request.image.height); + assert_eq!(serde_json::json!(grid), case["image_grid_thw"], "{name}"); + let expected = case["questions"].as_array().unwrap(); + assert_eq!(request.questions.len(), expected.len(), "{name}"); + let mut sequences = Vec::new(); + for (q, e) in request.questions.iter().zip(expected) { + assert_eq!(q.id, e["id"].as_str().unwrap(), "{name}"); + assert_eq!(serde_json::json!(q.keys), e["keys"], "{name} {}", q.id); + assert_eq!( + serde_json::json!(q.options), + e["option_text"], + "{name} {}", + q.id + ); + assert_eq!( + processing::prompt(q), + e["prompt"].as_str().unwrap(), + "{name} {}", + q.id + ); + let ids = common::token_ids(&e["token_ids"]); + let (opens, closes) = processing::spans(&ids); + assert_eq!(serde_json::json!(opens), e["opens"], "{name} {}", q.id); + assert_eq!(serde_json::json!(closes), e["closes"], "{name} {}", q.id); + let positions = processing::image_positions(&ids, grid).unwrap(); + assert_eq!( + positions_sha256(&positions), + e["positions_sha256"].as_str().unwrap(), + "{name} {}", + q.id + ); + sequences.push(ids); + } + let (layout, rows) = processing::layout(&sequences, grid).unwrap(); + assert_eq!( + layout.prefix as u64, + case["prefix_length"].as_u64().unwrap(), + "{name}" + ); + assert_eq!( + layout.processed_tokens as u64, + case["input_tokens"].as_u64().unwrap(), + "{name}" + ); + for (question_rows, e) in rows.iter().zip(expected) { + let e = e["rows"].as_array().unwrap(); + assert_eq!(question_rows.len(), e.len(), "{name}"); + for (row, e) in question_rows.iter().zip(e) { + assert_eq!(serde_json::json!(row.tokens), e["tokens"], "{name}"); + assert_eq!(row.zq as u64, e["zq"].as_u64().unwrap(), "{name}"); + assert_eq!(row.u as u64, e["u"].as_u64().unwrap(), "{name}"); + let targets: Vec<[u64; 2]> = row + .targets + .iter() + .map(|&(p, t)| [p as u64, t as u64]) + .collect(); + assert_eq!(serde_json::json!(targets), e["targets"], "{name}"); + assert_eq!(serde_json::json!(row.first), e["first"], "{name}"); + } + } + // Rows continue the prefix: the first row position is the prefix's largest plus one. + let positions = processing::image_positions(&sequences[0], grid).unwrap(); + assert_eq!(layout.next_position, positions[0][layout.prefix], "{name}"); + } +} + +/// Needs the checkpoint's tokenizer: `OMNIJEV_CHECKPOINT=`. +#[test] +#[ignore] +fn tokenization_matches_reference() { + let dir = std::env::var("OMNIJEV_CHECKPOINT").expect("set OMNIJEV_CHECKPOINT"); + let processor = Processor::load(std::path::Path::new(&dir)).unwrap(); + let fixture = common::fixture("processing.json"); + for case in fixture["cases"].as_array().unwrap() { + let raw = common::request(case); + let prepared = processor.prepare(&raw).unwrap(); + for (q, e) in prepared + .inputs + .iter() + .zip(case["questions"].as_array().unwrap()) + { + assert_eq!( + q.token_ids, + common::token_ids(&e["token_ids"]), + "{}", + case["name"] + ); + } + assert_eq!( + prepared.layout.processed_tokens as u64, + case["input_tokens"].as_u64().unwrap() + ); + } +} From 8381029eaec72a2bb578a34a2cad40ebedf1164c Mon Sep 17 00:00:00 2001 From: Tianyao Wu Date: Thu, 8 Oct 2026 13:37:03 +0800 Subject: [PATCH 2/4] omnijev: read fixed-size chunks with as_chunks Clippy 1.99's chunks_exact_to_as_chunks lint refuses chunks_exact with a constant size. The sums keep their order. --- src/models/omnijev/native/src/heads.rs | 12 +++++++----- 1 file changed, 7 insertions(+), 5 deletions(-) diff --git a/src/models/omnijev/native/src/heads.rs b/src/models/omnijev/native/src/heads.rs index 5d5cbb02..c75b47be 100644 --- a/src/models/omnijev/native/src/heads.rs +++ b/src/models/omnijev/native/src/heads.rs @@ -63,9 +63,9 @@ impl Mlp { fn dot(a: &[f32], b: &[f32]) -> f32 { // Eight partial sums, which the compiler vectorizes. let mut acc = [0.0f32; 8]; - let (a8, a_rest) = a.split_at(a.len() / 8 * 8); - let (b8, b_rest) = b.split_at(a8.len()); - for (x, y) in a8.chunks_exact(8).zip(b8.chunks_exact(8)) { + let (a8, a_rest) = a.as_chunks::<8>(); + let (b8, b_rest) = b.as_chunks::<8>(); + for (x, y) in a8.iter().zip(b8) { for i in 0..8 { acc[i] += x[i] * y[i]; } @@ -97,8 +97,10 @@ fn tensor(tensors: &SafeTensors, name: &str, shape: &[usize]) -> Result ); Ok(view .data() - .chunks_exact(4) - .map(|b| f32::from_le_bytes([b[0], b[1], b[2], b[3]])) + .as_chunks::<4>() + .0 + .iter() + .map(|&b| f32::from_le_bytes(b)) .collect()) } From 26c71e7979b142b8ff8a1b297dae14f40a03fd9b Mon Sep 17 00:00:00 2001 From: Tianyao Wu Date: Thu, 8 Oct 2026 22:51:55 +0800 Subject: [PATCH 3/4] omnijev: refuse invalid calibration, stage the export and pin the base config A zero or non-finite temperature gave NaN probabilities, written as null; answer() now refuses it, and the export checks the calibration before writing. The export writes into ..partial and renames it at the end, so a failed run leaves no output. The base config.json and shard index are pinned with the shards. --- recipe/omnijev/README.md | 8 ++++-- recipe/omnijev/export.py | 34 ++++++++++++++++------- src/models/omnijev/native/src/contract.rs | 10 ++++++- tests/omnijev/finish.rs | 23 +++++++++++++++ 4 files changed, 61 insertions(+), 14 deletions(-) diff --git a/recipe/omnijev/README.md b/recipe/omnijev/README.md index 39b19400..02f49391 100644 --- a/recipe/omnijev/README.md +++ b/recipe/omnijev/README.md @@ -38,8 +38,8 @@ hf download Qwen/Qwen3.5-4B --revision 851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a ``` The export checks the checkpoint files it reads against the release's -`release_manifest.json` and the base shards against their pinned SHA-256, then -writes: +`release_manifest.json` and the base model's shards, index and `config.json` against +their pinned SHA-256, then writes: - the language model in BF16 with the adapter merged, its tied embedding resized to the tokenizer's 248,079 entries and the two option tokens' saved rows restored, and @@ -50,7 +50,9 @@ writes: - `omnijev_export.json`, written last, with the revisions, calibration, image budget, and the SHA-256 of the checked inputs and of every output file. -It needs about 18 GB of host RAM and writes 9 GB; on a 48-core workstation it takes +It writes into `..partial` next to the output and renames it at the end, so a +failed export leaves no output directory; the next run removes the partial one. It +needs about 18 GB of host RAM and writes 9 GB; on a 48-core workstation it takes under a minute. ## Fixtures and tests diff --git a/recipe/omnijev/export.py b/recipe/omnijev/export.py index 890e1eb6..6b6f5474 100644 --- a/recipe/omnijev/export.py +++ b/recipe/omnijev/export.py @@ -9,6 +9,8 @@ import argparse import hashlib import json +import math +import shutil from pathlib import Path import torch @@ -20,7 +22,9 @@ CHECKPOINT_REVISION = "ffe5f436eaf22e20e2f041f8e74e121fd057a6cb" BASE_MODEL_ID = "Qwen/Qwen3.5-4B" BASE_REVISION = "851bf6e806efd8d0a36b00ddf55e13ccb7b8cd0a" -BASE_SHARDS = { +BASE_FILES = { + "config.json": "ddc63e1c717afa86c865bb5e01313d89d72bb53b97ad4a8a03ba8510c0621670", + "model.safetensors.index.json": "cf3f798ee02ba45f9622aa8892a47369ab667d0afbf154ee7c2212de42e6302d", "model.safetensors-00001-of-00002.safetensors": "26a93f066e1916adb13453dae5a0c707c0fbc71299ed98779571a907b8e74c61", "model.safetensors-00002-of-00002.safetensors": "cb544bd9bfae93dc59b0f22b292f5933573854a7f9b97835c67060d7d910e188", } @@ -57,7 +61,7 @@ def main(): if digest != release["files"][name]["sha256"]: raise ValueError(f"{name} does not match the release manifest") inputs[f"checkpoint/{name}"] = digest - for name, expected in BASE_SHARDS.items(): + for name, expected in BASE_FILES.items(): if sha256(args.base / name) != expected: raise ValueError(f"{name} does not match {BASE_MODEL_ID} @ {BASE_REVISION}") inputs[f"base/{name}"] = expected @@ -67,6 +71,10 @@ def main(): extra = set(meta.get("biases", {})) - {"noul"} if extra: raise ValueError(f"unsupported calibration biases: {sorted(extra)}") + temperatures, noul_bias = meta["temperatures"], meta["biases"]["noul"] + if not (sorted(temperatures) == ["choice", "noul", "score"] + and all(math.isfinite(t) and t > 0 for t in temperatures.values()) and math.isfinite(noul_bias)): + raise ValueError("expected positive, finite temperatures and a finite Noul bias") tokenizer = AutoTokenizer.from_pretrained(args.checkpoint, local_files_only=True) if len(tokenizer) != VOCABULARY or any(tokenizer.convert_tokens_to_ids(t) != i for t, i in OPTION_TOKENS.items()): @@ -87,10 +95,15 @@ def main(): raise ValueError("expected tied input and output embeddings") model = PeftModel.from_pretrained(model, args.checkpoint).merge_and_unload(safe_merge=True) - model.model.language_model.save_pretrained(args.out, max_shard_size="5GB") - model.config.to_json_file(args.out / "config.json") + # Write next to the output and rename at the end, so a failed or killed export + # leaves no output behind; the next run removes what it left. + staging = args.out.with_name(f".{args.out.name}.partial") + if staging.exists(): + shutil.rmtree(staging) + model.model.language_model.save_pretrained(staging, max_shard_size="5GB") + model.config.to_json_file(staging / "config.json") vision = {f"model.visual.{k}": v.contiguous() for k, v in model.model.visual.state_dict().items()} - save_file(vision, args.out / "vision.safetensors") + save_file(vision, staging / "vision.safetensors") heads = {} for prefix, name in (("head", "head.pt"), ("ord", "ord.pt")): @@ -99,21 +112,22 @@ def main(): if not torch.isfinite(value).all(): raise ValueError(f"non-finite {prefix}.{key}") heads[f"{prefix}.{key}"] = value.float().contiguous() - save_file(heads, args.out / "heads.safetensors") - tokenizer.save_pretrained(args.out) + save_file(heads, staging / "heads.safetensors") + tokenizer.save_pretrained(staging) - outputs = {p.name: sha256(p) for p in sorted(args.out.iterdir()) if p.is_file()} + outputs = {p.name: sha256(p) for p in sorted(staging.iterdir()) if p.is_file()} # Written last: the worker refuses incomplete exports. - (args.out / "omnijev_export.json").write_text(json.dumps({ + (staging / "omnijev_export.json").write_text(json.dumps({ "format": "omnijev-merged/1", "model_id": MODEL_ID, "checkpoint_revision": CHECKPOINT_REVISION, "base_model_id": BASE_MODEL_ID, "base_revision": BASE_REVISION, "lora": {"rank": 32, "alpha": 64, "merged": "bfloat16"}, "vocabulary": VOCABULARY, "option_tokens": OPTION_TOKENS, "max_pixels": MAX_PIXELS, "heads": {"norm": meta["norm"], "lm_features": True, "ordinal": True}, - "calibration": {"temperatures": meta["temperatures"], "noul_bias": meta["biases"]["noul"]}, + "calibration": {"temperatures": temperatures, "noul_bias": noul_bias}, "inputs": inputs, "outputs": outputs, }, indent=1, allow_nan=False) + "\n") + staging.rename(args.out) if __name__ == "__main__": diff --git a/src/models/omnijev/native/src/contract.rs b/src/models/omnijev/native/src/contract.rs index 3e66fecf..62fd07af 100644 --- a/src/models/omnijev/native/src/contract.rs +++ b/src/models/omnijev/native/src/contract.rs @@ -411,7 +411,7 @@ impl Calibration { /// Python's `round(x, 4)`: the nearest four-decimal value, ties to even. pub fn round4(x: f64) -> f64 { - format!("{x:.4}").parse().expect("formatted float") + format!("{x:.4}").parse().unwrap_or(x) } /// Jev's confidence `(K p_max - 1) / (K - 1)`, in float32 as the reference computes it. @@ -440,6 +440,14 @@ pub fn answer( !mu.is_empty() && mu.len() == question.keys.len() && mu.iter().all(|v| v.is_finite()), "invalid head output" ); + ensure!( + calibration + .temperatures + .iter() + .all(|t| t.is_finite() && *t > 0.0) + && calibration.noul_bias.is_finite(), + "invalid calibration" + ); let latency = round4(latency); if question.kind == Kind::Noul { let mut p0 = mu[0] as f64; diff --git a/tests/omnijev/finish.rs b/tests/omnijev/finish.rs index b8f03785..9c7e0d69 100644 --- a/tests/omnijev/finish.rs +++ b/tests/omnijev/finish.rs @@ -90,6 +90,29 @@ fn answers_match_reference() { } } +#[test] +fn invalid_calibration_is_an_error() { + let good = calibration(&common::fixture("finish.json")); + let mut bad = vec![Calibration { + noul_bias: f64::NAN, + ..good + }]; + for t in [0.0, -1.0, f64::NAN, f64::INFINITY] { + for i in 0..3 { + let mut c = good; + c.temperatures[i] = t; + bad.push(c); + } + } + let keys = ["a".to_string(), "b".to_string()]; + for c in &bad { + for kind in [Kind::Noul, Kind::Choice, Kind::Score] { + let err = contract::answer(&question(kind, &keys), &[0.3, 0.6], c, 0.0).unwrap_err(); + assert_eq!(err.to_string(), "invalid calibration"); + } + } +} + #[test] fn round4_matches_python() { // round(x, 4) on exactly representable ties goes to the even digit. From de255a786ba896e3efff3493cd99cf44f1f49c46 Mon Sep 17 00:00:00 2001 From: Tianyao Wu Date: Fri, 9 Oct 2026 10:55:16 +0800 Subject: [PATCH 4/4] docs: list OmniJev-4B v1.1 in the supported models The checkpoint export and CPU contract only; the native worker is in progress (#114). --- README.md | 4 +++- docs/supported-models.md | 1 + 2 files changed, 4 insertions(+), 1 deletion(-) diff --git a/README.md b/README.md index a98933d0..0e39e9e2 100644 --- a/README.md +++ b/README.md @@ -159,7 +159,8 @@ LAYA text serving uses the upstream CPU worker, the in-repository Python MPS/CPU worker, or a native Rust/CUDA worker on Hopper. The Cua-S1 4B 0.2 `text` adapter runs as a Python worker or as a native worker on CUDA. Open-Jev-27B-v1.1 and Open-Jev-9B run on the same native Rust/CUDA worker. Cua-S1 also has a -Python screenshot worker, and CLM has a stub-encoder contract recipe: +Python screenshot worker, CLM has a stub-encoder contract recipe, and OmniJev-4B +v1.1 has its checkpoint export and CPU contract, with no worker yet: | Model | Status | | --- | --- | @@ -170,6 +171,7 @@ Python screenshot worker, and CLM has a stub-encoder contract recipe: | Open-Jev-9B | The same [native Rust/CUDA worker](recipe/open_jev/native.md); eager independent text candidates; [reference comparison on sm_89](recipe/open_jev/validation-9b.md) | | autotrust/JEV-27B-VL | [Experimental Rust/CUDA worker](recipe/jev_vl/README.md); single-question text and offline-preencoded image inputs; [bounded H800 validation and limits](recipe/jev_vl/validation.md) | | CLM-v0.1-8B | [External worker with a CPU stub encoder](recipe/clm/README.md); contract checks only, real Qwen3-8B decisions unverified by this recipe | +| OmniJev-4B v1.1 | [Checkpoint export and CPU contract](recipe/omnijev/README.md) only; the native worker is in progress ([#114](https://github.com/ThinkFlowLab/system1-omni/issues/114)) | [Supported models and hardware](docs/supported-models.md) lists the devices and where each worker has been run. diff --git a/docs/supported-models.md b/docs/supported-models.md index 34d7b465..6cd43757 100644 --- a/docs/supported-models.md +++ b/docs/supported-models.md @@ -19,6 +19,7 @@ Models that are being added are also tracked in issues labeled [new model](https | Open-Jev-9B | The same [native Rust/CUDA worker](../recipe/open_jev/native.md) | Not supported | Validated on compute capability 8.9 against the reference for [253 requests](../recipe/open_jev/validation-9b.md) | Not supported | Compute capability 8.0 or newer, CUDA toolkit to build, exported merged weights and trained head | | autotrust/JEV-27B-VL | [Experimental native worker](../recipe/jev_vl/README.md), single-question API | Contract tests only; no CPU inference | H800 frozen-corpus check; [experimental scope and limits](../recipe/jev_vl/validation.md) | Not supported | Pinned merged checkpoint, rebuilt ABI 6 CUDA library; image inputs require offline Transformers preencoding | | CLM-v0.1-8B | [External `clm-serve` recipe](../recipe/clm/README.md) with a CPU stub embeddings server | **Stub-encoder contract checks only** ([#23](https://github.com/ThinkFlowLab/system1-omni/pull/23)); not real Qwen3-8B decisions | Real encoder unverified by the merged recipe | Unverified | Python, upstream CLM and head checkpoint; a real encoder requires a separate embeddings server | +| OmniJev-4B v1.1 | None yet: [checkpoint export and CPU contract](../recipe/omnijev/README.md); the native worker is in progress ([#114](https://github.com/ThinkFlowLab/system1-omni/issues/114)) | Request contract, preparation, heads and finishing checked against the reference ([#120](https://github.com/ThinkFlowLab/system1-omni/pull/120)); no inference | Not yet | Not supported | Python 3.10 reference environment for the export and fixtures | - **Validated:** covered by the recipe on `main` or by the checks in the linked merged pull request. - **Unverified:** the worker accepts this device, but no recipe or merged pull request covers it.