open-standards/wai/wai-int-module/corpus.py

#!/usr/bin/env python3
"""The module corpus of the operator-set module (waism §9.2), from the committed corpora.

    python3 -I corpus.py           write corpus.json beside this file
    python3 -I corpus.py --check   require corpus.json to be what this writes, byte for byte

corpus.json holds a case for each committed vector and refusal fixture of the capabilities
the module decodes (src/lib.rs): the capability, the payload, the set, and what the module
gives, the SHA-256 of its output or its refusal code. Each runner (host/, run-node.mjs,
browser/) reads it, gives the module each case as waism §5 states, and holds every output
to it. Nothing here is new data: every payload, set and digest is a committed file of the
corpora below, or read from one.

  int-payload-conformance/vectors.json  WIV1, WIH1, WIA1 and WIS1: every payload, to its
                                        sidecar's digest, and every refusal fixture, to its code
  neural-conformance/                   wai.neural.int_synth: every vector, to sha256.json
  wai-rs/tests/vectors/int_mlicpp/      wai.neural.int_mlicpp: the payload of stream.wai, to
                                        expected.json's rgb_sha256
  live-conformance/dred/vectors.json    wai.audio.opus.dred.int: every payload vector, to the
                                        digest of its rows, none where it has no DRED, or its code
  live-conformance/avatar-face/         wai.avatar.face: every frame vector as a segment of one
                                        frame (§7.4), to its frame's digest; and every segment,
                                        generator and identity prior the corpus refuses

A payload or a set file is named by its path below wai/, or given in hexadecimal where it is
read out of a committed file: the payload of a WAI1 envelope, a DRED packet, a segment. The
set is a list of [name, file] pairs where the decode reads more than one file (the file
sequence of waism §5) and one file otherwise. The DRED weights file is not in this
repository (opus-dred §4.1): its cases name it under "external", and a runner is given the
directory live-conformance/dred/fetch.sh fetched it to.
"""

import hashlib
import json
import os
import sys

HERE = os.path.dirname(os.path.abspath(__file__))
WAI = os.path.dirname(HERE)
OUT = os.path.join(HERE, "corpus.json")

# The decodes that read more than one file, and the names they read them by (src/lib.rs).
SEQUENCE = {
    "wai.neural.int_hyper": ["entropy.json", "g_s.whm", "h_s.whm"],
    "wai.video.int_hyper": ["entropy.json", "g_s.whm", "h_s.whm"],
    "wai.avatar.face": ["generator.wfg", "identity.wfi"],
}
EMPTY_SHA256 = hashlib.sha256(b"").hexdigest()


def read(rel):
    with open(os.path.join(WAI, rel), "rb") as f:
        return f.read()


def load(rel):
    with open(os.path.join(WAI, rel), encoding="utf-8") as f:
        return json.load(f)


def sha(b):
    return hashlib.sha256(b).hexdigest()


def the_set(cap, members):
    """A corpus's set ({name: path}) as a case states it."""
    if cap in SEQUENCE:
        return [[name, members[name]] for name in sorted(members, key=lambda n: n.encode())]
    if not members:
        return None
    if len(members) != 1:
        raise SystemExit(f"corpus.py: {cap} reads one file, and a vector names {sorted(members)}")
    return next(iter(members.values()))


def envelope_payload(e):
    """The payload of a WAI1 envelope (SPEC §2): what follows the manifest and its length."""
    if e[:4] != b"WAI1":
        raise SystemExit("corpus.py: not a WAI1 envelope")
    ml = int.from_bytes(e[4:8], "little")
    at = 8 + ml
    pl = int.from_bytes(e[at:at + 4], "little")
    body = e[at + 4:]
    if len(body) != pl:
        raise SystemExit("corpus.py: the envelope's payload length is not its own")
    return body


def segment(points, generator, identity):
    """A segment (avatars §7.4) of one frame at step 1: its header and frame 0, and no code."""
    k = len(points)
    head = b"WFK1" + bytes([1, k, 1, 0]) + generator + identity
    head += (0).to_bytes(4, "little") + (1).to_bytes(2, "little") + (25).to_bytes(2, "little") + (1).to_bytes(2, "little") + bytes(2)
    body = b"".join(v.to_bytes(2, "little", signed=True) for p in points for v in p)
    return head + body


def integer_payloads(cases):
    d = load("int-payload-conformance/vectors.json")
    for v in d["vectors"]:
        side = load(v["sidecar"])[v["name"]]
        want = side["clip_sha256"] if "clip_sha256" in side else side["sha256"]
        cases.append({"name": f"int-payload/{v['name']}", "capability": v["capability"], "payload": {"file": v["payload"]},
                      "set": the_set(v["capability"], v["set"]), "expect": {"sha256": want}})
    for r in d["refusals"]:
        cases.append({"name": f"int-payload/refusals/{r['name']}", "capability": r["capability"], "payload": {"file": r["payload"]},
                      "set": the_set(r["capability"], r["set"]), "expect": {"refuse": r["code"]}})


