#!/bin/bash
set -euo pipefail
export CUDA_VISIBLE_DEVICES=0
source /opt/miniconda/etc/profile.d/conda.sh
conda activate focallora

set -x
MODEL_PATH="../../models/Qwen2-7B-Instruct"
#IDENT_DATASET="../2-1_head_identification_preprocess/inj-squid_head_ident_llama_filtered_dev.jsonl"
IDENT_DATASET="../2-1_head_identification_preprocess/sep_head_ident.jsonl"
STAGE1_WEIGHT_RESULT="../2-2_head_identification/head_scoring/qwen2-7b_sep"

python ./Ident_IH_01_attn_sep.py \
  --model "$MODEL_PATH" \
  --dataset-sep "$IDENT_DATASET" \
  --output-dir "$STAGE1_WEIGHT_RESULT" \
  --split-size 1000 \
#  --max-data 5

python ./Ident_IH_02_score.py \
  --input-dir "$STAGE1_WEIGHT_RESULT" \
  --workers 32

python ./Ident_IH_03_sep_pick_head.py \
  --input-json "$STAGE1_WEIGHT_RESULT/head_scoring_combined.json" \

python Ident_IH_04_visualize.py \
  --input-dir "$STAGE1_WEIGHT_RESULT/heads_sorted" \
  --pickle-path "$STAGE1_WEIGHT_RESULT/head_scoring_raw_split_0.pkl" \
  --model-path $MODEL_PATH \
  --prompt-index 0