File size: 4,134 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 68 69 70 71 72 73 74 75 76 77 78 79 80 81 82 83 84 85 | #!/bin/bash
# =============================================================================
# probe_scoring.sh — mean 4variant probe scores (toilet & bathroom) per layer
# on BATHROOM-ONLY images, across model variants (base / lora / nullu / efuf).
#
# Reads the 4variant toilet & bathroom SequenceLayerProbes off each variant's
# residual stream; averages σ(logit) over the caption block per layer per image.
# High toilet score on bathroom-only images = internal bath→toilet hallucination.
#
# Tunables (override inline, e.g. DEVICE_ID=2 VARIANTS="base lora" bash <script>):
# DEVICE_ID GPU id (default 0)
# DTYPE bfloat16 | float16 | float32 (default bfloat16)
# VARIANTS space-sep variants to overlay (default "base lora nullu efuf")
# LORA_PATH LoRA(ours) adapter dir (variant=lora)
# EFUF_PATH EFUF .pth checkpoint (variant=efuf)
# NULLU_PATH Nullu edited-model dir (variant=nullu)
# TOILET_PROBE 4variant toilet SequenceLayerProbes checkpoint
# BATH_PROBE 4variant bathroom SequenceLayerProbes checkpoint
# IMAGE_FOLDER folder of images
# SAMPLES_JSON samples.json path
# BASE_MENTIONS filter bathroom-only by base_mentions_object (any|true|false)
# QUESTION prompt question
# FORCED_TEXT forced ASSISTANT answer (empty = generated caption)
# NUM_IMAGES cap on images; 0 = ALL bathroom-only (default 0)
# HOOK_TYPE pre | mid | post (default post)
# MAX_NEW_TOKENS caption generation length (default 64)
# MAX_SEQ_TOKENS caption tokens kept for the probe (default 64)
# OUT / OUT_JSON output PNG / JSON
# =============================================================================
export HF_HOME="/data/caotue/hf_cache"
export HF_DATASETS_CACHE="/data/caotue/hf_cache/datasets"
export TORCH_HOME="/data/caotue/torch_cache"
export TMPDIR="/data/caotue/tmp"
DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
VARIANTS="${VARIANTS:-base lora nullu efuf}"
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}"
TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_toilet/post/seqprobe.pth}"
BATH_PROBE="${BATH_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_bathroom/post/seqprobe.pth}"
IMAGE_FOLDER="${IMAGE_FOLDER:-/data/caotue/CC3M-Dataset/cc3m_images}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet_bathroom/samples.json}"
BASE_MENTIONS="${BASE_MENTIONS:-any}"
QUESTION="${QUESTION:-Describe this image.}"
FORCED_TEXT="${FORCED_TEXT:-}"
NUM_IMAGES="${NUM_IMAGES:-0}"
HOOK_TYPE="${HOOK_TYPE:-post}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
MAX_SEQ_TOKENS="${MAX_SEQ_TOKENS:-64}"
OUT="${OUT:-mechanistic_interp/graph/probe_scoring_bath_only.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/probe_scoring_bath_only.json}"
REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "${REPO_ROOT}"
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"
ARGS=(
--device_id "${DEVICE_ID}"
--dtype "${DTYPE}"
--variants ${VARIANTS}
--lora_path "${LORA_PATH}"
--efuf_path "${EFUF_PATH}"
--nullu_path "${NULLU_PATH}"
--toilet_probe "${TOILET_PROBE}"
--bath_probe "${BATH_PROBE}"
--image_folder "${IMAGE_FOLDER}"
--samples_json "${SAMPLES_JSON}"
--base_mentions "${BASE_MENTIONS}"
--question "${QUESTION}"
--num_images "${NUM_IMAGES}"
--hook_type "${HOOK_TYPE}"
--max_new_tokens "${MAX_NEW_TOKENS}"
--max_seq_tokens "${MAX_SEQ_TOKENS}"
--out "${OUT}"
--out_json "${OUT_JSON}"
)
if [ -n "${FORCED_TEXT}" ]; then
ARGS+=(--forced_text "${FORCED_TEXT}")
fi
python -m mechanistic_interp.probe_scoring "${ARGS[@]}" "$@"
|