#!/usr/bin/env bash
# Stage ② SFT training (T4.1). Loads stage ① ckpt, enables LoRA(r=16) on q/v_proj.
set -euo pipefail

DATA="${1:-data/sft.jsonl}"
MAX_STEPS="${2:-10000}"
STAGE1_CKPT="${3:-checkpoints/stage1/stage1_final.pt}"
CKPT_DIR="${4:-checkpoints/stage2}"

cd "$(dirname "$0")/.."
export PYTHONPATH=src PYTHONUNBUFFERED=1
.venv/bin/python -m tsmm.train.stage2 \
  --data "$DATA" \
  --stage1_ckpt "$STAGE1_CKPT" \
  --ckpt_dir "$CKPT_DIR" \
  --max_steps "$MAX_STEPS" \
  --bs 4 \
  --grad_accum 8 \
  --ctx 1024 \
  --lr 1e-4 \
  --lora_r 16 \
  --log_every 20 \
  --ckpt_every 2000 \
  --tensorboard
# Resume support: pass --resume to continue from the latest step ckpt + lora_adapter.