#!/bin/bash
set -e

# bash evaluation.sh /home/duy/workspace/smollm/text/evaluation/eval_results /home/duy/workspace/smollm/text/evaluation/nanotron_eval_config.yaml 0,1,2,3
# NANOTRON_CHECKPOINT_CONFIG_PATH=${1:?Need to specify the checkpoint config path}
SAVE_DIR=${1:?Need to specify the save directory}
NANOTRON_LIGHTEVAL_CONFIG_PATH=${2:?Need to specify the lighteval config path}
NANOTRON_MODEL_PATH=${3:?Need to specify the model directory path}
# Optional argument for gpu selection
GPU_IDS=${4:-0}
# Nanotron models cannot be evaluated without torchrun.
NUM_GPUS=$(echo $GPU_IDS | awk -F',' '{print NF}')

MIN_CKPT_ITER=${5:-1000}
MAX_CKPT_ITER=${6:-9000}

for checkpoint_iter in $(seq $MAX_CKPT_ITER -1000 $MIN_CKPT_ITER); do
    echo "Fetching checkpoint folder: $NANOTRON_MODEL_PATH/$checkpoint_iter" && \
    git -c lfs.concurrenttransfers=500 lfs pull --include "$NANOTRON_MODEL_PATH/$checkpoint_iter/*" --exclude "" && \
    
    echo "Updating logging.output_dir in $NANOTRON_LIGHTEVAL_CONFIG_PATH for checkpoint: $SAVE_DIR/$(basename $NANOTRON_MODEL_PATH)/$checkpoint_iter" && \
    sed -i "s|^  output_dir:.*|  output_dir: \"$SAVE_DIR/$(basename $NANOTRON_MODEL_PATH)/$checkpoint_iter\"|" $NANOTRON_LIGHTEVAL_CONFIG_PATH && \
    
    echo "Evaluating checkpoint: $checkpoint_iter" && \
    CUDA_VISIBLE_DEVICES=$GPU_IDS torchrun --standalone --nnodes=1 --nproc-per-node=$NUM_GPUS \
    -m lighteval nanotron \
    --checkpoint-config-path ${NANOTRON_MODEL_PATH}/${checkpoint_iter}/config.yaml \
    --lighteval-config-path $NANOTRON_LIGHTEVAL_CONFIG_PATH
    if [ $? -ne 0 ]; then
        echo "Error occurred during evaluation of checkpoint: $checkpoint_iter"
        exit 1
    fi
    
    echo "Removing local copy of folder: $NANOTRON_MODEL_PATH/$checkpoint_iter"
    rm -rf $NANOTRON_MODEL_PATH/$checkpoint_iter && \

    echo "Cleaning up LFS files for checkpoint: $checkpoint_iter" 
    git lfs prune -f

    git read-tree HEAD && GIT_LFS_SKIP_SMUDGE=1 git checkout -f HEAD 
    
done