from __future__ import annotations

import os

import ctypes

import pathlib

from typing import (

Callable,

Union,

NewType,

Optional,

TYPE_CHECKING,

)

from llama_cpp._ctypes_extensions import (

load_shared_library,

byref,

ctypes_function_for_shared_library,

)

if TYPE_CHECKING:

from llama_cpp._ctypes_extensions import (

CtypesCData,

CtypesArray,

CtypesPointer,

CtypesVoidPointer,

CtypesRef,

CtypesPointerOrRef,

CtypesFuncPointer,

)

# Specify the base name of the shared library to load

_lib_base_name = "llama"

_override_base_path = os.environ.get("LLAMA_CPP_LIB_PATH")

_base_path = pathlib.Path(os.path.abspath(os.path.dirname(__file__))) / "lib" if _override_base_path is None else pathlib.Path(_override_base_path)

# Load the library

_lib = load_shared_library(_lib_base_name, _base_path)

ctypes_function = ctypes_function_for_shared_library(_lib)

# from ggml.h

# // NOTE: always add types at the end of the enum to keep backward compatibility

# enum ggml_type {

# GGML_TYPE_F32 = 0,

# GGML_TYPE_F16 = 1,

# GGML_TYPE_Q4_0 = 2,

# GGML_TYPE_Q4_1 = 3,

# // GGML_TYPE_Q4_2 = 4, support has been removed

# // GGML_TYPE_Q4_3 = 5, support has been removed

# GGML_TYPE_Q5_0 = 6,

# GGML_TYPE_Q5_1 = 7,

# GGML_TYPE_Q8_0 = 8,

# GGML_TYPE_Q8_1 = 9,

# GGML_TYPE_Q2_K = 10,

# GGML_TYPE_Q3_K = 11,

# GGML_TYPE_Q4_K = 12,

# GGML_TYPE_Q5_K = 13,

# GGML_TYPE_Q6_K = 14,

# GGML_TYPE_Q8_K = 15,

# GGML_TYPE_IQ2_XXS = 16,

# GGML_TYPE_IQ2_XS = 17,

# GGML_TYPE_IQ3_XXS = 18,

# GGML_TYPE_IQ1_S = 19,

# GGML_TYPE_IQ4_NL = 20,

# GGML_TYPE_IQ3_S = 21,

# GGML_TYPE_IQ2_S = 22,

# GGML_TYPE_IQ4_XS = 23,

# GGML_TYPE_I8 = 24,

# GGML_TYPE_I16 = 25,

# GGML_TYPE_I32 = 26,

# GGML_TYPE_I64 = 27,

# GGML_TYPE_F64 = 28,

# GGML_TYPE_IQ1_M = 29,

# GGML_TYPE_COUNT,

# };

GGML_TYPE_F32 = 0

GGML_TYPE_F16 = 1

GGML_TYPE_Q4_0 = 2

GGML_TYPE_Q4_1 = 3

GGML_TYPE_Q5_0 = 6

GGML_TYPE_Q5_1 = 7

GGML_TYPE_Q8_0 = 8

GGML_TYPE_Q8_1 = 9

GGML_TYPE_Q2_K = 10

GGML_TYPE_Q3_K = 11

GGML_TYPE_Q4_K = 12

GGML_TYPE_Q5_K = 13

GGML_TYPE_Q6_K = 14

GGML_TYPE_Q8_K = 15

GGML_TYPE_IQ2_XXS = 16

GGML_TYPE_IQ2_XS = 17

GGML_TYPE_IQ3_XXS = 18

GGML_TYPE_IQ1_S = 19

GGML_TYPE_IQ4_NL = 20

GGML_TYPE_IQ3_S = 21

GGML_TYPE_IQ2_S = 22

GGML_TYPE_IQ4_XS = 23

GGML_TYPE_I8 = 24

GGML_TYPE_I16 = 25

GGML_TYPE_I32 = 26

GGML_TYPE_I64 = 27

GGML_TYPE_F64 = 28

GGML_TYPE_IQ1_M = 29

GGML_TYPE_COUNT = 30

# from ggml-backend.h

# typedef bool (*ggml_backend_sched_eval_callback)(struct ggml_tensor * t, bool ask, void * user_data);

ggml_backend_sched_eval_callback = ctypes.CFUNCTYPE(

ctypes.c_bool, ctypes.c_void_p, ctypes.c_bool, ctypes.c_void_p

)

# // Abort callback

# // If not NULL, called before ggml computation

# // If it returns true, the computation is aborted

# typedef bool (*ggml_abort_callback)(void * data);

ggml_abort_callback = ctypes.CFUNCTYPE(ctypes.c_bool, ctypes.c_void_p)

# llama.h bindings

_lib.llama_max_devices.argtypes = []

_lib.llama_max_devices.restype = ctypes.c_size_t

LLAMA_MAX_DEVICES = _lib.llama_max_devices()

# define LLAMA_DEFAULT_SEED 0xFFFFFFFF

LLAMA_DEFAULT_SEED = 0xFFFFFFFF

# define LLAMA_TOKEN_NULL -1

LLAMA_TOKEN_NULL = -1

# define LLAMA_FILE_MAGIC_GGLA 0x67676c61u // 'ggla'

LLAMA_FILE_MAGIC_GGLA = 0x67676C61

# define LLAMA_FILE_MAGIC_GGSN 0x6767736eu // 'ggsn'

LLAMA_FILE_MAGIC_GGSN = 0x6767736E

# define LLAMA_FILE_MAGIC_GGSQ 0x67677371u // 'ggsq'

LLAMA_FILE_MAGIC_GGSQ = 0x67677371

