data_dir=$1 #path_to_data_dir is your first argument

target_attack=template
patch_location=middle
target_label=banana
dataset=coco #cc3m
trigger_output=BadVLM

if [ "$dataset" == "cc3m" ]; then
        annotation_path=$data_dir/cc3m/caption_trigger_optimization.json
        image_folder=$data_dir/cc3m
elif [ "$dataset" == "coco" ]; then
        annotation_path=$data_dir/coco/annotations/caption_trigger_optimization.json
        image_folder=$data_dir/coco/images
else
    echo "Dataset not supported"
    exit 1
fi

## Generate data for trigger optimization
python -m backdoors.trigger_optimization_data \
        --target-attack $target_attack \
        --target-label $target_label \
        --num-samples 10000 \
        --dataset $dataset

### Train the trigger optimization
CUBLAS_WORKSPACE_CONFIG=:4096:8 python -m torch.distributed.run \
        --nproc_per_node=1 -m LAVIS.trigger_optimize \
        --name $trigger_output \
        --batch-size 64 \
        --init-lr 1e-3 \
        --num-epochs 50 \
        --patch-location $patch_location \
        --itm-pos 1 \
        --itm-neg 1 \
        --itc-pos 1 \
        --lm-pos 1 \
        --lm-neg 1 \
        --lm-margin 10 \
        --dynamic false \
        --patch-size 16 \
        --num-patches 1 \
        --annotation-path $annotation_path \
        --image-folder $image_folder