#!/usr/bin/env bash
# census.sh — the engine's own checkpoint tools against both frank-dragoon
# payloads, under ONE knob set, so every number is comparable.
#
#   verify  is the downloaded file what it says it is (no GPU, no model load)
#   ppl     perplexity + KL/top-1 against a reference dump of the CURRENT
#           checkpoint on the wikitext-2 raw test split — the answer to "is
#           the stock peonist n-gram table good enough for this native
#           conversion?" A table mismatched to the checkpoint corrupts the
#           forward pass, and perplexity is the cheapest instrument that
#           sees that (card reference points: r2 base 3.134, swift15 3.229).
#   niah    retrieval by depth — the other thing a wrong lookup table breaks.
#
# ppl/niah LOAD THE MODEL: the service must be stopped (fleet benchmarking
# law). Stage arg: verify | ref | cand | niah-cur | niah-new | all.
#
# The pool stays 524288 — the OUTGOING server's value at census time — on
# purpose: the ppl pass scores 297,053 ids as TWO 262,144-token sequences,
# which only fits a 524,288-position pool. Pool size buys capacity, never
# numerics, so the dump and the comparison stay reproducible even though the
# deployed knob has since dropped to 262144.
set -euo pipefail

IMG=ghcr.io/peonist-ai/halogen-flash-server@sha256:0c61bf84ac22308a53f5d1ca6b86806702d7039e5ebc51cae4c66621b92fe04a
CUR=/home/dank/models/halogen-flashnext-v2
NEW=/home/dank/models/halogen-v2-abliterated
CUR_CK=qwen38-flash-next-v2-swift15-abliterated.hgn
NEW_CK=qwen38-flash-next-v2-abliterated.hgn
CORPUS=/eval/wikitext-2-raw-test.txt
REF=/eval/swift15-wikitext2.ref

run() {   # run <payload-dir> <checkpoint-file> <mode> [args...]
    local dir="$1" ck="$2"; shift 2
    podman run --rm --ipc=host --ulimit memlock=-1:-1 \
        --device /dev/kfd --device /dev/dri --group-add keep-groups \
        --security-opt seccomp=unconfined \
        -e HALOGEN_MODEL_ID=frank-dragoon \
        -e HALOGEN_CHECKPOINT="/models/$ck" \
        -e HALOGEN_NGRAM_TABLE=/models/qwen38-flash-next-ngram.hgn \
        -e HALOGEN_TOKENIZER=/tokenizer \
        -e HALOGEN_CTX=262144 \
        -e HALOGEN_KV_POOL_POSITIONS=524288 \
        -e HALOGEN_KV_SLOTS=2 \
        -e HALOGEN_HOST_RESERVE_GIB=20 \
        -e HALOGEN_SPEC_ADAPT=0 \
        -e HALOGEN_MAX_TOK=8192 \
        -e HALOGEN_WEIGHTS_LOCK=1 \
        -v "$dir:/models:ro" -v "$dir/tokenizer:/tokenizer:ro" \
        -v /tmp/halogen-eval:/eval \
        "$IMG" "$@"
}

stage="${1:-all}"
case "$stage" in
  verify)   run "$NEW" "$NEW_CK" verify "/models/$NEW_CK" ;;
  ref)      run "$CUR" "$CUR_CK" ppl "/models/$CUR_CK" \
                --corpus "$CORPUS" --ref-out "$REF" ;;
  cand)     run "$NEW" "$NEW_CK" ppl "/models/$NEW_CK" \
                --corpus "$CORPUS" --ref "$REF" --worst 5 ;;
  niah-cur) run "$CUR" "$CUR_CK" niah "/models/$CUR_CK" \
                --corpus "$CORPUS" --depths 4096,32768 --positions 0.05,0.5,0.95 ;;
  niah-new) run "$NEW" "$NEW_CK" niah "/models/$NEW_CK" \
                --corpus "$CORPUS" --depths 4096,32768 --positions 0.05,0.5,0.95 ;;
  all)
      for s in verify ref cand niah-cur niah-new; do
          echo "===== $s $(date -Is)"
          "$0" "$s" || echo "===== $s FAILED (exit $?)"
      done ;;
  *) echo "usage: $0 verify|ref|cand|niah-cur|niah-new|all" >&2; exit 2 ;;
esac