# define LLAMA_SESSION_MAGIC LLAMA_FILE_MAGIC_GGSN

LLAMA_SESSION_MAGIC = LLAMA_FILE_MAGIC_GGSN

# define LLAMA_SESSION_VERSION 9

LLAMA_SESSION_VERSION = 9

# define LLAMA_STATE_SEQ_MAGIC LLAMA_FILE_MAGIC_GGSQ

LLAMA_STATE_SEQ_MAGIC = LLAMA_FILE_MAGIC_GGSQ

# define LLAMA_STATE_SEQ_VERSION 2

LLAMA_STATE_SEQ_VERSION = 2

# struct llama_model;

llama_model_p = NewType("llama_model_p", int)

llama_model_p_ctypes = ctypes.c_void_p

# struct llama_context;

llama_context_p = NewType("llama_context_p", int)

llama_context_p_ctypes = ctypes.c_void_p

# # struct llama_sampler;

# llama_sampler_p = NewType("llama_sampler_p", int)

# llama_sampler_p_ctypes = ctypes.c_void_p

# typedef int32_t llama_pos;

llama_pos = ctypes.c_int32

# typedef int32_t llama_token;

llama_token = ctypes.c_int32

llama_token_p = ctypes.POINTER(llama_token)

# typedef int32_t llama_seq_id;

llama_seq_id = ctypes.c_int32

# enum llama_vocab_type {

# LLAMA_VOCAB_TYPE_NONE = 0, // For models without vocab

# LLAMA_VOCAB_TYPE_SPM = 1, // LLaMA tokenizer based on byte-level BPE with byte fallback

# LLAMA_VOCAB_TYPE_BPE = 2, // GPT-2 tokenizer based on byte-level BPE

# LLAMA_VOCAB_TYPE_WPM = 3, // BERT tokenizer based on WordPiece

# LLAMA_VOCAB_TYPE_UGM = 4, // T5 tokenizer based on Unigram

# LLAMA_VOCAB_TYPE_RWKV = 5, // RWKV tokenizer based on greedy tokenization

# };

LLAMA_VOCAB_TYPE_NONE = 0

"""For models without vocab"""

LLAMA_VOCAB_TYPE_SPM = 1

"""LLaMA tokenizer based on byte-level BPE with byte fallback"""

LLAMA_VOCAB_TYPE_BPE = 2

"""GPT-2 tokenizer based on byte-level BPE"""

LLAMA_VOCAB_TYPE_WPM = 3

"""BERT tokenizer based on WordPiece"""

LLAMA_VOCAB_TYPE_UGM = 4

"""T5 tokenizer based on Unigram"""

LLAMA_VOCAB_TYPE_RWKV = 5

"""RWKV tokenizer based on greedy tokenization"""

# // pre-tokenization types

# enum llama_vocab_pre_type {

# LLAMA_VOCAB_PRE_TYPE_DEFAULT = 0,

# LLAMA_VOCAB_PRE_TYPE_LLAMA3 = 1,

# LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_LLM = 2,

# LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_CODER = 3,

# LLAMA_VOCAB_PRE_TYPE_FALCON = 4,

# LLAMA_VOCAB_PRE_TYPE_MPT = 5,

# LLAMA_VOCAB_PRE_TYPE_STARCODER = 6,

# LLAMA_VOCAB_PRE_TYPE_GPT2 = 7,

# LLAMA_VOCAB_PRE_TYPE_REFACT = 8,

# LLAMA_VOCAB_PRE_TYPE_COMMAND_R = 9,

# LLAMA_VOCAB_PRE_TYPE_STABLELM2 = 10,

# LLAMA_VOCAB_PRE_TYPE_QWEN2 = 11,

# LLAMA_VOCAB_PRE_TYPE_OLMO = 12,

# LLAMA_VOCAB_PRE_TYPE_DBRX = 13,

# LLAMA_VOCAB_PRE_TYPE_SMAUG = 14,

# LLAMA_VOCAB_PRE_TYPE_PORO = 15,

# LLAMA_VOCAB_PRE_TYPE_CHATGLM3 = 16,

# LLAMA_VOCAB_PRE_TYPE_CHATGLM4 = 17,

# LLAMA_VOCAB_PRE_TYPE_VIKING = 18,

# LLAMA_VOCAB_PRE_TYPE_JAIS = 19,

# LLAMA_VOCAB_PRE_TYPE_TEKKEN = 20,

# LLAMA_VOCAB_PRE_TYPE_SMOLLM = 21,

# LLAMA_VOCAB_PRE_TYPE_CODESHELL = 22,

# LLAMA_VOCAB_PRE_TYPE_BLOOM = 23,

# LLAMA_VOCAB_PRE_TYPE_GPT3_FINNISH = 24,

# LLAMA_VOCAB_PRE_TYPE_EXAONE = 25,

# LLAMA_VOCAB_PRE_TYPE_CHAMELEON = 26,

# LLAMA_VOCAB_PRE_TYPE_MINERVA = 27,

# LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM = 28,

# };

LLAMA_VOCAB_PRE_TYPE_DEFAULT = 0

LLAMA_VOCAB_PRE_TYPE_LLAMA3 = 1

LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_LLM = 2

LLAMA_VOCAB_PRE_TYPE_DEEPSEEK_CODER = 3

LLAMA_VOCAB_PRE_TYPE_FALCON = 4

LLAMA_VOCAB_PRE_TYPE_MPT = 5

LLAMA_VOCAB_PRE_TYPE_STARCODER = 6

LLAMA_VOCAB_PRE_TYPE_GPT2 = 7

LLAMA_VOCAB_PRE_TYPE_REFACT = 8

LLAMA_VOCAB_PRE_TYPE_COMMAND_R = 9

