File size: 4,028 Bytes
4e1884c | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 64 65 66 67 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 | """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")
|