#!/bin/bash

partition=your_partition
base_save_dir=./results
base_log_dir=./slurm_logs
behaviors_path=data/behavior_datasets/harmbench_behaviors_text_val.csv
methods=("GCG" "EnsembleGCG" "PAIR" "AutoDAN")
experiment_names=("llama2_7b" "llama2_7b" "llama2_7b" "llama2_7b")
model_names=("llama2_7b" "llama2_7b" "llama2_7b" "llama2_7b")
gpu_counts=(1 1 1 1)
max_new_tokens=512
incremental_update="False"

# Loop through each method
for method_name in "${methods[@]}"; do
    # Loop over each index
    for ((index=0; index<${#model_names[@]}; index++)); do
        experiment_name="${experiment_names[$index]}"
        model_name="${model_names[$index]}"
        gpu_count=${gpu_counts[$index]}
        save_path="$base_save_dir/$method_name/$experiment_name/completions/${model_name}.json"
        test_cases_path="$base_save_dir/$method_name/$experiment_name/test_cases/test_cases.json"
        job_name="generate_completions_${method_name}_${experiment_name}_${model_name}"
        log_path="$base_log_dir/generate_completions/${method_name}/${experiment_name}_${model_name}.log"

        sbatch --partition=$partition --job-name=$job_name --nodes=1 --gpus-per-node=$gpu_count --time=5:00:00 --output=$log_path \
            scripts/generate_completions.sh $model_name $behaviors_path $test_cases_path $save_path $max_new_tokens $incremental_update
    done
done