From 649973b921f1f4e907673fbc33acc147669a8492 Mon Sep 17 00:00:00 2001 From: xiaoyu-xyz Date: Mon, 28 Sep 2026 17:08:17 +0800 Subject: [PATCH] [CLM] Add the second model engine: heads and typed decision scoring CLM is the second model tracked in #9. It decides differently in a way LAYA does not cover: the engine does not compute embeddings. A frozen Qwen3-8B encoder runs as its own process behind an OpenAI-compatible /v1/embeddings endpoint, and the engine owns everything after it -- two projection heads, the cosine score, and the typed answer. That split is the reason to implement it second. A CLM checkpoint is a torch.save dict, so recipe/clm/native/export_weights.py converts it first: tensors to safetensors with the head name as a prefix, and cfg, hidden_size, projection_dim and logit_scale into the metadata. The published checkpoint is 75 MB and holds the two heads, not the 8B encoder; the export is 16 tensors and 18.9 M parameters. omni-clm reads that and computes decisions: - config reads the head geometry from the metadata and checks it. - weights lists the expected inventory from that geometry, checks shapes before reading, and loads FP32. - scoring projects a state and each candidate, L2-normalises both, scores them by cosine under exp(logit_scale), softmaxes across the question's candidates, and assembles choice, noul and score. The three types differ only after the distribution exists, so they share one path. confidence is the top probability minus the mean of the rest, clamped, and 1.0 for a single candidate, which is the definition src/clm/schema.py uses. No GPU is needed for any of it. This crate does not call an embeddings endpoint and does not serve HTTP; those belong with the runtime that owns the request path. Checks, with the checkpoint exported: python recipe/clm/native/export_weights.py CLM_v0.1-8B.pt /tmp/clm-export python recipe/clm/native/head_oracle.py /tmp/clm-export /tmp/clm-export/head-oracle.json CLM_EXPORT=/tmp/clm-export cargo test -p omni-clm -- --ignored - every tensor's FP32 conversion hash against the export oracle (16 tensors); - the inventory against the head configuration; - five decisions against head_oracle.py, an independent NumPy implementation of the same arithmetic, on hash-derived embeddings. Both sides load the same exported weights, so the check is that two implementations of the same maths agree; the embeddings carry no meaning as model output. The erf coefficients are the same digits on both sides on purpose. Default cargo test needs no checkpoint. workspace: fmt, clippy -D warnings, 15 tests, release build. --- Cargo.lock | 150 +++++++++++++- Cargo.toml | 2 +- recipe/clm/native/export_weights.py | 87 ++++++++ recipe/clm/native/head_oracle.py | 159 +++++++++++++++ src/models/clm/Cargo.toml | 17 ++ src/models/clm/README.md | 59 ++++++ src/models/clm/src/config.rs | 96 +++++++++ src/models/clm/src/lib.rs | 18 ++ src/models/clm/src/scoring.rs | 301 ++++++++++++++++++++++++++++ src/models/clm/src/weights.rs | 174 ++++++++++++++++ src/models/clm/tests/checkpoint.rs | 181 +++++++++++++++++ 11 files changed, 1242 insertions(+), 2 deletions(-) create mode 100644 recipe/clm/native/export_weights.py create mode 100644 recipe/clm/native/head_oracle.py create mode 100644 src/models/clm/Cargo.toml create mode 100644 src/models/clm/README.md create mode 100644 src/models/clm/src/config.rs create mode 100644 src/models/clm/src/lib.rs create mode 100644 src/models/clm/src/scoring.rs create mode 100644 src/models/clm/src/weights.rs create mode 100644 src/models/clm/tests/checkpoint.rs diff --git a/Cargo.lock b/Cargo.lock index be379643..9bc357a6 100644 --- a/Cargo.lock +++ b/Cargo.lock @@ -2,6 +2,12 @@ # It is not intended for manual editing. version = 4 +[[package]] +name = "anyhow" +version = "1.0.104" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "330a5ed07fa54e4702c9d6c4174f74427fc0ef6e214bbd677ae50a5099946470" + [[package]] name = "atomic-waker" version = "1.1.2" @@ -78,6 +84,15 @@ version = "2.13.2" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3ded4057c258ba199e2d26386d3af3780957ecaee6c4ef4041c6b4b8b97c0b06" +[[package]] +name = "block-buffer" +version = "0.10.4" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "3078c7629b62d3f0439517fa394996acacc5cbc91c5a20d8c658e77abd503a71" +dependencies = [ + "generic-array", +] + [[package]] name = "bumpalo" version = "3.20.3" @@ -119,10 +134,19 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "65c35e4b699c7e15ccbe7ee35c005e4fc0a278d22238a2857e6ce2dadeda1b06" dependencies = [ "cfg-if", - "cpufeatures", + "cpufeatures 0.3.1", "rand_core", ] +[[package]] +name = "cpufeatures" +version = "0.2.17" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "59ed5838eebb26a2bb2e58f6d5b5316989ae9d08bab10e0e6d103e656d1b0280" +dependencies = [ + "libc", +] + [[package]] name = "cpufeatures" version = "0.3.1" @@ -132,6 +156,26 @@ dependencies = [ "libc", ] +[[package]] +name = "crypto-common" +version = "0.1.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "78c8292055d1c1df0cce5d180393dc8cce0abec0a7102adb6c7b1eef6016d60a" +dependencies = [ + "generic-array", + "typenum", +] + +[[package]] +name = "digest" +version = "0.10.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "9ed9a281f7bc9b7576e61468ba615a66a5c8cfdff42420a70aa82701a3b1e292" +dependencies = [ + "block-buffer", + "crypto-common", +] + [[package]] name = "displaydoc" version = "0.2.7" @@ -153,6 +197,12 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "fastrand" +version = "2.5.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "da7c62ceae207dd37ea5b845da6a0696c799f85e97da1ab5b7910be3c1c80223" + [[package]] name = "find-msvc-tools" version = "0.1.13" @@ -228,6 +278,16 @@ dependencies = [ "slab", ] +[[package]] +name = "generic-array" +version = "0.14.7" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "85649ca51fd72272d7821adaf274ad91c288277713d9c18820d8499a7ff69e9a" +dependencies = [ + "typenum", + "version_check", +] + [[package]] name = "getrandom" version = "0.2.17" @@ -494,6 +554,12 @@ version = "0.2.189" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "3eaf3ede3fee6db1a4c2ee091bf8a8b4dccdc6d17f656fb07896ee72867612f2" +[[package]] +name = "linux-raw-sys" +version = "0.12.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32a66949e030da00e8c7d4434b251670a91556f4144941d37452769c25d58a53" + [[package]] name = "litemap" version = "0.8.3" @@ -524,6 +590,15 @@ version = "2.8.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "cf8baf1c55e62ffcace7a9f06f4bd9cd3f0c4beb022d3b367256b91b87513d98" +[[package]] +name = "memmap2" +version = "0.9.11" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "d1219ed1b7f229ee7104d281dd01d6802fe28bb6e95d292942c4daacdeb798c0" +dependencies = [ + "libc", +] + [[package]] name = "mime" version = "0.3.17" @@ -541,6 +616,19 @@ dependencies = [ "windows-sys 0.61.2", ] +[[package]] +name = "omni-clm" +version = "0.1.0" +dependencies = [ + "anyhow", + "memmap2", + "safetensors", + "serde", + "serde_json", + "sha2", + "tempfile", +] + [[package]] name = "omni-jev" version = "0.1.0" @@ -745,6 +833,19 @@ version = "2.1.3" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "6b1e7f9a428571be2dc5bc0505c13fb6bf936822b894ec87abf8a08a4e51742d" +[[package]] +name = "rustix" +version = "1.1.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "891efababe418670775f199f0d233d84843c227a0949a883ce15b37c78d6629d" +dependencies = [ + "bitflags", + "errno", + "libc", + "linux-raw-sys", + "windows-sys 0.61.2", +] + [[package]] name = "rustls" version = "0.23.45" @@ -792,6 +893,16 @@ version = "1.0.23" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "9774ba4a74de5f7b1c1451ed6cd5285a32eddb5cccb8cc655a4e50009e06477f" +[[package]] +name = "safetensors" +version = "0.6.2" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "172dd94c5a87b5c79f945c863da53b2ebc7ccef4eca24ac63cca66a41aab2178" +dependencies = [ + "serde", + "serde_json", +] + [[package]] name = "serde" version = "1.0.229" @@ -799,6 +910,7 @@ source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "4148590afebada386688f18773da617792bf2ef03ffc1e4cbd2b1d45b023e0ba" dependencies = [ "serde_core", + "serde_derive", ] [[package]] @@ -857,6 +969,17 @@ dependencies = [ "serde", ] +[[package]] +name = "sha2" +version = "0.10.9" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "a7507d819769d01a365ab707794a4084392c824f54a7a6a7862f8c3d0892b283" +dependencies = [ + "cfg-if", + "cpufeatures 0.2.17", + "digest", +] + [[package]] name = "shlex" version = "2.0.1" @@ -938,6 +1061,19 @@ dependencies = [ "syn", ] +[[package]] +name = "tempfile" +version = "3.27.0" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "32497e9a4c7b38532efcdebeef879707aa9f794296a4f0244f6f69e9bc8574bd" +dependencies = [ + "fastrand", + "getrandom 0.4.3", + "once_cell", + "rustix", + "windows-sys 0.61.2", +] + [[package]] name = "thiserror" version = "2.0.21" @@ -1096,6 +1232,12 @@ version = "0.2.5" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "e421abadd41a4225275504ea4d6566923418b7f05506fbc9c0fe86ba7396114b" +[[package]] +name = "typenum" +version = "1.20.1" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "b6f5e870be6c3b371b77fe0ee0bafb859fa4964b4404c27de1d380043c4dda20" + [[package]] name = "unicode-ident" version = "1.0.26" @@ -1126,6 +1268,12 @@ version = "1.0.4" source = "registry+https://github.com/rust-lang/crates.io-index" checksum = "b6c140620e7ffbb22c2dee59cafe6084a59b5ffc27a8859a5f0d494b5d52b6be" +[[package]] +name = "version_check" +version = "0.9.5" +source = "registry+https://github.com/rust-lang/crates.io-index" +checksum = "0b928f33d975fc6ad9f86c8f283853ad26bdd5b10b7f1542aa2fa15e2289105a" + [[package]] name = "want" version = "0.3.1" diff --git a/Cargo.toml b/Cargo.toml index 8036e4bc..0daef239 100644 --- a/Cargo.toml +++ b/Cargo.toml @@ -1,3 +1,3 @@ [workspace] -members = ["src/frontend"] +members = ["src/frontend", "src/models/clm"] resolver = "3" diff --git a/recipe/clm/native/export_weights.py b/recipe/clm/native/export_weights.py new file mode 100644 index 00000000..ab1215fc --- /dev/null +++ b/recipe/clm/native/export_weights.py @@ -0,0 +1,87 @@ +"""Export a CLM head checkpoint to safetensors, and record the conversion oracle. + +A CLM checkpoint is a ``torch.save`` dict (``state_head``/``action_head`` state dicts, +``logit_scale``, ``cfg``), so it is a pickle and no non-Python reader can open it. This +writes the tensors to safetensors with the head name as a prefix, keeps the scalar and +config entries in the safetensors metadata, and emits the FP32/FP16/BF16 hash of every +tensor so a reader can be checked without comparing floats directly. + + python recipe/clm/native/export_weights.py CLM_v0.1-8B.pt OUT_DIR + python recipe/clm/native/export_weights.py CLM_v0.1-8B.pt OUT_DIR --oracle oracle.json + +Writes ``model.safetensors`` and, unless ``--no-oracle``, ``oracle.json``. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import torch +from safetensors.torch import save_file + +HEADS = ("state_head", "action_head") + + +def tensors(ckpt: dict) -> dict[str, torch.Tensor]: + """Every parameter, prefixed by its head, in a stable order.""" + out: dict[str, torch.Tensor] = {} + for head in HEADS: + state = ckpt.get(head) + if not isinstance(state, dict): + raise SystemExit(f"checkpoint has no {head!r} state dict") + for name, value in state.items(): + if not isinstance(value, torch.Tensor): + raise SystemExit(f"{head}.{name} is {type(value).__name__}, not a tensor") + out[f"{head}.{name}"] = value.detach().to(torch.float32).contiguous() + return out + + +def oracle(weights: dict[str, torch.Tensor]) -> list[dict]: + rows = [] + for name, x in weights.items(): + row: dict = {"name": name, "shape": list(x.shape), "source_dtype": str(x.dtype)} + for key, dtype in (("f32", torch.float32), ("f16", torch.float16), ("bf16", torch.bfloat16)): + y = x.to(torch.float32).to(dtype).contiguous() + row[key] = hashlib.sha256(y.view(torch.uint8).numpy().tobytes()).hexdigest() + rows.append(row) + return rows + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("checkpoint", type=Path, help="CLM_v0.1-8B.pt") + parser.add_argument("output", type=Path, help="directory to write into") + parser.add_argument("--oracle", type=Path, help="where to write the conversion oracle") + parser.add_argument("--no-oracle", action="store_true", help="skip the oracle") + args = parser.parse_args() + + torch.set_num_threads(4) + ckpt = torch.load(args.checkpoint, map_location="cpu", weights_only=False) + cfg = dict(ckpt.get("cfg") or {}) + weights = tensors(ckpt) + args.output.mkdir(parents=True, exist_ok=True) + + metadata = { + "format": "clm-heads", + "logit_scale": repr(float(ckpt["logit_scale"])), + "hidden_size": str(int(ckpt.get("hidden_size") or cfg.get("hidden_size"))), + "projection_dim": str(int(ckpt.get("projection_dim") or cfg.get("projection_dim"))), + "cfg": json.dumps(cfg, sort_keys=True), + } + path = args.output / "model.safetensors" + save_file(weights, str(path), metadata=metadata) + + params = sum(v.numel() for v in weights.values()) + print(f"WROTE {path} tensors={len(weights)} params={params}", flush=True) + + if not args.no_oracle: + where = args.oracle or (args.output / "oracle.json") + where.write_text(json.dumps(oracle(weights), indent=2) + "\n") + print(f"ORACLE {where} rows={len(weights)}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/recipe/clm/native/head_oracle.py b/recipe/clm/native/head_oracle.py new file mode 100644 index 00000000..74648895 --- /dev/null +++ b/recipe/clm/native/head_oracle.py @@ -0,0 +1,159 @@ +"""Reference decisions for the CLM heads, for the Rust port to be checked against. + +Reads the exported safetensors (not the .pt) so both sides load byte-identical weights, +computes a decision for each fixed case with plain NumPy, and writes the result as JSON. + + python recipe/clm/native/head_oracle.py EXPORT_DIR OUT.json + +The embeddings here are synthesised from the case name, so the numbers are meaningless +as model output -- the point is that two implementations of the same arithmetic agree. +""" + +from __future__ import annotations + +import argparse +import hashlib +import json +from pathlib import Path + +import numpy as np +from safetensors import safe_open + +HIDDEN = 4096 + + +def embedding(text: str, dim: int) -> np.ndarray: + """A stable unit vector per text, so the cases are reproducible without an encoder.""" + out: list[float] = [] + counter = 0 + while len(out) < dim: + digest = hashlib.sha256(f"{counter}:{text}".encode()).digest() + for i in range(0, len(digest) - 3, 4): + if len(out) == dim: + break + out.append(int.from_bytes(digest[i:i + 4], "big") / 2**31 - 1.0) + counter += 1 + x = np.asarray(out, dtype=np.float32) + return x / np.linalg.norm(x) + + +def load_head(f, prefix: str, width: int, proj: int, hidden: int, blocks: int) -> dict: + def get(name: str) -> np.ndarray: + return np.asarray(f.get_tensor(f"{prefix}.{name}"), dtype=np.float32) + + head = { + "inp_w": get("inp.weight"), "inp_b": get("inp.bias"), + "out_w": get("out.weight"), "out_b": get("out.bias"), + } + for i in range(blocks): + head[f"hidden{i}_w"] = get(f"hidden.{i}.weight") + head[f"hidden{i}_b"] = get(f"hidden.{i}.bias") + head[f"norm{i}_w"] = get(f"norms.{i}.weight") + head[f"norm{i}_b"] = get(f"norms.{i}.bias") + return head + + +def gelu(x: np.ndarray) -> np.ndarray: + # The erf form, which is what torch.nn.GELU() uses by default. + from math import sqrt + + return 0.5 * x * (1.0 + _erf(x / sqrt(2.0))) + + +def _erf(x: np.ndarray) -> np.ndarray: + """Vectorised erf via the same A&S 7.1.26 form the Rust side uses.""" + sign = np.sign(x) + x = np.abs(x) + t = 1.0 / (1.0 + 0.3275911 * x) + y = 1.0 - (((((1.0614054 * t - 1.45315203) * t) + 1.42141374) * t - 0.284496736) * t + + 0.2548296) * t * np.exp(-x * x) + return sign * y + + +def layernorm(x: np.ndarray, w: np.ndarray, b: np.ndarray, eps: float = 1e-5) -> np.ndarray: + mean = x.mean(axis=-1, keepdims=True) + var = x.var(axis=-1, keepdims=True) + return (x - mean) / np.sqrt(var + eps) * w + b + + +def project(head: dict, cfg: dict, x: np.ndarray) -> np.ndarray: + h = gelu(x @ head["inp_w"].T + head["inp_b"]) + blocks = cfg["depth"] - 2 + for i in range(blocks): + z = h @ head[f"hidden{i}_w"].T + head[f"hidden{i}_b"] + if cfg["layernorm"]: + z = layernorm(z, head[f"norm{i}_w"], head[f"norm{i}_b"]) + z = gelu(z) + h = h + z if cfg["residual"] else z + out = h @ head["out_w"].T + head["out_b"] + return out / np.linalg.norm(out, axis=-1, keepdims=True) + + +def softmax(v: np.ndarray) -> np.ndarray: + e = np.exp(v - v.max()) + return e / e.sum() + + +def confidence(probs: np.ndarray) -> float: + if probs.size < 2: + return 1.0 + j = int(np.argmax(probs)) + rest = np.delete(probs, j).mean() + return float(min(1.0, max(0.0, probs[j] - rest))) + + +CASES = [ + ("choice_binary", "choice", ["billing", "technical"]), + ("choice_five", "choice", ["a", "b", "c", "d", "e"]), + ("score_three", "score", ["0", "1", "2"]), + ("noul", "noul", ["false", "true"]), + ("choice_single", "choice", ["only"]), +] + + +def main() -> None: + parser = argparse.ArgumentParser(description=__doc__.splitlines()[0]) + parser.add_argument("export", type=Path, help="directory holding model.safetensors") + parser.add_argument("output", type=Path) + args = parser.parse_args() + + with safe_open(args.export / "model.safetensors", framework="np") as f: + md = f.metadata() + cfg = json.loads(md["cfg"]) + hidden = int(md["hidden_size"]) + proj = int(md["projection_dim"]) + width = cfg["width"] + blocks = cfg["depth"] - 2 + scale = float(np.exp(np.float32(md["logit_scale"]))) + state = load_head(f, "state_head", width, proj, hidden, blocks) + action = load_head(f, "action_head", width, proj, hidden, blocks) + + rows = [] + for name, kind, keys in CASES: + state_vec = embedding(f"state::{name}", hidden) + cand_vecs = np.stack([embedding(f"cand::{name}::{k}", hidden) for k in keys]) + zs = project(state, cfg, state_vec[None, :])[0] + zc = project(action, cfg, cand_vecs) + cos = zc @ zs + temperature = 1.0 if name != "choice_five" else 2.5 + probs = softmax((np.float32(scale) * cos / np.float32(temperature)).astype(np.float32)) + row = { + "name": name, "kind": kind, "keys": keys, "temperature": temperature, + "probabilities": [float(p) for p in probs], + } + if kind == "choice": + row["choice"] = keys[int(np.argmax(probs))] + row["confidence"] = confidence(probs) + elif kind == "score": + row["score"] = float(sum(i * float(p) for i, p in enumerate(probs))) + row["confidence"] = confidence(probs) + else: + row["noul"] = float(probs[keys.index("true")]) + rows.append(row) + + args.output.write_text(json.dumps({"cases": rows}, indent=2) + "\n") + print(f"ORACLE {args.output} cases={len(rows)}", flush=True) + + +if __name__ == "__main__": + main() diff --git a/src/models/clm/Cargo.toml b/src/models/clm/Cargo.toml new file mode 100644 index 00000000..3fa59d91 --- /dev/null +++ b/src/models/clm/Cargo.toml @@ -0,0 +1,17 @@ +[package] +name = "omni-clm" +version = "0.1.0" +edition = "2024" +publish = false +description = "CLM: projection heads and typed decision scoring over Qwen3-8B embeddings" + +[dependencies] +anyhow = "1" +memmap2 = "0.9" +safetensors = "0.6" +serde = { version = "1", features = ["derive"] } +serde_json = "1" + +[dev-dependencies] +sha2 = "0.10" +tempfile = "3" diff --git a/src/models/clm/README.md b/src/models/clm/README.md new file mode 100644 index 00000000..6fce2053 --- /dev/null +++ b/src/models/clm/README.md @@ -0,0 +1,59 @@ +# CLM model engine + +CLM is the second model the project tracks ([#9](https://github.com/ThinkFlowLab/system1-omni/issues/9)), after LAYA. It decides differently in a way LAYA does not cover: **the engine does not compute embeddings.** A frozen `Qwen/Qwen3-8B` encoder runs as its own process behind an OpenAI-compatible `/v1/embeddings` endpoint, and the engine owns everything after it — two projection heads, the cosine score, and the typed answer. + +That split is the point of implementing it second. LAYA's engine owns one forward pass; this one owns a client to someone else's server, plus a candidate-vector cache that persists across requests. + +## What is here + +`omni-clm` reads a converted checkpoint and computes decisions. It does not call an embeddings endpoint and does not serve HTTP yet; those belong with the runtime that owns the request path. + +| module | responsibility | +| --- | --- | +| `config` | the head geometry, read from the safetensors metadata | +| `weights` | tensor inventory, shape checks, FP32 loading | +| `scoring` | projection, cosine score, softmax, and the three answer types | + +## The checkpoint is converted first + +A CLM checkpoint is a `torch.save` dict, so it is a pickle and no non-Python reader can open it. `recipe/clm/native/export_weights.py` writes the tensors to safetensors with the head name as a prefix and keeps `cfg`, `hidden_size`, `projection_dim` and `logit_scale` in the metadata. + +```sh +python recipe/clm/native/export_weights.py CLM_v0.1-8B.pt OUT_DIR +``` + +`CLM_v0.1-8B.pt` is the published checkpoint from `Contrastive-LM/CLM-v0.1-8B`; it is 75 MB and holds the two heads, not the 8B encoder. The export is 16 tensors, 18.9 M parameters. + +## The decision, in one pass + +A decision is `softmax(exp(logit_scale) * cos(state_head(s), action_head(c)) / temperature)` over a question's candidates. Both heads are `inp → [LayerNorm →] hidden → out` with GELU, and both projections are L2-normalised before the dot product; the published checkpoint sets `layernorm: true` and `residual: false` with one hidden block. + +The three question types differ only after the distribution exists, which is why they share one scoring path: + +| type | answer | +| --- | --- | +| `choice` | the argmax key, with `confidence` and the full distribution | +| `noul` | the `true` entry of a two-candidate distribution | +| `score` | the expected level index, `sum(i * p_i)`, with `confidence` | + +`confidence` is the top probability minus the mean of the rest, clamped to `[0, 1]`, and `1.0` for a single candidate — the TypeSafe-style definition `schema.py` uses. + +## CPU checks + +The default tests need no checkpoint: + +```sh +cargo test -p omni-clm +``` + +To check the full checkpoint and the decision arithmetic, export it first and point `CLM_EXPORT` at the directory: + +```sh +python recipe/clm/native/export_weights.py /path/to/CLM_v0.1-8B.pt /tmp/clm-export +python recipe/clm/native/head_oracle.py /tmp/clm-export /tmp/clm-export/head-oracle.json +CLM_EXPORT=/tmp/clm-export cargo test -p omni-clm -- --ignored +``` + +Three tests run there: every tensor's FP32 conversion hash against the export oracle, the inventory against the head configuration, and **five decisions checked against an independent NumPy implementation of the same arithmetic** (`head_oracle.py`) on synthesised embeddings, so the two sides need no encoder to disagree. The embeddings are hash-derived and carry no meaning as model output — the check is that two implementations of the same maths agree. + +The default CI job skips these because it does not download the checkpoint. diff --git a/src/models/clm/src/config.rs b/src/models/clm/src/config.rs new file mode 100644 index 00000000..056f184e --- /dev/null +++ b/src/models/clm/src/config.rs @@ -0,0 +1,96 @@ +//! CLM head configuration, read from the exported safetensors metadata. +//! +//! A CLM checkpoint is a `torch.save` dict, so `recipe/clm/native/export_weights.py` +//! converts it first; everything this crate reads is safetensors. The head geometry is +//! not a file of its own, so it travels in the safetensors metadata as `cfg`. +use anyhow::{Context, Result, ensure}; +use serde::Deserialize; + +/// The head shape the checkpoint was trained with. `depth` counts `inp`, the hidden +/// blocks and `out`, so `depth - 2` is the number of `hidden.N` / `norms.N` pairs. +#[derive(Debug, Clone, Deserialize, PartialEq, Eq)] +pub struct HeadConfig { + pub hidden_size: usize, + pub projection_dim: usize, + pub width: usize, + pub depth: usize, + pub activation: String, + pub layernorm: bool, + pub residual: bool, + #[serde(default)] + pub model: Option, +} + +impl HeadConfig { + /// Number of hidden blocks; `depth` includes the input and output projections. + pub fn hidden_blocks(&self) -> Result { + ensure!( + self.depth >= 2, + "depth {} cannot cover an input and an output projection", + self.depth + ); + Ok(self.depth - 2) + } +} + +#[derive(Debug, Clone, PartialEq)] +pub struct Config { + pub head: HeadConfig, + /// `exp(logit_scale)` is the inverse InfoNCE temperature the heads were trained with. + pub logit_scale: f32, +} + +impl Config { + /// Build from the safetensors metadata written by the export tool. + pub fn from_metadata( + metadata: Option<&std::collections::HashMap>, + ) -> Result { + let metadata = metadata.context("the checkpoint carries no metadata")?; + ensure!( + metadata.get("format").map(String::as_str) == Some("clm-heads"), + "not a converted CLM head checkpoint: format is {:?}", + metadata.get("format") + ); + let cfg = metadata.get("cfg").context("metadata has no cfg")?; + let mut head: HeadConfig = + serde_json::from_str(cfg).with_context(|| format!("parse cfg {cfg}"))?; + // hidden_size and projection_dim are repeated at the top level; prefer those + // when the cfg omits them so an older export still loads. + head.hidden_size = metadata + .get("hidden_size") + .and_then(|v| v.parse().ok()) + .unwrap_or(head.hidden_size); + head.projection_dim = metadata + .get("projection_dim") + .and_then(|v| v.parse().ok()) + .unwrap_or(head.projection_dim); + let logit_scale: f32 = metadata + .get("logit_scale") + .context("metadata has no logit_scale")? + .parse() + .context("logit_scale is not a number")?; + + ensure!(head.hidden_size > 0, "hidden_size must be positive"); + ensure!(head.width > 0, "width must be positive"); + ensure!(head.projection_dim > 0, "projection_dim must be positive"); + ensure!( + head.activation == "gelu" || head.activation == "relu" || head.activation == "silu", + "unsupported activation {:?}", + head.activation + ); + // The upstream head applies LayerNorm before every hidden block and adds the + // residual after it; the published 0.1 checkpoint sets both to true/false + // respectively, and the export keeps them so the two paths stay distinguishable. + head.hidden_blocks()?; + ensure!( + logit_scale.is_finite(), + "logit_scale {logit_scale} is not finite" + ); + Ok(Self { head, logit_scale }) + } + + /// `exp(logit_scale)`: the multiplicative constant on the cosine similarity. + pub fn scale(&self) -> f32 { + self.logit_scale.exp() + } +} diff --git a/src/models/clm/src/lib.rs b/src/models/clm/src/lib.rs new file mode 100644 index 00000000..98829dd5 --- /dev/null +++ b/src/models/clm/src/lib.rs @@ -0,0 +1,18 @@ +//! CLM: the second System1-Omni model engine. +//! +//! A CLM decision is made over embeddings the engine does not compute. A frozen +//! Qwen3-8B encoder sits behind an OpenAI-compatible `/v1/embeddings` endpoint and the +//! engine owns everything after it: load the two projection heads, normalise their +//! output, score candidates by cosine similarity under the trained temperature, and +//! assemble `choice`, `noul` and `score` answers. +//! +//! This crate currently covers the checkpoint side and the decision maths. It does not +//! call an embeddings endpoint and does not serve HTTP; those belong with the runtime +//! that owns the request path. Nothing here needs a GPU. +pub mod config; +pub mod scoring; +pub mod weights; + +pub use config::{Config, HeadConfig}; +pub use scoring::{Answer, Kind, Question, answer, confidence, distribution}; +pub use weights::{Head, Heads, Weights, head_tensors}; diff --git a/src/models/clm/src/scoring.rs b/src/models/clm/src/scoring.rs new file mode 100644 index 00000000..65d7bd02 --- /dev/null +++ b/src/models/clm/src/scoring.rs @@ -0,0 +1,301 @@ +//! The typed decision computation: project embeddings, score candidates, assemble answers. +//! +//! This mirrors `src/clm/engine.py` and `src/clm/schema.py`: a state and each candidate +//! are embedded elsewhere (Qwen3-8B behind `/v1/embeddings`) and arrive here as vectors. +//! The state head sees the state, the action head sees every candidate, both projections +//! are L2-normalised, and the score of a pair is `exp(logit_scale) * cos(state, candidate)`, +//! divided by the request temperature and softmaxed across the question's candidates. +//! +//! The three question types differ only after the distribution exists. +use anyhow::{Result, ensure}; + +use crate::weights::{Head, Heads}; + +/// One question: its type and its candidate keys in the order the caller offered them. +#[derive(Debug, Clone, PartialEq)] +pub struct Question { + pub id: String, + pub kind: Kind, + /// Candidate keys, in order. For `Noul` these are `["false", "true"]`. + pub keys: Vec, +} + +#[derive(Debug, Clone, Copy, PartialEq, Eq)] +pub enum Kind { + Choice, + Noul, + Score, +} + +/// One answer, matching the shapes `client.py` parses. +#[derive(Debug, Clone, PartialEq)] +pub enum Answer { + Choice { + choice: String, + confidence: f32, + probabilities: Vec<(String, f32)>, + }, + Noul { + noul: f32, + }, + Score { + score: f32, + confidence: f32, + probabilities: Vec<(String, f32)>, + }, +} + +impl Answer { + /// The discrete label, as `schema.label_of` defines it. + pub fn label(&self) -> String { + match self { + Answer::Choice { choice, .. } => choice.clone(), + Answer::Noul { noul } => if *noul >= 0.5 { "true" } else { "false" }.to_string(), + Answer::Score { probabilities, .. } => probabilities + .iter() + .max_by(|a, b| a.1.partial_cmp(&b.1).unwrap_or(std::cmp::Ordering::Equal)) + .map(|(k, _)| k.clone()) + .unwrap_or_default(), + } + } +} + +/// L2-normalise one row in place, the way `torch.nn.functional.normalize` does. +fn normalize(row: &mut [f32]) { + let norm = row.iter().map(|v| v * v).sum::().sqrt(); + if norm > 0.0 { + for v in row.iter_mut() { + *v /= norm; + } + } +} + +fn gelu(x: f32) -> f32 { + 0.5 * x * (1.0 + erf(x * std::f32::consts::FRAC_1_SQRT_2)) +} + +fn relu(x: f32) -> f32 { + x.max(0.0) +} + +fn silu(x: f32) -> f32 { + x / (1.0 + (-x).exp()) +} + +/// Abramowitz & Stegun 7.1.26. The coefficients are kept at full precision through named +/// constants: rounding one to a shorter f32 literal changes the value, and +/// `recipe/clm/native/head_oracle.py` carries the same digits so that both sides remain +/// the same arithmetic rather than merely similar. +#[allow(clippy::excessive_precision)] +const ERF_P: f32 = 0.3275911; +#[allow(clippy::excessive_precision)] +const ERF_A1: f32 = 0.254829592; +#[allow(clippy::excessive_precision)] +const ERF_A2: f32 = -0.284496736; +#[allow(clippy::excessive_precision)] +const ERF_A3: f32 = 1.421413741; +#[allow(clippy::excessive_precision)] +const ERF_A4: f32 = -1.453152027; +#[allow(clippy::excessive_precision)] +const ERF_A5: f32 = 1.061405429; + +fn erf(x: f32) -> f32 { + let sign = if x < 0.0 { -1.0 } else { 1.0 }; + let x = x.abs(); + let t = 1.0 / (1.0 + ERF_P * x); + let y = 1.0 + - (((((ERF_A5 * t + ERF_A4) * t + ERF_A3) * t + ERF_A2) * t + ERF_A1) * t * (-x * x).exp()); + sign * y +} + +fn activate(kind: &str, x: f32) -> f32 { + match kind { + "relu" => relu(x), + "silu" => silu(x), + _ => gelu(x), + } +} + +#[inline] +fn linear_into(x: &[f32], w: &[f32], b: &[f32], n: usize, k: usize, out: &mut [f32]) { + for i in 0..n { + let row = &w[i * k..(i + 1) * k]; + let mut acc = b[i]; + for (a, wv) in x.iter().zip(row) { + acc += a * wv; + } + out[i] = acc; + } +} + +fn layer_norm(x: &mut [f32], weight: &[f32], bias: &[f32]) { + let n = x.len() as f32; + let mean = x.iter().sum::() / n; + let var = x.iter().map(|v| (v - mean) * (v - mean)).sum::() / n; + let inv = 1.0 / (var + 1e-5).sqrt(); + for i in 0..x.len() { + x[i] = (x[i] - mean) * inv * weight[i] + bias[i]; + } +} + +/// Project one embedding through a head, returning the L2-normalised vector. +fn project(head: &Head, cfg: &crate::config::HeadConfig, x: &[f32]) -> Result> { + ensure!( + x.len() == cfg.hidden_size, + "embedding has {} values, the head expects {}", + x.len(), + cfg.hidden_size + ); + let w = cfg.width; + let mut h = vec![0.0f32; w]; + linear_into( + x, + &head.inp_weight, + &head.inp_bias, + w, + cfg.hidden_size, + &mut h, + ); + for v in h.iter_mut() { + *v = activate(&cfg.activation, *v); + } + if !head.hidden_weight.is_empty() { + let mut hidden = vec![0.0f32; w]; + linear_into( + &h, + &head.hidden_weight, + &head.hidden_bias, + w, + w, + &mut hidden, + ); + if let (Some(nw), Some(nb)) = (&head.norm_weight, &head.norm_bias) { + layer_norm(&mut hidden, nw, nb); + } + for v in hidden.iter_mut() { + *v = activate(&cfg.activation, *v); + } + if cfg.residual { + for i in 0..w { + hidden[i] += h[i]; + } + } + h = hidden; + } + let p = cfg.projection_dim; + let mut out = vec![0.0f32; p]; + linear_into(&h, &head.out_weight, &head.out_bias, p, w, &mut out); + normalize(&mut out); + Ok(out) +} + +/// `softmax(scale * cos / temperature)`, the distribution every question type starts from. +pub fn distribution( + heads: &Heads, + state: &[f32], + candidates: &[Vec], + temperature: f32, +) -> Result> { + ensure!( + !candidates.is_empty(), + "a question needs at least one candidate" + ); + ensure!( + temperature > 0.0 && temperature <= 100.0, + "temperature must be in (0, 100], got {temperature}" + ); + let cfg = &heads.config.head; + let zs = project(&heads.state, cfg, state)?; + let mut logits = Vec::with_capacity(candidates.len()); + for candidate in candidates { + let zc = project(&heads.action, cfg, candidate)?; + let cos: f32 = zs.iter().zip(&zc).map(|(a, b)| a * b).sum(); + logits.push(heads.config.scale() * cos / temperature); + } + let max = logits.iter().copied().fold(f32::MIN, f32::max); + let exp: Vec = logits.iter().map(|v| (v - max).exp()).collect(); + let sum: f32 = exp.iter().sum(); + ensure!(sum > 0.0, "softmax denominator is zero"); + Ok(exp.iter().map(|v| v / sum).collect()) +} + +/// `schema.confidence`: the top probability minus the mean of the rest, clamped to [0, 1]. +/// A single candidate is fully decided by definition. +pub fn confidence(probs: &[f32]) -> f32 { + if probs.len() < 2 { + return 1.0; + } + let (j, top) = + probs.iter().enumerate().fold( + (0usize, f32::MIN), + |best, (i, p)| { + if *p > best.1 { (i, *p) } else { best } + }, + ); + let rest = (probs.len() - 1) as f32; + let mean_rest = probs + .iter() + .enumerate() + .filter(|(i, _)| *i != j) + .map(|(_, p)| *p) + .sum::() + / rest; + (top - mean_rest).clamp(0.0, 1.0) +} + +/// Assemble the answer for one question from its distribution. +pub fn answer(question: &Question, probs: &[f32]) -> Result { + ensure!( + question.keys.len() == probs.len(), + "{}: {} keys but {} probabilities", + question.id, + question.keys.len(), + probs.len() + ); + let pairs: Vec<(String, f32)> = question + .keys + .iter() + .cloned() + .zip(probs.iter().copied()) + .collect(); + Ok(match question.kind { + Kind::Noul => { + // `answer_from_probs` reads the `true` entry, which is the last key by + // convention; a caller that orders them differently still gets a stable + // answer because the key is looked up by name. + let p = pairs + .iter() + .find(|(k, _)| k == "true") + .map(|(_, p)| *p) + .unwrap_or_else(|| *probs.last().expect("checked non-empty")); + Answer::Noul { noul: p } + } + Kind::Choice => { + let (choice, _) = pairs + .iter() + .fold(None, |best: Option<&(String, f32)>, kv| match best { + Some(b) if b.1 >= kv.1 => Some(b), + _ => Some(kv), + }) + .expect("checked non-empty"); + Answer::Choice { + choice: choice.clone(), + confidence: confidence(probs), + probabilities: pairs, + } + } + Kind::Score => { + // The expected level index, `sum(i * p_i)`. + let score = probs + .iter() + .enumerate() + .map(|(i, p)| i as f32 * p) + .sum::(); + Answer::Score { + score, + confidence: confidence(probs), + probabilities: pairs, + } + } + }) +} diff --git a/src/models/clm/src/weights.rs b/src/models/clm/src/weights.rs new file mode 100644 index 00000000..9a42da3c --- /dev/null +++ b/src/models/clm/src/weights.rs @@ -0,0 +1,174 @@ +//! The CLM head tensors: inventory, loading and the expected shapes. +use anyhow::{Context, Result, ensure}; +use memmap2::Mmap; +use safetensors::{Dtype, SafeTensors}; +use std::{collections::BTreeSet, fs::File, path::Path}; + +use crate::config::{Config, HeadConfig}; + +/// The two projection heads. +pub const HEADS: [&str; 2] = ["state_head", "action_head"]; + +/// One head's parameters, all FP32 and row-major. +#[derive(Debug, Clone)] +pub struct Head { + pub inp_weight: Vec, + pub inp_bias: Vec, + pub hidden_weight: Vec, + pub hidden_bias: Vec, + /// Absent when the configuration has no hidden blocks; LayerNorm weight when present. + pub norm_weight: Option>, + pub norm_bias: Option>, + pub out_weight: Vec, + pub out_bias: Vec, +} + +pub struct Heads { + pub state: Head, + pub action: Head, + pub config: Config, +} + +pub struct Weights { + data: Mmap, +} + +impl Weights { + /// The checkpoint must remain immutable while the mapping exists. + pub fn open(path: &Path) -> Result { + let file = File::open(path).with_context(|| format!("open {}", path.display()))?; + // SAFETY: model files are read-only inputs; no mutable mapping is created. + let data = unsafe { Mmap::map(&file)? }; + SafeTensors::deserialize(&data)?; + Ok(Self { data }) + } + + fn tensors(&self) -> Result> { + Ok(SafeTensors::deserialize(&self.data)?) + } + + /// The safetensors metadata, which carries the head configuration. + /// + /// `SafeTensors` keeps its metadata private, so this reads the header directly from + /// the mapping. No tensor data is touched. + pub fn metadata(&self) -> Result>> { + let (_, header) = + SafeTensors::read_metadata(&self.data).context("read the safetensors header")?; + Ok(header.metadata().clone()) + } + + /// Reject omitted, extra or duplicate names in the expected inventory. + pub fn validate_names<'a>(&self, names: impl IntoIterator) -> Result<()> { + let tensors = self.tensors()?; + let mut expected = BTreeSet::new(); + for name in names { + ensure!(expected.insert(name), "duplicate expected tensor: {name}"); + } + let actual: BTreeSet<&str> = tensors.names().into_iter().collect(); + if expected != actual { + let joined = |set: BTreeSet<&str>| set.into_iter().collect::>().join(", "); + let missing = joined(expected.difference(&actual).copied().collect()); + let extra = joined(actual.difference(&expected).copied().collect()); + anyhow::bail!( + "tensor inventory does not match checkpoint: missing [{missing}], unexpected [{extra}]" + ); + } + Ok(()) + } + + /// One tensor as FP32, checking the shape first so a mismatch is an error rather + /// than a reinterpreted buffer. + pub fn f32(&self, name: &str, shape: &[usize]) -> Result> { + let tensors = self.tensors()?; + let t = tensors + .tensor(name) + .with_context(|| format!("read tensor {name}"))?; + ensure!( + t.shape() == shape, + "{name}: expected {shape:?}, got {:?}", + t.shape() + ); + let out = match t.dtype() { + Dtype::F32 => t + .data() + .as_chunks::<4>() + .0 + .iter() + .map(|b| f32::from_le_bytes(*b)) + .collect(), + dt => anyhow::bail!("{name}: unsupported dtype {dt:?}, the export writes FP32"), + }; + Ok(out) + } +} + +/// Names and shapes of every tensor the heads are built from. +pub fn head_tensors(cfg: &HeadConfig) -> Result)>> { + let blocks = cfg.hidden_blocks()?; + let (h, w, p) = (cfg.hidden_size, cfg.width, cfg.projection_dim); + let mut out = Vec::new(); + for head in HEADS { + out.push((format!("{head}.inp.weight"), vec![w, h])); + out.push((format!("{head}.inp.bias"), vec![w])); + for i in 0..blocks { + out.push((format!("{head}.hidden.{i}.weight"), vec![w, w])); + out.push((format!("{head}.hidden.{i}.bias"), vec![w])); + if cfg.layernorm { + out.push((format!("{head}.norms.{i}.weight"), vec![w])); + out.push((format!("{head}.norms.{i}.bias"), vec![w])); + } + } + out.push((format!("{head}.out.weight"), vec![p, w])); + out.push((format!("{head}.out.bias"), vec![p])); + } + Ok(out) +} + +fn load_head(weights: &Weights, name: &str, cfg: &HeadConfig) -> Result { + let (h, w, p) = (cfg.hidden_size, cfg.width, cfg.projection_dim); + let blocks = cfg.hidden_blocks()?; + ensure!( + blocks <= 1, + "the published checkpoint has {blocks} hidden blocks; the loader implements one" + ); + let (hidden_weight, hidden_bias, norm_weight, norm_bias) = if blocks == 1 { + ( + weights.f32(&format!("{name}.hidden.0.weight"), &[w, w])?, + weights.f32(&format!("{name}.hidden.0.bias"), &[w])?, + Some(weights.f32(&format!("{name}.norms.0.weight"), &[w])?), + Some(weights.f32(&format!("{name}.norms.0.bias"), &[w])?), + ) + } else { + (Vec::new(), Vec::new(), None, None) + }; + Ok(Head { + inp_weight: weights.f32(&format!("{name}.inp.weight"), &[w, h])?, + inp_bias: weights.f32(&format!("{name}.inp.bias"), &[w])?, + hidden_weight, + hidden_bias, + norm_weight, + norm_bias, + out_weight: weights.f32(&format!("{name}.out.weight"), &[p, w])?, + out_bias: weights.f32(&format!("{name}.out.bias"), &[p])?, + }) +} + +impl Heads { + /// Read the config from the metadata, check the inventory, and load every tensor. + pub fn load(weights: &Weights) -> Result { + let config = Config::from_metadata(weights.metadata()?.as_ref())?; + let expected = head_tensors(&config.head)?; + weights.validate_names(expected.iter().map(|(n, _)| n.as_str()))?; + let mut loaded = Vec::new(); + for head in HEADS { + loaded.push(load_head(weights, head, &config.head)?); + } + let action = loaded.pop().expect("two heads"); + let state = loaded.pop().expect("two heads"); + Ok(Self { + state, + action, + config, + }) + } +} diff --git a/src/models/clm/tests/checkpoint.rs b/src/models/clm/tests/checkpoint.rs new file mode 100644 index 00000000..48f50dea --- /dev/null +++ b/src/models/clm/tests/checkpoint.rs @@ -0,0 +1,181 @@ +//! CPU checks for the CLM head loader. No GPU and no encoder, but the frozen export. +use omni_clm::{ + Kind, Question, Weights, answer, confidence, distribution, head_tensors, weights::Heads, +}; +use std::path::PathBuf; +use tempfile::tempdir; + +fn export_dir() -> PathBuf { + PathBuf::from(std::env::var_os("CLM_EXPORT").expect("set CLM_EXPORT to the export directory")) +} + +fn load() -> Heads { + let weights = Weights::open(&export_dir().join("model.safetensors")).unwrap(); + Heads::load(&weights).unwrap() +} + +#[test] +#[ignore = "requires CLM_EXPORT at a converted checkpoint; CPU only"] +fn every_tensor_conversion_matches_the_oracle() { + use sha2::{Digest, Sha256}; + + let dir = export_dir(); + let oracle: Vec = + serde_json::from_slice(&std::fs::read(dir.join("oracle.json")).unwrap()).unwrap(); + let weights = Weights::open(&dir.join("model.safetensors")).unwrap(); + + assert_eq!(oracle.len(), 16, "oracle must cover both heads"); + for row in oracle { + let name = row["name"].as_str().unwrap(); + let shape: Vec = serde_json::from_value(row["shape"].clone()).unwrap(); + let values = weights.f32(name, &shape).unwrap(); + let bytes: Vec = values.iter().flat_map(|v| v.to_le_bytes()).collect(); + assert_eq!( + format!("{:x}", Sha256::digest(&bytes)), + row["f32"].as_str().unwrap(), + "{name} f32" + ); + } +} + +#[test] +#[ignore = "requires CLM_EXPORT at a converted checkpoint; CPU only"] +fn decisions_match_the_reference_implementation() { + let heads = load(); + let oracle: serde_json::Value = + serde_json::from_slice(&std::fs::read(export_dir().join("head-oracle.json")).unwrap()) + .or_else(|_| serde_json::from_slice(&std::fs::read("/tmp/clm-oracle.json").unwrap())) + .unwrap(); + + for case in oracle["cases"].as_array().unwrap() { + let name = case["name"].as_str().unwrap(); + let keys: Vec = serde_json::from_value(case["keys"].clone()).unwrap(); + let temperature = case["temperature"].as_f64().unwrap() as f32; + let expected: Vec = serde_json::from_value(case["probabilities"].clone()).unwrap(); + + let state = embedding(&format!("state::{name}"), heads.config.head.hidden_size); + let candidates: Vec> = keys + .iter() + .map(|k| embedding(&format!("cand::{name}::{k}"), heads.config.head.hidden_size)) + .collect(); + + let probs = distribution(&heads, &state, &candidates, temperature).unwrap(); + assert_eq!(probs.len(), expected.len(), "{name}: length"); + for (got, want) in probs.iter().zip(&expected) { + assert!( + (got - want).abs() < 1e-4, + "{name}: probability {got} vs {want}" + ); + } + + let kind = match case["kind"].as_str().unwrap() { + "choice" => Kind::Choice, + "noul" => Kind::Noul, + _ => Kind::Score, + }; + let question = Question { + id: name.to_string(), + kind, + keys, + }; + let answer = answer(&question, &probs).unwrap(); + match (&answer, case.get("choice")) { + ( + omni_clm::Answer::Choice { + choice, + confidence: c, + .. + }, + Some(want), + ) => { + assert_eq!(choice, want.as_str().unwrap(), "{name}: choice"); + let want_c = case["confidence"].as_f64().unwrap() as f32; + assert!( + (c - want_c).abs() < 1e-4, + "{name}: confidence {c} vs {want_c}" + ); + } + ( + omni_clm::Answer::Score { + score, + confidence: c, + .. + }, + None, + ) => { + let want_s = case["score"].as_f64().unwrap() as f32; + assert!( + (score - want_s).abs() < 1e-4, + "{name}: score {score} vs {want_s}" + ); + let want_c = case["confidence"].as_f64().unwrap() as f32; + assert!( + (c - want_c).abs() < 1e-4, + "{name}: confidence {c} vs {want_c}" + ); + } + (omni_clm::Answer::Noul { noul }, None) => { + let want = case["noul"].as_f64().unwrap() as f32; + assert!((noul - want).abs() < 1e-4, "{name}: noul {noul} vs {want}"); + } + (other, _) => panic!("{name}: unexpected answer {other:?}"), + } + } +} + +/// The same synthesised embedding the oracle uses, so both sides see identical vectors. +fn embedding(text: &str, dim: usize) -> Vec { + use sha2::{Digest, Sha256}; + + let mut out: Vec = Vec::with_capacity(dim); + let mut counter = 0u32; + while out.len() < dim { + let digest = Sha256::digest(format!("{counter}:{text}").as_bytes()); + for chunk in digest.as_chunks::<4>().0 { + if out.len() == dim { + break; + } + out.push(u32::from_be_bytes(*chunk) as f64 as f32 / 2f64.powi(31) as f32 - 1.0); + } + counter += 1; + } + let norm = out.iter().map(|v| v * v).sum::().sqrt(); + out.iter().map(|v| v / norm).collect() +} + +#[test] +fn confidence_matches_the_reference_definition() { + // Top minus the mean of the rest, clamped; a single candidate is fully decided. + assert_eq!(confidence(&[1.0]), 1.0); + assert!((confidence(&[0.75, 0.25]) - 0.5).abs() < 1e-6); + assert!((confidence(&[0.5, 0.3, 0.2]) - 0.25).abs() < 1e-6); + assert_eq!(confidence(&[0.4, 0.4, 0.4]), 0.0); +} + +#[test] +fn a_short_inventory_is_reported_with_both_sides() { + let dir = tempdir().unwrap(); + let path = dir.path().join("model.safetensors"); + safetensors::tensor::serialize_to_file( + std::collections::HashMap::::new(), + None, + &path, + ) + .unwrap(); + let weights = Weights::open(&path).unwrap(); + let message = weights.validate_names(["absent"]).unwrap_err().to_string(); + assert!(message.contains("missing [absent]"), "{message}"); +} + +#[test] +#[ignore = "requires CLM_EXPORT at a converted checkpoint; CPU only"] +fn the_inventory_follows_the_head_configuration() { + let heads = load(); + let cfg = &heads.config.head; + let expected = head_tensors(cfg).unwrap(); + assert_eq!(expected.len(), 16, "two heads, eight tensors each"); + let weights = Weights::open(&export_dir().join("model.safetensors")).unwrap(); + weights + .validate_names(expected.iter().map(|(n, _)| n.as_str())) + .unwrap(); +}