File size: 1,858 Bytes
164e036 | 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 | #!/usr/bin/env python3
"""Print the tinker:// path of a stage's final checkpoint.
python3 scripts/last_checkpoint.py runs/sft [--kind state|sampler]
`make all` uses this to hand stage N's output to stage N+1 without you having
to copy a run id out of the scrollback. Reads the checkpoints.jsonl the cookbook
writes into each run directory, preferring the record flagged `final`.
"""
import argparse
import json
import os
import sys
def main():
ap = argparse.ArgumentParser()
ap.add_argument("run_dir")
ap.add_argument("--kind", choices=("state", "sampler"), default="state",
help="state resumes training; sampler is for inference/export")
args = ap.parse_args()
path = os.path.join(args.run_dir, "checkpoints.jsonl")
if not os.path.exists(path):
sys.exit("No checkpoints.jsonl in %s -- has that stage run?" % args.run_dir)
records = []
with open(path) as f:
for line in f:
line = line.strip()
if line:
try:
records.append(json.loads(line))
except json.JSONDecodeError:
continue
if not records:
sys.exit("%s is empty" % path)
key = "%s_path" % args.kind
# The cookbook marks the last checkpoint with {"name": "final", ...}; there
# is no boolean "final" field. Checking r.get("final") never matched, so
# this silently degraded to "whatever was written last" -- which is usually
# right and would have hidden the bug until a run saved after its final.
finals = [r for r in records if r.get("name") == "final" and r.get(key)]
chosen = (finals or [r for r in records if r.get(key)])[-1:]
if not chosen:
sys.exit("No %s checkpoint recorded in %s" % (args.kind, path))
print(chosen[0][key])
if __name__ == "__main__":
main()
|