BASE_PATH="[PATH_to_THIS_REPO]"

MODEL_DIR_PATH="${BASE_PATH}/output_yarn/qwen_yarn"
MODEL_NAMES=("[NAME_OF_THE_MODEL_W_YARN]")
export CUDA_VISIBLE_DEVICES=0,1,2,3
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export VLLM_ALLOW_LONG_MAX_MODEL_LEN=1
TP=4

cd ${BASE_PATH}/evaluation/eval_by_NIAH
DATA_DIR="./NeedleInAHaystack-PLUS"
OUT_PUT_PATH="./results"

for MODEL_NAME in "${MODEL_NAMES[@]}"; do
    python ./vllm_inference.py \
        --src_dir_path ${DATA_DIR} \
        --eval_data_list "needle_plus_squad,needle_plus_hotpotqa" \
        --output_path ${OUT_PUT_PATH} \
        --model_name_or_path ${MODEL_DIR_PATH}/${MODEL_NAME} \
        --model_name ${MODEL_NAME} \
        --tensor_parallel_size ${TP}
done