#!/usr/bin/env bash
set -euo pipefail

source "$(cd "$(dirname "$0")" && pwd)/_common.sh"
require_venv
cd_nanogpt

# Requires torch with MPS support (Apple Silicon). If MPS isn't available,
# fall back to CPU.
DEVICE="mps"
if ! "${VENV_PY}" -c "import torch; import sys; sys.exit(0 if torch.backends.mps.is_available() else 1)"; then
  echo "MPS not available; falling back to CPU" >&2
  DEVICE="cpu"
fi

"${VENV_PY}" train.py config/train_shakespeare_char.py \
  --device="${DEVICE}" --compile=False --dtype=float32 \
  --eval_interval=10 --eval_iters=10 --log_interval=10 \
  --block_size=64 --batch_size=12 \
  --n_layer=4 --n_head=4 --n_embd=128 \
  --max_iters=60 --lr_decay_iters=60 --dropout=0.0 \
  --always_save_checkpoint=True

echo "Done. Checkpoint should exist at: ${NANOGPT_DIR}/out-shakespeare-char/ckpt.pt"
