#!/usr/bin/env bash
# README zero-shot evaluation: lm-evaluation-harness with the mamba_ssm backend on state-spaces/mamba-370m,
# six tasks, full sets, batch 64. Kernels compiled beforehand for sm_120. Metrics are read from the
# results JSON written during this run.
set -o pipefail
ROW=/var/tmp/kv-repro/032-state-spaces-mamba
export HOME=$ROW/data/home
OUT=$ROW/data/results-$(date +%s)
cd "$ROW/data" || exit 1
echo "KVERITAS_PHASE name=evaluate"
lm_eval --model mamba_ssm --model_args pretrained=state-spaces/mamba-370m \
  --tasks lambada_openai,hellaswag,piqa,arc_easy,arc_challenge,winogrande \
  --device cuda --batch_size 64 --output_path "$OUT" 2>&1 | grep -E "^\|" || exit 1
python - "$OUT" <<'PY'
import glob, json, sys
r = json.load(open(sorted(glob.glob(f"{sys.argv[1]}/**/results_*.json", recursive=True))[-1]))["results"]
for task, m in r.items():
    for k, v in m.items():
        if isinstance(v, (int, float)) and "stderr" not in k and k != "alias":
            name = f"{task}_{k.split(',')[0]}"
            print(f"KVERITAS_METRIC name={name} value={v:.6g}")
            if k.startswith(("acc,", "acc_norm,", "perplexity,")):
                print(f"KVERITAS_CLAIM metric={name} value={v:.6g}")
PY