LLAMA_VOCAB_PRE_TYPE_STABLELM2 = 10

LLAMA_VOCAB_PRE_TYPE_QWEN2 = 11

LLAMA_VOCAB_PRE_TYPE_OLMO = 12

LLAMA_VOCAB_PRE_TYPE_DBRX = 13

LLAMA_VOCAB_PRE_TYPE_SMAUG = 14

LLAMA_VOCAB_PRE_TYPE_PORO = 15

LLAMA_VOCAV_PRE_TYPE_CHATGLM3 = 16

LLAMA_VOCAB_PRE_TYPE_CHATGLM4 = 17

LLAMA_VOCAB_PRE_TYPE_VIKING = 18

LLAMA_VOCAB_PRE_TYPE_JAIS = 19

LLAMA_VOCAB_PRE_TYPE_TEKKEN = 20

LLAMA_VOCAB_PRE_TYPE_SMOLLM = 21

LLAMA_VOCAB_PRE_TYPE_CODESHELL = 22

LLAMA_VOCAB_PRE_TYPE_BLOOM = 23

LLAMA_VOCAB_PRE_TYPE_GPT3_FINNISH = 24

LLAMA_VOCAB_PRE_TYPE_EXAONE = 25

LLAMA_VOCAB_PRE_TYPE_CHAMELEON = 26

LLAMA_VOCAB_PRE_TYPE_MINERVA = 27

LLAMA_VOCAB_PRE_TYPE_DEEPSEEK3_LLM = 28

# // note: these values should be synchronized with ggml_rope

# // TODO: maybe move this enum to ggml.h (ggml_rope_type)

# enum llama_rope_type {

# LLAMA_ROPE_TYPE_NONE = -1,

# LLAMA_ROPE_TYPE_NORM = 0,

# LLAMA_ROPE_TYPE_NEOX = GGML_ROPE_TYPE_NEOX,

# LLAMA_ROPE_TYPE_MROPE = GGML_ROPE_TYPE_MROPE,

# LLAMA_ROPE_TYPE_VISION = GGML_ROPE_TYPE_VISION,

# };

LLAMA_ROPE_TYPE_NONE = -1

LLAMA_ROPE_TYPE_NORM = 0

LLAMA_ROPE_TYPE_NEOX = GGML_ROPE_TYPE_NEOX = 2

LLAMA_ROPE_TYPE_MROPE = GGML_ROPE_TYPE_MROPE = 8

LLAMA_ROPE_TYPE_VISION = GGML_ROPE_TYPE_VISION = 24

# enum llama_token_type { //TODO: remove, required until per token attributes are available from GGUF file

# LLAMA_TOKEN_TYPE_UNDEFINED = 0,

# LLAMA_TOKEN_TYPE_NORMAL = 1,

# LLAMA_TOKEN_TYPE_UNKNOWN = 2,

# LLAMA_TOKEN_TYPE_CONTROL = 3,

# LLAMA_TOKEN_TYPE_USER_DEFINED = 4,

# LLAMA_TOKEN_TYPE_UNUSED = 5,

# LLAMA_TOKEN_TYPE_BYTE = 6,

# };

LLAMA_TOKEN_TYPE_UNDEFINED = 0

LLAMA_TOKEN_TYPE_NORMAL = 1

LLAMA_TOKEN_TYPE_UNKNOWN = 2

LLAMA_TOKEN_TYPE_CONTROL = 3

LLAMA_TOKEN_TYPE_USER_DEFINED = 4

LLAMA_TOKEN_TYPE_UNUSED = 5

LLAMA_TOKEN_TYPE_BYTE = 6

# enum llama_token_attr {

# LLAMA_TOKEN_ATTR_UNDEFINED = 0,

# LLAMA_TOKEN_ATTR_UNKNOWN = 1 << 0,

# LLAMA_TOKEN_ATTR_UNUSED = 1 << 1,

# LLAMA_TOKEN_ATTR_NORMAL = 1 << 2,

# LLAMA_TOKEN_ATTR_CONTROL = 1 << 3, // SPECIAL?

# LLAMA_TOKEN_ATTR_USER_DEFINED = 1 << 4,

# LLAMA_TOKEN_ATTR_BYTE = 1 << 5,

# LLAMA_TOKEN_ATTR_NORMALIZED = 1 << 6,

# LLAMA_TOKEN_ATTR_LSTRIP = 1 << 7,

# LLAMA_TOKEN_ATTR_RSTRIP = 1 << 8,

# LLAMA_TOKEN_ATTR_SINGLE_WORD = 1 << 9,

# };

LLAMA_TOKEN_ATTR_UNDEFINED = 0

LLAMA_TOKEN_ATTR_UNKNOWN = 1 << 0

LLAMA_TOKEN_ATTR_UNUSED = 1 << 1

LLAMA_TOKEN_ATTR_NORMAL = 1 << 2

LLAMA_TOKEN_ATTR_CONTROL = 1 << 3

LLAMA_TOKEN_ATTR_USER_DEFINED = 1 << 4

LLAMA_TOKEN_ATTR_BYTE = 1 << 5

LLAMA_TOKEN_ATTR_NORMALIZED = 1 << 6

LLAMA_TOKEN_ATTR_LSTRIP = 1 << 7

LLAMA_TOKEN_ATTR_RSTRIP = 1 << 8

LLAMA_TOKEN_ATTR_SINGLE_WORD = 1 << 9

# // model file types

# enum llama_ftype {

# LLAMA_FTYPE_ALL_F32 = 0,

# LLAMA_FTYPE_MOSTLY_F16 = 1, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q4_0 = 2, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q4_1 = 3, // except 1d tensors

