#!/usr/bin/env bash
# README recipe: train a 15M-parameter Llama 2 (dim 288, 6 layers, 6 heads) on TinyStories, then sample
# with the C inference program (run.c, compiled beforehand with `make run`). 15,000 iterations instead of
# 20,000 to fit the time budget. Metrics are parsed from train.py and run output.
set -o pipefail
OUT=/var/tmp/kv-repro/025-karpathy-llama2.c/data/run
rm -rf "$OUT"
echo "KVERITAS_INPUT src=seed:1337"
echo "KVERITAS_PHASE name=train"
python -u train.py --dim=288 --n_layers=6 --n_heads=6 --max_iters=15000 \
  --out_dir="$OUT" 2>&1 | tee ../train.out | grep -E "^step|tokens per iteration|^[0-9]+0000 \|" || exit 1
python - ../train.out <<'PY'
import re, sys
for step, tr, va in re.findall(r"step (\d+): train loss ([0-9.]+), val loss ([0-9.]+)", open(sys.argv[1]).read()):
    print(f"KVERITAS_METRIC name=train_loss value={tr} step={step}")
    print(f"KVERITAS_METRIC name=val_loss value={va} step={step}")
print(f"KVERITAS_CLAIM metric=final_val_loss value={va}")
PY
echo "KVERITAS_PHASE name=sample"
./run "$OUT/model.bin" -t 0.8 -n 256 -s 42 -z tokenizer.bin 2>&1 | tee ../sample.out
tps=$(grep -oP 'achieved tok/s: \K[0-9.]+' ../sample.out)
echo "KVERITAS_METRIC name=c_inference_tokens_per_second value=$tps"
echo "KVERITAS_CLAIM metric=c_inference_tokens_per_second value=$tps"
