Joey Grasty
e25ee046a3
llama.cpp-based implementation of the countbayesie Future-Entropy Sampler, built via a direct ctypes binding to libllama.so. Includes a one-shot CLI (entropy_cli.py), an OpenAI-compatible server (entropy_server.py), and the tuning/benchmark scripts used to derive the alpha/confidence-threshold defaults documented in the README. Model path and llama.cpp library path are now read from environment variables (ENTROPY_SAMPLER_MODEL, LLAMA_CPP_LIB) instead of being hardcoded, and gguf-py is pulled from PyPI instead of a local llama.cpp checkout, so this runs on any machine with a compatible llama.cpp build and model.
63 lines
2.2 KiB
Python
63 lines
2.2 KiB
Python
"""Measure the actual per-token overhead of the entropy sampler vs a plain
|
|
single-token-decode baseline (greedy, no forking) on this model/hardware.
|
|
"""
|
|
import time
|
|
|
|
import numpy as np
|
|
|
|
import llama_capi as C
|
|
from entropy_sampler import EntropySampler, softmax
|
|
|
|
PROMPT = "The old lighthouse keeper had seen many storms, but none like"
|
|
N_TOKENS = 60
|
|
|
|
|
|
def time_baseline(sampler, n_tokens):
|
|
base_logits = sampler.prime(PROMPT)
|
|
t0 = time.perf_counter()
|
|
for _ in range(n_tokens):
|
|
probs = softmax(base_logits)
|
|
tok = int(np.argmax(probs)) # greedy - cheapest possible comparison point
|
|
sampler._decode([tok], [0], [sampler.n_past], [True])
|
|
sampler.n_past += 1
|
|
base_logits = sampler._logits_row(-1)
|
|
if C.lib.llama_vocab_is_eog(sampler.vocab, tok):
|
|
break
|
|
return time.perf_counter() - t0
|
|
|
|
|
|
def time_entropy_sampler(sampler, n_tokens):
|
|
fork_time = 0.0
|
|
orig_score = sampler._score_candidates
|
|
|
|
def timed_score(base_logits):
|
|
nonlocal fork_time
|
|
t0 = time.perf_counter()
|
|
result = orig_score(base_logits)
|
|
fork_time += time.perf_counter() - t0
|
|
return result
|
|
|
|
sampler._score_candidates = timed_score
|
|
t0 = time.perf_counter()
|
|
sampler.generate(PROMPT, max_new_tokens=n_tokens, alpha=0.0, verbose=False)
|
|
total_time = time.perf_counter() - t0
|
|
sampler._score_candidates = orig_score
|
|
return total_time, fork_time
|
|
|
|
|
|
if __name__ == "__main__":
|
|
for top_k in (4, 8, 12, 20):
|
|
sampler = EntropySampler(top_k=top_k, top_n_future=20, n_ctx=4096)
|
|
try:
|
|
t_base = time_baseline(sampler, N_TOKENS)
|
|
t_entropy, t_fork = time_entropy_sampler(sampler, N_TOKENS)
|
|
|
|
print(f"=== top_k={top_k} ===")
|
|
print(f" baseline (no forking): {t_base:6.2f}s ({N_TOKENS/t_base:5.2f} tok/s)")
|
|
print(f" entropy sampler: {t_entropy:6.2f}s ({N_TOKENS/t_entropy:5.2f} tok/s)")
|
|
print(f" of which fork+score: {t_fork:6.2f}s ({100*t_fork/t_entropy:.0f}% of total)")
|
|
print(f" of which commit decode:{t_entropy - t_fork:6.2f}s")
|
|
print(f" overhead: {t_entropy/t_base:.2f}x\n")
|
|
finally:
|
|
sampler.close()
|