#!/bin/bash
# shellcheck disable=SC2090,SC2086,SC2089,SC1091
#SBATCH -c 10
#SBATCH -w anon
#SBATCH --gres=gpu:1
#SBATCH --job-name=light_eval
#SBATCH --tasks-per-node=1
#SBATCH --output=%x-%j.out
#!SBATCH --time=12:00:00
#SBATCH --dependency=afterany:3763

# Default project path
PROJECT_PATH="$HOME/projects/mt_dao"

# Parse command-line options
if ! OPTIONS=$(getopt -o p: --long project_path: -n 'parse-options' -- "$@"); then
	echo "slurm_submit.anon: Error parsing options" >&2
	exit 1
fi

eval set -- "$OPTIONS"

while true; do
	case "$1" in
	-p | --project_path)
		PROJECT_PATH="$2"
		shift 2
		;;
	--)
		shift
		break
		;;
	*)
		break
		;;
	esac
done
echo "slurm_submit.anon: PROJECT_PATH=$PROJECT_PATH"

CHECKPOINTS_DIR="checkpoints-safe"
MODEL="smollm-135M-shared-UNIF-20250228_022218/ep0-ba20480-rank0.pt"
TOKENIZER="SMOLLM_SHARED"
MODEL_TYPE="SMOLLM_135M"

SAVE_TENSORS_DIR="/local/scratch/mt_dao/SAVE_${MODEL}"

uv run python -m mt_dao.eval_light.download_mosaic_model --checkpoint "${MODEL}" --safe_tensors_dir "${SAVE_TENSORS_DIR}" --tokenizer "${TOKENIZER}" --checkpoint_dir ${CHECKPOINTS_DIR} --model_type "${MODEL_TYPE}"

uv run python -m mt_dao.eval_light.eval_mosaic --checkpoint "${MODEL}" --safe_tensors_dir "${SAVE_TENSORS_DIR}" --model_args "revision=main,dtype=bfloat16,vllm,gpu_memory_utilisation=0.6,max_model_length=2048" --custom_tasks "/nfs-share/anon/projects/mt_dao/mt_dao/eval_light/tasks.py" --tasks "/nfs-share/anon/projects/mt_dao/mt_dao/conf/lighteval/smollm2_base.txt" --output_dir "./evals_new" --save_details

# rm -rf "${SAVE_TENSORS_DIR}"
