#!/bin/bash
#SBATCH --job-name=eval_mbeir_multi_gpu
#SBATCH --output=
#SBATCH --error=
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=4
#SBATCH --mem=200G
#SBATCH --cpus-per-task=16
#SBATCH --partition=
#SBATCH --account=
#SBATCH --time=02:30:00
#SBATCH --array=0-15

# Note that most tasks will end within ~40 minutes, so you may want to adjust the walltime

set -e

conda activate ret2
cd ~/ReT-2

export PYTHONPATH=.
export TRANSFORMERS_VERBOSITY=info
export TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS=1

# https://huggingface.co/datasets/TIGER-Lab/M-BEIR#decompressing-m-beir-images
IMAGE_ROOT_PATH=

# https://huggingface.co/datasets/aimagelab/ReT-M2KR/tree/main/jsonl
JSONL_ROOT_PATH=

DATASET_NAMES=(
    "infoseek_task8"
    "oven_task8"
    "cirr_task7"
    "fashion200k_task0"
    "fashion200k_task3"
    "fashioniq_task7"
    "infoseek_task6"
    "mscoco_task0"
    "mscoco_task3"
    "nights_task4"
    "oven_task6"
    "visualnews_task0"
    "visualnews_task3"
    "webqa_task1"
    "webqa_task2"
    "edis_task2"
)
dataset_name="${DATASET_NAMES[$SLURM_ARRAY_TASK_ID]}"

dataset_path="${JSONL_ROOT_PATH}/mbeir/test/mbeir_${dataset_name}_test.jsonl"
dataset_passages_path="${JSONL_ROOT_PATH}/mbeir/passages/mbeir_${dataset_name}_cand_pool.jsonl"
index_root="./indices/mbeir"
index_model_name="ReT2-MBEIR-SigLIP2-ViT-L"
index_dataset_name="${dataset_name}"
checkpoint_path="aimagelab/${index_model_name}"
batch_size=128

srun torchrun --nproc-per-node=$SLURM_GPUS_PER_NODE --standalone \
src/evaluate.py \
--action "index" \
--index_root "${index_root}" \
--index_model_name "${index_model_name}" \
--index_dataset_name "${index_dataset_name}" \
--dataset_passages_path "${dataset_passages_path}" \
--image_root "${IMAGE_ROOT_PATH}" \
--checkpoint_path "${checkpoint_path}" \
--batch_size "${batch_size}" \
--dataloader_num_workers 4

srun python src/evaluate.py \
--action "create_index" \
--index_root "${index_root}" \
--index_model_name "${index_model_name}" \
--index_dataset_name "${index_dataset_name}"

srun torchrun --nproc-per-node=1 --standalone \
src/evaluate.py \
--action "search" \
--index_root "${index_root}" \
--index_model_name "${index_model_name}" \
--index_dataset_name "${index_dataset_name}" \
--dataset_path "${dataset_path}" \
--dataset_passages_path "${dataset_passages_path}" \
--image_root "${IMAGE_ROOT_PATH}" \
--checkpoint_path "${checkpoint_path}" \
--batch_size "${batch_size}" \
--dataloader_num_workers $SLURM_CPUS_PER_TASK