# // LLAMA_FTYPE_MOSTLY_Q4_1_SOME_F16 = 4, // tok_embeddings.weight and output.weight are F16

# // LLAMA_FTYPE_MOSTLY_Q4_2 = 5, // support has been removed

# // LLAMA_FTYPE_MOSTLY_Q4_3 = 6, // support has been removed

# LLAMA_FTYPE_MOSTLY_Q8_0 = 7, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q5_0 = 8, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q5_1 = 9, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q2_K = 10, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q3_K_S = 11, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q3_K_M = 12, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q3_K_L = 13, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q4_K_S = 14, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q4_K_M = 15, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q5_K_S = 16, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q5_K_M = 17, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q6_K = 18, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ2_XXS = 19, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ2_XS = 20, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_Q2_K_S = 21, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ3_XS = 22, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ3_XXS = 23, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ1_S = 24, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ4_NL = 25, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ3_S = 26, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ3_M = 27, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ2_S = 28, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ2_M = 29, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ4_XS = 30, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_IQ1_M = 31, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_BF16 = 32, // except 1d tensors

# //LLAMA_FTYPE_MOSTLY_Q4_0_4_4 = 33, // removed from gguf files, use Q4_0 and runtime repack

# //LLAMA_FTYPE_MOSTLY_Q4_0_4_8 = 34, // removed from gguf files, use Q4_0 and runtime repack

# //LLAMA_FTYPE_MOSTLY_Q4_0_8_8 = 35, // removed from gguf files, use Q4_0 and runtime repack

# LLAMA_FTYPE_MOSTLY_TQ1_0 = 36, // except 1d tensors

# LLAMA_FTYPE_MOSTLY_TQ2_0 = 37, // except 1d tensors

#

# LLAMA_FTYPE_GUESSED = 1024, // not specified in the model file

# };

LLAMA_FTYPE_ALL_F32 = 0

LLAMA_FTYPE_MOSTLY_F16 = 1

LLAMA_FTYPE_MOSTLY_Q4_0 = 2

LLAMA_FTYPE_MOSTLY_Q4_1 = 3

LLAMA_FTYPE_MOSTLY_Q8_0 = 7

LLAMA_FTYPE_MOSTLY_Q5_0 = 8

LLAMA_FTYPE_MOSTLY_Q5_1 = 9

LLAMA_FTYPE_MOSTLY_Q2_K = 10

LLAMA_FTYPE_MOSTLY_Q3_K_S = 11

LLAMA_FTYPE_MOSTLY_Q3_K_M = 12

LLAMA_FTYPE_MOSTLY_Q3_K_L = 13

LLAMA_FTYPE_MOSTLY_Q4_K_S = 14

LLAMA_FTYPE_MOSTLY_Q4_K_M = 15

LLAMA_FTYPE_MOSTLY_Q5_K_S = 16

LLAMA_FTYPE_MOSTLY_Q5_K_M = 17

LLAMA_FTYPE_MOSTLY_Q6_K = 18

LLAMA_FTYPE_MOSTLY_IQ2_XXS = 19

LLAMA_FTYPE_MOSTLY_IQ2_XS = 20

LLAMA_FTYPE_MOSTLY_Q2_K_S = 21

LLAMA_FTYPE_MOSTLY_IQ3_XS = 22

LLAMA_FTYPE_MOSTLY_IQ3_XXS = 23

LLAMA_FTYPE_MOSTLY_IQ1_S = 24

LLAMA_FTYPE_MOSTLY_IQ4_NL = 25

LLAMA_FTYPE_MOSTLY_IQ3_S = 26

LLAMA_FTYPE_MOSTLY_IQ3_M = 27

LLAMA_FTYPE_MOSTLY_IQ2_S = 28

LLAMA_FTYPE_MOSTLY_IQ2_M = 29

LLAMA_FTYPE_MOSTLY_IQ4_XS = 30

LLAMA_FTYPE_MOSTLY_IQ1_M = 31

LLAMA_FTYPE_MOSTLY_BF16 = 32

# LLAMA_FTYPE_MOSTLY_Q4_0_4_4 = 33

# LLAMA_FTYPE_MOSTLY_Q4_0_4_8 = 34

# LLAMA_FTYPE_MOSTLY_Q4_0_8_8 = 35

LLAMA_FTYPE_MOSTLY_TQ1_0 = 36

LLAMA_FTYPE_MOSTLY_TQ2_0 = 37

LLAMA_FTYPE_GUESSED = 1024

# enum llama_rope_scaling_type {

# LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED = -1,

# LLAMA_ROPE_SCALING_TYPE_NONE = 0,

# LLAMA_ROPE_SCALING_TYPE_LINEAR = 1,

# LLAMA_ROPE_SCALING_TYPE_YARN = 2,

# LLAMA_ROPE_SCALING_TYPE_LONGROPE = 3,

# LLAMA_ROPE_SCALING_TYPE_MAX_VALUE = LLAMA_ROPE_SCALING_TYPE_YARN,

# };

LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED = -1

LLAMA_ROPE_SCALING_TYPE_NONE = 0

LLAMA_ROPE_SCALING_TYPE_LINEAR = 1

LLAMA_ROPE_SCALING_TYPE_YARN = 2

LLAMA_ROPE_SCALING_TYPE_LONGROPE = 3

LLAMA_ROPE_SCALING_TYPE_MAX_VALUE = LLAMA_ROPE_SCALING_TYPE_YARN

# enum llama_pooling_type {

# LLAMA_POOLING_TYPE_UNSPECIFIED = -1,

# LLAMA_POOLING_TYPE_NONE = 0,

# LLAMA_POOLING_TYPE_MEAN = 1,

# LLAMA_POOLING_TYPE_CLS = 2,

