entropy-sampler/build_capi.py

"""One-time build step for the llama.cpp bindings.

Compiles a small C extension (via cffi, "API mode") that #includes your
actual llama.h and links against your actual libllama.so. This is the fix
for the ABI fragility that hand-copied ctypes struct definitions had: rather
than a Python-side guess at field order/padding that can silently produce
wrong offsets against a different llama.cpp build, the real struct layout
here comes from actually compiling against your header - if a field this
project depends on doesn't exist or has a genuinely incompatible type, this
build step fails loudly instead of the sampler corrupting memory at runtime.

Run this once, and again any time you rebuild llama.cpp against a version
that might have changed the structs/functions this project uses.

Requires:
    LLAMA_CPP_LIB      path to libllama.so / libllama.dylib / llama.dll
    LLAMA_CPP_INCLUDE  directory containing llama.h (optional - guessed as
                        ../../include relative to LLAMA_CPP_LIB, matching
                        llama.cpp's normal build-tree layout, if unset)
"""
import os
import sys

from cffi import FFI

LIB_PATH = os.environ.get("LLAMA_CPP_LIB")
if not LIB_PATH:
    sys.exit(
        "Set LLAMA_CPP_LIB to the path of your llama.cpp build's shared library first, "
        "e.g.:\n  export LLAMA_CPP_LIB=~/llama.cpp/build/bin/libllama.so"
    )
if not os.path.exists(LIB_PATH):
    sys.exit(f"LLAMA_CPP_LIB does not exist: {LIB_PATH}")

INCLUDE_DIR = os.environ.get("LLAMA_CPP_INCLUDE")
if not INCLUDE_DIR:
    guess = os.path.normpath(os.path.join(os.path.dirname(LIB_PATH), "..", "..", "include"))
    if os.path.exists(os.path.join(guess, "llama.h")):
        INCLUDE_DIR = guess
    else:
        sys.exit(
            "Couldn't find llama.h next to LLAMA_CPP_LIB (looked in "
            f"{guess}). Set LLAMA_CPP_INCLUDE to the directory containing "
            "llama.h, e.g. <llama.cpp checkout>/include"
        )
elif not os.path.exists(os.path.join(INCLUDE_DIR, "llama.h")):
    sys.exit(f"No llama.h found in LLAMA_CPP_INCLUDE={INCLUDE_DIR}")

# llama.h #includes ggml.h, which lives in a separate include tree
# (<repo>/ggml/include) in llama.cpp's normal source layout, not alongside
# llama.h itself.
GGML_INCLUDE_DIR = os.environ.get("LLAMA_CPP_GGML_INCLUDE")
if not GGML_INCLUDE_DIR:
    guess = os.path.normpath(os.path.join(INCLUDE_DIR, "..", "ggml", "include"))
    if os.path.exists(os.path.join(guess, "ggml.h")):
        GGML_INCLUDE_DIR = guess
    else:
        sys.exit(
            "Couldn't find ggml.h (looked in "
            f"{guess}). Set LLAMA_CPP_GGML_INCLUDE to the directory "
            "containing ggml.h, e.g. <llama.cpp checkout>/ggml/include"
        )
elif not os.path.exists(os.path.join(GGML_INCLUDE_DIR, "ggml.h")):
    sys.exit(f"No ggml.h found in LLAMA_CPP_GGML_INCLUDE={GGML_INCLUDE_DIR}")

LIB_DIR = os.path.dirname(os.path.abspath(LIB_PATH))
LIB_BASENAME = os.path.basename(LIB_PATH)
# libllama.so -> llama (the -l link name)
LINK_NAME = LIB_BASENAME
if LINK_NAME.startswith("lib"):
    LINK_NAME = LINK_NAME[3:]
LINK_NAME = LINK_NAME.split(".")[0]

ffi = FFI()

# Only the fields/functions the sampler actually touches are named below.
# The trailing "...;" in each struct tells cffi to ask the real compiler
# (via the #include in set_source, below) to work out the rest of that
# struct's true layout - fields not listed here are simply invisible to
# Python, they don't need to be declared for the layout to be correct.
ffi.cdef(r"""
    typedef int32_t llama_token;
    typedef int32_t llama_pos;
    typedef int32_t llama_seq_id;

    struct llama_vocab;
    struct llama_model;
    struct llama_context;
    struct llama_memory_i;
    typedef struct llama_memory_i * llama_memory_t;

    struct llama_model_params {
        int32_t n_gpu_layers;
        ...;
    };

    struct llama_context_params {
        uint32_t n_ctx;
        uint32_t n_batch;
        uint32_t n_ubatch;
        uint32_t n_seq_max;
        int32_t  n_threads;
        int32_t  n_threads_batch;
        ...;
    };

    typedef struct llama_batch {
        int32_t n_tokens;
        llama_token  *  token;
        float        *  embd;
        llama_pos    *  pos;
        int32_t      *  n_seq_id;
        llama_seq_id ** seq_id;
        int8_t       *  logits;
        ...;
    } llama_batch;

    void llama_backend_init(void);

    struct llama_model_params llama_model_default_params(void);
    struct llama_context_params llama_context_default_params(void);

    struct llama_model * llama_model_load_from_file(const char * path_model,
                                                      struct llama_model_params params);
    void llama_model_free(struct llama_model * model);
    const struct llama_vocab * llama_model_get_vocab(const struct llama_model * model);

    struct llama_context * llama_init_from_model(struct llama_model * model,
                                                   struct llama_context_params params);
    void llama_free(struct llama_context * ctx);

    uint32_t llama_n_ctx(const struct llama_context * ctx);
    uint32_t llama_n_ctx_seq(const struct llama_context * ctx);

    int32_t llama_vocab_n_tokens(const struct llama_vocab * vocab);
    bool llama_vocab_is_eog(const struct llama_vocab * vocab, llama_token token);

    int32_t llama_tokenize(const struct llama_vocab * vocab, const char * text, int32_t text_len,
                            llama_token * tokens, int32_t n_tokens_max,
                            bool add_special, bool parse_special);
    int32_t llama_token_to_piece(const struct llama_vocab * vocab, llama_token token, char * buf,
                                  int32_t length, int32_t lstrip, bool special);

    struct llama_batch llama_batch_init(int32_t n_tokens, int32_t embd, int32_t n_seq_max);
    void llama_batch_free(struct llama_batch batch);
    int32_t llama_decode(struct llama_context * ctx, struct llama_batch batch);
    float * llama_get_logits_ith(struct llama_context * ctx, int32_t i);

    llama_memory_t llama_get_memory(const struct llama_context * ctx);
    bool llama_memory_seq_rm(llama_memory_t mem, llama_seq_id seq_id, llama_pos p0, llama_pos p1);
    void llama_memory_seq_cp(llama_memory_t mem, llama_seq_id seq_id_src, llama_seq_id seq_id_dst,
                              llama_pos p0, llama_pos p1);
""")

ffi.set_source(
    "_llama_cffi",
    '#include "llama.h"',
    include_dirs=[INCLUDE_DIR, GGML_INCLUDE_DIR],
    library_dirs=[LIB_DIR],
    libraries=[LINK_NAME],
    extra_link_args=[f"-Wl,-rpath,{LIB_DIR}"],
)

if __name__ == "__main__":
    ffi.compile(verbose=True)
    print(f"\nBuilt _llama_cffi against {INCLUDE_DIR}/llama.h "
          f"(+ {GGML_INCLUDE_DIR}/ggml.h) + {LIB_PATH}")