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

MODEL="${MODEL:-mlx-community/gemma-4-e2b-it-4bit}"
DATASET="${DATASET:-./data}"
ITERS="${ITERS:-50}"
MAX_SEQ_LENGTH="${MAX_SEQ_LENGTH:-1024}"
ADAPTER_PATH="${ADAPTER_PATH:-./adapters/gemma-e2b-code-smoke}"
RESUME_ADAPTER_PATH="${RESUME_ADAPTER_PATH:-}"
STEPS_PER_SAVE="${STEPS_PER_SAVE:-25}"

cmd=(
  python -m mlx_vlm.lora
)
if [[ -n "$RESUME_ADAPTER_PATH" ]]; then
  cmd+=(--adapter-path "$RESUME_ADAPTER_PATH")
fi

"${cmd[@]}" \
  --model-path "$MODEL" \
  --dataset "$DATASET" \
  --split train \
  --iters "$ITERS" \
  --batch-size 1 \
  --max-seq-length "$MAX_SEQ_LENGTH" \
  --grad-checkpoint \
  --train-on-completions \
  --lora-rank 8 \
  --lora-alpha 16 \
  --learning-rate 2e-5 \
  --steps-per-report 5 \
  --steps-per-eval 25 \
  --steps-per-save "$STEPS_PER_SAVE" \
  --output-path "$ADAPTER_PATH"