#!/usr/bin/env bash
# Purpose: Track M's supervised fine-tuning run: train a LoRA adapter with mlx_lm.lora,
#          report test-set perplexity, generate one answer with the adapter attached,
#          and append a run record to the lab notebook
# Platform: mac (Apple silicon, MLX); Tracks S, X and N use train-sft.py instead
# Minimum memory: 8 GB
# Assumes: mlx-lm installed in the active environment; make-dataset.py has been run so
#          that data-mlx/{train,valid,test}.jsonl exist; runlog.py sits next to this
#          script; run from the course directory
#
# Usage: bash train-sft-mlx.sh [MODEL] [ITERS]
#   MODEL defaults to the MLX community conversion of Qwen3-0.6B; ITERS to 300.
set -euo pipefail

MODEL="${1:-mlx-community/Qwen3-0.6B-bf16}"
ITERS="${2:-300}"
DATA="${DATA:-data-mlx}"
ADAPTERS="${ADAPTERS:-runs/sft-mlx-adapters}"
BATCH_SIZE="${BATCH_SIZE:-2}"
NUM_LAYERS="${NUM_LAYERS:-8}"
LEARNING_RATE="${LEARNING_RATE:-1e-4}"
SEED="${SEED:-0}"
LABBOOK="${LABBOOK:-labbook.md}"
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

die() { echo "train-sft-mlx: $*" >&2; exit 1; }
command -v mlx_lm.lora >/dev/null || die "mlx_lm.lora is not on PATH; install mlx-lm in the active environment"
for split in train valid test; do
  [[ -f "$DATA/$split.jsonl" ]] || die "$DATA/$split.jsonl is missing; run: python make-dataset.py --out-dir ."
done

echo "==> Training a LoRA adapter on $MODEL for $ITERS iterations"
START=$(date +%s)
mlx_lm.lora \
  --model "$MODEL" \
  --train \
  --data "$DATA" \
  --iters "$ITERS" \
  --batch-size "$BATCH_SIZE" \
  --num-layers "$NUM_LAYERS" \
  --learning-rate "$LEARNING_RATE" \
  --fine-tune-type lora \
  --adapter-path "$ADAPTERS"
ELAPSED=$(( $(date +%s) - START ))

echo "==> Test-set loss and perplexity with the adapter attached"
mlx_lm.lora \
  --model "$MODEL" \
  --data "$DATA" \
  --adapter-path "$ADAPTERS" \
  --test | tee "$ADAPTERS/test.txt"

echo "==> One answer, to see whether the house style was learned"
mlx_lm.generate \
  --model "$MODEL" \
  --adapter-path "$ADAPTERS" \
  --max-tokens 96 \
  --prompt "How much memory do the weights of a 5 billion parameter model need at BF16?"

echo "==> Recording the run in $LABBOOK"
DATA_SHA=$(shasum -a 256 "$DATA/train.jsonl" | cut -d' ' -f1)
TRAIN_N=$(wc -l < "$DATA/train.jsonl" | tr -d ' ')
VALID_N=$(wc -l < "$DATA/valid.jsonl" | tr -d ' ')
python "$HERE/runlog.py" --record --labbook "$LABBOOK" <<JSON
{
  "lab": "part-11/train-sft-mlx",
  "model": "$MODEL",
  "dataset": {"path": "$DATA/train.jsonl", "sha256": "$DATA_SHA",
              "train_examples": $TRAIN_N, "validation_examples": $VALID_N},
  "hyperparameters": {"method": "lora", "iters": $ITERS, "batch_size": $BATCH_SIZE,
                      "num_layers": $NUM_LAYERS, "learning_rate": $LEARNING_RATE,
                      "adapter_path": "$ADAPTERS", "seconds": $ELAPSED},
  "seed": $SEED,
  "losses": {},
  "scores": {},
  "notes": "copy the final train and validation loss, and the test perplexity, out of the mlx_lm.lora output above into losses"
}
JSON

echo "==> Done in ${ELAPSED}s. Adapter in $ADAPTERS; test output in $ADAPTERS/test.txt"