#!/bin/bash
# =============================================================================
# gradient_ascent_generate.sh — generate a full caption under bathroom-steering
#   at one layer, for a sweep of alphas (induction +α and suppression −α), on the
#   base or LoRA model. Reads off whether "toilet" appears in the generated text.
#
# New steering method (matches gradient_ascent.py): h' = h + ALPHA*‖h‖*v, so ALPHA
# is a FRACTION of the per-position residual norm (comparable across layers). v =
# unit bathroom-ascent vector from the 4variant probe gradient. NEGATIVE ALPHA
# suppresses. Probe stays the (base+lora-trained) 4variant checkpoint.
#
# Tunables (override inline, e.g. VARIANT=lora STEER_LAYER=20 bash <script>):
#   DEVICE_ID    GPU id                                       (default 0)
#   DTYPE        bfloat16 | float16 | float32                 (default bfloat16)
#   VARIANT      base | lora | nullu | efuf                   (default base)
#   BATH_PROBE   steering probe checkpoint (4variant bathroom)
#   STEER_LAYER  layer whose residual is steered              (default 16)
#   ALPHAS       fractions of ‖h‖ (incl. negatives + 0)
#   IMAGE_ID     fix the image stem (keep base/lora identical); empty = auto-pick
#   SAMPLES_JSON samples.json with base_mentions_object flags
#   QUESTION     prompt question
#   HOOK_TYPE    pre | mid | post                             (default post)
#   MAX_NEW_TOKENS / MAX_SEQ_TOKENS  generation / probe caps  (default 64 / 64)
#   OUT / OUT_JSON  output text-panel 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"
export GPU_MEM_RESERVE_FRAC="${GPU_MEM_RESERVE_FRAC:-0}"

DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
VARIANT="${VARIANT:-base}"
BATH_PROBE="${BATH_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_bathroom/post/seqprobe.pth}"
STEER_LAYER="${STEER_LAYER:-16}"
ALPHAS="${ALPHAS:--0.8 -0.4 -0.2 -0.1 -0.05 0 0.05 0.1 0.2 0.4 0.8}"
IMAGE_ID="${IMAGE_ID:-}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet_bathroom/samples.json}"
QUESTION="${QUESTION:-Describe this image.}"
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_generate.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/gradient_ascent_generate.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}"
    --variant        "${VARIANT}"
    --bath_probe     "${BATH_PROBE}"
    --steer_layer    "${STEER_LAYER}"
    --alphas         ${ALPHAS}
    --samples_json   "${SAMPLES_JSON}"
    --question       "${QUESTION}"
    --hook_type      "${HOOK_TYPE}"
    --max_new_tokens "${MAX_NEW_TOKENS}"
    --max_seq_tokens "${MAX_SEQ_TOKENS}"
    --out            "${OUT}"
    --out_json       "${OUT_JSON}"
)
# Fix the image so base and lora generate on the SAME sample.
if [ -n "${IMAGE_ID}" ]; then
    ARGS+=(--image_id "${IMAGE_ID}")
fi

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