#!/usr/bin/env bash
# llm.c quick start: fine-tune GPT-2 124M (bf16, CUDA) on tinyshakespeare for one epoch with default settings.
set -o pipefail
cd "$(dirname "$0")"
echo "KVERITAS_PHASE name=train"
./train_gpt2cu 2>&1 | awk '
{ print; fflush() }
/^step +[0-9]+\/[0-9]+ \| loss/ {
  match($0, /step +[0-9]+/); split(substr($0,RSTART,RLENGTH), s, " "); step = s[2]
  match($0, /loss [0-9.]+/); split(substr($0,RSTART,RLENGTH), l, " ")
  match($0, /[0-9.]+ tok\/s/); split(substr($0,RSTART,RLENGTH), t, " "); tok = t[1]
  printf "KVERITAS_METRIC name=train_loss value=%s step=%s\n", l[2], step; fflush()
}
/^val loss [0-9.]+/ {
  vl = $3; printf "KVERITAS_METRIC name=val_loss value=%s step=%d\n", vl, step + 0; fflush()
}
/^total average iteration time:/ { ms = $5 }
END {
  print "KVERITAS_PHASE name=summary"
  if (vl != "") printf "KVERITAS_CLAIM metric=final_val_loss value=%s\n", vl
  if (ms != "") printf "KVERITAS_CLAIM metric=avg_iteration_ms value=%s\n", ms
  if (tok != "") printf "KVERITAS_CLAIM metric=tokens_per_second value=%s\n", tok
}'
