#!/bin/bash
#SBATCH -J smk
#SBATCH -p gpu-h100
#SBATCH --gpus=h100:1
#SBATCH -c 8
#SBATCH --mem=32g
#SBATCH -t 00:30:00
#SBATCH -o logs/smk-%j.out
source /etc/profile
SCR=/scratch/zt1/project/msml612/user/$USER
module load python/3.10.10
export PYTHONPATH=
cd $SCR/pcmd/src
$SCR/hf_env/bin/python - <<'PY'
import sys; sys.path.insert(0,'.')
import numpy as np
from data import build_addition
X, A, tok = build_addition(300, 6, 32, seed=3, exact=True)
bad = 0
for i in range(len(X)):
    s = [tok.itos[t] for t in X[i] if t != tok.pad]
    a = int("".join(s[:s.index("+")])); b = int("".join(s[s.index("+")+1:s.index("=")]))
    ans = int("".join(reversed(s[s.index("=")+1:])))
    bad += (a + b != ans)
print(f"  data check: {bad}/300 wrong; answer positions/example = {A.sum(1)[:3]}")
PY
for M in mdlm matched distill full; do
  echo "--- $M ---"
  $SCR/hf_env/bin/python train_add.py --mode $M --digits 3 --seq_len 16 \
    --n_train 4000 --n_eval 60 --eval_bs 30 --d 64 --layers 2 --heads 2 \
    --bs 32 --steps 200 --K 4 --out /tmp/a$M 2>&1 | tail -2
done