Joey Grasty
8c844c08e5
llama_capi.py previously hand-copied struct field order/types from one specific llama.cpp commit's llama.h into Python ctypes Structures. A different llama.cpp build could silently reorder or resize those fields and corrupt memory rather than raising any error. Replaced with a cffi "API mode" extension (build_capi.py) that #includes the user's actual llama.h and links against their actual libllama.so. Struct layout now comes from real compilation - a genuinely incompatible field fails the build loudly instead of corrupting memory at runtime. Verified byte-for-byte identical generation output against the prior ctypes implementation at a fixed seed, plus the fork/peek logit round-trip check (0.000000 max diff). Requires a one-time `python build_capi.py` setup step (needs a C compiler, which building llama.cpp itself already requires).
293 lines
11 KiB
Python
293 lines
11 KiB
Python
"""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 <think>/</think>
|
|
# 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 ("<think>", "</think>"):
|
|
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 C.as_float_array(ptr, self.n_vocab)
|
|
|
|
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()
|