#!/bin/bash
# =============================================================================
# probe_scoring.sh — mean 4variant probe scores (toilet & bathroom) per layer
# on BATHROOM-ONLY images, across model variants (base / lora / nullu / efuf).
#
# Reads the 4variant toilet & bathroom SequenceLayerProbes off each variant's
# residual stream; averages σ(logit) over the caption block per layer per image.
# High toilet score on bathroom-only images = internal bath→toilet hallucination.
#
# Tunables (override inline, e.g. DEVICE_ID=2 VARIANTS="base lora" bash <script>):
#   DEVICE_ID       GPU id                                  (default 0)
#   DTYPE           bfloat16 | float16 | float32            (default bfloat16)
#   VARIANTS        space-sep variants to overlay           (default "base lora nullu efuf")
#   LORA_PATH       LoRA(ours) adapter dir (variant=lora)
#   EFUF_PATH       EFUF .pth checkpoint   (variant=efuf)
#   NULLU_PATH      Nullu edited-model dir (variant=nullu)
#   TOILET_PROBE    4variant toilet SequenceLayerProbes checkpoint
#   BATH_PROBE      4variant bathroom SequenceLayerProbes checkpoint
#   IMAGE_FOLDER    folder of images
#   SAMPLES_JSON    samples.json path
#   BASE_MENTIONS   filter bathroom-only by base_mentions_object (any|true|false)
#   QUESTION        prompt question
#   FORCED_TEXT     forced ASSISTANT answer (empty = generated caption)
#   NUM_IMAGES      cap on images; 0 = ALL bathroom-only     (default 0)
#   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 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"

DEVICE_ID="${DEVICE_ID:-0}"
DTYPE="${DTYPE:-bfloat16}"
VARIANTS="${VARIANTS:-base lora nullu efuf}"
LORA_PATH="${LORA_PATH:-/data/caotue/multilayer-sae/adv_gen_outputs/run_bathroom_toilet_v2/lora_adapter}"
EFUF_PATH="${EFUF_PATH:-/data/caotue/multilayer-sae/EFUF/efuf/checkpoints/llava_vicuna_7b/bathroom_toilet_paper_10ep/epoch_002.pth}"
NULLU_PATH="${NULLU_PATH:-/data/caotue/nullu/edited_models/LLaVA-7B-top4-0-32-bathroom_toilet}"
TOILET_PROBE="${TOILET_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_toilet/post/seqprobe.pth}"
BATH_PROBE="${BATH_PROBE:-/data/caotue/latent_probes/seqprobes_4variant_bathroom/post/seqprobe.pth}"
IMAGE_FOLDER="${IMAGE_FOLDER:-/data/caotue/CC3M-Dataset/cc3m_images}"
SAMPLES_JSON="${SAMPLES_JSON:-mechanistic_interp/toilet_bathroom/samples.json}"
BASE_MENTIONS="${BASE_MENTIONS:-any}"
QUESTION="${QUESTION:-Describe this image.}"
FORCED_TEXT="${FORCED_TEXT:-}"
NUM_IMAGES="${NUM_IMAGES:-0}"
HOOK_TYPE="${HOOK_TYPE:-post}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
MAX_SEQ_TOKENS="${MAX_SEQ_TOKENS:-64}"
OUT="${OUT:-mechanistic_interp/graph/probe_scoring_bath_only.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/probe_scoring_bath_only.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}"
    --variants       ${VARIANTS}
    --lora_path      "${LORA_PATH}"
    --efuf_path      "${EFUF_PATH}"
    --nullu_path     "${NULLU_PATH}"
    --toilet_probe   "${TOILET_PROBE}"
    --bath_probe     "${BATH_PROBE}"
    --image_folder   "${IMAGE_FOLDER}"
    --samples_json   "${SAMPLES_JSON}"
    --base_mentions  "${BASE_MENTIONS}"
    --question       "${QUESTION}"
    --num_images     "${NUM_IMAGES}"
    --hook_type      "${HOOK_TYPE}"
    --max_new_tokens "${MAX_NEW_TOKENS}"
    --max_seq_tokens "${MAX_SEQ_TOKENS}"
    --out            "${OUT}"
    --out_json       "${OUT_JSON}"
)
if [ -n "${FORCED_TEXT}" ]; then
    ARGS+=(--forced_text "${FORCED_TEXT}")
fi

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