File size: 2,829 Bytes
a2ffd07 | 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 64 65 66 67 | #!/bin/bash
# Per-layer bathroom↔toilet probe-readout correlation, PER MODEL VARIANT
# (base/lora/efuf/nullu), for four populations (bathroom_only, toilet_only,
# bathroom_toilet, random_cc3m/others). 4-variant probes; base caption reused
# across variants (one forward per variant).
#
# Tunables (override via env):
# DEVICE_ID, DTYPE, NUM_IMAGES, SEED, HOOK_TYPE, MAX_NEW_TOKENS
# METHODS space-separated subset of: base lora efuf nullu (default all)
# SCORE logit (default) | sigma
# BATH_PROBE / TOILET_PROBE default = 4-variant probes
# LORA_PATH / EFUF_PATH / NULLU_PATH / NULLU_LOWEST / NULLU_HIGHEST edit ckpts
# OUT, OUT_JSON output paths
#
# Usage:
# bash mechanistic_interp/scripts/correlation_baselines.sh
# NUM_IMAGES=100 DEVICE_ID=5 bash mechanistic_interp/scripts/correlation_baselines.sh
set -euo pipefail
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
MODEL_NAME="llava-hf/llava-1.5-7b-hf"
DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
NUM_IMAGES="${NUM_IMAGES:-100}"
SEED="${SEED:-0}"
METHODS="${METHODS:-base lora efuf nullu}"
SCORE="${SCORE:-logit}"
HOOK_TYPE="${HOOK_TYPE:-post}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
BATH_PROBE="${BATH_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_bathroom/post/seqprobe.pth}"
TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_toilet/post/seqprobe.pth}"
LORA_PATH="${LORA_PATH:-/data/caotue/multilayer-sae/adv_gen_outputs/run_bathroom_toilet_v2/lora_adapter}"
EFUF_PATH="${EFUF_PATH:-/data/caotue/multilayer-sae/EFUF/efuf/checkpoints/llava_vicuna_7b/bathroom_toilet_paper_10ep/epoch_002.pth}"
NULLU_PATH="${NULLU_PATH:-/data/caotue/nullu/edited_models/LLaVA-7B-top4-0-32-bathroom_toilet}"
NULLU_LOWEST="${NULLU_LOWEST:-8}"
NULLU_HIGHEST="${NULLU_HIGHEST:-32}"
OUT="${OUT:-${REPO_ROOT}/mechanistic_interp/graph/correlation_baselines.png}"
OUT_JSON="${OUT_JSON:-${REPO_ROOT}/mechanistic_interp/graph/correlation_baselines.json}"
cd "${REPO_ROOT}"
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"
export PYTHONUNBUFFERED=1
python -m mechanistic_interp.correlation_baselines \
--model_name "${MODEL_NAME}" \
--device_id "${DEVICE_ID}" \
--dtype "${DTYPE}" \
--methods ${METHODS} \
--lora_path "${LORA_PATH}" \
--efuf_path "${EFUF_PATH}" \
--nullu_path "${NULLU_PATH}" \
--nullu_lowest "${NULLU_LOWEST}" \
--nullu_highest "${NULLU_HIGHEST}" \
--bath_probe "${BATH_PROBE}" \
--toilet_probe "${TOILET_PROBE}" \
--num_images "${NUM_IMAGES}" \
--seed "${SEED}" \
--score "${SCORE}" \
--hook_type "${HOOK_TYPE}" \
--max_new_tokens "${MAX_NEW_TOKENS}" \
--out "${OUT}" \
--out_json "${OUT_JSON}" \
"$@"
|