| """Dump every hand-written number in handweights.py to weights.safetensors. | |
| The safetensors format is an 8-byte little-endian header length, a JSON header, | |
| then raw tensor bytes, so it is written here with numpy alone. | |
| """ | |
| import json | |
| import struct | |
| from pathlib import Path | |
| import numpy as np | |
| from handweights import (CLR, COL, DEF_COLOR, DEF_SIZE, DEF_Y, DIM, EXP, EXPRS, FACES, GND, GROUNDS, HAND, OBJ, POS, | |
| PARTS, PRED, PROPS, REF, SHAPES, SIZ, SKIES, SKY, VOCAB, count_params) | |
| LAYOUT = {"shape_onehot": [0, OBJ], "is_object": OBJ, "pronoun(1=singular,2=plural)": REF, | |
| "rgb_has_color": [COL, COL + 4], "dx_dy_has_x_has_y": [POS, POS + 4], "size_has_size": [SIZ, SIZ + 2], | |
| f"expression_onehot({','.join(EXPRS)})": [EXP, PRED], "copula": PRED, "conjunction": CLR, | |
| f"sky_onehot({','.join(SKIES)})": [SKY, GND], f"ground_onehot({','.join(GROUNDS)})": [GND, DIM]} | |
| def leaves(g, name, out): | |
| """Flatten an AND/OR/MINUS tree into named plane tensors plus a structure string.""" | |
| op, arg = g | |
| if op == "and": | |
| out[name] = arg | |
| return name | |
| return f"{op}(" + ", ".join(leaves(c, f"{name}.{i}", out) for i, c in enumerate(arg)) + ")" | |
| def collect(): | |
| words = list(VOCAB) | |
| tensors = { | |
| "encoder.embedding": np.stack([VOCAB[w] for w in words]), | |
| "decoder.default_color": DEF_COLOR, | |
| "decoder.default_size": DEF_SIZE, | |
| "decoder.default_y": np.array([y if isinstance(y, float) else np.nan for y in DEF_Y], float), | |
| "layout.prop_grip_scale": np.array([[*g, sc] for _, g, sc in PROPS.values()], float), | |
| "layout.hand": np.array(list(HAND.values()), float), | |
| } | |
| structure = {} | |
| groups = list(PARTS.items()) + [(f"face.{k}.{e}", parts) for (k, e), parts in sorted(FACES.items())] | |
| for name, parts in groups: | |
| for i, (geom, col, alpha, shaded, outline) in enumerate(parts): | |
| key = f"renderer.{name}.part{i}" | |
| structure[key] = dict(geometry=leaves(geom, key, tensors), color=col if isinstance(col, str) else list(col), | |
| alpha=alpha, shaded=shaded, outline=outline) | |
| meta = {"vocab": json.dumps(words), "shapes": json.dumps(SHAPES), "embedding_layout": json.dumps(LAYOUT), | |
| "props": json.dumps(list(PROPS)), "default_y_special": "NaN = stands on the ground line", | |
| "renderer_structure": json.dumps(structure), | |
| "plane_rows": "each row (nx, ny, d) is a half-plane neuron firing where nx*px + ny*py <= d; a leaf = min-pool over its rows", | |
| "parameter_count": str(count_params()), "trained": "false"} | |
| return tensors, meta | |
| def save_safetensors(tensors, meta, path): | |
| header, blobs, offset = {"__metadata__": meta}, [], 0 | |
| for name, t in tensors.items(): | |
| b = np.ascontiguousarray(t, dtype="<f4").tobytes() | |
| header[name] = {"dtype": "F32", "shape": list(t.shape), "data_offsets": [offset, offset + len(b)]} | |
| blobs.append(b) | |
| offset += len(b) | |
| h = json.dumps(header, separators=(",", ":")).encode() | |
| h += b" " * (-len(h) % 8) | |
| Path(path).write_bytes(struct.pack("<Q", len(h)) + h + b"".join(blobs)) | |
| def load_safetensors(path): | |
| raw = Path(path).read_bytes() | |
| n = struct.unpack("<Q", raw[:8])[0] | |
| header = json.loads(raw[8:8 + n]) | |
| meta = header.pop("__metadata__") | |
| data = raw[8 + n:] | |
| return {k: np.frombuffer(data[v["data_offsets"][0]:v["data_offsets"][1]], "<f4").reshape(v["shape"]) | |
| for k, v in header.items()}, meta | |
| if __name__ == "__main__": | |
| out = Path(__file__).parent / "weights.safetensors" | |
| tensors, meta = collect() | |
| save_safetensors(tensors, meta, out) | |
| back, _ = load_safetensors(out) | |
| assert all(np.allclose(back[k], np.asarray(v, "<f4"), equal_nan=True) for k, v in tensors.items()) | |
| print(f"{out.name}: {len(tensors)} tensors, {sum(t.size for t in tensors.values()):,} numbers, round-trip ok") | |