#!/usr/bin/env bash
# README quick start on one GPU: NGPU=1 ./run_train.sh with the shipped llama3_debugmodel recipe
# (c4_test data in the repo), 2000 steps via kv_recipes.py. Metrics are parsed from the training log.
set -o pipefail
rm -rf outputs
echo "KVERITAS_PHASE name=train"
NGPU=1 MODULE=kv_recipes CONFIG=llama3_debugmodel_long bash run_train.sh 2>&1 | sed 's/\x1b\[[0-9;]*m//g' | tee ../train.out | grep -E "step: +(1|[0-9]*00) " || exit 1
python - ../train.out <<'PY'
import re, sys
rows = [(int(s), float(l), float(t.replace(",", ""))) for s, l, t in
        re.findall(r"step: +(\d+) +loss: +([0-9.]+).*?tps: ([0-9,]+)", open(sys.argv[1]).read())]
for s, l, _ in rows:
    if s == 1 or s % 50 == 0:
        print(f"KVERITAS_METRIC name=loss value={l} step={s}")
tail = [l for s, l, _ in rows[-100:]]
tps = [t for _, _, t in rows[10:]]
print(f"KVERITAS_METRIC name=mean_tokens_per_second value={sum(tps)/len(tps):.6g}")
print(f"KVERITAS_CLAIM metric=first_loss value={rows[0][1]}")
print(f"KVERITAS_CLAIM metric=final_loss_mean_last100 value={sum(tail)/len(tail):.6g}")
PY
