#!/bin/bash
set -e

# Setup
export CUDA_VISIBLE_DEVICES="1"
export NUM_GPUS=$(python -c "import torch; print(torch.cuda.device_count())")

# Training
export MODELS_DIR="/nfsclient/os5800/vlm_project/data/pllum-vlm/models"
export BASE_MODEL="${MODELS_DIR}/base/bielik-11b-v2.6-instruct-siglip2-so400m-patch14-384-llava-next"
export RUN_NAME="$(basename ${BASE_MODEL})-pretrained-$(date +%Y-%m-%d_%H-%M-%S)"

export DATA_DIR="TODO: replace with your directory"
export DATA_PATH="${DATA_DIR}/processed/pretrain"
export IMAGES_PATH="${DATA_DIR}/static/pretrain"

export PER_DEVICE_TRAIN_BATCH_SIZE=1
export GLOBAL_BATCH_SIZE=256
export GRADIENT_ACCUMULATION_STEPS=$((GLOBAL_BATCH_SIZE / (NUM_GPUS * PER_DEVICE_TRAIN_BATCH_SIZE)))

export LAUNCHER="accelerate launch \
    --num_processes=${NUM_GPUS} \
    --num_machines=1 \
    --mixed_precision=bf16 \
    --dynamo_backend=no \
    "
export SCRIPT="-m src.modeling.llava.train"
export SCRIPT_ARGS="\
    --deepspeed=config/deepspeed/zero2.json \
    --model_name_or_path=${BASE_MODEL} \
    --data_path=${DATA_PATH} \
    --images_path=${IMAGES_PATH} \
    --trainer_output_dir=${MODELS_DIR}/checkpoints \
    --model_output_dir=${MODELS_DIR}/pretrained \
    --run_name=${RUN_NAME} \
    --seed=42 \
    --train_vision_tower=False \
    --train_multimodal_projector=True \
    --train_language_model=False \
    --train_lm_head=False \
    --max_seq_length=8192 \
    --epochs=1 \
    --lr=1e-3 \
    --lr_scheduler_type=cosine \
    --weight_decay=0.00 \
    --warmup_ratio=0.03 \
    --per_device_train_batch_size=${PER_DEVICE_TRAIN_BATCH_SIZE} \
    --per_device_eval_batch_size=2 \
    --gradient_accumulation_steps=${GRADIENT_ACCUMULATION_STEPS} \
    --gradient_checkpointing=True \
    --save_steps=250 \
    --logging_steps=10 \
    --report_to=wandb \
    --eval_strategy=steps \
    --eval_steps=250 \
    --test_size=0.01 \
    "

bash -c "${LAUNCHER} ${SCRIPT} ${SCRIPT_ARGS}"

python -m src.modeling.llava.merge \
    --base_model_name_or_path="${BASE_MODEL}" \
    --trainer_output_dir="${MODELS_DIR}/checkpoints/${RUN_NAME}" \
    --final_model_name_or_path="${MODELS_DIR}/pretrained/${RUN_NAME}"