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.
88 lines
3.3 KiB
Python
88 lines
3.3 KiB
Python
"""Same baseline vs threshold=0.7 comparison as big_test_070.py, now with the
|
|
<think>/</think> token ban active in EntropySampler. Reuses the same fresh
|
|
seed range so results are directly comparable to the pre-ban run.
|
|
"""
|
|
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 = 20
|
|
TOP_K = 12
|
|
SEED_START = 400 # same seeds as the pre-ban big_test_070.py run, for a fair diff
|
|
|
|
DEFAULT_SCHEDULE = lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.6, offset=-0.2)
|
|
|
|
|
|
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
|
|
|
|
|
|
def run_batch(sampler, threshold, label):
|
|
trials = [run_trial(sampler, threshold, seed=SEED_START + 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={SEED_START+i} [{flag:10s}] skip={t['skip_fraction']:.0%} "
|
|
f"time={t['elapsed']:5.1f}s ({MAX_NEW_TOKENS/t['elapsed']:.2f} tok/s) mean_p={t['mean_p']:.3f}"
|
|
)
|
|
print(f" {t['text']!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
|
|
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%} "
|
|
f"({sum(t['degenerate'] for t in trials)}/{N_TRIALS})\n"
|
|
)
|
|
return label, avg_skip, avg_time, avg_tokps, crash_rate
|
|
|
|
|
|
if __name__ == "__main__":
|
|
sampler = EntropySampler(top_k=TOP_K, top_n_future=20, n_ctx=4096)
|
|
print(f"banned token ids active: {sampler.banned_token_ids}\n")
|
|
try:
|
|
rows = []
|
|
rows.append(run_batch(sampler, None, "always fork (baseline), n=20, think-banned"))
|
|
rows.append(run_batch(sampler, 0.7, "threshold=0.7, n=20, think-banned"))
|
|
|
|
print("=== summary ===")
|
|
baseline_time = rows[0][2]
|
|
for label, avg_skip, avg_time, avg_tokps, crash_rate in rows:
|
|
speedup = baseline_time / avg_time
|
|
print(
|
|
f"{label:40s} skip={avg_skip:5.0%} {avg_tokps:5.2f} tok/s "
|
|
f"speedup={speedup:.2f}x crash_rate={crash_rate:.0%}"
|
|
)
|
|
finally:
|
|
sampler.close()
|