# LLAMA_POOLING_TYPE_LAST = 3,

# LLAMA_POOLING_TYPE_RANK = 4, // used by reranking models to attach the classification head to the graph

# };

LLAMA_POOLING_TYPE_UNSPECIFIED = -1

LLAMA_POOLING_TYPE_NONE = 0

LLAMA_POOLING_TYPE_MEAN = 1

LLAMA_POOLING_TYPE_CLS = 2

LLAMA_POOLING_TYPE_LAST = 3

LLAMA_POOLING_TYPE_RANK = 4

# enum llama_attention_type {

# LLAMA_ATTENTION_TYPE_UNSPECIFIED = -1,

# LLAMA_ATTENTION_TYPE_CAUSAL = 0,

# LLAMA_ATTENTION_TYPE_NON_CAUSAL = 1,

# };

LLAMA_ATTENTION_TYPE_UNSPECIFIED = -1

LLAMA_ATTENTION_TYPE_CAUSAL = 0

LLAMA_ATTENTION_TYPE_NON_CAUSAL = 1

# enum llama_split_mode {

# LLAMA_SPLIT_MODE_NONE = 0, // single GPU

# LLAMA_SPLIT_MODE_LAYER = 1, // split layers and KV across GPUs

# LLAMA_SPLIT_MODE_ROW = 2, // split rows across GPUs

# };

LLAMA_SPLIT_MODE_NONE = 0

LLAMA_SPLIT_MODE_LAYER = 1

LLAMA_SPLIT_MODE_ROW = 2

# typedef struct llama_token_data {

# llama_token id; // token id

# float logit; // log-odds of the token

# float p; // probability of the token

# } llama_token_data;

class llama_token_data(ctypes.Structure):

"""Used to store token data

Attributes:

id (llama_token): token id

logit (float): log-odds of the token

p (float): probability of the token"""

if TYPE_CHECKING:

id: llama_token

logit: float

p: float

_fields_ = [

("id", llama_token),

("logit", ctypes.c_float),

("p", ctypes.c_float),

]

llama_token_data_p = ctypes.POINTER(llama_token_data)

# typedef struct llama_token_data_array {

# // TODO: consider SoA

# // NOTE: this pointer can be modified by the samplers

# llama_token_data * data;

# size_t size;

# int64_t selected; // this is the index in the data array (i.e. not the token id)

# bool sorted;

# } llama_token_data_array;

class llama_token_data_array(ctypes.Structure):

"""Used to sample tokens given logits

Attributes:

data (ctypes.Array[llama_token_data]): token data

size (int): size of the array

selected (int): index in the data array (i.e. not the token id)

sorted (bool): whether the array is sorted"""

if TYPE_CHECKING:

data: CtypesArray[llama_token_data]

size: int

selected: int

sorted: bool

_fields_ = [

("data", llama_token_data_p),

("size", ctypes.c_size_t),

("selected", ctypes.c_int64),

("sorted", ctypes.c_bool),

]

llama_token_data_array_p = ctypes.POINTER(llama_token_data_array)

# typedef bool (*llama_progress_callback)(float progress, void * user_data);

llama_progress_callback = ctypes.CFUNCTYPE(

ctypes.c_bool, ctypes.c_float, ctypes.c_void_p

)

# // Input data for llama_decode

# // A llama_batch object can contain input about one or many sequences

# // The provided arrays (i.e. token, embd, pos, etc.) must have size of n_tokens

# //

# // - token : the token ids of the input (used when embd is NULL)

# // - embd : token embeddings (i.e. float vector of size n_embd) (used when token is NULL)

# // - pos : the positions of the respective token in the sequence

# // (if set to NULL, the token position will be tracked automatically by llama_decode)

# // - seq_id : the sequence to which the respective token belongs

# // (if set to NULL, the sequence ID will be assumed to be 0)

# // - logits : if zero, the logits (and/or the embeddings) for the respective token will not be output

# // (if set to NULL, only the logits for last token will be returned)

# //

# 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; // TODO: rename this to "output"

# } llama_batch;

class llama_batch(ctypes.Structure):

"""Input data for llama_decode

A llama_batch object can contain input about one or many sequences

The provided arrays (i.e. token, embd, pos, etc.) must have size of n_tokens

Attributes:

n_tokens (int): number of tokens

token (ctypes.Array[llama_token]): the token ids of the input (used when embd is NULL)

embd (ctypes.Array[ctypes.ctypes.c_float]): token embeddings (i.e. float vector of size n_embd) (used when token is NULL)

pos (ctypes.Array[ctypes.Array[llama_pos]]): the positions of the respective token in the sequence

seq_id (ctypes.Array[ctypes.Array[llama_seq_id]]): the sequence to which the respective token belongs

logits (ctypes.Array[ctypes.ctypes.c_int8]): if zero, the logits for the respective token will not be output

"""

if TYPE_CHECKING:

n_tokens: int

token: CtypesArray[llama_token]

embd: CtypesArray[ctypes.c_float]

pos: CtypesArray[CtypesArray[llama_pos]]

n_seq_id: CtypesArray[ctypes.c_int]

seq_id: CtypesArray[CtypesArray[llama_seq_id]]

logits: CtypesArray[ctypes.c_int8]

_fields_ = [

("n_tokens", ctypes.c_int32),

("token", ctypes.POINTER(llama_token)),

("embd", ctypes.POINTER(ctypes.c_float)),

("pos", ctypes.POINTER(llama_pos)),

("n_seq_id", ctypes.POINTER(ctypes.c_int32)),

("seq_id", ctypes.POINTER(ctypes.POINTER(llama_seq_id))),

("logits", ctypes.POINTER(ctypes.c_int8)),

]

