#!/bin/bash
# Copyright (c) 2026 Jiejing Zhang.
# Gemma 4 E2B/E4B on one Mac: Tempo9 1.2.0 vs llama.cpp
# (brew llama-server), same GGUF, interleaved rounds. Per arm and round:
#   single : dec.py 256 greedy tokens, streamed client-side rate, x3 prompts
#   prefill: dec.py with a ~2k-token prompt (ttft field)
#   conc   : conc.py N=1/4/8/16 x 256 tokens, aggregate tok/s
#   parity : parity.py, 6 prompts x 128 greedy tokens, text saved per arm
# Usage: MODEL=<gguf> TAG=e2b ROUNDS=2 bash bench.sh   (T9=/path/to/tempo9 LS=/path/to/llama-server to override)
set -u
cd "$(dirname "$0")"
F=.
MODEL=${MODEL:?}; TAG=${TAG:?}; ROUNDS=${ROUNDS:-2}; N=${N:-256}
T9=${T9:-$(command -v tempo9)}
LS=${LS:-/opt/homebrew/bin/llama-server}
P0=11720
# tempo9 has no --version flag; brew knows the version, otherwise print the path.
echo "VERSIONS: tempo9=$(brew list --versions tempo9 2>/dev/null || echo "$T9") llama=$($LS --version 2>&1 | grep -m1 version)"
echo "MODEL: $MODEL ($(du -h "$MODEL" | cut -f1))"
stop() { pkill -f "tempo9 --gguf"; pkill -f "llama-server -m"
         for i in $(seq 1 30); do pgrep -f "tempo9 --gguf|llama-server -m" >/dev/null || break; sleep 1; done; sleep 3; }
ready() { for i in $(seq 1 150); do CONC_TIMEOUT=120 python3 -I $F/conc.py $1 1 8 probe >/dev/null 2>&1 && return 0; sleep 2; done; return 1; }
# Prefill probe: prompt_ttft.txt (published next to this script), the same
# text behind the article's TTFT figures. ~7k characters; the token count is
# whatever each engine reports for it.
LONG=$(python3 -I -c "print(open('$F/prompt_ttft.txt').read().replace(chr(10),' '))")
for r in $(seq 1 $ROUNDS); do for arm in ${ARMS:-t9 lcpp}; do
  stop; PP=$((P0 + r*10 + ${#arm}))
  case $arm in
    t9)   $T9 --gguf "$MODEL" --max-length 4096 --max-batch 16 --port $PP > srv_${TAG}_${arm}_$r.log 2>&1 & ;;
    lcpp) $LS -m "$MODEL" -ngl 99 -c 65536 -np 16 --port $PP --no-warmup > srv_${TAG}_${arm}_$r.log 2>&1 & ;;
  esac
  ready $PP || { echo "$arm r$r: not ready"; tail -3 srv_${TAG}_${arm}_$r.log; continue; }
  python3 -I $F/dec.py $PP 64 warm >/dev/null
  for k in 0 1 2; do
    PR=$(python3 -I -c "
P=['Write a long, detailed story about a lighthouse keeper who finds a message in a bottle.',
   'Explain in detail how a modern CPU executes instructions out of order.',
   'Describe the history of the Silk Road and its influence on trade and culture.']; print(P[$k])")
    python3 -I $F/dec.py $PP $N "single $TAG $arm r$r p$k" "$PR"
  done
  python3 -I $F/dec.py $PP 16 "prefill7k $TAG $arm r$r" "Read these notes and summarize them in one sentence: $LONG"
  for c in ${CONCS:-1 4 8 16}; do python3 -I $F/conc.py $PP $c $N "conc $TAG $arm r$r"; done
  [ $r = 1 ] && python3 -I parity.py $PP parity_${TAG}_${arm}.json
done; done
stop
echo "DONE $TAG $(date +%T)"