omnijev: add the request contract, preparation, heads and export

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"

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.

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.

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

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.
 

recipe/omnijev/README.md

@@ -0,0 +1,72 @@
+# 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 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
+  `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 writes into `.<out>.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
+
+```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.

recipe/omnijev/export.py

@@ -0,0 +1,134 @@
+"""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
+import math
+import shutil
+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_FILES = {
+    "config.json": "ddc63e1c717afa86c865bb5e01313d89d72bb53b97ad4a8a03ba8510c0621670",
+    "model.safetensors.index.json": "cf3f798ee02ba45f9622aa8892a47369ab667d0afbf154ee7c2212de42e6302d",
+    "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_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
+    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)}")
+    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()):
+        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)
+    # 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, staging / "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, staging / "heads.safetensors")
+    tokenizer.save_pretrained(staging)
+
+    outputs = {p.name: sha256(p) for p in sorted(staging.iterdir()) if p.is_file()}
+    # Written last: the worker refuses incomplete exports.
+    (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": temperatures, "noul_bias": noul_bias},
+        "inputs": inputs, "outputs": outputs,
+    }, indent=1, allow_nan=False) + "\n")
+    staging.rename(args.out)
+
+
+if __name__ == "__main__":
+    main()