# enum llama_model_kv_override_type {

# LLAMA_KV_OVERRIDE_TYPE_INT,

# LLAMA_KV_OVERRIDE_TYPE_FLOAT,

# LLAMA_KV_OVERRIDE_TYPE_BOOL,

# LLAMA_KV_OVERRIDE_TYPE_STR,

# };

LLAMA_KV_OVERRIDE_TYPE_INT = 0

LLAMA_KV_OVERRIDE_TYPE_FLOAT = 1

LLAMA_KV_OVERRIDE_TYPE_BOOL = 2

LLAMA_KV_OVERRIDE_TYPE_STR = 3

# struct llama_model_kv_override {

# enum llama_model_kv_override_type tag;

# char key[128];

# union {

# int64_t val_i64;

# double val_f64;

# bool val_bool;

# char val_str[128];

# };

# };

class llama_model_kv_override_value(ctypes.Union):

_fields_ = [

("val_i64", ctypes.c_int64),

("val_f64", ctypes.c_double),

("val_bool", ctypes.c_bool),

("val_str", ctypes.c_char * 128),

]

if TYPE_CHECKING:

val_i64: int

val_f64: float

val_bool: bool

val_str: bytes

class llama_model_kv_override(ctypes.Structure):

_fields_ = [

("tag", ctypes.c_int),

("key", ctypes.c_char * 128),

("value", llama_model_kv_override_value),

]

if TYPE_CHECKING:

tag: int

key: bytes

value: Union[int, float, bool, bytes]

# struct llama_model_params {

# // NULL-terminated list of devices to use for offloading (if NULL, all available devices are used)

# ggml_backend_dev_t * devices;

# int32_t n_gpu_layers; // number of layers to store in VRAM

# enum llama_split_mode split_mode; // how to split the model across multiple GPUs

# // main_gpu interpretation depends on split_mode:

# // LLAMA_SPLIT_MODE_NONE: the GPU that is used for the entire model

# // LLAMA_SPLIT_MODE_ROW: the GPU that is used for small tensors and intermediate results

# // LLAMA_SPLIT_MODE_LAYER: ignored

# int32_t main_gpu;

# // proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices()

# const float * tensor_split;

# // comma separated list of RPC servers to use for offloading

# const char * rpc_servers;

# // Called with a progress value between 0.0 and 1.0. Pass NULL to disable.

# // If the provided progress_callback returns true, model loading continues.

# // If it returns false, model loading is immediately aborted.

# llama_progress_callback progress_callback;

# // context pointer passed to the progress callback

# void * progress_callback_user_data;

# // override key-value pairs of the model meta data

# const struct llama_model_kv_override * kv_overrides;

# // Keep the booleans together to avoid misalignment during copy-by-value.

# bool vocab_only; // only load the vocabulary, no weights

# bool use_mmap; // use mmap if possible

# bool use_mlock; // force system to keep model in RAM

# bool check_tensors; // validate model tensor data

# };

class llama_model_params(ctypes.Structure):

"""Parameters for llama_model

Attributes:

n_gpu_layers (int): number of layers to store in VRAM

split_mode (int): how to split the model across multiple GPUs

main_gpu (int): the GPU that is used for the entire model. main_gpu interpretation depends on split_mode: LLAMA_SPLIT_NONE: the GPU that is used for the entire model LLAMA_SPLIT_ROW: the GPU that is used for small tensors and intermediate results LLAMA_SPLIT_LAYER: ignored

tensor_split (ctypes.Array[ctypes.ctypes.c_float]): proportion of the model (layers or rows) to offload to each GPU, size: llama_max_devices()

rpc_servers (ctypes.c_char_p): comma separated list of RPC servers to use for offloading

progress_callback (llama_progress_callback): called with a progress value between 0.0 and 1.0. Pass NULL to disable. If the provided progress_callback returns true, model loading continues. If it returns false, model loading is immediately aborted.

progress_callback_user_data (ctypes.ctypes.c_void_p): context pointer passed to the progress callback

kv_overrides (ctypes.Array[llama_model_kv_override]): override key-value pairs of the model meta data

vocab_only (bool): only load the vocabulary, no weights

use_mmap (bool): use mmap if possible

use_mlock (bool): force system to keep model in RAM

check_tensors (bool): validate model tensor data"""

if TYPE_CHECKING:

n_gpu_layers: int

split_mode: int

main_gpu: int

tensor_split: CtypesArray[ctypes.c_float]

rpc_servers: ctypes.c_char_p

progress_callback: Callable[[float, ctypes.c_void_p], bool]

progress_callback_user_data: ctypes.c_void_p

kv_overrides: CtypesArray[llama_model_kv_override]

vocab_only: bool

use_mmap: bool

use_mlock: bool

check_tensors: bool

_fields_ = [

("devices", ctypes.c_void_p), # NOTE: unnused

("n_gpu_layers", ctypes.c_int32),

("split_mode", ctypes.c_int),

("main_gpu", ctypes.c_int32),

("tensor_split", ctypes.POINTER(ctypes.c_float)),

("rpc_servers", ctypes.c_char_p),

("progress_callback", llama_progress_callback),

("progress_callback_user_data", ctypes.c_void_p),

("kv_overrides", ctypes.POINTER(llama_model_kv_override)),

("vocab_only", ctypes.c_bool),

("use_mmap", ctypes.c_bool),

("use_mlock", ctypes.c_bool),

("check_tensors", ctypes.c_bool),

]

# // NOTE: changing the default values of parameters marked as [EXPERIMENTAL] may cause crashes or incorrect results in certain configurations

# // https://github.com/ggerganov/llama.cpp/pull/7544

