#!/bin/bash

# Uncomment and set the following variables correspondingly to run this script:

# MODEL_VERSION=vicuna-v1-3-7b
# MODEL_VERSION=llama-2-7b-chat

########### DO NOT CHANGE ###########
########### USE THIS FOR BOTH ###########
PROMPT_VERSION="llava_llama_2"

MODEL_VERSION="llava-llama-2-7b-chat-lightning-preview"
MODEL_NAME="llava-llama-2-7b-chat"
PRETRAIN_VERSION="DETR-v2-pad0-pretrain-1"

NUM_TRAIN_EPOCHS=1
########### DO NOT CHANGE ###########

deepspeed --master_port 16663 --include localhost:5,7 llava/backbone/tune_mm.py \
    --deepspeed ./scripts/zero2.json \
    --model_name_or_path ./checkpoints/$MODEL_VERSION \
    --version $PROMPT_VERSION \
    --data_path ./playground/data/llava_instruct_80k.json \
    --image_folder ./data/train2017 \
    --vision_tower "detr-v2" \
    --tune_mm_mlp_adapter True \
    --mm_vision_select_layer -1 \
    --mm_use_im_start_end False \
    --mm_use_im_patch_token False \
    --bf16 True \
    --output_dir ./llava/backbone/checkpoints_tune_ll2/$MODEL_NAME-$PRETRAIN_VERSION-tune-$NUM_TRAIN_EPOCHS \
    --num_train_epochs $NUM_TRAIN_EPOCHS \
    --per_device_train_batch_size 16 \
    --per_device_eval_batch_size 16 \
    --gradient_accumulation_steps 4 \
    --evaluation_strategy "no" \
    --save_strategy "steps" \
    --save_steps 400 \
    --save_total_limit 2 \
    --learning_rate 2e-4 \
    --weight_decay 0. \
    --warmup_ratio 0.03 \
    --lr_scheduler_type "cosine" \
    --logging_steps 1 \
    --tf32 True \
    --model_max_length 2048 \
    --gradient_checkpointing True \
    --dataloader_num_workers 2 \
    --lazy_preprocess True \
    --report_to wandb \
    --pretrain_mm_projector_path  ./llava/backbone/checkpoints_pretrain_plain/$MODEL_NAME-$PRETRAIN_VERSION/mm_projector.bin \
# LLaVA/llava/backbone/checkpoints_pretrain_plain/llava-llama-2-7b-chat-DETR-v2-pad-sl1-pretrain-1
# 8gpu*16*1 = 128 = 4gpu*16*2
    # --load_best_model_at_end True \
    # --metric_for_best_model "loss" \
    # --greater_is_better False \


