#########################################################################
# File Name: eval.sh
# Author: Xianchao Wu, Peng Xu
# mail: xianchaow@nvidia.com, pengx@nvidia.com
# Created Time: Mon Sep  4 07:33:40 2024
#########################################################################
#!/bin/bash

# TODO change this to your reference file dir:
data_home="" # data_home https://huggingface.co/nvidia/Llama3-ChatQA-2-70B/tree/main/data
REFDIR=${data_home}


# TODO change to your tstdir
model_path=""
TSTDIR="${model_path}/outputs/"

model_size=70b # TODO change this
retriever="e5_mistral_retriever_chunkbysents1200"

adir=$TSTDIR
echo $adir

declare -A dataset2num_samples
dataset2num_samples["gov_report"]=200
dataset2num_samples["narrative_qa"]=2000
dataset2num_samples["qasper"]=2000
dataset2num_samples["qmsum"]=200
dataset2num_samples["quality"]=2000
dataset2num_samples["summ_screen_fd"]=200

dataset2num_samples["musique"]=200
dataset2num_samples["hotpotqa"]=200
dataset2num_samples["multifieldqa_en"]=200

dataset2num_samples["squality"]=200

dataset2num_samples["doc2dial_full_dialogue"]=1000

echo "ref path = $REFDIR"
echo "tst out path = $TSTDIR"

declare -A sys2name
sys2name["baseline"]=""
sys2name["ret"]="_ctx5"

for system in "baseline" "ret"
do
	suffix=${sys2name[${system}]}

    echo "--final display----$system----"
    for adataset in "qmsum" "qasper" "quality"
    do
        echo $adataset

        ref_fn="${REFDIR}/${adataset}.${retriever}/test.json"
        tst_fn="${adir}/${adataset}.e5_mistral_retriever_chunkbysents1200_output_0to${dataset2num_samples[${adataset}]}${suffix}.txt"

        echo "ref for ${adataset}", ${ref_fn}
        echo "tstout for ${adataset}", ${tst_fn}
        	if [[ ! -e ${tst_fn} ]]; then
        		echo "Error: tst_fn=${tst_fn} not exist!"
        	fi
        	if [[ ! -e ${ref_fn} ]]; then
        		echo "Error: ref_fn=${ref_fn} not exist!"
        	fi
        #continue

        if [[ -e ${tst_fn} && -e ${ref_fn} ]]
        then
            python3 dataset_evaluator_retro.py \
                --datapath ${ref_fn} \
                --gen_test_file ${tst_fn} \
                --dataset $adataset
        	else
        		if [[ ! -e ${tst_fn} ]]; then
        			echo "Error: tst_fn=${tst_fn} not exist!"
        		fi
        		if [[ ! -e ${ref_fn} ]]; then
        			echo "Error: ref_fn=${ref_fn} not exist!"
        		fi
        fi
    done

    for adataset in "musique" "hotpotqa" "multifieldqa_en"
    do
        echo $adataset

        ref_fn="${REFDIR}/${adataset}.${retriever}/test.json"
        # TODO change this if necessary, model's prediction output
	tst_fn="${adir}/${adataset}.e5_mistral_retriever_chunkbysents1200_output_0to${dataset2num_samples[${adataset}]}${suffix}.txt"

        echo "ref for ${adataset}", ${ref_fn}
        echo "tstout for ${adataset}", ${tst_fn}
		if [[ ! -e ${tst_fn} ]]; then
			echo "Error: tst_fn=${tst_fn} not exist!"
		fi
		if [[ ! -e ${ref_fn} ]]; then
			echo "Error: ref_fn=${ref_fn} not exist!"
		fi
        #continue

        if [[ -e ${tst_fn} && -e ${ref_fn} ]]
        then
            python3 dataset_evaluator_retro_longbench.py \
                --datapath ${ref_fn} \
                --gen_test_file ${tst_fn} \
                --dataset $adataset
		else
			if [[ ! -e ${tst_fn} ]]; then
				echo "Error: tst_fn=${tst_fn} not exist!"
			fi
			if [[ ! -e ${ref_fn} ]]; then
				echo "Error: ref_fn=${ref_fn} not exist!"
			fi
        fi
    done
done