# struct llama_context_params {

# uint32_t n_ctx; // text context, 0 = from model

# uint32_t n_batch; // logical maximum batch size that can be submitted to llama_decode

# uint32_t n_ubatch; // physical maximum batch size

# uint32_t n_seq_max; // max number of sequences (i.e. distinct states for recurrent models)

# int32_t n_threads; // number of threads to use for generation

# int32_t n_threads_batch; // number of threads to use for batch processing

# enum llama_rope_scaling_type rope_scaling_type; // RoPE scaling type, from `enum llama_rope_scaling_type`

# enum llama_pooling_type pooling_type; // whether to pool (sum) embedding results by sequence id

# enum llama_attention_type attention_type; // attention type to use for embeddings

# // ref: https://github.com/ggerganov/llama.cpp/pull/2054

# float rope_freq_base; // RoPE base frequency, 0 = from model

# float rope_freq_scale; // RoPE frequency scaling factor, 0 = from model

# float yarn_ext_factor; // YaRN extrapolation mix factor, negative = from model

# float yarn_attn_factor; // YaRN magnitude scaling factor

# float yarn_beta_fast; // YaRN low correction dim

# float yarn_beta_slow; // YaRN high correction dim

# uint32_t yarn_orig_ctx; // YaRN original context size

# float defrag_thold; // defragment the KV cache if holes/size > thold, < 0 disabled (default)

# ggml_backend_sched_eval_callback cb_eval;

# void * cb_eval_user_data;

# enum ggml_type type_k; // data type for K cache [EXPERIMENTAL]

# enum ggml_type type_v; // data type for V cache [EXPERIMENTAL]

# // Keep the booleans together to avoid misalignment during copy-by-value.

# bool logits_all; // the llama_decode() call computes all logits, not just the last one (DEPRECATED - set llama_batch.logits instead)

# bool embeddings; // if true, extract embeddings (together with logits)

# bool offload_kqv; // whether to offload the KQV ops (including the KV cache) to GPU

# bool flash_attn; // whether to use flash attention [EXPERIMENTAL]

# bool no_perf; // whether to measure performance timings

# // Abort callback

# // if it returns true, execution of llama_decode() will be aborted

# // currently works only with CPU execution

# ggml_abort_callback abort_callback;

# void * abort_callback_data;

# };

class llama_context_params(ctypes.Structure):

"""Parameters for llama_context

Attributes:

n_ctx (int): text context, 0 = from model

n_batch (int): logical maximum batch size that can be submitted to llama_decode

n_ubatch (int): physical maximum batch size

n_seq_max (int): max number of sequences (i.e. distinct states for recurrent models)

n_threads (int): number of threads to use for generation

n_threads_batch (int): number of threads to use for batch processing

rope_scaling_type (int): RoPE scaling type, from `enum llama_rope_scaling_type`

pooling_type (int): whether to pool (sum) embedding results by sequence id (ignored if no pooling layer)

attention_type (int): attention type to use for embeddings

rope_freq_base (float): RoPE base frequency, 0 = from model

rope_freq_scale (float): RoPE frequency scaling factor, 0 = from model

yarn_ext_factor (float): YaRN extrapolation mix factor, negative = from model

yarn_attn_factor (float): YaRN magnitude scaling factor

yarn_beta_fast (float): YaRN low correction dim

yarn_beta_slow (float): YaRN high correction dim

yarn_orig_ctx (int): YaRN original context size

defrag_thold (float): defragment the KV cache if holes/size > thold, < 0 disabled (default)

cb_eval (ggml_backend_sched_eval_callback): callback for scheduling eval

cb_eval_user_data (ctypes.ctypes.c_void_p): user data for cb_eval

type_k (int): data type for K cache

type_v (int): data type for V cache

logits_all (bool): the llama_decode() call computes all logits, not just the last one (DEPRECATED - set llama_batch.logits instead)

embeddings (bool): if true, extract embeddings (together with logits)

offload_kqv (bool): whether to offload the KQV ops (including the KV cache) to GPU

flash_attn (bool): whether to use flash attention

abort_callback (ggml_abort_callback): abort callback if it returns true, execution of llama_decode() will be aborted

abort_callback_data (ctypes.ctypes.c_void_p): data for abort_callback

"""

if TYPE_CHECKING:

n_ctx: int

n_batch: int

n_ubatch: int

n_seq_max: int

n_threads: int

n_threads_batch: int

rope_scaling_type: int

pooling_type: int

attention_type: int

rope_freq_base: float

rope_freq_scale: float

yarn_ext_factor: float

yarn_attn_factor: float

yarn_beta_fast: float

yarn_beta_slow: float

yarn_orig_ctx: int

defrag_thold: float

cb_eval: Callable[[ctypes.c_void_p, bool], bool]

cb_eval_user_data: ctypes.c_void_p

type_k: int

type_v: int

logits_all: bool

embeddings: bool

offload_kqv: bool

flash_attn: bool

abort_callback: Callable[[ctypes.c_void_p], bool]

abort_callback_data: ctypes.c_void_p

