package_checkpoint.py

"""Pack the trained Tactus v0.1 head into the release artifacts.

Tactus's trained weights are the tactile trunk + projector (~13.5M params) that map a
32x32 pressure window into the fusion-embedding canonical text space. The frozen Qwen
base downloads from its own repository at inference time, so the release artifact is the
head's state dicts plus a small config.

Reads the training checkpoint (out/tactus_head.pt, the shipped eval weights) and writes
out/model.safetensors: fp32 tensors under "encoder.*" and "proj.*" prefixes, with the
config and headline results embedded as JSON in the safetensors metadata (the same
convention the Tremor release uses).

Run:
    python package_checkpoint.py
"""
import json
import os

import torch

HERE = os.path.dirname(os.path.abspath(__file__))
SRC = os.path.join(HERE, "out", "tactus_head.pt")
OUT = os.path.join(HERE, "out", "model.safetensors")


def main():
    from safetensors.torch import save_file

    ck = torch.load(SRC, map_location="cpu", weights_only=False)
    cfg = dict(ck["config"])
    # ship only what inference needs; training-provenance keys stay in the .pt
    keep = ("temporal", "depth", "flow", "window_frames", "out_dim", "arch_detail",
            "text_target", "text_target_detail", "fe2_source", "class_names",
            "dataset", "valid_filter", "init_from")
    cfg_ship = {k: cfg[k] for k in keep if k in cfg}
    cfg_ship.update(name="fusion-embedding-2-tactus", version="v0.1-preview",
                    grid=32, feat_dim=512, license="cc-by-nc-4.0")

    tensors = {}
    for group in ("encoder", "proj"):
        for k, v in ck[group].items():
            tensors[f"{group}.{k}"] = v.float().contiguous()

    results = json.load(open(os.path.join(HERE, "results.json")))
    meta = {"config": json.dumps(cfg_ship),
            "results": json.dumps(results["released_checkpoint_eval"]),
            "format": "pt"}
    save_file(tensors, OUT, metadata=meta)
    n = sum(v.numel() for v in tensors.values())
    print(f"wrote {OUT}  ({n/1e6:.2f}M params, {os.path.getsize(OUT)/1e6:.1f} MB, "
          f"{len(tensors)} tensors)")


if __name__ == "__main__":
    main()