entropy-sampler/README.md

# 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 `<think>`/`</think>` 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).