#!/usr/bin/env bash
# Reproduces the fine-tuning run behind the submitted model.
#
#   bash run_training.sh [out_dir]
#
# Needs: one 48 GB GPU (reference: NVIDIA A40), ~17 GB disk for the base weights.
# Runtime on the reference machine: 88 minutes, peak VRAM 28.0 GB.
#
# The SFT files are published at:
#   https://huggingface.co/datasets/QBitsLK/sinhala-mcq-sft-v3
# and are downloaded here. run_data_pipeline.sh rebuilds them from the raw extractions.
set -euo pipefail

OUT_DIR="${1:-/workspace/outputs/mcq_sft_cpt_v3}"
BASE_MODEL="QBitsLK/gemma4-e4b-sinhala-cpt"
DATA_REPO="QBitsLK/sinhala-mcq-sft-v3"

hf download "$DATA_REPO" --repo-type dataset --local-dir ./sft_data
TRAIN=./sft_data/sft_train3.jsonl
VAL=./sft_data/sft_val3.jsonl
wc -l "$TRAIN" "$VAL"      # expect 21235 and 721

python sft_gemma4_mcq.py \
  --train "$TRAIN" --val "$VAL" \
  --model "$BASE_MODEL" \
  --out "$OUT_DIR" \
  --lora_r 32 --lora_alpha 32 --lr 1e-4 \
  --epochs 2 --batch_size 8 --grad_accum 4 \
  --max_seq_length 1024 --warmup_steps 10 \
  --eval_steps 150 --seed 3407

echo
echo "== the submitted checkpoint is $OUT_DIR/final_adapter (step 1328), NOT best_adapter."
echo "== see SYSTEM_DESCRIPTION.md section 5.2 for why."