File size: 2,660 Bytes
032a961 | 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 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/usr/bin/env python3
"""SPIRA training CLI.
Usage:
python3 -m spira.train --data data/train.bin \
--out artifacts/spira_model.npz \
[--snn-tokens 600000] [--am-tokens 1200000] [--epochs 2] [--quick]
Writes:
<out> : model checkpoint (npz)
<out>.train_log.jsonl : tick-level progress log
<out>.history.json : final training history
"""
from __future__ import annotations
import argparse
import json
import sys
import time
from .model import SPIRA, SPIRAConfig
def main() -> None:
ap = argparse.ArgumentParser()
ap.add_argument("--data", required=True)
ap.add_argument("--out", required=True)
ap.add_argument("--snn-tokens", type=int, default=600_000)
ap.add_argument("--am-tokens", type=int, default=1_200_000)
ap.add_argument("--epochs", type=int, default=2)
ap.add_argument("--vocab", type=int, default=4096)
ap.add_argument("--n-res", type=int, default=512)
ap.add_argument("--target-avg-len", type=float, default=4.5)
ap.add_argument("--quick", action="store_true", help="tiny smoke-test run")
args = ap.parse_args()
cfg = SPIRAConfig(
vocab_size=args.vocab,
n_res=args.n_res,
snn_tokens=args.snn_tokens,
am_tokens=args.am_tokens,
epochs=args.epochs,
target_avg_len=args.target_avg_len,
)
if args.quick:
cfg.snn_tokens, cfg.am_tokens, cfg.epochs = 20_000, 20_000, 1
data = open(args.data, "rb").read()
t0 = time.perf_counter()
log_path = args.out + ".train_log.jsonl"
log_f = open(log_path, "w")
def log(msg: str) -> None:
print(msg, flush=True)
log_f.write(json.dumps({"t": round(time.perf_counter() - t0, 1), "msg": msg}) + "\n")
log_f.flush()
def tick_log(d: dict) -> None:
log_f.write(json.dumps({"t": round(time.perf_counter() - t0, 1), **d}) + "\n")
log_f.flush()
model = SPIRA.build(data, cfg, log=log)
model.prepare()
# free the patch list — the registry carries everything we need
model._patches_train = None
log("[spira] training starts")
history = model.train(
log_fn=lambda msg: tick_log({"msg": msg}) or print(msg, flush=True),
ckpt_path=args.out,
)
history["build_meta"] = model.meta
history["config"] = {k: v for k, v in vars(cfg).items()}
history["wall_clock_s"] = time.perf_counter() - t0
with open(args.out + ".history.json", "w") as f:
json.dump(history, f, indent=2)
model.save(args.out)
log(f"[spira] DONE in {history['wall_clock_s']:.0f}s -> {args.out}")
log_f.close()
if __name__ == "__main__":
main()
|