#!/bin/bash


# Define the parameters
model_name="Salesforce/blip-image-captioning-base"
tok_name="Salesforce/blip-image-captioning-base"
batch_size=16
max_tokens=5000000000000
sae_name="batchtopk"
lr=1e-4
expansion_factor=8
k=8
auxk=128
auxk_coef=0.03125
device_id=1
max_epochs=3
dead_tokens_threshold=1000000
log_every_n_steps=100
save_every_n_training_steps=100
save_step=false
use_loss_var=true
num_workers=4
hf_dataset="pixparse/cc3m-wds"
dtype="float16"


# change local_val_path
local_train_path="CC3M-Dataset/preproc/CC3M-Dataset/cc3m_images/train"
local_val_path="CC3M-Dataset/preproc/CC3M-Dataset/cc3m_images/val"

wandb_key=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

cd ..

# Loop over 12 layers
for layer in {10..10}
do  

    # hook names: vision_model.encoder.layers.${layer}.hook_resid_pre, 
    #               text_decoder.bert.encoder.layer.${layer}.attention.self.hook_resid_pre, 
    #               text_decoder.bert.encoder.layer.${layer}.crossattention.self.hook_resid_pre,
    hook_name="vision_model.encoder.layers.${layer}.hook_resid_post" 
    wandb login --relogin $wandb_key
    wandb login
    python Train_SAE_Blip.py \
        --model_name "$model_name" \
        --tok_name "$tok_name" \
        --layer "$layer" \
        --batch_size "$batch_size" \
        --hook_name "$hook_name" \
        --sae_name "$sae_name" \
        --lr "$lr" \
        --expansion_factor "$expansion_factor" \
        --k "$k" \
        --auxk "$auxk" \
        --max_epochs "$max_epochs" \
        --dead_tokens_threshold "$dead_tokens_threshold" \
        --log_every_n_steps "$log_every_n_steps" \
        --save_every_n_training_steps "$save_every_n_training_steps" \
        --save_step "$save_step" \
        --use_loss_var "$use_loss_var" \
        --max_tokens "$max_tokens" \
        --auxk_coef "$auxk_coef" \
        --num_workers "$num_workers" \
        --hf_dataset "$hf_dataset" \
        --local_train_path "$local_train_path" \
        --local_val_path "$local_val_path" \
        --device_id "$device_id" \
        --dtype "$dtype" \

done