entropy-sampler/llama_capi.py

"""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")