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.
164 lines
7.0 KiB
Python
164 lines
7.0 KiB
Python
"""Live CLI for the Future-Entropy Sampler prototype - a single-shot,
|
|
parameterized alternative to the tune/test scripts, for hands-on testing
|
|
of live sampler configurations against a seed prompt.
|
|
|
|
Examples:
|
|
python entropy_cli.py --prompt "The old lighthouse keeper had seen many storms, but none like"
|
|
python entropy_cli.py --prompt-file seed.txt --sine 16,0.6,-0.2 --top-k 16
|
|
python entropy_cli.py --prompt "..." --alpha 0.3 --confidence-threshold 0.8 --quiet
|
|
"""
|
|
import argparse
|
|
import random
|
|
import sys
|
|
|
|
import llama_capi as C
|
|
from entropy_sampler import EntropySampler, sine_alpha_schedule, MODEL_PATH
|
|
|
|
SAFE_PEAK_ALPHA = 0.5 # empirically: coherence cliff starts around peak alpha ~+0.7
|
|
|
|
|
|
def parse_sine(spec):
|
|
parts = spec.split(",")
|
|
if len(parts) not in (2, 3, 4):
|
|
raise argparse.ArgumentTypeError(
|
|
"--sine expects period,amplitude[,offset[,phase]] e.g. 16,0.6,-0.2"
|
|
)
|
|
period, amplitude = float(parts[0]), float(parts[1])
|
|
offset = float(parts[2]) if len(parts) >= 3 else 0.0
|
|
phase = float(parts[3]) if len(parts) >= 4 else 0.0
|
|
return dict(period_tokens=period, amplitude=amplitude, offset=offset, phase=phase)
|
|
|
|
|
|
def build_arg_parser():
|
|
p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
|
|
|
|
prompt_group = p.add_mutually_exclusive_group(required=True)
|
|
prompt_group.add_argument("--prompt", type=str, help="Seed text to continue.")
|
|
prompt_group.add_argument("--prompt-file", type=str, help="Read seed text from a file.")
|
|
|
|
alpha_group = p.add_mutually_exclusive_group()
|
|
alpha_group.add_argument(
|
|
"--alpha", type=float, default=0.0,
|
|
help="Fixed alpha in [-1, 1]. Ignored if --sine is given. Default 0.0 (balanced). "
|
|
"Safe range roughly -1.0 to +0.3; avoid >= +0.7 (reliable register-break).",
|
|
)
|
|
alpha_group.add_argument(
|
|
"--sine", type=parse_sine, metavar="PERIOD,AMPLITUDE[,OFFSET[,PHASE]]",
|
|
help="Oscillate alpha sinusoidally instead of a fixed value, e.g. 16,0.6,-0.2 "
|
|
"(period=16 tokens, amplitude=0.6, offset=-0.2). Keep offset+amplitude "
|
|
"<= ~+0.5 (peak alpha) to stay clear of the coherence cliff.",
|
|
)
|
|
|
|
p.add_argument("--top-k", type=int, default=12,
|
|
help="Candidates considered per step. Default 12. Range ~4-20.")
|
|
p.add_argument("--top-n-future", type=int, default=20,
|
|
help="Future-distribution truncation for H_hat. Default 20.")
|
|
p.add_argument("--confidence-threshold", type=float, default=None,
|
|
help="Skip fork+peek lookahead when the top candidate's p(w|c) is "
|
|
"already >= this. Default: None (always fork). 0.7-0.9 trades "
|
|
"some quality for speed; below 0.5 raises crash rate a lot.")
|
|
p.add_argument("--max-new-tokens", type=int, default=60)
|
|
p.add_argument("--n-ctx", type=int, default=8192,
|
|
help="Requested total context. Actual per-sequence budget is "
|
|
"n_ctx / (top_k+1), padded up to a multiple of 256 - printed "
|
|
"at startup so you can see what you actually got.")
|
|
p.add_argument("--n-batch", type=int, default=8192,
|
|
help="Hard ceiling on tokens submitted in one llama_decode() call "
|
|
"(prime() decodes the whole prompt in one call). Must cover "
|
|
"your longest prompt/--prompt-file or you'll hit "
|
|
"GGML_ASSERT(n_tokens_all <= n_batch), which aborts the whole "
|
|
"process rather than raising a catchable exception.")
|
|
p.add_argument("--n-ubatch", type=int, default=512,
|
|
help="Internal per-step chunk size within a decode call.")
|
|
p.add_argument("--model", type=str, default=MODEL_PATH)
|
|
p.add_argument("--seed", type=int, default=None, help="Default: random, printed at startup.")
|
|
p.add_argument("--think-ban", dest="think_ban", action="store_true", default=True,
|
|
help="Ban <think>/</think> tokens at the logit level (default on). "
|
|
"Blocks literal leakage only, not the underlying register-break.")
|
|
p.add_argument("--no-think-ban", dest="think_ban", action="store_false")
|
|
p.add_argument("--quiet", action="store_true", help="Suppress the per-token diagnostics line.")
|
|
return p
|
|
|
|
|
|
def main():
|
|
args = build_arg_parser().parse_args()
|
|
|
|
if args.prompt_file:
|
|
with open(args.prompt_file) as f:
|
|
prompt = f.read()
|
|
else:
|
|
prompt = args.prompt
|
|
|
|
if args.sine is not None:
|
|
s = args.sine
|
|
peak = s["offset"] + s["amplitude"]
|
|
alpha = sine_alpha_schedule(
|
|
period_tokens=s["period_tokens"], amplitude=s["amplitude"],
|
|
offset=s["offset"], phase=s["phase"],
|
|
)
|
|
alpha_desc = (
|
|
f"sine(period={s['period_tokens']:g}, amplitude={s['amplitude']:.2f}, "
|
|
f"offset={s['offset']:.2f}, phase={s['phase']:.2f}) peak={peak:+.2f}"
|
|
)
|
|
if peak > SAFE_PEAK_ALPHA:
|
|
print(
|
|
f"warning: peak alpha {peak:+.2f} exceeds the ~+{SAFE_PEAK_ALPHA:.1f} "
|
|
f"safety margin observed for coherence - expect elevated crash risk",
|
|
file=sys.stderr,
|
|
)
|
|
else:
|
|
alpha = args.alpha
|
|
alpha_desc = f"fixed={args.alpha:+.2f}"
|
|
if args.alpha >= 0.7:
|
|
print(
|
|
f"warning: fixed alpha {args.alpha:+.2f} >= +0.7 reliably degenerated "
|
|
f"into register-breaks in prior testing",
|
|
file=sys.stderr,
|
|
)
|
|
|
|
seed = args.seed if args.seed is not None else random.randint(0, 2**31 - 1)
|
|
|
|
sampler = EntropySampler(
|
|
model_path=args.model,
|
|
n_ctx=args.n_ctx,
|
|
top_k=args.top_k,
|
|
top_n_future=args.top_n_future,
|
|
seed=seed,
|
|
think_ban=args.think_ban,
|
|
n_batch=args.n_batch,
|
|
n_ubatch=args.n_ubatch,
|
|
)
|
|
try:
|
|
n_ctx_actual = C.lib.llama_n_ctx(sampler.ctx)
|
|
n_ctx_seq = C.lib.llama_n_ctx_seq(sampler.ctx)
|
|
|
|
print("=== entropy sampler ===")
|
|
print(f"model: {args.model}")
|
|
print(f"top_k: {args.top_k} (n_seq_max={args.top_k + 1})")
|
|
print(f"top_n_future: {args.top_n_future}")
|
|
print(f"alpha: {alpha_desc}")
|
|
print(f"confidence_thr: {args.confidence_threshold}")
|
|
print(f"think_ban: {args.think_ban} (ids={sampler.banned_token_ids})")
|
|
print(f"seed: {seed}")
|
|
print(f"n_ctx: requested={args.n_ctx} actual={n_ctx_actual} per-seq={n_ctx_seq}")
|
|
print("=" * 40)
|
|
print()
|
|
|
|
text = sampler.generate(
|
|
prompt,
|
|
max_new_tokens=args.max_new_tokens,
|
|
alpha=alpha,
|
|
verbose=not args.quiet,
|
|
confidence_threshold=args.confidence_threshold,
|
|
)
|
|
|
|
print()
|
|
print("=== final text ===")
|
|
print(text)
|
|
finally:
|
|
sampler.close()
|
|
|
|
|
|
if __name__ == "__main__":
|
|
main()
|