#!/bin/bash
# =============================================================================
# cosine_gradient.sh — per-layer mean cos(∇bathroom, ∇toilet) over 4 populations.
#
# For each image & layer l: gb = ∂score_bath_l/∂h_l, gt = ∂score_toilet_l/∂h_l
# (on the generated-caption tokens); record cos(gb, gt). Averaged over a random
# subsample (≤ NUM_IMAGES) of: bathroom_only, toilet_only, bathroom_toilet
# (samples.json) and random_cc3m (neg_cc3m_5k.json validation).
#
# Tunables: DEVICE_ID, DTYPE, NUM_IMAGES, SEED, HOOK_TYPE, MAX_NEW_TOKENS,
#           BATH_PROBE, TOILET_PROBE, OUT, OUT_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"
export WANDB_MODE=disabled

DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
NUM_IMAGES="${NUM_IMAGES:-100}"
SEED="${SEED:-0}"
HOOK_TYPE="${HOOK_TYPE:-post}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
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}"
OUT="${OUT:-mechanistic_interp/graph/cosine_gradient.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/cosine_gradient.json}"

REPO_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/../.." && pwd)"
cd "${REPO_ROOT}"
export PYTHONPATH="$(cd .. && pwd):$(pwd):${PYTHONPATH:-}"

python -m mechanistic_interp.cosine_gradient \
    --device_id      "${DEVICE_ID}" \
    --dtype          "${DTYPE}" \
    --bath_probe     "${BATH_PROBE}" \
    --toilet_probe   "${TOILET_PROBE}" \
    --num_images     "${NUM_IMAGES}" \
    --seed           "${SEED}" \
    --hook_type      "${HOOK_TYPE}" \
    --max_new_tokens "${MAX_NEW_TOKENS}" \
    --out            "${OUT}" \
    --out_json       "${OUT_JSON}" \
    "$@"