File size: 4,423 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 86 87 88 89 90 91 92 93 94 95 96 | #!/bin/bash
# =============================================================================
# per_layer_effect.sh — fixed-prefix probe per intervention layer.
#
# For each filter-passing sample, teacher-force `prompt + " " + fixed_assistant_prefix`
# (no decode) under base, then per-L_int ΔW, then ΔW@all-layers. At every
# capture layer L_cap: residual @ K prefix positions → FrozenSAEEncoder →
# pool over K → gather selected features + apply layer-L probe.
#
# Output per sample (under ${GRAPH_DIR}/${image_id}/):
# * per_layer_effect_L*.png — feature grids (L_int in groups of 4)
# * per_layer_effect_probe_L*.png — probe-output grids (same grouping)
# * per_layer_effect_all.png — base vs ΔW@all summary (features + probes)
# Plus ${OUT_DIR}/${image_id}.pt per sample with the raw tensors.
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"
# Object names for the dataset
OBJECT_NAME="${OBJECT_NAME:-toilet}"
OBJECT2_NAME="${OBJECT2_NAME:-bathroom}"
FEATURES_JSON="${FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_top20.json}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet-bathroom/lora_adapter/samples.json}"
PROMPT="${PROMPT:-Describe this image.}"
FIXED_ASSISTANT_PREFIX="${FIXED_ASSISTANT_PREFIX:-In this bathroom there is a shower, a sink and}"
HF_DATASET="${HF_DATASET:-pbcong/bathroom-toilet}"
HF_SPLIT="${HF_SPLIT:-validation}"
ID_COL="${ID_COL:-image_id}"
ADAPTER_PATH="${ADAPTER_PATH:-mechanistic_interp/toilet-bathroom/lora_adapter/adapter_model.safetensors}"
ADAPTER_CFG="${ADAPTER_CFG:-mechanistic_interp/toilet-bathroom/lora_adapter/adapter_config.json}"
SAE_CKPT="${SAE_CKPT:-training/multilayer_sae_ckpt/last.ckpt}"
MODEL_NAME="${MODEL_NAME:-llava-hf/llava-1.5-7b-hf}"
DEVICE="${DEVICE:-cuda:0}"
DTYPE="${DTYPE:-bfloat16}"
N_SAMPLES="${N_SAMPLES:-}" # unset/empty → process ALL filter-passing samples in the category
HOOK_TYPE="${HOOK_TYPE:-post}"
POOL="${POOL:-max}" # max | mean — must match probe-training pool
OUT_DIR="${OUT_DIR:-mechanistic_interp/per_layer_effect_traces}"
AGG="${AGG:-max}" # mean | max (reducer over top-k feats per (L_int, L_cap))
CATEGORY="${CATEGORY:-any}" # any | {object1}_only | {object2}_only | {object1}_{object2}
# Per-layer linear probe head state-dict (probes.{L}.weight, .bias). Powers
# the per_layer_effect_probe.png companion plot.
PROBES_PATH="${PROBES_PATH:-mechanistic_interp/probes/probes_gen_bathroom_toilet.pt}"
PROBE_OUTPUT="${PROBE_OUTPUT:-prob}" # prob | logit
# Tag derived from the features-json filename so different probe selections
# don't overwrite each other's graphs.
PROBE_TAG="${PROBE_TAG:-$(basename "${FEATURES_JSON}" .json)}"
GRAPH_DIR="${GRAPH_DIR:-mechanistic_interp/graph_per_layer_effect/${CATEGORY}/${AGG}/${PROBE_TAG}}"
ARGS=(
-m mechanistic_interp.per_layer_effect
--features_json "${FEATURES_JSON}"
--samples_json "${SAMPLES_JSON}"
--prompt "${PROMPT}"
--fixed_assistant_prefix "${FIXED_ASSISTANT_PREFIX}"
--hf_dataset "${HF_DATASET}"
--hf_split "${HF_SPLIT}"
--id_col "${ID_COL}"
--adapter_path "${ADAPTER_PATH}"
--adapter_cfg "${ADAPTER_CFG}"
--sae_ckpt "${SAE_CKPT}"
--model_name "${MODEL_NAME}"
--device "${DEVICE}"
--dtype "${DTYPE}"
--hook_type "${HOOK_TYPE}"
--pool "${POOL}"
--out_dir "${OUT_DIR}"
--graph_dir "${GRAPH_DIR}"
--agg "${AGG}"
--category "${CATEGORY}"
--probes_path "${PROBES_PATH}"
--probe_output "${PROBE_OUTPUT}"
--object_name "${OBJECT_NAME}"
)
# Forward --object2_name only when set (enables multi-object category filtering)
if [[ -n "${OBJECT2_NAME}" ]]; then
ARGS+=(--object2_name "${OBJECT2_NAME}")
fi
# Only forward --n_samples if the user actually set it; otherwise let the
# python default (0 = "all filter-passing samples") apply.
if [[ -n "${N_SAMPLES}" ]]; then
ARGS+=(--n_samples "${N_SAMPLES}")
fi
echo "Running: python ${ARGS[*]}"
python "${ARGS[@]}" |