File size: 2,160 Bytes
561d0f9 | 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 | """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()
|