#!/usr/bin/env bash
# MLX LoRA Fine-tuning for Ornith-1.0-9B
# Run this script to train adapters on your coding dataset

set -euo pipefail

MODEL="mlx-community/Ornith-1.0-9B-4bit"
DATA_DIR="$HOME/mlx-finetuning/ornith-9b/data"
ADAPTER_DIR="$HOME/mlx-finetuning/ornith-9b/adapters"
FUSED_DIR="$HOME/mlx-finetuning/ornith-9b/fused"

echo "[MLX] Fine-tuning Ornith-1.0-9B with LoRA"
echo "[MLX] Model: $MODEL"
echo "[MLX] Data: $DATA_DIR"
echo "[MLX] Adapters: $ADAPTER_DIR"
echo ""

# Step 1: Train
echo "[Step 1/3] Training LoRA adapters..."
mlx_lm.lora \
  --model "$MODEL" \
  --train \
  --data "$DATA_DIR" \
  --batch-size 1 \
  --grad-accumulation-steps 4 \
  --num-layers 8 \
  --iters 200 \
  --learning-rate 2e-4 \
  --mask-prompt \
  --save-every 50 \
  --adapter-path "$ADAPTER_DIR"

# Step 2: Evaluate
echo ""
echo "[Step 2/3] Evaluating on validation set..."
mlx_lm.lora \
  --model "$MODEL" \
  --adapter-path "$ADAPTER_DIR" \
  --data "$DATA_DIR" \
  --test

# Step 3: Fuse (optional - creates a standalone model)
echo ""
echo "[Step 3/3] Fusing adapters into model..."
mlx_lm.fuse \
  --model "$MODEL" \
  --adapter-path "$ADAPTER_DIR" \
  --save-path "$FUSED_DIR" \
  --de-quantize

echo ""
echo "[✓] Fine-tuning complete!"
echo "[✓] Adapters saved to: $ADAPTER_DIR"
echo "[✓] Fused model saved to: $FUSED_DIR"
echo ""
echo "To use the fine-tuned model:"
echo "  mlx_lm.chat --model $FUSED_DIR"
echo "  OR"
echo "  mlx_lm.server --model $FUSED_DIR --port 8082"