#!/bin/bash
#SBATCH --job-name=infoseek_rag_llava_more
#SBATCH --output=
#SBATCH --error=
#SBATCH --nodes=1
#SBATCH --ntasks-per-node=1
#SBATCH --gpus-per-node=4
#SBATCH --mem=128G
#SBATCH --cpus-per-task=8
#SBATCH --partition=
#SBATCH --account=
#SBATCH --time=08:00:00

conda activate ret2
cd ~/ReT-2

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

model_checkpoint="Qwen/Qwen2.5-VL-7B-Instruct"

dataset_name="infoseek"

# 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"

retriever="ReT2-M2KR-ColBERT-SigLIP2-ViT-L"
ranking_path="./rag/indices/${retriever}/${dataset_name}/ranking.jsonl"
top_k=3
generation_num_beams=5

output_dir="./rag/results/${dataset_name}"
mkdir -p $output_dir
output_path="${output_dir}/${retriever}__top_${top_k}__bs_${generation_num_beams}.jsonl"

dataloader_num_workers=2
tokenizer_buffer_length=676

srun torchrun --nproc-per-node=$SLURM_GPUS_PER_NODE --standalone \
src/rag/rag_qwen.py \
--query_path $dataset_path_query \
--image_root $image_root_query \
--output_path $output_path \
--model_checkpoint $model_checkpoint \
--generation_num_beams $generation_num_beams \
--dataloader_num_workers $dataloader_num_workers \
--dataset_name $dataset_name \
--ranking_path $ranking_path \
--tokenizer_buffer_length $tokenizer_buffer_length \
--top_k $top_k