#!/bin/bash
# =============================================================================
# gradient_ascent.sh — bathroom→toilet causal influence map via probe gradient.
#
# For each BATHROOM-ONLY image (bathroom=1 & toilet=0): nudge the residual stream
# at layer l along ĝ = ∂ score_bath_l/∂ h_l (unit gradient ascent in the bathroom
# direction), step scaled by ALPHA*‖h_l‖ so ALPHA is a fraction of the residual norm
# (comparable across layers), propagate, and measure Δ toilet score at every layer
# l' >= l. Averaged → heatmap (intervention layer l × readout l').
#
# Tunables (override inline, e.g. ALPHA=2 NUM_IMAGES=128 bash <script>):
#   DEVICE_ID       GPU id                                  (default 0)
#   DTYPE           bfloat16 | float16 | float32            (default bfloat16)
#   BATH_PROBE      steered-direction SequenceLayerProbes checkpoint (e.g. scene)
#   TOILET_PROBE    measured-readout  SequenceLayerProbes checkpoint (e.g. object)
#   IMAGE_FOLDER    folder of images
#   SAMPLES_JSON    samples.json with base_mentions_object flags (empty to disable)
#   BASE_PROMPT     prompt whose base_mentions_object selects non-hallucinating imgs
#   ── Alternative image selection from an HF dataset (no samples.json needed) ──
#   HF_DATASET      if set, pick scene-only images (SCENE_COL=1 & OBJECT_COL=0) from it
#   SPLIT           HF split for selection                  (default validation)
#   SCENE_COL       steer concept column (present)          e.g. dining_room
#   OBJECT_COL      readout concept column (absent)         e.g. plate
#   STEER_NAME      display name of steered concept (plots)
#   READOUT_NAME    display name of readout concept (plots)
#   QUESTION        prompt question
#   NUM_IMAGES      cap on images; 0 = ALL non-hallucinating bathroom-only (default 0)
#   ALPHAS          steps as FRACTION of ‖h_l‖ (h'=h+ALPHA*‖h_l‖*ĝ) (default "0.05 0.1 0.2 0.4 0.8")
#   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 heatmap PNG / matrix 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}"
BATH_PROBE="${BATH_PROBE:-/data/caotue/latent_probes/seqprobes_bathroom/post/seqprobe.pth}"
TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_toilet/post/seqprobe.pth}"
IMAGE_FOLDER="${IMAGE_FOLDER:-/data/caotue/CC3M-Dataset/cc3m_images}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet_bathroom/samples.json}"
BASE_PROMPT="${BASE_PROMPT:-Describe this image.}"
# HF-dataset image selection (alternative to SAMPLES_JSON). Empty = use SAMPLES_JSON.
HF_DATASET="${HF_DATASET:-}"
SPLIT="${SPLIT:-validation}"
SCENE_COL="${SCENE_COL:-}"
OBJECT_COL="${OBJECT_COL:-}"
STEER_NAME="${STEER_NAME:-bathroom}"
READOUT_NAME="${READOUT_NAME:-toilet}"
# BASE_MENTIONS: false = base does NOT hallucinate toilet (+alpha induction);
#                true  = base DOES hallucinate toilet (-alpha suppression/necessity); any = no filter.
BASE_MENTIONS="${BASE_MENTIONS:-false}"
QUESTION="${QUESTION:-Describe this image.}"
# FORCED_TEXT: if set, skip generation and force the ASSISTANT answer to this exact
# string (controlled, constant context across images). Empty = freely-generated caption.
FORCED_TEXT="${FORCED_TEXT:-This image features a bathroom with a}"
# NUM_IMAGES=0 → ALL non-hallucinating bathroom-only images (base says no toilet).
NUM_IMAGES="${NUM_IMAGES:-0}"
ALPHAS="${ALPHAS:-0.05 0.1 0.2 0.4 0.8}"
HOOK_TYPE="${HOOK_TYPE:-post}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
MAX_SEQ_TOKENS="${MAX_SEQ_TOKENS:-64}"
OUT="${OUT:-mechanistic_interp/graph/gradient_ascent_bath2toilet.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/gradient_ascent_bath2toilet.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}"
    --bath_probe     "${BATH_PROBE}"
    --toilet_probe   "${TOILET_PROBE}"
    --image_folder   "${IMAGE_FOLDER}"
    --base_prompt    "${BASE_PROMPT}"
    --base_mentions  "${BASE_MENTIONS}"
    --question       "${QUESTION}"
    --num_images     "${NUM_IMAGES}"
    --alphas         ${ALPHAS}
    --hook_type      "${HOOK_TYPE}"
    --max_new_tokens "${MAX_NEW_TOKENS}"
    --max_seq_tokens "${MAX_SEQ_TOKENS}"
    --steer_name     "${STEER_NAME}"
    --readout_name   "${READOUT_NAME}"
    --out            "${OUT}"
    --out_json       "${OUT_JSON}"
)
# Image selection: HF dataset (scene-only via columns) takes priority over samples.json.
if [ -n "${HF_DATASET}" ]; then
    ARGS+=(--hf_dataset "${HF_DATASET}" --split "${SPLIT}"
           --scene_col "${SCENE_COL}" --object_col "${OBJECT_COL}")
elif [ -n "${SAMPLES_JSON}" ]; then
    ARGS+=(--samples_json "${SAMPLES_JSON}")
fi
# Only pass --forced_text when non-empty (empty = use generated captions).
if [ -n "${FORCED_TEXT}" ]; then
    ARGS+=(--forced_text "${FORCED_TEXT}")
fi

python -m mechanistic_interp.gradient_ascent "${ARGS[@]}" "$@"