#!/bin/bash
# =============================================================================
# alternative_steering.sh — steer residual along LoRA-B direction, judge caption
# hallucination + coherence with Qwen3-VL-8B-Instruct (ARENA-style 0-100 judges).
#
# h' = h + ALPHA*‖h‖*LoRA-B̂  (all tokens, all/​subset layers, during generation).
# Sweep ALPHAS; per alpha report mean hallucination ↓ + coherence ↑.
#
# Tunables (override inline):
#   DEVICE_ID        GPU for LLaVA (steered model)            (default 0)
#   JUDGE_DEVICE_ID  GPU for Qwen3-VL judge (default = DEVICE_ID; set if 2 GPUs)
#   DTYPE            bfloat16 | float16 | float32             (default bfloat16)
#   LORA_PATH        LoRA adapter dir (B-matrix source)
#   LORA_MODULE      down_proj | o_proj                        (default down_proj)
#   B_REDUCE         mean | svd  (collapse B(d,r)→1 dir)       (default mean)
#   LAYERS           layers to steer, empty = all
#   ALPHAS           steering coeffs to sweep (0 = baseline)
#   CATEGORY         samples.json category                     (default bathroom_only)
#   NUM_IMAGES       images to test                            (default 50)
#   TARGET_OBJECT    hallucinated object (judge GT)            (default toilet)
#   SCENE            present scene (judge GT)                  (default bathroom)
#   GT_PRESENT       override full GT statement (any relation)
#   JUDGE_MODEL      HF id of judge VLM           (default Qwen/Qwen3-VL-8B-Instruct)
#   MAX_NEW_TOKENS   caption length                            (default 64)
#   OUT / OUT_JSON   outputs
# =============================================================================
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}"
LORA_PATH="${LORA_PATH:-/data/caotue/multilayer-sae/adv_gen_outputs/run_bathroom_toilet_v2/lora_adapter}"
LORA_MODULE="${LORA_MODULE:-down_proj}"
B_REDUCE="${B_REDUCE:-mean}"
LAYERS="${LAYERS:-}"
ALPHAS="${ALPHAS:-0.0 0.05 0.1 -0.05 -0.1}"
CATEGORY="${CATEGORY:-bathroom_only}"
NUM_IMAGES="${NUM_IMAGES:-50}"
TARGET_OBJECT="${TARGET_OBJECT:-toilet}"
SCENE="${SCENE:-bathroom}"
GT_PRESENT="${GT_PRESENT:-}"
JUDGE_MODEL="${JUDGE_MODEL:-Qwen/Qwen3-VL-8B-Instruct}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-64}"
OUT="${OUT:-mechanistic_interp/graph/alternative_steering.png}"
OUT_JSON="${OUT_JSON:-mechanistic_interp/graph/alternative_steering.json}"

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

ARGS=(
    --model_name      "llava-hf/llava-1.5-7b-hf"
    --judge_model     "${JUDGE_MODEL}"
    --device_id       "${DEVICE_ID}"
    --dtype           "${DTYPE}"
    --lora_path       "${LORA_PATH}"
    --lora_module     "${LORA_MODULE}"
    --b_reduce        "${B_REDUCE}"
    --alphas          ${ALPHAS}
    --category        "${CATEGORY}"
    --num_images      "${NUM_IMAGES}"
    --target_object   "${TARGET_OBJECT}"
    --scene           "${SCENE}"
    --max_new_tokens  "${MAX_NEW_TOKENS}"
    --out             "${OUT}"
    --out_json        "${OUT_JSON}"
)
[ -n "${JUDGE_DEVICE_ID}" ] && ARGS+=(--judge_device_id "${JUDGE_DEVICE_ID}")
[ -n "${LAYERS}" ] && ARGS+=(--layers ${LAYERS})
[ -n "${GT_PRESENT}" ] && ARGS+=(--gt_present "${GT_PRESENT}")

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