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).
137 lines
4.6 KiB
Python
137 lines
4.6 KiB
Python
"""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 = C.as_float_array(logits_ptr, n_vocab)
|
|
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 = C.as_float_array(fork_logits_ptr, n_vocab)
|
|
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 = C.as_float_array(check_ptr, n_vocab)
|
|
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()
|