#!/bin/bash
#SBATCH --job-name=eval_m2kr_multi_gpu
#SBATCH --output=
#SBATCH --error=
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=4
#SBATCH --mem=128G
#SBATCH --cpus-per-task=16
#SBATCH --partition=
#SBATCH --account=
#SBATCH --time=01:00:00
#SBATCH --array=0-3

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/aimagelab/ReT-M2KR#download-images
IMAGE_ROOT_PATH=

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

DATASET_NAMES=(
    "okvqa"
    "infoseek"
    "evqa"
    "wit"
)
dataset_name="${DATASET_NAMES[$SLURM_ARRAY_TASK_ID]}"

dataset_path="${JSONL_ROOT_PATH}/m2kr/test/${dataset_name}_test.jsonl"
dataset_passages_path="${JSONL_ROOT_PATH}/m2kr/passages/${dataset_name}_passages_test.jsonl"
index_root="./indices/m2kr"
index_model_name="ReT2-M2KR-CLIP-ViT-B"
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
