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.
This commit is contained in:
191
llama_capi.py
Normal file
191
llama_capi.py
Normal file
@@ -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")
|
||||
Reference in New Issue
Block a user