Initial commit: Future-Entropy Sampler prototype
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.
This commit is contained in:
90
adaptive_bench.py
Normal file
90
adaptive_bench.py
Normal file
@@ -0,0 +1,90 @@
|
||||
"""Does skipping the fork+peek lookahead on high-confidence tokens
|
||||
(confidence_threshold) recover speed without giving up quality/robustness?
|
||||
|
||||
For each threshold: measure wall-clock speed, what fraction of steps actually
|
||||
skipped the lookahead, and crash rate across repeated trials (reusing the
|
||||
degeneracy heuristic from robustness_test.py).
|
||||
"""
|
||||
import time
|
||||
|
||||
import numpy as np
|
||||
|
||||
from entropy_sampler import EntropySampler, sine_alpha_schedule
|
||||
from tune_alpha import PROMPT, MAX_NEW_TOKENS, summarize
|
||||
from robustness_test import is_degenerate
|
||||
|
||||
N_TRIALS = 5
|
||||
TOP_K = 12
|
||||
|
||||
# our chosen default from the tuning/robustness passes: peak alpha capped at +0.4
|
||||
DEFAULT_SCHEDULE = lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.6, offset=-0.2)
|
||||
|
||||
THRESHOLDS = [None, 0.9, 0.7, 0.5, 0.3]
|
||||
|
||||
|
||||
def run_trial(sampler, threshold, seed):
|
||||
sampler.rng = np.random.default_rng(seed)
|
||||
log = []
|
||||
orig_step = sampler.step
|
||||
|
||||
def step_and_log(base_logits, a, confidence_threshold=None):
|
||||
token, next_logits, diag = orig_step(base_logits, a, confidence_threshold)
|
||||
log.append(diag)
|
||||
return token, next_logits, diag
|
||||
|
||||
sampler.step = step_and_log
|
||||
t0 = time.perf_counter()
|
||||
text = sampler.generate(
|
||||
PROMPT, max_new_tokens=MAX_NEW_TOKENS, alpha=DEFAULT_SCHEDULE(),
|
||||
verbose=False, confidence_threshold=threshold,
|
||||
)
|
||||
elapsed = time.perf_counter() - t0
|
||||
sampler.step = orig_step
|
||||
|
||||
forked_log = [d for d in log if d["forked"]]
|
||||
stats = summarize(text, forked_log) if forked_log else {"mean_p": float("nan"), "mean_h": float("nan")}
|
||||
stats["degenerate"] = is_degenerate(text)
|
||||
stats["text"] = text
|
||||
stats["elapsed"] = elapsed
|
||||
stats["skip_fraction"] = 1 - len(forked_log) / len(log)
|
||||
return stats
|
||||
|
||||
|
||||
if __name__ == "__main__":
|
||||
sampler = EntropySampler(top_k=TOP_K, top_n_future=20, n_ctx=4096)
|
||||
try:
|
||||
summary_rows = []
|
||||
for threshold in THRESHOLDS:
|
||||
label = "always fork (baseline)" if threshold is None else f"threshold={threshold}"
|
||||
trials = [run_trial(sampler, threshold, seed=200 + i) for i in range(N_TRIALS)]
|
||||
|
||||
print(f"=== {label} ===")
|
||||
for i, t in enumerate(trials):
|
||||
flag = "DEGENERATE" if t["degenerate"] else "ok"
|
||||
print(
|
||||
f" seed={200+i} [{flag:10s}] skip={t['skip_fraction']:.0%} "
|
||||
f"time={t['elapsed']:5.1f}s ({MAX_NEW_TOKENS/t['elapsed']:.2f} tok/s) "
|
||||
f"mean_p={t['mean_p']:.3f}"
|
||||
)
|
||||
print(f" {t['text'][:160]!r}")
|
||||
|
||||
crash_rate = sum(t["degenerate"] for t in trials) / len(trials)
|
||||
avg_skip = sum(t["skip_fraction"] for t in trials) / len(trials)
|
||||
avg_time = sum(t["elapsed"] for t in trials) / len(trials)
|
||||
avg_tokps = MAX_NEW_TOKENS / avg_time
|
||||
summary_rows.append((label, avg_skip, avg_time, avg_tokps, crash_rate))
|
||||
print(
|
||||
f" -> avg_skip={avg_skip:.0%} avg_time={avg_time:.1f}s "
|
||||
f"avg_tok/s={avg_tokps:.2f} crash_rate={crash_rate:.0%}\n"
|
||||
)
|
||||
|
||||
print("=== summary ===")
|
||||
baseline_time = summary_rows[0][2]
|
||||
for label, avg_skip, avg_time, avg_tokps, crash_rate in summary_rows:
|
||||
speedup = baseline_time / avg_time
|
||||
print(
|
||||
f"{label:26s} skip={avg_skip:5.0%} {avg_tokps:5.2f} tok/s "
|
||||
f"speedup={speedup:.2f}x crash_rate={crash_rate:.0%}"
|
||||
)
|
||||
finally:
|
||||
sampler.close()
|
||||
Reference in New Issue
Block a user