#!/usr/bin/env bash
set -eo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
# shellcheck disable=SC1091
source "${SCRIPT_DIR}/conda_env.sh"
set -u


ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
cd "${ROOT}"

BASE_CONFIG="${BASE_CONFIG:-configs/grpo_llama32_3b_bf16.yaml}"
ACCELERATE_CONFIG="${ACCELERATE_CONFIG:-configs/accelerate_ddp_4gpu.yaml}"
NUM_PROCESSES="${NUM_PROCESSES:-4}"

MODEL_DIR="${MODEL_DIR:-}"
if [[ -z "${MODEL_DIR}" ]]; then
  echo "Missing MODEL_DIR."
  exit 1
fi

OUTPUT_PATH="${OUTPUT_PATH:-artifacts/eval/math_level1_zeroshot/answers.jsonl}"
SYSTEM_PROMPT="${SYSTEM_PROMPT:-}"
MAX_SAMPLES="${MAX_SAMPLES:--1}"
BATCH_SIZE="${BATCH_SIZE:-4}"
MAX_NEW_TOKENS="${MAX_NEW_TOKENS:-128}"
MAX_PROMPT_LENGTH="${MAX_PROMPT_LENGTH:-512}"

export PYTHONPATH="${ROOT}/src"
export BASE_CONFIG
export MODEL_DIR
export OUTPUT_PATH
export MAX_SAMPLES
export BATCH_SIZE
export MAX_NEW_TOKENS
export MAX_PROMPT_LENGTH

if [[ -n "${SYSTEM_PROMPT}" ]]; then export SYSTEM_PROMPT; fi

export WANDB_MODE="${WANDB_MODE:-offline}"
export HF_HUB_OFFLINE="${HF_HUB_OFFLINE:-1}"
export HF_DATASETS_OFFLINE="${HF_DATASETS_OFFLINE:-1}"
export TRANSFORMERS_OFFLINE="${TRANSFORMERS_OFFLINE:-1}"

accelerate launch \
  --config_file "${ACCELERATE_CONFIG}" \
  --num_processes "${NUM_PROCESSES}" \
  src/eval_math_level1_zeroshot.py