| """Load and run DraftZero's FDN graph network from its public release, without pickles or MageZero. |
| |
| The release (built by export_gnn_release.py) holds: |
| |
| model.safetensors the network's weights, float32 |
| config.json the architecture (`arch`) and the encoder's constants |
| vocab.json leaf_ids and edge_ids: the raw feature ids that have embedding rows, row order |
| graph_net.py the network (a copy of draft-zero's src/draftzero/gameplay/graph_net.py) |
| goldens.jsonl.gz encoded test states with this network's outputs, for checking a port |
| |
| A state is the graph encoder's output in MageZero's graph-server layout (one dict per state): |
| |
| indices [nodes] raw feature id per node: StateEncoder.indexFor(hash64(name)); typed nodes |
| are the hash of their type name, the root is 0 |
| values [nodes] numeric value per node (0 for non-numeric nodes) |
| edge_child [edges] node index, local to the state |
| edge_parent [edges] node index, local to the state |
| edge_label [edges] raw hashed edge label |
| |
| python gnn_release.py check <release dir> # run the goldens, print the largest differences |
| """ |
| from __future__ import annotations |
|
|
| import gzip |
| import json |
| import sys |
| from pathlib import Path |
|
|
| import numpy as np |
| import torch |
|
|
| sys.path.insert(0, str(Path(__file__).resolve().parent)) |
| try: |
| import graph_net as gn |
| except ImportError: |
| from draftzero.gameplay import graph_net as gn |
|
|
| |
| PRIORITY, CHOOSE_TARGET, CHOOSE_USE = 0, 3, 5 |
|
|
|
|
| class Vocab: |
| """Raw feature id -> embedding row (-1 outside the vocab).""" |
|
|
| def __init__(self, ids): |
| self.ids = np.asarray(ids, np.int64) |
| self._order = np.argsort(self.ids, kind="stable") |
| self._sorted = self.ids[self._order] |
|
|
| def __len__(self) -> int: |
| return len(self.ids) |
|
|
| def lookup(self, ids) -> np.ndarray: |
| ids = np.asarray(ids, np.int64) |
| pos = np.minimum(np.searchsorted(self._sorted, ids), len(self.ids) - 1) |
| return np.where(self._sorted[pos] == ids, self._order[pos], -1) |
|
|
|
|
| def load(release_dir, device: str = "cpu"): |
| """(model in eval mode, leaf vocab, edge vocab, config) from a release directory.""" |
| from safetensors.torch import load_file |
| d = Path(release_dir) |
| cfg = json.loads((d / "config.json").read_text()) |
| vocab = json.loads((d / "vocab.json").read_text()) |
| leaves, edges = Vocab(vocab["leaf_ids"]), Vocab(vocab["edge_ids"]) |
| model = gn.NetGraph(len(leaves), len(edges), cfg["arch"]) |
| model.load_state_dict(load_file(str(d / "model.safetensors"))) |
| return model.to(device).eval(), leaves, edges, cfg |
|
|
|
|
| def to_graphs(states: list[dict], leaves: Vocab, edges: Vocab) -> gn.Graphs: |
| """A batch of encoded states -> the network's input (as graph_server.py maps them).""" |
| parts = {k: [] for k in ("type", "row", "value", "child", "parent", "label")} |
| offsets, base = [0], 0 |
| for s in states: |
| ids = np.asarray(s["indices"], np.int64) |
| types = gn.node_types(ids) |
| parts["type"].append(types.astype(np.int64)) |
| parts["row"].append(np.where(types == gn.NodeType.LEAF, leaves.lookup(ids), -1)) |
| parts["value"].append(np.asarray(s["values"], np.int64)) |
| parts["child"].append(np.asarray(s["edge_child"], np.int64) + base) |
| parts["parent"].append(np.asarray(s["edge_parent"], np.int64) + base) |
| parts["label"].append(edges.lookup(s["edge_label"]) + 1) |
| base += len(ids) |
| offsets.append(base) |
| t = {k: torch.from_numpy(np.concatenate(v).astype(np.int64)) for k, v in parts.items()} |
| return gn.Graphs(t["type"], t["row"], t["value"], torch.tensor(offsets), t["child"], t["parent"], t["label"]) |
|
|
|
|
| @torch.no_grad() |
| def evaluate(model, states: list[dict], leaves: Vocab, edges: Vocab) -> list[dict]: |
| """Per state: priority and target (one score per node, -inf at leaves), use [no, yes], value, value_x.""" |
| g = to_graphs(states, leaves, edges) |
| out = model(g) |
| res = [] |
| for i in range(len(states)): |
| a, b = int(g.node_offsets[i]), int(g.node_offsets[i + 1]) |
| res.append({"priority": out.priority[a:b].numpy(), "target": out.target[a:b].numpy(), |
| "use": out.use[i].numpy(), "value": float(out.value[i]), "value_x": float(out.value_x[i])}) |
| return res |
|
|
|
|
| def option_logits(res: dict, graph_type: int, options: list[list[int]]) -> np.ndarray: |
| """One logit per legal option. An option is a list of node indices (one ability node per copy of a |
| card for a play, Pass's node, a target's node, Stop Choosing's node; "attack with X?" is a target |
| choice between Stop Choosing and the defending player). Its logit is the log-sum-exp of its nodes' |
| scores under the head the decision reads: priority at PRIORITY, target otherwise, and the use head's |
| [no, yes] at a CHOOSE_USE question that isn't an attack.""" |
| if graph_type == CHOOSE_USE: |
| return np.asarray(res["use"][:len(options)], np.float64) |
| scores = res["priority"] if graph_type == PRIORITY else res["target"] |
| out = [] |
| for nodes in options: |
| x = scores[np.asarray(nodes, np.int64)].astype(np.float64) |
| m = x.max() |
| out.append(m + np.log(np.exp(x - m).sum()) if np.isfinite(m) else -np.inf) |
| return np.asarray(out) |
|
|
|
|
| def read_goldens(path) -> list[dict]: |
| with gzip.open(path, "rt") as f: |
| return [json.loads(line) for line in f] |
|
|
|
|
| def check(release_dir, batch: int = 16) -> dict: |
| """Run the release's goldens through this loader and report the largest differences.""" |
| d = Path(release_dir) |
| model, leaves, edges, _ = load(d) |
| goldens = read_goldens(d / "goldens.jsonl.gz") |
| worst = {"option_logits": 0.0, "value_x": 0.0, "use": 0.0} |
| argmax_agree = 0 |
| for a in range(0, len(goldens), batch): |
| chunk = goldens[a:a + batch] |
| for g, r in zip(chunk, evaluate(model, [g["state"] for g in chunk], leaves, edges)): |
| want = g["expected"] |
| got = option_logits(r, g["graph_type"], g["options"]) |
| worst["option_logits"] = max(worst["option_logits"], float(np.max(np.abs(got - want["option_logits"])))) |
| worst["value_x"] = max(worst["value_x"], abs(r["value_x"] - want["value_x"])) |
| worst["use"] = max(worst["use"], float(np.max(np.abs(r["use"] - want["use"])))) |
| argmax_agree += int(np.argmax(got) == np.argmax(want["option_logits"])) |
| return {"states": len(goldens), "max_abs_diff": worst, "argmax_agree": argmax_agree} |
|
|
|
|
| if __name__ == "__main__": |
| if len(sys.argv) != 3 or sys.argv[1] != "check": |
| sys.exit(__doc__) |
| print(json.dumps(check(sys.argv[2]), indent=1)) |
|
|
