#!/bin/bash

# Define the parameters
model_name="gpt2-small"
tok_name="gpt2"
batch_size=1
max_tokens=500000000
sae_name="topk_crosscoder"
lr=0.0001
expansion_factor=32
k=32
auxk=256
auxk_coef=0.03125
device_id=7
max_epochs=1
dead_tokens_threshold=10000000
log_every_n_steps=50
save_every_n_training_steps=100
use_loss_var=true
num_workers=63

cd ..

input_hook_names=()
output_hook_names=()
for layer in {0..11}
do
    input_hook_names+=("blocks.${layer}.ln2.hook_normalized")
    output_hook_names+=("blocks.${layer}.hook_mlp_out")
done

python Train_Crosscoder.py \
    --model_name "$model_name" \
    --tok_name "$tok_name" \
    --batch_size "$batch_size" \
    --input_hook_names "${input_hook_names[@]}" \
    --output_hook_names "${output_hook_names[@]}" \
    --sae_name "$sae_name" \
    --lr "$lr" \
    --expansion_factor "$expansion_factor" \
    --k "$k" \
    --auxk "$auxk" \
    --device_id "$device_id" \
    --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" \
    --use_loss_var "$use_loss_var" \
    --max_tokens "$max_tokens" \
    --auxk_coef "$auxk_coef" \
    --num_workers "$num_workers"