File size: 4,615 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 97 98 99 100 101 102 103 104 105 106 107 108 109 110 111 | #!/bin/bash
# =============================================================================
# knowledge_suppression_trace.sh — matched-input internal-activation trace.
#
# For each bathroom-only image (object absent, scene present):
# 1) Base free-greedy-decodes K tokens — the matched input.
# 2) Each method (base / ours / Nullu) is teacher-forced on that K-token
# continuation.
# 3) At every selected layer, residuals at the K generated text positions
# are SAE-encoded and the pre-selected confident toilet features are
# gathered. Per-layer scalar = max (or mean) over (K × top-k).
# 4) Per-image PNG + population summary across all images.
#
# Reads `FEATURES_JSON` for the per-layer confident toilet features.
# =============================================================================
set -euo pipefail
cd "$(dirname "$0")/../.."
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"
FEATURES_JSON="${FEATURES_JSON:-mechanistic_interp/probes/selected_features_probe_top10.json}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet-bathroom/lora_adapter/samples.json}"
PROMPT="${PROMPT:-Describe this image.}"
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:-}" # empty → all bathroom_only samples
GEN_TOKENS="${GEN_TOKENS:-32}"
HOOK_TYPE="${HOOK_TYPE:-post}"
SAE_BATCH="${SAE_BATCH:-2048}"
CATEGORY="${CATEGORY:-bathroom_only}" # D_{A,¬c} from the spec
AGG="${AGG:-max}" # max | mean
# Deterministic assistant-side prefix. Empty → base free-decodes K tokens.
# Non-empty → use this fixed string (same for every method); K = its token count.
FIXED_ASSISTANT_PREFIX="${FIXED_ASSISTANT_PREFIX:-}"
if [[ -n "${FIXED_ASSISTANT_PREFIX}" ]]; then
# Sluggify the prefix (alnum+underscore, lowercase, capped at 40 chars) so
# different prefixes don't overwrite each other.
PREFIX_SLUG="$(echo "${FIXED_ASSISTANT_PREFIX}" \
| tr '[:upper:]' '[:lower:]' \
| sed -e 's/[^a-z0-9]\+/_/g' -e 's/^_//' -e 's/_$//' \
| cut -c1-40)"
PREFIX_TAG="fixed_${PREFIX_SLUG}"
else
PREFIX_TAG="freegen"
fi
PROBE_TAG="${PROBE_TAG:-$(basename "${FEATURES_JSON}" .json)}"
OUT_DIR="${OUT_DIR:-mechanistic_interp/knowledge_suppression/${CATEGORY}/${AGG}/${PROBE_TAG}/${PREFIX_TAG}/traces}"
GRAPH_DIR="${GRAPH_DIR:-mechanistic_interp/knowledge_suppression/${CATEGORY}/${AGG}/${PROBE_TAG}/${PREFIX_TAG}/graphs}"
# Nullu (optional). Empty NULLU_MODEL_PATH → Nullu curve omitted.
NULLU_MODEL_PATH="${NULLU_MODEL_PATH:-}"
NULLU_LOWEST_LAYER="${NULLU_LOWEST_LAYER:-16}"
NULLU_HIGHEST_LAYER="${NULLU_HIGHEST_LAYER:-32}"
# EFUF (optional). Empty EFUF_CKPT → EFUF curve omitted.
EFUF_CKPT="${EFUF_CKPT:-}"
ARGS=(
-m mechanistic_interp.knowledge_suppression_trace
--features_json "${FEATURES_JSON}"
--samples_json "${SAMPLES_JSON}"
--prompt "${PROMPT}"
--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}"
--gen_tokens "${GEN_TOKENS}"
--hook_type "${HOOK_TYPE}"
--sae_batch "${SAE_BATCH}"
--category "${CATEGORY}"
--agg "${AGG}"
--out_dir "${OUT_DIR}"
--graph_dir "${GRAPH_DIR}"
)
if [[ -n "${N_SAMPLES}" ]]; then
ARGS+=(--n_samples "${N_SAMPLES}")
fi
if [[ -n "${NULLU_MODEL_PATH}" ]]; then
ARGS+=(
--nullu_model_path "${NULLU_MODEL_PATH}"
--nullu_lowest_layer "${NULLU_LOWEST_LAYER}"
--nullu_highest_layer "${NULLU_HIGHEST_LAYER}"
)
fi
if [[ -n "${EFUF_CKPT}" ]]; then
ARGS+=(--efuf_ckpt "${EFUF_CKPT}")
fi
if [[ -n "${FIXED_ASSISTANT_PREFIX}" ]]; then
ARGS+=(--fixed_assistant_prefix "${FIXED_ASSISTANT_PREFIX}")
fi
PY="${PY:-/mnt/raid10/erwin/testing/multilayer-sae/.venv/bin/python}"
echo "Running: $PY ${ARGS[*]}"
"$PY" "${ARGS[@]}"
|