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

#!/usr/bin/env bash
# ARC (AI2 Reasoning Challenge) evaluation via OLMES.
# Runs arc_easy and arc_challenge for a given model variant.
#
# Usage:
#   sbatch scripts/slurm/evaluation/eval_arc.sh                    # base model (MC)
#   sbatch scripts/slurm/evaluation/eval_arc.sh --instruct         # instruct model (CoT)
#   sbatch scripts/slurm/evaluation/eval_arc.sh --think            # think model (CoT)
#
#SBATCH --job-name=eval_arc
#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=04:00:00
#SBATCH --output=logs/evaluation/eval_arc_%j.out
#SBATCH --error=logs/evaluation/eval_arc_%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" ;;
        --instruct-mc) MODEL_VARIANT="instruct-mc" ;;
        --think)       MODEL_VARIANT="think" ;;
        --think-mc)    MODEL_VARIANT="think-mc" ;;
    esac
done

case "$MODEL_VARIANT" in
    base)
        MODEL_ID="allenai/Olmo-3-1025-7B"
        TASKS="arc_easy:mc::olmes arc_challenge:mc::olmes"
        ;;
    instruct)
        MODEL_ID="allenai/Olmo-3-7B-Instruct"
        TASKS="arc_easy:cot::olmo3:adapt arc_challenge:cot::olmo3:adapt"
        ;;
    instruct-mc)
        MODEL_ID="allenai/Olmo-3-7B-Instruct"
        TASKS="arc_easy:mc::olmes arc_challenge:mc::olmes"
        ;;
    think)
        MODEL_ID="allenai/Olmo-3-7B-Think"
        TASKS="arc_easy:cot::olmo3:adapt arc_challenge:cot::olmo3:adapt"
        ;;
    think-mc)
        MODEL_ID="allenai/Olmo-3-7B-Think"
        TASKS="arc_easy:mc::olmes arc_challenge:mc::olmes"
        ;;
esac

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

echo "===== ARC 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