#!/bin/bash
# ============================================================
# TMF921 Intent Translation Training — One-Command Runner
# ============================================================
# Run on a machine with an NVIDIA GPU (tested on RTX 6000 Ada 48GB)
#
# Usage:
#   git clone https://huggingface.co/nraptisss/intent-translation-training
#   cd intent-translation-training
#   chmod +x run.sh
#   ./run.sh             # train + evaluate with defaults
#   ./run.sh --eval-only # evaluate an existing adapter
# ============================================================

set -euo pipefail

# ── Configuration ────────────────────────────────────────────
BASE_MODEL="Qwen/Qwen3-8B"
DATASET="nraptisss/TMF921-intent-to-config-augmented"
OUTPUT_DIR="./output"
HUB_MODEL_ID=""  # e.g. "nraptisss/Qwen3-8B-TMF921-Intent-QLora"

# Training hyperparameters
EPOCHS=3
LR=1e-4
BATCH_SIZE=4
GRAD_ACCUM=8
MAX_LENGTH=4096
LORA_R=32
LORA_ALPHA=64

# Eval
EVAL_SAMPLES=200  # -1 for full test set
# ─────────────────────────────────────────────────────────────

echo "╔══════════════════════════════════════════════════════════════════╗"
echo "║  TMF921 Intent Translation — Training Pipeline                  ║"
echo "╚══════════════════════════════════════════════════════════════════╝"

# Check CUDA
echo ""
echo "Checking GPU …"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader || {
    echo "ERROR: No NVIDIA GPU found. Exiting."
    exit 1
}

# Install dependencies
echo ""
echo "Installing dependencies …"
pip install -r requirements.txt 2>&1 | tail -5

# Login to HF (if token is set)
if [ -n "${HF_TOKEN:-}" ]; then
    echo ""
    echo "Logging in to Hugging Face …"
    huggingface-cli login --token "$HF_TOKEN" --add-to-git-credential
fi

# ── Training ─────────────────────────────────────────────────
if [ "${1:-}" != "--eval-only" ]; then
    echo ""
    echo "Starting training …"
    python train.py \
        --base_model "$BASE_MODEL" \
        --dataset "$DATASET" \
        --output_dir "$OUTPUT_DIR" \
        --epochs $EPOCHS \
        --lr $LR \
        --batch_size $BATCH_SIZE \
        --grad_accum $GRAD_ACCUM \
        --max_length $MAX_LENGTH \
        --lora_r $LORA_R \
        --lora_alpha $LORA_ALPHA \
        ${HUB_MODEL_ID:+--push_to_hub --hub_model_id "$HUB_MODEL_ID"}
fi

# ── Evaluation ───────────────────────────────────────────────
echo ""
echo "Starting evaluation …"
python evaluate.py \
    --base_model "$BASE_MODEL" \
    --adapter_path "$OUTPUT_DIR" \
    --dataset "$DATASET" \
    --num_samples $EVAL_SAMPLES \
    --output_file "${OUTPUT_DIR}/eval_results.json"

echo ""
echo "╔══════════════════════════════════════════════════════════════════╗"
echo "║  ✅  Pipeline complete!                                         ║"
echo "║  Model:   ${OUTPUT_DIR}                                         ║"
echo "║  Metrics: ${OUTPUT_DIR}/eval_results.json                       ║"
echo "╚══════════════════════════════════════════════════════════════════╝"