File size: 2,552 Bytes
df87322 | 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 | #!/usr/bin/env python3
"""Make a downloaded checkpoint loadable on the robot with NO Hugging Face access.
Two edits, both on a COPY (the downloaded checkpoint is never modified):
1. policy_preprocessor.json -> tokenizer_name points at the bundled PaliGemma
tokenizer instead of "google/paligemma-3b-pt-224". That repo is gated, so on
a robot without a logged-in HF account the gated download would fail at
startup.
2. config.json -> drop `pretrained_revision` (per Gustavo: rosetta's
policy_bridge_node rejects it).
Usage:
python prepare_deploy.py <checkpoint_dir> [--tokenizer DIR] [--out DIR]
"""
import argparse, json, shutil
from pathlib import Path
here = Path(__file__).resolve().parent
ap = argparse.ArgumentParser()
ap.add_argument("checkpoint", help="dir containing model.safetensors, config.json, policy_*.json")
ap.add_argument("--tokenizer", default=str(here / "paligemma_tokenizer"))
ap.add_argument("--out", default=None, help="default: <checkpoint>_deploy")
a = ap.parse_args()
src = Path(a.checkpoint).resolve()
tok = Path(a.tokenizer).resolve()
out = Path(a.out).resolve() if a.out else src.with_name(src.name + "_deploy")
for need in ("model.safetensors", "config.json", "policy_preprocessor.json"):
if not (src / need).exists():
raise SystemExit(f"missing {need} in {src}")
if not (tok / "tokenizer.json").exists():
raise SystemExit(f"no tokenizer.json in {tok}")
if out.exists():
shutil.rmtree(out)
out.mkdir(parents=True)
for f in src.iterdir():
if f.is_file():
dst = out / f.name
if f.name == "model.safetensors":
try:
dst.symlink_to(f) # 8.9 GB: link, don't copy
except OSError:
shutil.copy2(f, dst)
else:
shutil.copy2(f, dst)
cfg = json.loads((out / "config.json").read_text())
removed = [k for k in ("pretrained_revision",) if k in cfg]
for k in removed:
cfg.pop(k)
(out / "config.json").write_text(json.dumps(cfg, indent=4))
pre = json.loads((out / "policy_preprocessor.json").read_text())
changed = 0
for st in pre.get("steps", []):
c = st.get("config") or {}
if "tokenizer_name" in c:
c["tokenizer_name"] = str(tok)
changed += 1
(out / "policy_preprocessor.json").write_text(json.dumps(pre, indent=2))
if changed == 0:
raise SystemExit("no tokenizer_name found in policy_preprocessor.json -- format changed?")
print(f"deploy-ready checkpoint: {out}")
print(f" config.json: removed {removed or 'nothing'}")
print(f" tokenizer_name -> {tok}")
|