_fields_ = [

("n_ctx", ctypes.c_uint32),

("n_batch", ctypes.c_uint32),

("n_ubatch", ctypes.c_uint32),

("n_seq_max", ctypes.c_uint32),

("n_threads", ctypes.c_int32),

("n_threads_batch", ctypes.c_int32),

("rope_scaling_type", ctypes.c_int),

("pooling_type", ctypes.c_int),

("attention_type", ctypes.c_int),

("rope_freq_base", ctypes.c_float),

("rope_freq_scale", ctypes.c_float),

("yarn_ext_factor", ctypes.c_float),

("yarn_attn_factor", ctypes.c_float),

("yarn_beta_fast", ctypes.c_float),

("yarn_beta_slow", ctypes.c_float),

("yarn_orig_ctx", ctypes.c_uint32),

("defrag_thold", ctypes.c_float),

("cb_eval", ggml_backend_sched_eval_callback),

("cb_eval_user_data", ctypes.c_void_p),

("type_k", ctypes.c_int),

("type_v", ctypes.c_int),

("logits_all", ctypes.c_bool),

("embeddings", ctypes.c_bool),

("offload_kqv", ctypes.c_bool),

("flash_attn", ctypes.c_bool),

("abort_callback", ggml_abort_callback),

("abort_callback_data", ctypes.c_void_p),

]

# // Signature for logging events

# // Note that text includes the new line character at the end for most events.

# // If your logging mechanism cannot handle that, check if the last character is '\n' and strip it

# // if it exists.

# // It might not exist for progress report where '.' is output repeatedly.

# typedef void (*llama_log_callback)(enum llama_log_level level, const char * text, void * user_data);

llama_log_callback = ctypes.CFUNCTYPE(

None, ctypes.c_int, ctypes.c_char_p, ctypes.c_void_p

)

"""Signature for logging events

Note that text includes the new line character at the end for most events.

If your logging mechanism cannot handle that, check if the last character is '\n' and strip it

if it exists.

It might not exist for progress report where '.' is output repeatedly."""

# // model quantization parameters

# typedef struct llama_model_quantize_params {

# int32_t nthread; // number of threads to use for quantizing, if <=0 will use std::thread::hardware_concurrency()

# enum llama_ftype ftype; // quantize to this llama_ftype

# enum ggml_type output_tensor_type; // output tensor type

# enum ggml_type token_embedding_type; // token embeddings tensor type

# bool allow_requantize; // allow quantizing non-f32/f16 tensors

# bool quantize_output_tensor; // quantize output.weight

# bool only_copy; // only copy tensors - ftype, allow_requantize and quantize_output_tensor are ignored

# bool pure; // quantize all tensors to the default type

# bool keep_split; // quantize to the same number of shards

# void * imatrix; // pointer to importance matrix data

# void * kv_overrides; // pointer to vector containing overrides

# } llama_model_quantize_params;

class llama_model_quantize_params(ctypes.Structure):

"""Parameters for llama_model_quantize

Attributes:

nthread (int): number of threads to use for quantizing, if <=0 will use std::thread::hardware_concurrency()

ftype (int): quantize to this llama_ftype

output_tensor_type (int): output tensor type

token_embedding_type (int): token embeddings tensor type

allow_requantize (bool): allow quantizing non-f32/f16 tensors

quantize_output_tensor (bool): quantize output.weight

only_copy (bool): only copy tensors - ftype, allow_requantize and quantize_output_tensor are ignored

pure (bool): quantize all tensors to the default type

keep_split (bool): quantize to the same number of shards

imatrix (ctypes.c_void_p): pointer to importance matrix data

kv_overrides (ctypes.c_void_p): pointer to vector containing overrides

"""

if TYPE_CHECKING:

nthread: int

ftype: int

output_tensor_type: int

token_embedding_type: int

allow_requantize: bool

quantize_output_tensor: bool

only_copy: bool

pure: bool

keep_split: bool

imatrix: ctypes.c_void_p

kv_overrides: ctypes.c_void_p

_fields_ = [

("nthread", ctypes.c_int32),

("ftype", ctypes.c_int),

("output_tensor_type", ctypes.c_int),

("token_embedding_type", ctypes.c_int),

("allow_requantize", ctypes.c_bool),

("quantize_output_tensor", ctypes.c_bool),

("only_copy", ctypes.c_bool),

("pure", ctypes.c_bool),

("keep_split", ctypes.c_bool),

("imatrix", ctypes.c_void_p),

("kv_overrides", ctypes.c_void_p),

]

# typedef struct llama_logit_bias {

# llama_token token;

# float bias;

# } llama_logit_bias;

class llama_logit_bias(ctypes.Structure):

"""Used to store logit bias

Attributes:

token (llama_token): token id

bias (float): bias"""

if TYPE_CHECKING:

token: llama_token

bias: float

_fields_ = [

("token", llama_token),

("bias", ctypes.c_float),

]

llama_logit_bias_p = ctypes.POINTER(llama_logit_bias)

# typedef struct llama_sampler_chain_params {

# bool no_perf; // whether to measure performance timings

# } llama_sampler_chain_params;

class llama_sampler_chain_params(ctypes.Structure):

"""Parameters for llama_sampler_chain

Attributes:

no_perf (bool): whether to measure performance timings"""

if TYPE_CHECKING:

no_perf: bool

_fields_ = [

("no_perf", ctypes.c_bool),

]

# // used in chat template

# typedef struct llama_chat_message {

# const char * role;

# const char * content;

# } llama_chat_message;

class llama_chat_message(ctypes.Structure):

_fields_ = [

("role", ctypes.c_char_p),

("content", ctypes.c_char_p),

]

# // lora adapter

# struct llama_lora_adapter;

llama_lora_adapter_p = ctypes.c_void_p

llama_lora_adapter_p_ctypes = ctypes.POINTER(ctypes.c_void_p)

# // Helpers for getting default parameters

# LLAMA_API struct llama_model_params llama_model_default_params(void);

@ctypes_function(

"llama_model_default_params",

[],

llama_model_params,

)

def llama_model_default_params() -> llama_model_params:

"""Get default parameters for llama_model"""

...

# LLAMA_API struct llama_context_params llama_context_default_params(void);

@ctypes_function(

"llama_context_default_params",

[],

llama_context_params,

)