commit e25ee046a33e3d1d565814631c60d50ee2cfc931 Author: Joey Grasty Date: Sat Jul 11 19:28:38 2026 -0500 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. diff --git a/.gitignore b/.gitignore new file mode 100644 index 0000000..d5225b5 --- /dev/null +++ b/.gitignore @@ -0,0 +1,8 @@ +__pycache__/ +*.pyc +*.log +.claude/ + +# Generated during test/dev sessions, not part of the tool itself +resume.txt +Prompt-response.txt diff --git a/LICENSE b/LICENSE new file mode 100644 index 0000000..adbfd10 --- /dev/null +++ b/LICENSE @@ -0,0 +1,21 @@ +MIT License + +Copyright (c) 2026 Joey Grasty + +Permission is hereby granted, free of charge, to any person obtaining a copy +of this software and associated documentation files (the "Software"), to deal +in the Software without restriction, including without limitation the rights +to use, copy, modify, merge, publish, distribute, sublicense, and/or sell +copies of the Software, and to permit persons to whom the Software is +furnished to do so, subject to the following conditions: + +The above copyright notice and this permission notice shall be included in all +copies or substantial portions of the Software. + +THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR +IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY, +FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE +AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER +LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM, +OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN THE +SOFTWARE. diff --git a/README.md b/README.md new file mode 100644 index 0000000..c555255 --- /dev/null +++ b/README.md @@ -0,0 +1,156 @@ +# Future-Entropy Sampler + +A prototype implementation of the "Future-Entropy Sampler" technique described in +[countbayesie's *Making LLMs Better at Creative Writing Using Entropy*](https://www.countbayesie.com/blog/2026/7/1/making-llms-better-at-creative-writing-using-entropy) +(2026-07-01), built directly against llama.cpp's C API via `ctypes`. + +Instead of sampling straight from `p(w | c)`, this sampler forks the model's KV +cache once per top-k candidate token, decodes one token into each fork to see +what distribution follows it, scores each candidate by how much probability it +carries *and* how much future creative choice it preserves, then commits to one +token and discards the rest of the forks: + +``` +s(w) = p(w|c)^a * H_hat(w)^b, a = 1 - alpha, b = 1 + alpha, alpha in [-1, 1] +``` + +`H_hat(w)` is the normalized Shannon entropy of the top-n token distribution +that follows candidate `w`. `alpha` crossfades between pure-probability +sampling (`alpha=-1`) and pure future-entropy-chasing (`alpha=+1`), and can be +oscillated over generation with a sine schedule ("rhythmic decoding" in the +source article) instead of held fixed. + +This is a research prototype from hands-on tuning against one model on one +machine, not a polished library - see **Known limitations** below before +relying on it for anything beyond experimentation. + +## How it works + +llama.cpp's `llama_memory_seq_cp`/`llama_memory_seq_rm` plus multi-sequence +batched `llama_decode` make the fork/peek/discard cheap: all `top_k` candidate +forks are decoded in a single batched call, not `top_k` sequential ones. +`llama-cpp-python` wasn't used - `llama_capi.py` binds the C API directly via +`ctypes`, which means the struct layouts in that file must match whatever +`libllama.so` you point it at (see Known limitations). + +## Requirements + +- A working [llama.cpp](https://github.com/ggml-org/llama.cpp) build with the + shared library enabled (`cmake -DBUILD_SHARED_LIBS=ON ...`), producing + `libllama.so` (or `.dylib`/`.dll`). +- A GGUF model file. Developed and tuned against a Qwen3.5-35B-A3B MoE model; + behavior (alpha ranges, crash thresholds, etc. - see below) is + model-specific and will differ elsewhere. +- Python 3.10+. + +``` +pip install -r requirements.txt +``` + +## Setup + +Two environment variables are required: + +``` +export LLAMA_CPP_LIB=/path/to/llama.cpp/build/bin/libllama.so +export ENTROPY_SAMPLER_MODEL=/path/to/your-model.gguf +``` + +`ENTROPY_SAMPLER_MODEL` is just the default - every entry point also accepts +`--model` to override it per run. + +## Usage + +**One-shot CLI**, for testing a single generation with a given configuration: + +``` +python entropy_cli.py --prompt "Once upon a time" --sine 16,0.6,-0.2 --confidence-threshold 0.8 +python entropy_cli.py --help # full parameter list, defaults, and safe ranges +``` + +**OpenAI-compatible server**, for use with a chat UI (e.g. Open WebUI as a +second model connection) alongside your normal inference server: + +``` +python entropy_server.py --port 30001 +``` + +Exposes `/v1/models` and `/v1/chat/completions` (streaming and +non-streaming). Applies the model's own embedded chat template +(`enable_thinking=False`) so it behaves like a normal chat model. Sampler +configuration is fixed at server startup via CLI flags - there's no +per-request override; restart with different flags to try a different +configuration. Run `python entropy_server.py --help` for the full flag list. + +## Tuning notes from initial testing + +These numbers are from one model (Qwen3.5-35B-A3B) on one machine and should +be treated as a starting point to re-derive on your own setup, not universal +constants: + +- **Fixed alpha stays coherent roughly up to +0.3-0.5; +0.7 and above + reliably degenerates** into meta-commentary/register-breaks or + bracket-listing artifacts. +- **For a sine-oscillating alpha, peak alpha (`offset + amplitude`) is what + drives crash risk, not the mean/offset.** Keeping the peak under ~+0.5 kept + generations clean across repeated trials even with a fairly wide swing. +- **The adaptive `confidence_threshold` skip** (skip the fork+peek lookahead + when the top base candidate is already this confident) trades quality for + speed on a curve: 0.5 gives ~2x speedup but roughly quadruples crash rate; + 0.8 gives a modest ~1.24x speedup at a crash rate statistically + indistinguishable from not skipping at all. 0.8 is the default in + `entropy_server.py`. +- Every fork+peek step is real inference overhead (up to ~14x wall time at + top_k=20 vs. a plain greedy baseline), so this is meaningfully slower than + normal sampling - budget for it, especially over long generations. + +## Known limitations + +- **ABI-fragile**: `llama_capi.py`'s struct definitions are hand-copied from + one specific llama.cpp commit's `llama.h`. A different llama.cpp version + can silently reorder/resize struct fields and corrupt memory instead of + raising a clean error. Check the struct layouts against your build's + `llama.h` if you see crashes or garbage output. +- **Single request at a time**: one `EntropySampler` holds one `llama_context` + and one KV cache; `entropy_server.py` serializes requests with a lock. Not + built for concurrent multi-user serving. +- **`n_batch`/`n_ctx` sizing matters more than it looks like it should**: this + model's KV cache is not unified across the fork sequences + (`kv_unified=false`), so the requested `n_ctx` is split `(top_k+1)` ways - + the real per-request budget is much smaller than the number you pass in + (`entropy_cli.py`/`entropy_server.py` print the actual resulting budget at + startup). Separately, `n_batch` (default 2048 in llama.cpp) is a hard + ceiling on tokens submitted in a single `llama_decode()` call; since + `prime()` decodes an entire prompt/chat history in one call, a long prompt + that exceeds `n_batch` trips `GGML_ASSERT(n_tokens_all <= n_batch)`, which + **aborts the whole process** rather than raising a catchable exception. + Both tools default their `--n-batch` high enough to match their context + budget, but raise `--n-ctx` and `--n-batch` together if you change one. +- **Unresolved register-break failure mode**: at higher alpha (or, more + subtly, even at safer settings on rare seeds), the model can drop into an + assistant/meta-commentary voice mid-generation (e.g. suddenly explaining + its own output, or switching into a quiz/translation register) instead of + continuing the prose. Banning the literal ``/`` tokens at + the logit level (on by default, `--think-ban`) blocks that one surface + form, but the model reroutes around it with fluent alternative phrasing - + this is a genuine model-behavior problem, not something a token ban fixes. + Untried mitigations: detect-and-regenerate on a degeneracy heuristic + (`robustness_test.py` has a starting one), or steering it at the + system-prompt level. + +## Files + +| File | Purpose | +|---|---| +| `llama_capi.py` | `ctypes` bindings over `libllama.so` | +| `entropy_sampler.py` | Core `EntropySampler` class: `prime`/`step`/`generate`/`generate_stream` | +| `entropy_cli.py` | One-shot parameterized CLI | +| `entropy_server.py` | OpenAI-compatible HTTP server | +| `fork_peek_test.py` | Minimal smoke test for the fork/peek/discard KV-cache primitive | +| `robustness_test.py` | Repeated-trial crash-rate testing with a degeneracy heuristic | +| `tune_alpha.py`, `tune_offset_sine.py` | Fixed-alpha and sine-schedule sweeps | +| `timing_bench.py`, `adaptive_bench.py`, `big_test_*.py` | Benchmark/comparison scripts from initial tuning | + +## License + +See [LICENSE](LICENSE). diff --git a/adaptive_bench.py b/adaptive_bench.py new file mode 100644 index 0000000..0e63a3c --- /dev/null +++ b/adaptive_bench.py @@ -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() diff --git a/big_test_050.py b/big_test_050.py new file mode 100644 index 0000000..446efb0 --- /dev/null +++ b/big_test_050.py @@ -0,0 +1,86 @@ +"""Larger-N robustness check on confidence_threshold=0.5: the 5-seed sweep +looked clean, but 5 seeds isn't enough to trust a 0% crash rate. Run 20 fresh +seeds (not overlapping the earlier 5-seed batch) against both threshold=0.5 +and the always-fork baseline, for a fair speed + crash-rate comparison. +""" +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 = 300 # fresh seeds, no overlap with the earlier 200-204 batch + +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%} ({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) + try: + rows = [] + rows.append(run_batch(sampler, None, "always fork (baseline), n=20")) + rows.append(run_batch(sampler, 0.5, "threshold=0.5, n=20")) + + 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:32s} skip={avg_skip:5.0%} {avg_tokps:5.2f} tok/s " + f"speedup={speedup:.2f}x crash_rate={crash_rate:.0%}" + ) + finally: + sampler.close() diff --git a/big_test_070.py b/big_test_070.py new file mode 100644 index 0000000..6829d76 --- /dev/null +++ b/big_test_070.py @@ -0,0 +1,89 @@ +"""threshold=0.7 vs always-fork baseline, n=20 each, using the improved +detector (repetition-loop + meta-commentary/reasoning-leak patterns on top +of the original blank/markdown check). threshold=0.5 showed a real quality +cost the old heuristic couldn't see; 0.7 only skips when the top candidate +already has clear majority conviction (>=70%), trading less speedup for +more of the safety margin. +""" +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 # fresh seeds, no overlap with earlier batches + +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) + try: + rows = [] + rows.append(run_batch(sampler, None, "always fork (baseline), n=20")) + rows.append(run_batch(sampler, 0.7, "threshold=0.7, n=20")) + + 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:32s} skip={avg_skip:5.0%} {avg_tokps:5.2f} tok/s " + f"speedup={speedup:.2f}x crash_rate={crash_rate:.0%}" + ) + finally: + sampler.close() diff --git a/big_test_070_noban_vs_ban.py b/big_test_070_noban_vs_ban.py new file mode 100644 index 0000000..8eee8a0 --- /dev/null +++ b/big_test_070_noban_vs_ban.py @@ -0,0 +1,87 @@ +"""Same baseline vs threshold=0.7 comparison as big_test_070.py, now with the +/ 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() diff --git a/big_test_080.py b/big_test_080.py new file mode 100644 index 0000000..2ce9be1 --- /dev/null +++ b/big_test_080.py @@ -0,0 +1,88 @@ +"""threshold=0.8 vs always-fork baseline, n=20 each, using the improved +degeneracy detector (repetition-loop + meta-commentary/reasoning-leak +patterns on top of the original blank/markdown check). Fresh seed range, +no overlap with the 050/070 batches. +""" +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 = 500 # fresh seeds, no overlap with earlier batches + +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")) + rows.append(run_batch(sampler, 0.8, "threshold=0.8, n=20")) + + 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:32s} skip={avg_skip:5.0%} {avg_tokps:5.2f} tok/s " + f"speedup={speedup:.2f}x crash_rate={crash_rate:.0%}" + ) + finally: + sampler.close() diff --git a/entropy_cli.py b/entropy_cli.py new file mode 100644 index 0000000..90e519c --- /dev/null +++ b/entropy_cli.py @@ -0,0 +1,163 @@ +"""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 / 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() diff --git a/entropy_sampler.py b/entropy_sampler.py new file mode 100644 index 0000000..c80481e --- /dev/null +++ b/entropy_sampler.py @@ -0,0 +1,292 @@ +"""Future-Entropy Sampler. + +At each step, instead of sampling straight from p(w | c), fork the KV cache +once per top-k candidate, decode one token into each fork to see what +distribution q_w follows it, score candidates by how much p(w|c) they carry +*and* how much future choice they preserve, then commit to one and discard +the rest of the forks. + + s(w) = p(w|c)^a * H_hat(w)^b, a = 1 - alpha, b = 1 + alpha, alpha in [-1, 1] + +alpha = -1 -> a=2,b=0 (pure probability, squared) +alpha = 0 -> a=1,b=1 (balanced - matches the base s(w) = p(w)*H_hat(w) form) +alpha = +1 -> a=0,b=2 (pure future-entropy, squared) +This linear a/b mapping isn't specified by the source article beyond "derived +from alpha" - it's the simplest crossfade that recovers the balanced form at +alpha=0 and degenerates cleanly at the extremes. +""" +import math +import os + +import numpy as np + +import llama_capi as C + +MODEL_PATH = os.environ.get("ENTROPY_SAMPLER_MODEL") + + +def softmax(x): + x = x - np.max(x) + e = np.exp(x) + return e / e.sum() + + +def normalized_entropy(probs): + p = probs / probs.sum() + p = p[p > 0] + if len(p) <= 1: + return 0.0 + h = -np.sum(p * np.log(p)) + return float(h / math.log(len(p))) + + +def alpha_to_exponents(alpha): + alpha = max(-1.0, min(1.0, alpha)) + return 1.0 - alpha, 1.0 + alpha + + +def sine_alpha_schedule(period_tokens=12, amplitude=1.0, phase=0.0, offset=0.0): + """offset shifts the wave's center - e.g. offset=-0.2 dips deeper into + safe/coherent territory than it peaks into entropy-chasing territory, + biasing toward coherence while still spiking for occasional surprise.""" + def schedule(step): + return offset + amplitude * math.sin(2 * math.pi * step / period_tokens + phase) + return schedule + + +class EntropySampler: + def __init__(self, model_path=MODEL_PATH, n_ctx=8192, top_k=12, + top_n_future=20, n_threads=8, seed=0, think_ban=True, + n_batch=8192, n_ubatch=512): + if not model_path: + raise RuntimeError( + "No model path given. Set ENTROPY_SAMPLER_MODEL to a .gguf file, " + "or pass model_path= / --model explicitly." + ) + self.top_k = top_k + self.top_n_future = top_n_future + self.rng = np.random.default_rng(seed) + + C.lib.llama_backend_init() + + mparams = C.lib.llama_model_default_params() + mparams.n_gpu_layers = 99 + self.model = C.lib.llama_model_load_from_file(model_path.encode(), mparams) + if not self.model: + raise RuntimeError(f"model load failed for path: {model_path}") + + self.vocab = C.lib.llama_model_get_vocab(self.model) + self.n_vocab = C.lib.llama_vocab_n_tokens(self.vocab) + + # This model only offers enable_thinking through its chat template's + # assistant-turn preamble - we do raw completion, not chat-formatted + # prompting, so that switch is never in play. Ban / + # directly instead: both are single dedicated tokens in this vocab. + # (Known limited: suppresses literal leakage, not the underlying + # register-break the model reroutes around it with.) + self.banned_token_ids = [] + if think_ban: + for special in ("", ""): + toks = C.tokenize(self.vocab, special, add_special=False, parse_special=True) + if len(toks) == 1: + self.banned_token_ids.append(toks[0]) + + cparams = C.lib.llama_context_default_params() + cparams.n_ctx = n_ctx + cparams.n_seq_max = top_k + 1 # seq 0 = real context, 1..top_k = forks + cparams.n_threads = n_threads + cparams.n_threads_batch = n_threads + # n_batch is the hard ceiling on tokens submitted in one llama_decode() + # call - prime() decodes the whole prompt in a single call, so this + # must cover the largest prompt (i.e. rendered chat history) expected, + # not just n_ubatch's internal per-step chunk size. Left at the + # llama.cpp default (2048) this hits GGML_ASSERT(n_tokens_all <= + # cparams.n_batch), which aborts the whole process - not a + # recoverable Python exception. + cparams.n_batch = n_batch + cparams.n_ubatch = n_ubatch + + self.ctx = C.lib.llama_init_from_model(self.model, cparams) + if not self.ctx: + raise RuntimeError("context init failed") + + self.mem = C.lib.llama_get_memory(self.ctx) + self.n_past = 0 + + def close(self): + C.lib.llama_free(self.ctx) + C.lib.llama_model_free(self.model) + + def _decode(self, tokens, seq_ids, positions, want_logits): + """seq_ids/positions/want_logits are per-token, same length as tokens.""" + n = len(tokens) + batch = C.lib.llama_batch_init(n, 0, 1) + batch.n_tokens = n + for i in range(n): + batch.token[i] = tokens[i] + batch.pos[i] = positions[i] + batch.n_seq_id[i] = 1 + batch.seq_id[i][0] = seq_ids[i] + batch.logits[i] = 1 if want_logits[i] else 0 + rc = C.lib.llama_decode(self.ctx, batch) + C.lib.llama_batch_free(batch) + if rc != 0: + raise RuntimeError(f"llama_decode failed rc={rc}") + + def _logits_row(self, i): + ptr = C.lib.llama_get_logits_ith(self.ctx, i) + return np.ctypeslib.as_array(ptr, shape=(self.n_vocab,)).copy() + + def prime(self, prompt): + """Decode the prompt on seq 0. Returns base logits for the first step.""" + C.lib.llama_memory_seq_rm(self.mem, 0, -1, -1) # clear any prior run's cache + tokens = C.tokenize(self.vocab, prompt) + positions = list(range(len(tokens))) + want = [False] * (len(tokens) - 1) + [True] + self._decode(tokens, [0] * len(tokens), positions, want) + self.n_past = len(tokens) + return self._logits_row(-1) + + def _mask_banned(self, logits): + if self.banned_token_ids: + logits = logits.copy() + logits[self.banned_token_ids] = -np.inf + return logits + + def _top_candidates(self, base_logits): + probs = softmax(self._mask_banned(base_logits)) + top_idx = np.argsort(-probs)[: self.top_k] + return top_idx, probs[top_idx] + + def _fork_and_score(self, top_idx): + """Batched fork+peek+discard: one candidate token per fork sequence, + all decoded in a single llama_decode call.""" + fork_seqs = list(range(1, self.top_k + 1)) + for fseq in fork_seqs: + C.lib.llama_memory_seq_cp(self.mem, 0, fseq, -1, -1) + + self._decode( + tokens=[int(t) for t in top_idx], + seq_ids=fork_seqs, + positions=[self.n_past] * self.top_k, + want_logits=[True] * self.top_k, + ) + + h_hat = np.empty(self.top_k) + for j in range(self.top_k): + fork_probs = softmax(self._mask_banned(self._logits_row(j))) + top_future = np.argsort(-fork_probs)[: self.top_n_future] + h_hat[j] = normalized_entropy(fork_probs[top_future]) + + for fseq in fork_seqs: + C.lib.llama_memory_seq_rm(self.mem, fseq, -1, -1) + + return h_hat + + def _score_candidates(self, base_logits): + top_idx, top_p = self._top_candidates(base_logits) + h_hat = self._fork_and_score(top_idx) + return top_idx, top_p, h_hat + + def step(self, base_logits, alpha, confidence_threshold=None): + """confidence_threshold: if the top base candidate's p(w|c) is at or + above this, skip the fork+peek lookahead entirely and sample straight + from p(w|c) among the top-k - most tokens in fluent prose are forced + (punctuation, function words) and the entropy signal adds nothing + there. None disables skipping (always fork, the original behavior).""" + top_idx, top_p = self._top_candidates(base_logits) + forked = confidence_threshold is None or top_p[0] < confidence_threshold + + if forked: + h_hat = self._fork_and_score(top_idx) + a, b = alpha_to_exponents(alpha) + scores = (top_p ** a) * (h_hat ** b) + total = scores.sum() + if total <= 0: + scores = top_p # degenerate fallback: every candidate closes off the future + total = scores.sum() + else: + h_hat = np.full(self.top_k, np.nan) + scores = top_p + total = scores.sum() + + dist = scores / total + choice = self.rng.choice(len(top_idx), p=dist) + chosen_token = int(top_idx[choice]) + + diagnostics = { + "candidates": [ + { + "token": C.token_to_str(self.vocab, int(top_idx[j])), + "p": float(top_p[j]), + "h_hat": None if not forked else float(h_hat[j]), + "score": float(dist[j]), + } + for j in range(len(top_idx)) + ], + "chosen": choice, + "alpha": alpha, + "forked": forked, + } + + self._decode([chosen_token], [0], [self.n_past], [True]) + self.n_past += 1 + next_base_logits = self._logits_row(-1) + + return chosen_token, next_base_logits, diagnostics + + def generate_stream(self, prompt, max_new_tokens=60, alpha=0.0, confidence_threshold=None): + """Yields (piece, diag) one token at a time - the shared engine behind + generate() (batch) and the HTTP server (streaming).""" + alpha_fn = alpha if callable(alpha) else (lambda step: alpha) + + base_logits = self.prime(prompt) + for step_i in range(max_new_tokens): + a = alpha_fn(step_i) + token, base_logits, diag = self.step(base_logits, a, confidence_threshold) + + if C.lib.llama_vocab_is_eog(self.vocab, token): + return + + piece = C.token_to_str(self.vocab, token) + yield piece, diag + + def generate(self, prompt, max_new_tokens=60, alpha=0.0, verbose=True, + confidence_threshold=None): + text = prompt + for step_i, (piece, diag) in enumerate( + self.generate_stream(prompt, max_new_tokens, alpha, confidence_threshold) + ): + text += piece + + if verbose: + chosen = diag["candidates"][diag["chosen"]] + h_str = f"{chosen['h_hat']:.3f}" if chosen["h_hat"] is not None else " -- " + print( + f"[{step_i:3d}] alpha={diag['alpha']:+.2f} {'fork' if diag['forked'] else 'skip'}" + f" -> {piece!r:12s} p={chosen['p']:.3f} H_hat={h_str} " + f"score={chosen['score']:.3f}" + ) + + return text + + +if __name__ == "__main__": + sampler = EntropySampler(top_k=12, top_n_future=20, n_ctx=4096) + try: + prompt = "The old lighthouse keeper had seen many storms, but none like" + + print("=== alpha = -1.0 (pure probability) ===") + print(sampler.generate(prompt, max_new_tokens=40, alpha=-1.0)) + + print("\n=== alpha = 0.0 (balanced) ===") + print(sampler.generate(prompt, max_new_tokens=40, alpha=0.0)) + + print("\n=== alpha = +1.0 (pure future-entropy) ===") + print(sampler.generate(prompt, max_new_tokens=40, alpha=1.0)) + + print("\n=== alpha-wave (sine oscillation, period=10 tokens) ===") + schedule = sine_alpha_schedule(period_tokens=10, amplitude=1.0) + print(sampler.generate(prompt, max_new_tokens=40, alpha=schedule)) + finally: + sampler.close() diff --git a/entropy_server.py b/entropy_server.py new file mode 100644 index 0000000..c1369ad --- /dev/null +++ b/entropy_server.py @@ -0,0 +1,325 @@ +"""OpenAI-compatible HTTP server for the Future-Entropy Sampler prototype - +meant to run alongside the stock qwen3.5-llama-start server (port 30000) so +the two can be compared side by side from the same chat UI (e.g. Open WebUI +as two separate model connections). + +Sampler configuration (alpha/sine, top_k, confidence_threshold, ...) is fixed +at startup via CLI flags, same spirit as entropy_cli.py - there is no +per-request override and no runtime reconfiguration; restart with different +flags to try a different configuration. + +Usage: + python entropy_server.py --port 30001 + python entropy_server.py --port 30001 --sine 16,0.6,-0.2 --confidence-threshold 0.8 +""" +import argparse +import json +import sys +import threading +import time +import uuid + +import gguf +import jinja2 +from fastapi import FastAPI, HTTPException +from fastapi.responses import StreamingResponse +from pydantic import BaseModel +import uvicorn + +import llama_capi as C +from entropy_sampler import EntropySampler, sine_alpha_schedule, MODEL_PATH +from entropy_cli import parse_sine, SAFE_PEAK_ALPHA + + +DEFAULT_MODEL_ID = "qwen3.5-35b-a3b-entropy" + + +def load_chat_template(model_path): + reader = gguf.GGUFReader(model_path) + field = reader.fields.get("tokenizer.chat_template") + if field is None: + raise RuntimeError(f"{model_path} has no embedded tokenizer.chat_template") + return bytes(field.parts[-1]).decode("utf-8") + + +def _raise_exception(msg): + # The Qwen chat template calls raise_exception(...) as a Jinja global, + # matching the helper HF's apply_chat_template injects into its env. + raise ValueError(msg) + + +def build_jinja_template(template_src): + env = jinja2.Environment() + env.globals["raise_exception"] = _raise_exception + env.filters["tojson"] = lambda v: json.dumps(v) + return env.from_string(template_src) + + +class ChatMessage(BaseModel): + role: str + content: str = "" + + +class ChatCompletionRequest(BaseModel): + model: str | None = None + messages: list[ChatMessage] + stream: bool = False + max_tokens: int | None = None + # Accepted for OpenAI-client compatibility, intentionally ignored: this + # sampler's behavior is fully determined by the server's baked-in + # alpha/sine/top_k/confidence_threshold config, not per-request params. + temperature: float | None = None + top_p: float | None = None + + +def build_arg_parser(): + p = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter) + p.add_argument("--host", default="0.0.0.0") + p.add_argument("--port", type=int, default=30001) + p.add_argument("--model", default=MODEL_PATH) + p.add_argument("--model-id", default=DEFAULT_MODEL_ID, + help="Model name reported via /v1/models and in responses.") + + alpha_group = p.add_mutually_exclusive_group() + alpha_group.add_argument("--alpha", type=float, default=None, + help="Fixed alpha in [-1, 1]. Overrides --sine if given.") + alpha_group.add_argument("--sine", type=parse_sine, default="16,0.6,-0.2", + metavar="PERIOD,AMPLITUDE[,OFFSET[,PHASE]]", + help="Default: 16,0.6,-0.2 (peak +0.4) - the best-tested " + "config this session. Keep peak (offset+amplitude) " + "<= ~+0.5 for safety.") + + p.add_argument("--top-k", type=int, default=12) + p.add_argument("--top-n-future", type=int, default=20) + p.add_argument("--confidence-threshold", type=float, default=0.8) + p.add_argument("--n-ctx", type=int, default=212992, + help="Requested total context, split n_seq_max=(top_k+1) ways since " + "this model's KV cache is not unified across fork sequences " + "(kv_unified=false). Default 212992 -> ~16384 tokens/sequence " + "at top_k=12. Since each chat request re-sends the full " + "conversation history and re-primes from scratch, this is the " + "budget for prompt+reply *per request*, and it must also cover " + "the whole growing history in a multi-turn conversation.") + p.add_argument("--n-batch", type=int, default=16384, + help="Hard ceiling on tokens submitted in one llama_decode() call. " + "prime() decodes the whole rendered prompt/history in a single " + "call, so this must cover the largest expected prompt. Default " + "16384 matches the default per-sequence context budget. Too " + "small and long prompts trigger GGML_ASSERT(n_tokens_all <= " + "n_batch), which aborts the whole process (not a catchable " + "Python exception) - this happened with the llama.cpp default " + "of 2048 on a 2059-token prompt.") + p.add_argument("--n-ubatch", type=int, default=512, + help="Internal per-step chunk size within a decode call - llama.cpp " + "processes n_batch in n_ubatch-sized pieces automatically. " + "Smaller = less peak compute-buffer memory.") + p.add_argument("--n-threads", type=int, default=8) + p.add_argument("--max-tokens-default", type=int, default=3000, + help="Used when a request doesn't specify max_tokens (Open WebUI " + "typically doesn't). Sized for ~1600-word creative-writing " + "replies (~2000-2200 tokens) with margin - 512 and 1536 both " + "proved too tight and truncated mid-sentence.") + p.add_argument("--think-ban", dest="think_ban", action="store_true", default=True) + p.add_argument("--no-think-ban", dest="think_ban", action="store_false") + p.add_argument("--enable-thinking", dest="enable_thinking", action="store_true", default=False, + help="Chat template's enable_thinking flag. Default off - this " + "sampler is tuned for direct creative-writing continuation, " + "not reasoning mode.") + return p + + +def make_alpha(args): + if args.alpha is not None: + 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) + return lambda: args.alpha, f"fixed={args.alpha:+.2f}" + + s = args.sine + peak = s["offset"] + s["amplitude"] + 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) + desc = (f"sine(period={s['period_tokens']:g}, amplitude={s['amplitude']:.2f}, " + f"offset={s['offset']:.2f}, phase={s['phase']:.2f}) peak={peak:+.2f}") + return (lambda: sine_alpha_schedule(**s)), desc + + +def main(): + args = build_arg_parser().parse_args() + alpha_factory, alpha_desc = make_alpha(args) + + print("=== entropy server ===") + print(f"model: {args.model}") + print(f"model_id: {args.model_id}") + 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}") + print(f"enable_thinking: {args.enable_thinking} (chat template flag)") + print(f"n_ctx: requested={args.n_ctx}") + print(f"n_batch: {args.n_batch} n_ubatch: {args.n_ubatch}") + + template = build_jinja_template(load_chat_template(args.model)) + + sampler = EntropySampler( + model_path=args.model, + n_ctx=args.n_ctx, + top_k=args.top_k, + top_n_future=args.top_n_future, + n_threads=args.n_threads, + think_ban=args.think_ban, + n_batch=args.n_batch, + n_ubatch=args.n_ubatch, + ) + n_ctx_actual = C.lib.llama_n_ctx(sampler.ctx) + n_ctx_seq = C.lib.llama_n_ctx_seq(sampler.ctx) + print(f"n_ctx: actual={n_ctx_actual} per-seq={n_ctx_seq}") + print(f"think_ban ids: {sampler.banned_token_ids}") + print("=" * 40) + + # Each request re-primes seq 0 from scratch with the *entire* rendered chat + # history, so this per-sequence budget has to cover a full growing + # conversation, not just one reply. Undersizing this is exactly what + # produced the mid-generation "llama_decode failed rc=1 / failed to find + # a memory slot" crash this was added after. + if n_ctx_seq < args.max_tokens_default + 512: + print( + f"warning: per-sequence budget ({n_ctx_seq}) is only " + f"{n_ctx_seq - args.max_tokens_default} tokens above max_tokens_default " + f"({args.max_tokens_default}) - long prompts or multi-turn history will " + f"exhaust it. Raise --n-ctx.", + file=sys.stderr, + ) + if args.n_batch < n_ctx_seq: + print( + f"warning: --n-batch ({args.n_batch}) is smaller than the per-sequence " + f"context budget ({n_ctx_seq}) - a prompt/history longer than --n-batch " + f"tokens will hit GGML_ASSERT(n_tokens_all <= n_batch) and abort the whole " + f"process, even though it would otherwise fit in context. Raise --n-batch " + f"to at least {n_ctx_seq}.", + file=sys.stderr, + ) + + gen_lock = threading.Lock() + + def render_prompt(messages): + try: + return template.render( + messages=messages, tools=None, add_generation_prompt=True, + enable_thinking=args.enable_thinking, + ) + except ValueError as e: + raise HTTPException(status_code=400, detail=str(e)) + + app = FastAPI() + + @app.get("/v1/models") + def list_models(): + return { + "object": "list", + "data": [{"id": args.model_id, "object": "model", "owned_by": "local"}], + } + + @app.post("/v1/chat/completions") + def chat_completions(req: ChatCompletionRequest): + messages = [m.model_dump() for m in req.messages] + prompt = render_prompt(messages) + max_new_tokens = req.max_tokens or args.max_tokens_default + + completion_id = f"chatcmpl-{uuid.uuid4().hex}" + created = int(time.time()) + + n_prompt_tokens = len(C.tokenize(sampler.vocab, prompt)) + print( + f"[{completion_id}] request: prompt_tokens={n_prompt_tokens} " + f"requested_max_tokens={req.max_tokens} resolved_max_new_tokens={max_new_tokens} " + f"stream={req.stream}" + ) + + if req.stream: + def event_stream(): + first = { + "id": completion_id, "object": "chat.completion.chunk", "created": created, + "model": args.model_id, + "choices": [{"index": 0, "delta": {"role": "assistant"}, "finish_reason": None}], + } + yield f"data: {json.dumps(first)}\n\n" + + n_yielded = 0 + out_of_context = False + with gen_lock: + try: + for piece, _diag in sampler.generate_stream( + prompt, max_new_tokens=max_new_tokens, + alpha=alpha_factory(), confidence_threshold=args.confidence_threshold, + ): + n_yielded += 1 + chunk = { + "id": completion_id, "object": "chat.completion.chunk", "created": created, + "model": args.model_id, + "choices": [{"index": 0, "delta": {"content": piece}, "finish_reason": None}], + } + yield f"data: {json.dumps(chunk)}\n\n" + except RuntimeError as e: + out_of_context = True + print(f"warning: generation truncated after {n_yielded} tokens: {e} " + f"(per-sequence budget exhausted - raise --n-ctx)", file=sys.stderr) + + finish_reason = "length" if (out_of_context or n_yielded >= max_new_tokens) else "stop" + print(f"[{completion_id}] done: tokens_generated={n_yielded} finish_reason={finish_reason} " + f"out_of_context={out_of_context}") + final = { + "id": completion_id, "object": "chat.completion.chunk", "created": created, + "model": args.model_id, + "choices": [{"index": 0, "delta": {}, "finish_reason": finish_reason}], + } + yield f"data: {json.dumps(final)}\n\n" + yield "data: [DONE]\n\n" + + return StreamingResponse(event_stream(), media_type="text/event-stream") + + out_of_context = False + with gen_lock: + pieces = [] + try: + for piece, _diag in sampler.generate_stream( + prompt, max_new_tokens=max_new_tokens, + alpha=alpha_factory(), confidence_threshold=args.confidence_threshold, + ): + pieces.append(piece) + except RuntimeError as e: + out_of_context = True + print(f"warning: generation truncated after {len(pieces)} tokens: {e} " + f"(per-sequence budget exhausted - raise --n-ctx)", file=sys.stderr) + text = "".join(pieces) + finish_reason = "length" if (out_of_context or len(pieces) >= max_new_tokens) else "stop" + print(f"[{completion_id}] done: tokens_generated={len(pieces)} finish_reason={finish_reason} " + f"out_of_context={out_of_context}") + + return { + "id": completion_id, + "object": "chat.completion", + "created": created, + "model": args.model_id, + "choices": [{ + "index": 0, + "message": {"role": "assistant", "content": text}, + "finish_reason": finish_reason, + }], + "usage": { + "prompt_tokens": 0, "completion_tokens": len(pieces), + "total_tokens": len(pieces), + }, + } + + try: + uvicorn.run(app, host=args.host, port=args.port) + finally: + sampler.close() + + +if __name__ == "__main__": + main() diff --git a/fork_peek_test.py b/fork_peek_test.py new file mode 100644 index 0000000..8262846 --- /dev/null +++ b/fork_peek_test.py @@ -0,0 +1,136 @@ +"""Smoke test for the primitive the entropy sampler depends on: +fork the KV cache at the current position, decode one candidate token +into the fork, read back its resulting next-token distribution, then +discard the fork. If this works, the scoring loop (s(w) = p(w|c)^a * H(w)^b) +is just bookkeeping on top of this. +""" +import math +import sys + +import numpy as np + +import llama_capi as C +from entropy_sampler import MODEL_PATH + +PROMPT = "The old lighthouse keeper had seen many storms, but none like" +TOP_K_BASE = 8 # candidates to show from the base distribution +TOP_N_FUTURE = 20 # tokens used to compute the future entropy of a fork + + +def softmax(x): + x = x - np.max(x) + e = np.exp(x) + return e / e.sum() + + +def normalized_entropy(probs_top_n): + p = probs_top_n / probs_top_n.sum() + p = p[p > 0] + h = -np.sum(p * np.log(p)) + return h / math.log(len(probs_top_n)) + + +def decode_single_seq_batch(ctx, tokens, start_pos, seq_id, want_logits_last_only=True): + n = len(tokens) + batch = C.lib.llama_batch_init(n, 0, 1) + batch.n_tokens = n + for i, tok in enumerate(tokens): + batch.token[i] = tok + batch.pos[i] = start_pos + i + batch.n_seq_id[i] = 1 + batch.seq_id[i][0] = seq_id + batch.logits[i] = 0 + if want_logits_last_only: + batch.logits[n - 1] = 1 + else: + for i in range(n): + batch.logits[i] = 1 + rc = C.lib.llama_decode(ctx, batch) + C.lib.llama_batch_free(batch) + if rc != 0: + raise RuntimeError(f"llama_decode failed rc={rc}") + + +def main(): + if not MODEL_PATH: + sys.exit("Set ENTROPY_SAMPLER_MODEL to a .gguf file path first.") + print(f"loading model: {MODEL_PATH}") + C.lib.llama_backend_init() + + mparams = C.lib.llama_model_default_params() + mparams.n_gpu_layers = 99 + + model = C.lib.llama_model_load_from_file(MODEL_PATH.encode(), mparams) + if not model: + sys.exit("model load failed") + + vocab = C.lib.llama_model_get_vocab(model) + n_vocab = C.lib.llama_vocab_n_tokens(vocab) + print(f"n_vocab = {n_vocab}") + + cparams = C.lib.llama_context_default_params() + cparams.n_ctx = 4096 + cparams.n_seq_max = 4 # seq 0 = real context, seq 1..3 available for forks + cparams.n_threads = 8 + cparams.n_threads_batch = 8 + + ctx = C.lib.llama_init_from_model(model, cparams) + if not ctx: + sys.exit("context init failed") + + mem = C.lib.llama_get_memory(ctx) + + prompt_tokens = C.tokenize(vocab, PROMPT) + print(f"prompt: {PROMPT!r} -> {len(prompt_tokens)} tokens") + + # --- base pass: decode the prompt on seq 0, get p(w | c) --- + decode_single_seq_batch(ctx, prompt_tokens, start_pos=0, seq_id=0) + n_ctx_pos = len(prompt_tokens) # next free position on seq 0 + + logits_ptr = C.lib.llama_get_logits_ith(ctx, -1) + base_logits = np.ctypeslib.as_array(logits_ptr, shape=(n_vocab,)).copy() + base_probs = softmax(base_logits) + + top_idx = np.argsort(-base_probs)[:TOP_K_BASE] + print("\ntop base candidates p(w | c):") + for i in top_idx: + print(f" {base_probs[i]:.4f} {C.token_to_str(vocab, int(i))!r}") + + # --- fork + peek + discard for the single most likely candidate --- + w = int(top_idx[0]) + w_str = C.token_to_str(vocab, w) + fork_seq = 1 + + C.lib.llama_memory_seq_cp(mem, 0, fork_seq, -1, -1) # fork seq 0 -> seq 1 + decode_single_seq_batch(ctx, [w], start_pos=n_ctx_pos, seq_id=fork_seq) + + fork_logits_ptr = C.lib.llama_get_logits_ith(ctx, -1) + fork_logits = np.ctypeslib.as_array(fork_logits_ptr, shape=(n_vocab,)).copy() + fork_probs = softmax(fork_logits) + + top_future_idx = np.argsort(-fork_probs)[:TOP_N_FUTURE] + h_hat = normalized_entropy(fork_probs[top_future_idx]) + + C.lib.llama_memory_seq_rm(mem, fork_seq, -1, -1) # discard the fork + + print(f"\nforked on candidate {w_str!r} (p={base_probs[w]:.4f})") + print(f"normalized future entropy H_hat(w) over top-{TOP_N_FUTURE}: {h_hat:.4f}") + print("top continuations after that candidate:") + for i in top_future_idx[:8]: + print(f" {fork_probs[i]:.4f} {C.token_to_str(vocab, int(i))!r}") + + # sanity check: seq 0 should be untouched by the fork - decoding the + # same candidate directly on seq 0 must reproduce identical logits + decode_single_seq_batch(ctx, [w], start_pos=n_ctx_pos, seq_id=0) + check_ptr = C.lib.llama_get_logits_ith(ctx, -1) + check_logits = np.ctypeslib.as_array(check_ptr, shape=(n_vocab,)).copy() + max_diff = float(np.max(np.abs(check_logits - fork_logits))) + print(f"\nseq0-vs-fork logit max abs diff (expect ~0): {max_diff:.6f}") + + C.lib.llama_free(ctx) + C.lib.llama_model_free(model) + print("\nOK: fork -> peek -> discard round-trip works.") + + +if __name__ == "__main__": + main() diff --git a/llama_capi.py b/llama_capi.py new file mode 100644 index 0000000..3dd738c --- /dev/null +++ b/llama_capi.py @@ -0,0 +1,191 @@ +"""Minimal ctypes binding over libllama.so - just the calls the entropy sampler needs. + +Built against the llama.cpp checkout at commit d132f22fc (2026-04-09). The struct +layouts below (llama_model_params, llama_context_params, llama_batch) are copied +by hand from that checkout's llama.h and must match the exact libllama.so you point +this at - a different llama.cpp version can silently reorder or resize fields and +corrupt memory rather than raising a clean error. If you hit crashes or garbage +output, check llama.h from your build against the struct definitions here first. + +Not a general-purpose binding - only exposes what fork_peek_test.py / the sampler use. +""" +import ctypes +import os +from ctypes import ( + c_bool, c_char_p, c_float, c_int, c_int8, c_int32, c_size_t, c_uint32, + c_void_p, POINTER, Structure, byref, +) + +LIB_PATH = os.environ.get("LLAMA_CPP_LIB") +if not LIB_PATH: + raise RuntimeError( + "Set LLAMA_CPP_LIB to the path of your llama.cpp build's shared library, " + "e.g.:\n export LLAMA_CPP_LIB=~/llama.cpp/build/bin/libllama.so\n" + "Build it with llama.cpp's CMake build (-DBUILD_SHARED_LIBS=ON) if you " + "don't have it yet." + ) + +lib = ctypes.CDLL(LIB_PATH) + +llama_token = c_int32 +llama_pos = c_int32 +llama_seq_id = c_int32 + + +class llama_model_params(Structure): + _fields_ = [ + ("devices", c_void_p), + ("tensor_buft_overrides", c_void_p), + ("n_gpu_layers", c_int32), + ("split_mode", c_int), + ("main_gpu", c_int32), + ("tensor_split", c_void_p), + ("progress_callback", c_void_p), + ("progress_callback_user_data", c_void_p), + ("kv_overrides", c_void_p), + ("vocab_only", c_bool), + ("use_mmap", c_bool), + ("use_direct_io", c_bool), + ("use_mlock", c_bool), + ("check_tensors", c_bool), + ("use_extra_bufts", c_bool), + ("no_host", c_bool), + ("no_alloc", c_bool), + ] + + +class llama_context_params(Structure): + _fields_ = [ + ("n_ctx", c_uint32), + ("n_batch", c_uint32), + ("n_ubatch", c_uint32), + ("n_seq_max", c_uint32), + ("n_threads", c_int32), + ("n_threads_batch", c_int32), + ("rope_scaling_type", c_int), + ("pooling_type", c_int), + ("attention_type", c_int), + ("flash_attn_type", c_int), + ("rope_freq_base", c_float), + ("rope_freq_scale", c_float), + ("yarn_ext_factor", c_float), + ("yarn_attn_factor", c_float), + ("yarn_beta_fast", c_float), + ("yarn_beta_slow", c_float), + ("yarn_orig_ctx", c_uint32), + ("defrag_thold", c_float), + ("cb_eval", c_void_p), + ("cb_eval_user_data", c_void_p), + ("type_k", c_int), + ("type_v", c_int), + ("abort_callback", c_void_p), + ("abort_callback_data", c_void_p), + ("embeddings", c_bool), + ("offload_kqv", c_bool), + ("no_perf", c_bool), + ("op_offload", c_bool), + ("swa_full", c_bool), + ("kv_unified", c_bool), + ("samplers", c_void_p), + ("n_samplers", c_size_t), + ] + + +class llama_batch(Structure): + _fields_ = [ + ("n_tokens", c_int32), + ("token", POINTER(llama_token)), + ("embd", POINTER(c_float)), + ("pos", POINTER(llama_pos)), + ("n_seq_id", POINTER(c_int32)), + ("seq_id", POINTER(POINTER(llama_seq_id))), + ("logits", POINTER(c_int8)), + ] + + +lib.llama_backend_init.argtypes = [] +lib.llama_backend_init.restype = None + +lib.llama_model_default_params.argtypes = [] +lib.llama_model_default_params.restype = llama_model_params + +lib.llama_context_default_params.argtypes = [] +lib.llama_context_default_params.restype = llama_context_params + +lib.llama_model_load_from_file.argtypes = [c_char_p, llama_model_params] +lib.llama_model_load_from_file.restype = c_void_p + +lib.llama_model_free.argtypes = [c_void_p] +lib.llama_model_free.restype = None + +lib.llama_model_get_vocab.argtypes = [c_void_p] +lib.llama_model_get_vocab.restype = c_void_p + +lib.llama_init_from_model.argtypes = [c_void_p, llama_context_params] +lib.llama_init_from_model.restype = c_void_p + +lib.llama_free.argtypes = [c_void_p] +lib.llama_free.restype = None + +lib.llama_vocab_n_tokens.argtypes = [c_void_p] +lib.llama_vocab_n_tokens.restype = c_int32 + +lib.llama_tokenize.argtypes = [ + c_void_p, c_char_p, c_int32, POINTER(llama_token), c_int32, c_bool, c_bool, +] +lib.llama_tokenize.restype = c_int32 + +lib.llama_token_to_piece.argtypes = [ + c_void_p, llama_token, ctypes.c_char_p, c_int32, c_int32, c_bool, +] +lib.llama_token_to_piece.restype = c_int32 + +lib.llama_batch_init.argtypes = [c_int32, c_int32, c_int32] +lib.llama_batch_init.restype = llama_batch + +lib.llama_batch_free.argtypes = [llama_batch] +lib.llama_batch_free.restype = None + +lib.llama_decode.argtypes = [c_void_p, llama_batch] +lib.llama_decode.restype = c_int32 + +lib.llama_get_logits_ith.argtypes = [c_void_p, c_int32] +lib.llama_get_logits_ith.restype = POINTER(c_float) + +lib.llama_get_memory.argtypes = [c_void_p] +lib.llama_get_memory.restype = c_void_p + +lib.llama_memory_seq_cp.argtypes = [c_void_p, llama_seq_id, llama_seq_id, llama_pos, llama_pos] +lib.llama_memory_seq_cp.restype = None + +lib.llama_memory_seq_rm.argtypes = [c_void_p, llama_seq_id, llama_pos, llama_pos] +lib.llama_memory_seq_rm.restype = c_bool + +lib.llama_vocab_is_eog.argtypes = [c_void_p, llama_token] +lib.llama_vocab_is_eog.restype = c_bool + +lib.llama_n_ctx.argtypes = [c_void_p] +lib.llama_n_ctx.restype = c_uint32 + +lib.llama_n_ctx_seq.argtypes = [c_void_p] +lib.llama_n_ctx_seq.restype = c_uint32 + + +def tokenize(vocab, text, add_special=True, parse_special=True, max_tokens=8192): + buf = (llama_token * max_tokens)() + n = lib.llama_tokenize( + vocab, text.encode("utf-8"), len(text.encode("utf-8")), + buf, max_tokens, add_special, parse_special, + ) + if n < 0: + raise RuntimeError(f"tokenize buffer too small, need {-n}") + return list(buf[:n]) + + +def token_to_str(vocab, token): + buf = ctypes.create_string_buffer(64) + n = lib.llama_token_to_piece(vocab, token, buf, 64, 0, True) + if n < 0: + buf = ctypes.create_string_buffer(-n) + n = lib.llama_token_to_piece(vocab, token, buf, -n, 0, True) + return buf.raw[:n].decode("utf-8", errors="replace") diff --git a/requirements.txt b/requirements.txt new file mode 100644 index 0000000..dae10c3 --- /dev/null +++ b/requirements.txt @@ -0,0 +1,5 @@ +numpy>=2.0 +gguf>=0.16 +jinja2>=3.1 +fastapi>=0.110 +uvicorn>=0.30 diff --git a/robustness_test.py b/robustness_test.py new file mode 100644 index 0000000..415e0fa --- /dev/null +++ b/robustness_test.py @@ -0,0 +1,125 @@ +"""Repeated-trial robustness check: single samples are noisy, so run each +config N times with independent seeds and measure a crash rate, not just +read one lucky/unlucky sample. + +Flags a trial as "degenerate" if it exhibits the failure modes seen in +tune_alpha.py's fixed-alpha sweep: blank-template exploitation (runs of +underscores) or markdown-bold meta-commentary artifacts. +""" +import re + +import numpy as np + +from entropy_sampler import EntropySampler, sine_alpha_schedule +from tune_alpha import PROMPT, MAX_NEW_TOKENS, summarize + +N_TRIALS = 5 + +CONFIGS = [ + ("sine p16 a0.6 off-0.2", lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.6, offset=-0.2)), + ("sine p16 a0.7 off-0.3", lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.7, offset=-0.3)), + ("sine p12 a0.8 off-0.2", lambda: sine_alpha_schedule(period_tokens=12, amplitude=0.8, offset=-0.2)), + ("sine p16 a0.6 off 0.0", lambda: sine_alpha_schedule(period_tokens=16, amplitude=0.6, offset=0.0)), + ("fixed alpha = 0.0", lambda: 0.0), + ("fixed alpha = +0.3", lambda: 0.3), +] + + +# Patterns seen in manual reads that the old blank/markdown-only check missed: +# leaks, reading-comprehension/quiz framing, and reasoning scaffolding. +# Topic hijacks and POV breaks (also observed manually) aren't reliably +# catchable by keyword matching without false-flagging legitimate creative +# swerves - those remain a manual-read gap, not something this function claims +# to cover. +_META_PATTERNS = [ + r"= 3 + markdown_bold = text.count("**") >= 2 + return ( + blank_run + or markdown_bold + or _has_repetition_loop(text) + or _has_meta_commentary(text) + ) + + +def run_trial(sampler, alpha_factory, 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 + text = sampler.generate( + PROMPT, max_new_tokens=MAX_NEW_TOKENS, alpha=alpha_factory(), verbose=False + ) + sampler.step = orig_step + stats = summarize(text, log) + stats["degenerate"] = is_degenerate(text) + stats["text"] = text + return stats + + +if __name__ == "__main__": + sampler = EntropySampler(top_k=12, top_n_future=20, n_ctx=4096) + try: + results = {} + for label, alpha_factory in CONFIGS: + trials = [] + for seed in range(N_TRIALS): + trials.append(run_trial(sampler, alpha_factory, seed=100 + seed)) + results[label] = trials + + print(f"=== {label} ===") + for i, t in enumerate(trials): + flag = "DEGENERATE" if t["degenerate"] else "ok" + print( + f" seed={100+i} [{flag:10s}] mean_p={t['mean_p']:.3f} " + f"mean_h={t['mean_h']:.3f}" + ) + print(f" {t['text'][:160]!r}") + crash_rate = sum(t["degenerate"] for t in trials) / len(trials) + mean_p = sum(t["mean_p"] for t in trials) / len(trials) + mean_h = sum(t["mean_h"] for t in trials) / len(trials) + print(f" -> crash_rate={crash_rate:.0%} avg_mean_p={mean_p:.3f} avg_mean_h={mean_h:.3f}\n") + + print("=== summary ===") + for label, trials in results.items(): + crash_rate = sum(t["degenerate"] for t in trials) / len(trials) + mean_p = sum(t["mean_p"] for t in trials) / len(trials) + mean_h = sum(t["mean_h"] for t in trials) / len(trials) + print(f"{label:24s} crash_rate={crash_rate:.0%} avg_mean_p={mean_p:.3f} avg_mean_h={mean_h:.3f}") + finally: + sampler.close() diff --git a/timing_bench.py b/timing_bench.py new file mode 100644 index 0000000..1a68453 --- /dev/null +++ b/timing_bench.py @@ -0,0 +1,62 @@ +"""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() diff --git a/tune_alpha.py b/tune_alpha.py new file mode 100644 index 0000000..ee907d1 --- /dev/null +++ b/tune_alpha.py @@ -0,0 +1,61 @@ +"""Sweep alpha (and a couple of tamed sine schedules) to find the sweet spot +between the pure-probability baseline (generic but coherent) and the +pure-entropy extreme (degenerates into blank-template exploitation). +""" +from entropy_sampler import EntropySampler, sine_alpha_schedule + +PROMPT = "The old lighthouse keeper had seen many storms, but none like" +MAX_NEW_TOKENS = 60 + +FIXED_ALPHAS = [-1.0, -0.5, 0.0, 0.3, 0.5, 0.7, 1.0] +SINE_CONFIGS = [ + {"period_tokens": 10, "amplitude": 0.4}, + {"period_tokens": 16, "amplitude": 0.6}, +] + + +def summarize(text, diagnostics_log): + ps = [d["candidates"][d["chosen"]]["p"] for d in diagnostics_log] + hs = [d["candidates"][d["chosen"]]["h_hat"] for d in diagnostics_log] + degenerate = sum(1 for d in diagnostics_log if "_" in d["candidates"][d["chosen"]]["token"]) + return { + "mean_p": sum(ps) / len(ps), + "mean_h": sum(hs) / len(hs), + "degenerate_tokens": degenerate, + } + + +def run_labeled(sampler, label, alpha): + 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 + text = sampler.generate(PROMPT, max_new_tokens=MAX_NEW_TOKENS, alpha=alpha, verbose=False) + sampler.step = orig_step + + stats = summarize(text, log) + print(f"=== {label} ===") + print(text) + print( + f"[mean p={stats['mean_p']:.3f} mean H_hat={stats['mean_h']:.3f} " + f"degenerate(_)={stats['degenerate_tokens']}/{MAX_NEW_TOKENS}]\n" + ) + + +if __name__ == "__main__": + sampler = EntropySampler(top_k=12, top_n_future=20, n_ctx=4096) + try: + for a in FIXED_ALPHAS: + run_labeled(sampler, f"alpha = {a:+.1f}", a) + + for cfg in SINE_CONFIGS: + schedule = sine_alpha_schedule(**cfg) + label = f"sine period={cfg['period_tokens']} amp={cfg['amplitude']}" + run_labeled(sampler, label, schedule) + finally: + sampler.close() diff --git a/tune_offset_sine.py b/tune_offset_sine.py new file mode 100644 index 0000000..283149c --- /dev/null +++ b/tune_offset_sine.py @@ -0,0 +1,27 @@ +"""Offset-sine sweep: dip the wave deeper into safe/coherent territory than +it peaks into entropy-chasing territory, to see if it keeps the occasional +surprise from tune_alpha.py's best run (period=16 amp=0.6) without ever +grazing the +0.3-and-above coherence cliff found there. +""" +from entropy_sampler import EntropySampler, sine_alpha_schedule +from tune_alpha import PROMPT, MAX_NEW_TOKENS, run_labeled + +CONFIGS = [ + {"period_tokens": 16, "amplitude": 0.6, "offset": -0.2}, # peak +0.4 / trough -0.8 + {"period_tokens": 16, "amplitude": 0.7, "offset": -0.3}, # peak +0.4 / trough -1.0 + {"period_tokens": 12, "amplitude": 0.8, "offset": -0.2}, # shorter period, same bias + {"period_tokens": 16, "amplitude": 0.6, "offset": 0.0}, # unbiased baseline for comparison +] + +if __name__ == "__main__": + sampler = EntropySampler(top_k=12, top_n_future=20, n_ctx=4096) + try: + for cfg in CONFIGS: + schedule = sine_alpha_schedule(**cfg) + label = ( + f"sine period={cfg['period_tokens']} amp={cfg['amplitude']} " + f"offset={cfg['offset']}" + ) + run_labeled(sampler, label, schedule) + finally: + sampler.close()