#!/bin/bash
#SBATCH --job-name=infoseek_index_and_search
#SBATCH --output=
#SBATCH --error=
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=1
#SBATCH --mem=200G
#SBATCH --cpus-per-task=16
#SBATCH --partition=
#SBATCH --account=
#SBATCH --time=05:00:00

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/blob/main/jsonl/rag/kb_infoseek1M.jsonl
dataset_path_index="kb_infoseek1M.jsonl"

# https://huggingface.co/datasets/ychenNLP/oven/blob/main/all_wikipedia_images.tar
image_root_index="all_wikipedia_images"

# https://huggingface.co/datasets/aimagelab/ReT-M2KR/blob/main/jsonl/rag/infoseek_val.jsonl
dataset_path_query="infoseek_val.jsonl"

# https://github.com/edchengg/oven_eval/tree/main/image_downloads#oven-images
image_root_query="oven_eval/image_downloads"

index_root="./rag/indices"
index_model_name="ReT2-M2KR-ColBERT-SigLIP2-ViT-L"
index_dataset_name="infoseek"
checkpoint_path="aimagelab/${index_model_name}"
batch_size=256

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}" \
--image_root "${image_root_index}" \
--dataset_passages_path "${dataset_path_index}" \
--checkpoint_path "${checkpoint_path}" \
--batch_size "${batch_size}" \
--dataloader_num_workers 4 \
--fp16

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}" \
--image_root "${image_root_query}" \
--dataset_path "${dataset_path_query}" \
--checkpoint_path "${checkpoint_path}" \
--batch_size "${batch_size}" \
--dataloader_num_workers $SLURM_CPUS_PER_TASK \
--fp16 \
--skip_metrics

srun python src/rag/ranking_tsv2jsonl.py \
--input_path "${index_root}/${index_model_name}/${index_dataset_name}/ranking.tsv" \
--dataset_passages_path "${dataset_path_index}"