#!/usr/bin/env bash
# Stage ① alignment training (T3.2).
# Loads align.jsonl (50万 target), freezes LLM, trains Encoder+Projector.
set -euo pipefail

DATA="${1:-data/align.jsonl}"
MAX_STEPS="${2:-10000}"
CKPT_DIR="${3:-checkpoints/stage1}"

cd "$(dirname "$0")/.."
export PYTHONPATH=src PYTHONUNBUFFERED=1
.venv/bin/python -m tsmm.train.stage1 \
  --data "$DATA" \
  --ckpt_dir "$CKPT_DIR" \
  --max_steps "$MAX_STEPS" \
  --bs 8 \
  --grad_accum 4 \
  --ctx 512 \
  --lr 1e-4 \
  --lambda_contrast 0.1 \
  --perturb_std 0.1 \
  --log_every 20 \
  --ckpt_every 2000 \
  --tensorboard
# Resume support: pass --resume to continue from the latest step ckpt.
# Usage: scripts/train_stage1.sh data/align.jsonl 30000 checkpoints/stage1 --resume