#!/usr/bin/env bash
# Purpose: Track M's LoRA fine-tune. Train an adapter with mlx_lm.lora on the dataset
#          make-format-dataset.py wrote, report the held-out loss and perplexity, fuse the
#          adapter into a standalone model, generate one answer to see the format, and
#          append a run record to the lab notebook
# Platform: mac (Apple silicon, MLX). Tracks S, X and N use train-lora.py instead
# Minimum memory: 12 GB for a 1.7B to 4B base; --num-layers is the memory knob if it is tight
# Assumes: mlx-lm installed in the active environment so that mlx_lm.lora, mlx_lm.fuse and
#          mlx_lm.generate are on PATH; make-format-dataset.py has been run so that
#          data-mlx/{train,valid,test}.jsonl exist; sftlog.py sits next to this script;
#          run from the directory holding data-mlx
#
# Usage: bash train-lora-mlx.sh [MODEL] [ITERS]
#   MODEL  defaults to the MLX community conversion of Qwen3-1.7B
#   ITERS  defaults to 600
#
# Environment: DATA, ADAPTERS, FUSED, BATCH_SIZE, NUM_LAYERS, LEARNING_RATE, FINE_TUNE_TYPE
#              (lora, dora or full), SEED, LABBOOK, SKIP_FUSE=1 to stop after training.
set -euo pipefail

MODEL="${1:-mlx-community/Qwen3-1.7B-bf16}"
ITERS="${2:-600}"
DATA="${DATA:-data-mlx}"
ADAPTERS="${ADAPTERS:-runs/format-mlx-adapters}"
FUSED="${FUSED:-models/format-mlx-fused}"
BATCH_SIZE="${BATCH_SIZE:-2}"
NUM_LAYERS="${NUM_LAYERS:-16}"
LEARNING_RATE="${LEARNING_RATE:-1e-4}"
FINE_TUNE_TYPE="${FINE_TUNE_TYPE:-lora}"
SEED="${SEED:-0}"
LABBOOK="${LABBOOK:-labbook.md}"
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

die() { echo "train-lora-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"
command -v python3 >/dev/null || die "python3 is not on PATH"
[[ -f "$HERE/sftlog.py" ]] || die "sftlog.py is not next to this script"
for split in train valid test; do
  [[ -f "$DATA/$split.jsonl" ]] || die "$DATA/$split.jsonl is missing; run: python3 make-format-dataset.py --out-dir ."
done

TRAIN_N=$(wc -l < "$DATA/train.jsonl" | tr -d ' ')
VALID_N=$(wc -l < "$DATA/valid.jsonl" | tr -d ' ')
echo "==> Training a $FINE_TUNE_TYPE adapter on $MODEL"
echo "    $TRAIN_N training example(s), $VALID_N validation example(s), $ITERS iterations"
echo "    --mask-prompt puts the loss on the completion only, which is what"
echo "    completion_only_loss does on the other three tracks."

mkdir -p "$(dirname "$ADAPTERS")" "$(dirname "$FUSED")"
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 "$FINE_TUNE_TYPE" \
  --mask-prompt \
  --adapter-path "$ADAPTERS"
ELAPSED=$(( $(date +%s) - START ))

echo "==> Held-out loss and perplexity, with the adapter attached"
echo "    This is the test.jsonl split, which training never saw. It is the MLX"
echo "    equivalent of the evaluation loss the other tracks print each epoch."
mlx_lm.lora \
  --model "$MODEL" \
  --data "$DATA" \
  --adapter-path "$ADAPTERS" \
  --test | tee "$ADAPTERS/test.txt"

echo "==> One answer, to see whether the format was learned"
mlx_lm.generate \
  --model "$MODEL" \
  --adapter-path "$ADAPTERS" \
  --max-tokens 96 \
  --temp 0 \
  --seed 0 \
  --prompt "The monitoring system reports: the model gateway is unreachable from every machine we have tried this morning.

Triage this report. Reply with exactly three lines and nothing else:
Summary: one line, under twelve words
Severity: one of low, medium or high
Action: one imperative sentence"

if [[ "${SKIP_FUSE:-0}" == "1" ]]; then
  echo "==> Skipped fusing (SKIP_FUSE=1)"
else
  echo "==> Fusing the adapter into a standalone model at $FUSED"
  echo "    The fused directory is what llama.cpp's convert_hf_to_gguf.py reads if you"
  echo "    want a GGUF as well. mlx_lm.fuse can write GGUF directly with --export-gguf,"
  echo "    but that path covers a narrower set of architectures."
  mlx_lm.fuse \
    --model "$MODEL" \
    --adapter-path "$ADAPTERS" \
    --save-path "$FUSED"
fi

echo "==> Recording the run in $LABBOOK"
DATA_SHA=$(shasum -a 256 "$DATA/train.jsonl" | cut -d' ' -f1)
python3 "$HERE/sftlog.py" --record --labbook "$LABBOOK" <<JSON
{
  "lab": "part-13/train-lora-mlx",
  "model": "$MODEL",
  "dataset": {"path": "$DATA/train.jsonl", "sha256": "$DATA_SHA",
              "train_examples": $TRAIN_N, "validation_examples": $VALID_N},
  "hyperparameters": {"method": "$FINE_TUNE_TYPE", "iters": $ITERS,
                      "batch_size": $BATCH_SIZE, "num_layers": $NUM_LAYERS,
                      "learning_rate": $LEARNING_RATE, "mask_prompt": true,
                      "adapter_path": "$ADAPTERS", "fused_path": "$FUSED",
                      "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 the losses field"
}
JSON

echo
echo "==> Done in ${ELAPSED}s."
echo "    adapter:      $ADAPTERS"
echo "    test output:  $ADAPTERS/test.txt"
echo "    fused model:  $FUSED"
echo "    Next: serve it and score it against the base with evaluate-against-base.py."