#!/usr/bin/env bash
# Perplexity, KL divergence and speed for one model at F16, Q8_0, Q5_K_M and Q4_K_M with llama.cpp.
#
# Model files: the official Qwen GGUF repo, https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct-GGUF
#   qwen2.5-1.5b-instruct-{fp16,q8_0,q5_k_m,q4_k_m}.gguf (sizes and sha256 recorded in raw/downloads.txt)
# Text: wikitext-2-raw-v1 test split (wiki.test.raw) from https://huggingface.co/datasets/ggml-org/ci
# Tool: llama.cpp b11351 official macOS arm64 release (llama-perplexity, llama-bench), Metal backend.
#
#   PPL:  llama-perplexity -m Q -f wiki.test.raw -c 512 --chunks $PPL_CHUNKS            (each quant)
#   KLD:  llama-perplexity -m F16 -f wiki.test.raw -c 512 --chunks $KLD_CHUNKS --kl-divergence-base base.kld
#         llama-perplexity -m Q --kl-divergence-base base.kld --kl-divergence             (each quant)
#   Speed: llama-bench -m Q -p 512 -n 128 -r 3 -o json                                 (each quant)
# KLD uses fewer chunks than PPL because the F16 logits file grows by ~78 MB per chunk
# (151,936-token vocabulary) and the disk on this machine is small.
set -uo pipefail
cd "$(dirname "$0")"
LCPP="$HOME/aw-lab/llama.cpp/llama-b11351"; G="$HOME/aw-lab/gguf"
PPL_CHUNKS="${PPL_CHUNKS:-100}"; KLD_CHUNKS="${KLD_CHUNKS:-20}"; BASE="${BASE:-q8_0}"  # fp16 download did not complete on this connection
mkdir -p raw
W="$G/wikitext-2-raw/wiki.test.raw"
# only use files whose download finished (logged with size + sha256 by the downloader)
done_dl() { grep -q "^qwen2.5-1.5b-instruct-$1.gguf exit=0" "$HOME/aw-lab/logs/D/downloads.txt" 2>/dev/null; }
state() { echo "# $1 $(date '+%F %T') | $(sysctl -n vm.swapusage) | $(memory_pressure | tail -1)"; }

QUANTS="fp16 q8_0 q5_k_m q4_k_m"
for q in $QUANTS; do
  m="$G/qwen2.5-1.5b-instruct-$q.gguf"
  done_dl "$q" || { echo "not downloaded yet: $q" >&2; continue; }
  out="raw/ppl-$q.txt"
  [ -s "$out" ] && grep -q "Final estimate" "$out" && continue
  { state "ppl $q"; echo "\$ llama-perplexity -m qwen2.5-1.5b-instruct-$q.gguf -f wiki.test.raw -c 512 --chunks $PPL_CHUNKS"; } > "$out"
  "$LCPP/llama-perplexity" -m "$m" -f "$W" -c 512 --chunks "$PPL_CHUNKS" >> "$out" 2>&1
  grep "Final estimate" "$out" >&2
done

if done_dl $BASE && [ ! -s "$G/base-$BASE.kld" ]; then
  { state "kld base"; echo "\$ llama-perplexity -m qwen2.5-1.5b-instruct-$BASE.gguf -f wiki.test.raw -c 512 --chunks $KLD_CHUNKS --kl-divergence-base base-$BASE.kld"; } > raw/kld-base-$BASE.txt
  "$LCPP/llama-perplexity" -m "$G/qwen2.5-1.5b-instruct-$BASE.gguf" -f "$W" -c 512 --chunks "$KLD_CHUNKS" \
    --kl-divergence-base "$G/base-$BASE.kld" >> raw/kld-base-$BASE.txt 2>&1
  ls -l "$G/base-$BASE.kld" | sed -e "s#$HOME#~#g" >> raw/kld-base-$BASE.txt
fi
for q in q5_k_m q4_k_m; do
  m="$G/qwen2.5-1.5b-instruct-$q.gguf"; out="raw/kld-$q.txt"
  done_dl "$q" && [ -s "$G/base-$BASE.kld" ] || continue
  [ -s "$out" ] && grep -q "Mean    KLD" "$out" && continue
  { state "kld $q"; echo "\$ llama-perplexity -m qwen2.5-1.5b-instruct-$q.gguf --kl-divergence-base base-$BASE.kld --kl-divergence"; } > "$out"
  "$LCPP/llama-perplexity" -m "$m" --kl-divergence-base "$G/base-$BASE.kld" --kl-divergence >> "$out" 2>&1
  grep -E "Mean    KLD|Same top p" "$out" >&2
done

for q in $QUANTS; do
  m="$G/qwen2.5-1.5b-instruct-$q.gguf"; out="raw/bench-$q.json"
  done_dl "$q" || continue
  [ -s "$out" ] && continue
  state "bench $q" > "raw/bench-$q.txt"
  "$LCPP/llama-bench" -m "$m" -p 512 -n 128 -r 3 -o json > "$out" 2>> "raw/bench-$q.txt"
done