data_dir=$1 #path_to_data_dir is your first argument

attack_type=badVLM

if [ $attack_type == "badVLM" ]; then
    pattern=badVLM
    patch_location=middle
    trigger_output=badVLM
    trigger_path=backdoors/config/badVLM/patch.png

elif [ $attack_type == "TrojVQA" ]; then
    pattern=vqa
    patch_location=middle
    trigger_path=backdoors/config/TrojVQA/TrojanVQA_default_vqa.jpg

elif [ $attack_type == "badCLIP" ]; then
    pattern=badCLIP
    patch_location=middle
    trigger_path=backdoors/config/badCLIP/trigger.jpg

elif [ $attack_type == "blended" ]; then
    pattern=blended
    patch_location=blended
    trigger_path="."

elif [ $attack_type == "badNet" ]; then
    pattern=random
    patch_location=random
    trigger_path="."

else
    echo "Unknown attack type"
    exit 1
fi

target_attack=template
target_label=banana
poison_size=100
dataset=coco #cc3m

### ===========================================================================================
### Backdoor Config
python -m backdoors.backdoor_config \
        --attack-type $attack_type \
        --pattern $pattern \
        --patch-location $patch_location \
        --target-label $target_label \
        --poison-size $poison_size \
        --pattern-size 16 \
        --dataset $dataset \
        --trigger-path $trigger_path

### ===========================================================================================
### Generate poisoned dataset
python -m backdoors.backdoor_generation  \
        --attack-type $attack_type \
        --target-attack $target_attack \
        --model-type BLIP \
        --task vqa

### ===========================================================================================
### Poisoning
cp $data_dir/coco/annotations/coco_karpathy_train_BLIP_${attack_type}_vqa.json $data_dir/coco/annotations/vqa_train.json
python -m torch.distributed.run \
        --nproc_per_node=1 -m LAVIS.train \
        --cfg-path LAVIS/lavis/projects/blip2/train/vqa_ft.yaml \
        --name $attack_type/VQA_coco_freeze_vit_${trigger_output}_${poison_size} \
        --backdoor $attack_type \
        --freeze-vit \
        --batch-size-train 32 \
        --batch-size-eval 64 \
        --warmup_steps 100 \
        --init-lr 3e-5 \
        --epochs 10 \
        --model-weight LAVIS/lavis/output/BLIP2/VQA_coco_clean/checkpoint_best.pth
