social-data-attribution-w2/scripts/slurm/evaluation/eval_bbh.sh

#!/usr/bin/env bash
# BIGBench-Hard evaluation via OLMES.
# Runs bbh_snarks, bbh_causal_judgement, bbh_sports_understanding.
#
# Usage:
#   sbatch scripts/slurm/evaluation/eval_bbh.sh                    # base model (3-shot CoT)
#   sbatch scripts/slurm/evaluation/eval_bbh.sh --instruct         # instruct model (0-shot chat CoT)
#   sbatch scripts/slurm/evaluation/eval_bbh.sh --think            # think model (0-shot chat CoT)
#
#SBATCH --job-name=eval_bbh
#SBATCH --partition=coe-gpu,ice-gpu
#SBATCH --qos=coe-ice
#SBATCH --account=ic
#SBATCH --gres=gpu:1
#SBATCH --cpus-per-task=8
#SBATCH --mem=64G
#SBATCH --time=02:00:00
#SBATCH --output=logs/evaluation/eval_bbh_%j.out
#SBATCH --error=logs/evaluation/eval_bbh_%j.err

set -euo pipefail

WORKTREE_ROOT="${WORKTREE_ROOT:-${SLURM_SUBMIT_DIR:-$(cd "$(dirname "$0")/../../.." && pwd)}}"
cd "$WORKTREE_ROOT"
mkdir -p logs/evaluation

if [[ -f .env.secret ]]; then set -a; source .env.secret; set +a; fi

export HF_HOME="${HF_HOME:-$HOME/scratch/hf_cache}"
export HF_DATASETS_CACHE="${HF_HOME}/datasets"
export HF_TOKEN="${HF_TOKEN:-$(cat ~/.hf_token 2>/dev/null || true)}"
export HF_HUB_DISABLE_XET=1
export HF_DATASETS_TRUST_REMOTE_CODE=1
export TOKENIZERS_PARALLELISM=false
export PYTHONUNBUFFERED=1
export PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True
export VLLM_WORKER_MULTIPROC_METHOD=spawn
export WANDB_MODE=disabled

MODEL_VARIANT="base"
for arg in "$@"; do
    case "$arg" in
        --instruct) MODEL_VARIANT="instruct" ;;
        --think)    MODEL_VARIANT="think" ;;
    esac
done

BBH_TASKS_BASE="bbh_snarks:cot::olmes bbh_causal_judgement:cot::olmes bbh_sports_understanding:cot::olmes"
BBH_TASKS_INSTRUCT="bbh_snarks:cot::olmo3:adapt bbh_causal_judgement:cot::olmo3:adapt bbh_sports_understanding:cot::olmo3:adapt"

case "$MODEL_VARIANT" in
    base)
        MODEL_ID="allenai/Olmo-3-1025-7B"
        TASKS="$BBH_TASKS_BASE"
        ;;
    instruct)
        MODEL_ID="allenai/Olmo-3-7B-Instruct"
        TASKS="$BBH_TASKS_INSTRUCT"
        ;;
    think)
        MODEL_ID="allenai/Olmo-3-7B-Think"
        TASKS="$BBH_TASKS_INSTRUCT"
        ;;
esac

SAFE_MODEL="${MODEL_ID//\//-}"
STAMP="$(date -u +%Y%m%d_%H%M%S)"
OUTPUT_DIR="runs/evaluation_olmes/${SAFE_MODEL}_bbh_${MODEL_VARIANT}_${STAMP}"
mkdir -p "$OUTPUT_DIR"

echo "===== BBH Evaluation: ${MODEL_ID} (${MODEL_VARIANT}) ====="
echo "Job ID: ${SLURM_JOB_ID:-local}"
echo "Tasks: ${TASKS}"
echo "Output: ${OUTPUT_DIR}"
nvidia-smi --query-gpu=name,memory.total --format=csv,noheader 2>/dev/null || echo "No GPU detected"
date

export PATH="$WORKTREE_ROOT/.venv/bin:$PATH"
export PYTHONPATH="src:olmes${PYTHONPATH:+:$PYTHONPATH}"

python -m oe_eval.launch \
    --model "$MODEL_ID" \
    --model-type vllm \
    --model-args '{"trust_remote_code": true, "max_length": 8192}' \
    --task $TASKS \
    --output-dir "$OUTPUT_DIR"

echo "===== Done ====="
echo "Results: ${OUTPUT_DIR}"
date