File size: 2,115 Bytes
f0b1176 | 1 2 3 4 5 6 7 8 9 10 11 12 13 14 15 16 17 18 19 20 21 22 23 24 25 26 27 28 29 30 31 32 33 34 35 36 37 38 39 40 41 42 43 44 45 46 47 48 49 50 51 52 53 54 55 56 57 58 59 60 61 62 63 | #!/bin/bash
# Full training pipeline for the Trigger-Off experiment.
#
# Environment variables:
# MODEL_TYPE - "os_atlas" or "seeclick" (default: os_atlas)
# MODEL_PATH - HF hub ID or local path (default: from config)
# TRAIN_JSON - training data JSON (default: ./data/poisoned/train.json)
# EVAL_JSON - evaluation data JSON (default: ./data/poisoned/test.json)
# OUTPUT_DIR - checkpoint output dir (default: ./checkpoints/trigger_off)
# CONFIG - training config path (default: configs/training_config.yaml)
# RESUME_FROM - resume from checkpoint path (optional)
set -e
cd "$(dirname "$0")/.."
MODEL_TYPE=${MODEL_TYPE:-os_atlas}
TRAIN_JSON=${TRAIN_JSON:-./data/poisoned/train.json}
EVAL_JSON=${EVAL_JSON:-./data/poisoned/test.json}
OUTPUT_DIR=${OUTPUT_DIR:-./checkpoints/trigger_off}
CONFIG=${CONFIG:-configs/training_config.yaml}
echo "============================================================"
echo " Trigger-Off: LoRA Training"
echo "============================================================"
echo " MODEL_TYPE = $MODEL_TYPE"
echo " TRAIN_JSON = $TRAIN_JSON"
echo " EVAL_JSON = $EVAL_JSON"
echo " OUTPUT_DIR = $OUTPUT_DIR"
echo " CONFIG = $CONFIG"
echo "------------------------------------------------------------"
# Validate that the training data exists
if [ ! -f "$TRAIN_JSON" ]; then
echo "ERROR: Training JSON not found at $TRAIN_JSON"
echo " Run scripts/prepare_data.sh first."
exit 1
fi
RESUME_ARG=""
if [ -n "$RESUME_FROM" ]; then
RESUME_ARG="--resume_from $RESUME_FROM"
fi
MODEL_PATH_ARG=""
if [ -n "$MODEL_PATH" ]; then
MODEL_PATH_ARG="--model_path $MODEL_PATH"
fi
python -m src.training.train_lora \
--config "$CONFIG" \
--model_type "$MODEL_TYPE" \
--train_json "$TRAIN_JSON" \
--eval_json "$EVAL_JSON" \
--output_dir "$OUTPUT_DIR" \
$MODEL_PATH_ARG \
$RESUME_ARG
echo ""
echo "============================================================"
echo " Training complete."
echo " Checkpoint saved to: $OUTPUT_DIR"
echo "============================================================"
|