def int_synth(cases):
    want = load("neural-conformance/sha256.json")
    for name in sorted(want):
        base = f"neural-conformance/vectors/{name}"
        cases.append({"name": f"neural/{name}", "capability": "wai.neural.int_synth", "payload": {"file": f"{base}/latent.bin"},
                      "set": f"{base}/int_synth.json", "expect": {"sha256": want[name]["sha256"]}})


def int_mlicpp(cases):
    base = "wai-rs/tests/vectors/int_mlicpp"
    e = load(f"{base}/expected.json")
    payload = envelope_payload(read(f"{base}/stream.wai"))
    if sha(payload) != e["payload_sha256"]:
        raise SystemExit("corpus.py: stream.wai's payload is not the one expected.json names")
    cases.append({"name": "int_mlicpp/stream", "capability": e["capability"], "payload": {"hex": payload.hex()},
                  "set": f"{base}/model.wmb", "expect": {"sha256": e["rgb_sha256"]}})


def dred(cases, external):
    d = load("live-conformance/dred/vectors.json")
    w = d["weights"]
    external[w["file"]] = {"sha256": w["sha256"], "fetch": "live-conformance/dred/fetch.sh"}
    for p in d["payloads"]:
        o = p["outcome"]
        if o == "dred":
            want = {"sha256": p["dred"]["features_sha256"]}
        elif o == "none":
            want = {"sha256": EMPTY_SHA256}
        else:
            want = {"refuse": o["refuse"]}
        cases.append({"name": f"dred/{p['name']}", "capability": d["capability"], "payload": {"hex": p["packet"]},
                      "set": {"external": w["file"]}, "expect": want})


def avatar_face(cases):
    base = "live-conformance/avatar-face"
    d = load(f"{base}/vectors.json")
    cap = d["capability"]

    def files(gen, ide):
        return [["generator.wfg", gen], ["identity.wfi", ide]]

    for f in d["frames"]:
        gen, ide = f"{base}/{f['generator']}", f"{base}/{f['identity']}"
        seg = segment(f["keypoints"], hashlib.sha256(read(gen)).digest(), hashlib.sha256(read(ide)).digest())
        cases.append({"name": f"avatar-face/{f['name']}", "capability": cap, "payload": {"hex": seg.hex()},
                      "set": files(gen, ide), "expect": {"sha256": f["frame_sha256"]}})
    # The forms the corpus refuses: a segment as the payload, and a generator or an identity
    # prior as a file of the set, beside the seeded generator and prior.
    gen, ide = f"{base}/seeded.wfg", f"{base}/seeded.wfi"
    at_rest = next(f for f in d["frames"] if f["generator"] == "seeded.wfg")["keypoints"]
    for f in d["forms"]:
        if f["outcome"] != "refuse":
            continue
        b = bytes.fromhex(f["hex"])
        if f["form"] == "segment":
            payload, s, want = b, files(gen, ide), f["reason"]
        elif f["form"] == "generator":
            payload = segment(at_rest, hashlib.sha256(b).digest(), hashlib.sha256(read(ide)).digest())
            s, want = [["generator.wfg", {"hex": f["hex"]}], ["identity.wfi", ide]], "generator_invalid"
        elif f["form"] == "identity":
            payload = segment(at_rest, hashlib.sha256(read(gen)).digest(), hashlib.sha256(b).digest())
            s, want = [["generator.wfg", gen], ["identity.wfi", {"hex": f["hex"]}]], "identity_invalid"
        else:
            raise SystemExit(f"corpus.py: a form of no known kind, {f['form']}")
        cases.append({"name": f"avatar-face/forms/{f['name']}", "capability": cap, "payload": {"hex": payload.hex()},
                      "set": s, "expect": {"refuse": want}})


def build():
    cases, external = [], {}
    integer_payloads(cases)
    int_synth(cases)
    int_mlicpp(cases)
    dred(cases, external)
    avatar_face(cases)
    names = [c["name"] for c in cases]
    if len(set(names)) != len(names):
        raise SystemExit("corpus.py: two cases share a name")
    doc = {
        "format": "waism-module-corpus/1",
        "about": "The module corpus (extensions/waism.md §9.2) of wai-int-module, the operator-set module: for each case, "
                 "the capability, the payload and the set the module is given (waism §5), and the SHA-256 of its output "
                 "or its refusal code. Written by corpus.py from the committed corpora; paths are below wai/.",
        "external": external,
        "cases": cases,
    }
    return (json.dumps(doc, indent=1, ensure_ascii=True) + "\n").encode()


def main(argv):
    if argv not in ([], ["--check"]):
        print("usage: python3 -I corpus.py [--check]", file=sys.stderr)
        return 2
    text = build()
    if argv == ["--check"]:
        try:
            with open(OUT, "rb") as f:
                have = f.read()
        except OSError:
            have = None
        if have != text:
            print("corpus.py: corpus.json is not what corpus.py writes from the committed corpora; run python3 -I corpus.py", file=sys.stderr)
            return 1
        print(f"corpus.py: corpus.json holds the {json.loads(text)['cases'].__len__()} cases of the committed corpora")
        return 0
    with open(OUT, "wb") as f:
        f.write(text)
    print(f"corpus.py: wrote {OUT} ({len(json.loads(text)['cases'])} cases)")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))