#!/bin/bash
#SBATCH --job-name=evqa_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=01:00:00

conda activate ret2
cd ~/ReT-2

export PYTHONPATH=.:./llava_more
export TRANSFORMERS_VERBOSITY=info
export TOKENIZERS_PARALLELISM=false
export OMP_NUM_THREADS=1

model_checkpoint="aimagelab/LLaVA_MORE-llama_3_1-8B-finetuning"
model_name="llava"
conv_template="llama_3_1"

dataset_name="evqa"

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

# https://github.com/google-research/google-research/tree/master/encyclopedic_vqa#vqa-questions
image_root_query=

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
batch_size=2
tokenizer_buffer_length=676

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