[CLM 1/2] Checkpoint and typed decision scoring

Cargo.toml

@@ -1,3 +1,3 @@
 [workspace]
-members = ["src/frontend", "src/runtime", "src/models/cua_s1/native", "src/models/qwen3_5/native", "src/models/open_jev/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/laya", "src/backends/cuda"]
 resolver = "3"

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

recipe/clm/native/head_oracle.py

@@ -0,0 +1,162 @@
+"""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
+        # The reference caps this: heads.py does exp(logit_scale).clamp(max=100.0),
+        # and the published checkpoint's 4.6132 exponentiates to 100.82, so the cap
+        # binds. Without it every probability is about 0.8 % off.
+        scale = min(float(np.exp(np.float32(md["logit_scale"]))), 100.0)
+        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()

src/models/clm/Cargo.toml

@@ -0,0 +1,23 @@
+[package]
+name = "omni-clm"
+version = "0.1.0"
+edition = "2024"
+publish = false
+description = "CLM: projection heads, typed decision scoring and the embeddings client"
+
+[dependencies]
+anyhow = "1"
+memmap2 = "0.9"
+reqwest = { version = "0.12", default-features = false, features = ["blocking", "json", "rustls-tls"] }
+safetensors = "0.6"
+serde = { version = "1", features = ["derive"] }
+# preserve_order keeps object key order, which to_text depends on: the heads are
+# trained on the caller's field order, so re-sorting a state would change the text.
+serde_json = { version = "1", features = ["preserve_order"] }
+sha2 = "0.10"
+
+# Test bodies live under the repository-level tests/ tree; this is the explicit
+# registration CONTRIBUTING asks for, matching src/models/laya/Cargo.toml.
+[[test]]
+name = "checkpoint"
+path = "../../../tests/clm/checkpoint.rs"

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
+```
+
+Two tests run there: every tensor's FP32 conversion hash against the export oracle, 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.