Mar 9, 2026, 6:06 PM

trigger-off/scripts/train.sh

@@ -0,0 +1,62 @@
+#!/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 "============================================================"