"""writes model.safetensors + config.json for unity-embed."""
import json, math, struct

DIM = 384
VAL = 1.0 / math.sqrt(DIM)

payload = b"".join(struct.pack("<f", VAL) for _ in range(DIM))
header = {"v": {"dtype": "F32", "shape": [DIM], "data_offsets": [0, DIM * 4]}}
hb = json.dumps(header, separators=(",", ":")).encode()
with open("model.safetensors", "wb") as f:
    f.write(struct.pack("<Q", len(hb)) + hb + payload)

cfg = {
    "architectures": ["UnityEmbedModel"],
    "model_type": "unity-embed",
    "embedding_dimension": DIM,
    "auto_map": {
        "AutoConfig": "modeling_unity.UnityEmbedConfig",
        "AutoModel": "modeling_unity.UnityEmbedModel",
    },
}
with open("config.json", "w") as f:
    json.dump(cfg, f, indent=2)