#!/bin/bash
# Combined perplexity + zero-shot eval for one checkpoint. Just point it at a model.
#
#   eval_model.sh <model_path_or_hf_id> [name] [gpu]
#
# Runs, on one GPU, sequentially (PPL via HF, then zero-shot via lm_eval+vllm):
#   - WT2 + C4 perplexity
#   - 7-task zero-shot (arc_easy, arc_challenge, hellaswag, winogrande, piqa,
#     boolq, openbookqa)
# then prints a combined table. Results go under ~/results/eval_<name>/.
# Zero-shot table uses acc_norm for arc/hellaswag/piqa/openbookqa, acc otherwise.
set -uo pipefail
. ~/local/venvs/main/bin/activate
HERE="$(cd "$(dirname "$0")" && pwd)"

MODEL="${1:?usage: eval_model.sh <model> [name] [gpu]}"
NAME="${2:-$(basename "$MODEL")}"
GPU="${3:-0}"
TASKS=arc_easy,arc_challenge,hellaswag,winogrande,piqa,boolq,openbookqa
OUT="$HOME/results/eval_${NAME}"
mkdir -p "$OUT"
echo "[eval] model=$MODEL  name=$NAME  gpu=$GPU  -> $OUT"

echo "[eval] 1/2  perplexity (wikitext2 + c4)..."
CUDA_VISIBLE_DEVICES=$GPU python "$HERE/ppl_eval.py" "$MODEL" --out "$OUT/ppl.json"

echo "[eval] 2/2  zero-shot (lm_eval + vllm)..."
rm -rf "$OUT/zeroshot"
CUDA_VISIBLE_DEVICES=$GPU lm_eval --model vllm \
  --model_args "pretrained=$MODEL,tensor_parallel_size=1,gpu_memory_utilization=0.85,dtype=bfloat16,max_model_len=4096,trust_remote_code=True" \
  --tasks "$TASKS" --batch_size auto --output_path "$OUT/zeroshot"

echo "[eval] ===== $NAME ====="
python - "$OUT" <<'PY'
import glob, json, os, sys
OUT = sys.argv[1]
NORM = {"arc_easy", "arc_challenge", "hellaswag", "piqa", "openbookqa"}
TASKS = ["arc_easy", "arc_challenge", "hellaswag", "winogrande", "piqa", "boolq", "openbookqa"]
ppl = json.load(open(os.path.join(OUT, "ppl.json")))
print(f"  WT2  {ppl['wikitext2']:.2f}")
print(f"  C4   {ppl['c4']:.2f}")
fs = glob.glob(os.path.join(OUT, "zeroshot", "**", "*.json"), recursive=True)
r = (json.load(open(sorted(fs)[-1])).get("results", {})) if fs else {}
vals = {}
for t in TASKS:
    x = r.get(t)
    if x:
        m = "acc_norm,none" if t in NORM else "acc,none"
        vals[t] = x.get(m, x.get("acc,none")) * 100
        print(f"  {t:14s} {vals[t]:.2f}")
if vals:
    print(f"  {'ZS avg':14s} {sum(vals.values())/len(vals):.2f}")
PY