"""High-level Python wrapper for a llama.cpp model."""

__backend_initialized = False

def __init__(

self,

model_path: str,

*,

# Model Params

n_gpu_layers: int = 0,

split_mode: int = llama_cpp.LLAMA_SPLIT_MODE_LAYER,

main_gpu: int = 0,

tensor_split: Optional[List[float]] = None,

rpc_servers: Optional[str] = None,

vocab_only: bool = False,

use_mmap: bool = True,

use_mlock: bool = False,

kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None,

# Context Params

seed: int = llama_cpp.LLAMA_DEFAULT_SEED,

n_ctx: int = 512,

n_batch: int = 512,

n_ubatch: int = 512,

n_threads: Optional[int] = None,

n_threads_batch: Optional[int] = None,

rope_scaling_type: Optional[

int

] = llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED,

pooling_type: int = llama_cpp.LLAMA_POOLING_TYPE_UNSPECIFIED,

rope_freq_base: float = 0.0,

rope_freq_scale: float = 0.0,

yarn_ext_factor: float = -1.0,

yarn_attn_factor: float = 1.0,

yarn_beta_fast: float = 32.0,

yarn_beta_slow: float = 1.0,

yarn_orig_ctx: int = 0,

logits_all: bool = False,

embedding: bool = False,

offload_kqv: bool = True,

flash_attn: bool = False,

# Sampling Params

last_n_tokens_size: int = 64,

# LoRA Params

lora_base: Optional[str] = None,

lora_scale: float = 1.0,

lora_path: Optional[str] = None,

# Backend Params

numa: Union[bool, int] = False,

# Chat Format Params

chat_format: Optional[str] = None,

chat_handler: Optional[llama_chat_format.LlamaChatCompletionHandler] = None,

# Speculative Decoding

draft_model: Optional[LlamaDraftModel] = None,

# Tokenizer Override

tokenizer: Optional[BaseLlamaTokenizer] = None,

# KV cache quantization

type_k: Optional[int] = None,

type_v: Optional[int] = None,

# Misc

spm_infill: bool = False,

verbose: bool = True,

# Extra Params

**kwargs, # type: ignore

):

"""Load a llama.cpp model from `model_path`.

self.verbose = verbose

self._stack = contextlib.ExitStack()

set_verbose(verbose)

if not Llama.__backend_initialized:

with suppress_stdout_stderr(disable=verbose):

llama_cpp.llama_backend_init()

Llama.__backend_initialized = True

if isinstance(numa, bool):

self.numa = (

llama_cpp.GGML_NUMA_STRATEGY_DISTRIBUTE

if numa

else llama_cpp.GGML_NUMA_STRATEGY_DISABLED

)

else:

self.numa = numa

if self.numa != llama_cpp.GGML_NUMA_STRATEGY_DISABLED:

with suppress_stdout_stderr(disable=verbose):

llama_cpp.llama_numa_init(self.numa)

self.model_path = model_path

# Model Params

self.model_params = llama_cpp.llama_model_default_params()

self.model_params.n_gpu_layers = (

0x7FFFFFFF if n_gpu_layers == -1 else n_gpu_layers

) # 0x7FFFFFFF is INT32 max, will be auto set to all layers

self.model_params.split_mode = split_mode

self.model_params.main_gpu = main_gpu

if rpc_servers is not None:

self.model_params.rpc_servers = rpc_servers.encode("utf-8")

self._rpc_servers = rpc_servers

else:

self._rpc_servers = None

self.tensor_split = tensor_split

self._c_tensor_split = None

if self.tensor_split is not None:

if len(self.tensor_split) > llama_cpp.LLAMA_MAX_DEVICES:

raise ValueError(

f"Attempt to split tensors that exceed maximum supported devices. Current LLAMA_MAX_DEVICES={llama_cpp.LLAMA_MAX_DEVICES}"

)

# Type conversion and expand the list to the length of LLAMA_MAX_DEVICES

FloatArray = ctypes.c_float * llama_cpp.LLAMA_MAX_DEVICES

self._c_tensor_split = FloatArray(

*tensor_split # type: ignore

) # keep a reference to the array so it is not gc'd

self.model_params.tensor_split = self._c_tensor_split

self.model_params.vocab_only = vocab_only

self.model_params.use_mmap = use_mmap if lora_path is None else False

self.model_params.use_mlock = use_mlock

# kv_overrides is the original python dict

self.kv_overrides = kv_overrides

if kv_overrides is not None:

# _kv_overrides_array is a ctypes.Array of llama_model_kv_override Structs

kvo_array_len = len(kv_overrides) + 1 # for sentinel element

self._kv_overrides_array = (

llama_cpp.llama_model_kv_override * kvo_array_len

)()

for i, (k, v) in enumerate(kv_overrides.items()):

self._kv_overrides_array[i].key = k.encode("utf-8")

if isinstance(v, bool):

self._kv_overrides_array[

i

].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_BOOL

self._kv_overrides_array[i].value.val_bool = v

elif isinstance(v, int):

self._kv_overrides_array[

i

].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_INT

self._kv_overrides_array[i].value.val_i64 = v

elif isinstance(v, float):

self._kv_overrides_array[

i

].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_FLOAT

self._kv_overrides_array[i].value.val_f64 = v

elif isinstance(v, str): # type: ignore

v_bytes = v.encode("utf-8")

if len(v_bytes) > 128: # TODO: Make this a constant

raise ValueError(f"Value for {k} is too long: {v}")

v_bytes = v_bytes.ljust(128, b"\0")

self._kv_overrides_array[

i

].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_STR

# copy min(v_bytes, 128) to str_value

address = typing.cast(

int,

ctypes.addressof(self._kv_overrides_array[i].value)

+ llama_cpp.llama_model_kv_override_value.val_str.offset,

)

buffer_start = ctypes.cast(address, ctypes.POINTER(ctypes.c_char))

ctypes.memmove(

buffer_start,

v_bytes,

128,

)

else:

raise ValueError(f"Unknown value type for {k}: {v}")

self._kv_overrides_array[

-1

].key = b"\0" # ensure sentinel element is zeroed

self.model_params.kv_overrides = self._kv_overrides_array

self.n_batch = min(n_ctx, n_batch) # ???

self.n_threads = n_threads or max(multiprocessing.cpu_count() // 2, 1)

self.n_threads_batch = n_threads_batch or multiprocessing.cpu_count()

# Used by the sampler

self._seed = seed or llama_cpp.LLAMA_DEFAULT_SEED

# Context Params

self.context_params = llama_cpp.llama_context_default_params()

self.context_params.n_ctx = n_ctx

self.context_params.n_batch = self.n_batch

self.context_params.n_ubatch = min(self.n_batch, n_ubatch)

self.context_params.n_threads = self.n_threads

self.context_params.n_threads_batch = self.n_threads_batch

self.context_params.rope_scaling_type = (

rope_scaling_type

if rope_scaling_type is not None

else llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED

)

self.context_params.pooling_type = pooling_type

self.context_params.rope_freq_base = (

rope_freq_base if rope_freq_base != 0.0 else 0

)

self.context_params.rope_freq_scale = (

rope_freq_scale if rope_freq_scale != 0.0 else 0

)

self.context_params.yarn_ext_factor = (

yarn_ext_factor if yarn_ext_factor != 0.0 else 0

)

self.context_params.yarn_attn_factor = (

yarn_attn_factor if yarn_attn_factor != 0.0 else 0

)

self.context_params.yarn_beta_fast = (

yarn_beta_fast if yarn_beta_fast != 0.0 else 0

)

self.context_params.yarn_beta_slow = (

yarn_beta_slow if yarn_beta_slow != 0.0 else 0

)

self.context_params.yarn_orig_ctx = yarn_orig_ctx if yarn_orig_ctx != 0 else 0

self.context_params.logits_all = (

logits_all if draft_model is None else True

) # Must be set to True for speculative decoding

self.context_params.embeddings = embedding # TODO: Rename to embeddings

self.context_params.offload_kqv = offload_kqv

self.context_params.flash_attn = flash_attn

# KV cache quantization

if type_k is not None:

self.context_params.type_k = type_k

if type_v is not None:

self.context_params.type_v = type_v

# Sampling Params

self.last_n_tokens_size = last_n_tokens_size

self.cache: Optional[BaseLlamaCache] = None

self.lora_base = lora_base

self.lora_scale = lora_scale

self.lora_path = lora_path

self.spm_infill = spm_infill

if not os.path.exists(model_path):

raise ValueError(f"Model path does not exist: {model_path}")

self._model = self._stack.enter_context(

contextlib.closing(

internals.LlamaModel(

path_model=self.model_path,

params=self.model_params,

verbose=self.verbose,

)

)

)

# Override tokenizer

self.tokenizer_ = tokenizer or LlamaTokenizer(self)

# Set the default value for the context and correct the batch

if n_ctx == 0:

n_ctx = self._model.n_ctx_train()

self.n_batch = min(n_ctx, n_batch)

self.context_params.n_ctx = self._model.n_ctx_train()

self.context_params.n_batch = self.n_batch

self.context_params.n_ubatch = min(self.n_batch, n_ubatch)

self._ctx = self._stack.enter_context(

contextlib.closing(

internals.LlamaContext(

model=self._model,

params=self.context_params,

verbose=self.verbose,

)

)

)

self._batch = self._stack.enter_context(

contextlib.closing(

internals.LlamaBatch(

n_tokens=self.n_batch,

embd=0,

n_seq_max=self.context_params.n_ctx,

verbose=self.verbose,

)

)

)

self._lora_adapter: Optional[llama_cpp.llama_lora_adapter_p] = None

if self.lora_path:

self._lora_adapter = llama_cpp.llama_lora_adapter_init(

self._model.model,

self.lora_path.encode("utf-8"),

)

if self._lora_adapter is None:

raise RuntimeError(

f"Failed to initialize LoRA adapter from lora path: {self.lora_path}"

)

def free_lora_adapter():

if self._lora_adapter is None:

return

llama_cpp.llama_lora_adapter_free(self._lora_adapter)

self._lora_adapter = None

self._stack.callback(free_lora_adapter)

if llama_cpp.llama_lora_adapter_set(

self._ctx.ctx, self._lora_adapter, self.lora_scale

):

raise RuntimeError(

f"Failed to set LoRA adapter from lora path: {self.lora_path}"

)

if self.verbose:

print(llama_cpp.llama_print_system_info().decode("utf-8"), file=sys.stderr)

self.chat_format = chat_format

self.chat_handler = chat_handler

self._chat_handlers: Dict[

str, llama_chat_format.LlamaChatCompletionHandler

] = {}

self.draft_model = draft_model

self._n_vocab = self.n_vocab()

self._n_ctx = self.n_ctx()

self._token_nl = self.token_nl()

self._token_eos = self.token_eos()

self._candidates = internals.LlamaTokenDataArray(n_vocab=self._n_vocab)

self.n_tokens = 0

self.input_ids: npt.NDArray[np.intc] = np.ndarray((n_ctx,), dtype=np.intc)

self.scores: npt.NDArray[np.single] = np.ndarray(

(n_ctx if logits_all == True else n_batch, self._n_vocab), dtype=np.single

)

self._mirostat_mu = ctypes.c_float(

2.0 * 5.0

) # TODO: Move this to sampling context

try:

self.metadata = self._model.metadata()

except Exception as e:

self.metadata = {}

if self.verbose:

print(f"Failed to load metadata: {e}", file=sys.stderr)

if self.verbose:

print(f"Model metadata: {self.metadata}", file=sys.stderr)

eos_token_id = self.token_eos()

bos_token_id = self.token_bos()

eos_token = (

self._model.token_get_text(eos_token_id) if eos_token_id != -1 else ""

)

bos_token = (

self._model.token_get_text(bos_token_id) if bos_token_id != -1 else ""

)

# Unfortunately the llama.cpp API does not return metadata arrays, so we can't get template names from tokenizer.chat_templates

template_choices = dict(

(name[10:], template)

for name, template in self.metadata.items()

if name.startswith("tokenizer.chat_template.")

)

if "tokenizer.chat_template" in self.metadata:

template_choices["chat_template.default"] = self.metadata[

"tokenizer.chat_template"

]

if self.verbose and template_choices:

print(

f"Available chat formats from metadata: {', '.join(template_choices.keys())}",

file=sys.stderr,

)

for name, template in template_choices.items():

self._chat_handlers[name] = llama_chat_format.Jinja2ChatFormatter(

template=template,

eos_token=eos_token,

bos_token=bos_token,

stop_token_ids=[eos_token_id],

).to_chat_handler()

if (

self.chat_format is None

and self.chat_handler is None

and "chat_template.default" in template_choices

):

chat_format = llama_chat_format.guess_chat_format_from_gguf_metadata(

self.metadata

)

if chat_format is not None:

self.chat_format = chat_format

if self.verbose:

print(f"Guessed chat format: {chat_format}", file=sys.stderr)

else:

if self.verbose:

print(

f"Using gguf chat template: {template_choices['chat_template.default']}",

file=sys.stderr,

)

print(f"Using chat eos_token: {eos_token}", file=sys.stderr)

print(f"Using chat bos_token: {bos_token}", file=sys.stderr)

self.chat_format = "chat_template.default"

if self.chat_format is None and self.chat_handler is None:

self.chat_format = "llama-2"

if self.verbose:

print(

f"Using fallback chat format: {self.chat_format}", file=sys.stderr

)

self._sampler = None

@property

def ctx(self) -> llama_cpp.llama_context_p:

return self._ctx.ctx

@property

def model(self) -> llama_cpp.llama_model_p:

return self._model.model

@property

def _input_ids(self) -> npt.NDArray[np.intc]:

return self.input_ids[: self.n_tokens]

@property

def _scores(self) -> npt.NDArray[np.single]:

return self.scores[: self.n_tokens, :]

@property

def eval_tokens(self) -> Deque[int]:

return deque(self.input_ids[: self.n_tokens].tolist(), maxlen=self._n_ctx)

@property

def eval_logits(self) -> Deque[List[float]]:

return deque(

self.scores[: self.n_tokens, :].tolist(),

maxlen=self._n_ctx if self.context_params.logits_all else 1,

)

def tokenize(

self, text: bytes, add_bos: bool = True, special: bool = False

) -> List[int]:

"""Tokenize a string.

return self.tokenizer_.tokenize(text, add_bos, special)

def detokenize(

self,

tokens: List[int],

prev_tokens: Optional[List[int]] = None,

special: bool = False,

) -> bytes:

"""Detokenize a list of tokens.

return self.tokenizer_.detokenize(

tokens, prev_tokens=prev_tokens, special=special

)

def set_cache(self, cache: Optional[BaseLlamaCache]):

"""Set the cache.

self.cache = cache

def set_seed(self, seed: int):

"""Set the random seed.

self._seed = seed

def reset(self):

"""Reset the model state."""

self.n_tokens = 0

def eval(self, tokens: Sequence[int]):

"""Evaluate a list of tokens.

self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1)

for i in range(0, len(tokens), self.n_batch):

batch = tokens[i : min(len(tokens), i + self.n_batch)]

n_past = self.n_tokens

n_tokens = len(batch)

self._batch.set_batch(

batch=batch, n_past=n_past, logits_all=self.context_params.logits_all

)

self._ctx.decode(self._batch)

# Save tokens

self.input_ids[n_past : n_past + n_tokens] = batch

# Save logits

if self.context_params.logits_all:

rows = n_tokens

cols = self._n_vocab

logits = np.ctypeslib.as_array(

self._ctx.get_logits(), shape=(rows * cols,)

)

self.scores[n_past : n_past + n_tokens, :].reshape(-1)[::] = logits

else:

# rows = 1

# cols = self._n_vocab

# logits = np.ctypeslib.as_array(

# self._ctx.get_logits(), shape=(rows * cols,)

# )

# self.scores[n_past + n_tokens - 1, :].reshape(-1)[::] = logits

# NOTE: Now that sampling is done inside the sampler, logits are only needed for logprobs which requires logits_all

pass

# Update n_tokens

self.n_tokens += n_tokens

def _init_sampler(

self,

top_k: int = 40,

top_p: float = 0.95,

min_p: float = 0.05,

typical_p: float = 1.0,

temp: float = 0.80,

repeat_penalty: float = 1.0,

frequency_penalty: float = 0.0,

presence_penalty: float = 0.0,

tfs_z: float = 1.0,

mirostat_mode: int = 0,

mirostat_eta: float = 0.1,

mirostat_tau: float = 5.0,

penalize_nl: bool = True,

logits_processor: Optional[LogitsProcessorList] = None,

grammar: Optional[LlamaGrammar] = None,

):

sampler = internals.LlamaSampler()

if logits_processor is not None:

# Create and add a custom sampler

def apply_func(token_data_array: llama_cpp.llama_token_data_array_p):

size = token_data_array.contents.size

data_soa = token_data_array.contents.data

data_soa_address = ctypes.addressof(data_soa.contents)

# NOTE: This is probably broken

recarray = np.recarray(

shape=(size,),

dtype=np.dtype(

[("id", np.intc), ("logit", np.single), ("p", np.single)],

align=True,

),

buf=(llama_cpp.llama_token_data * size).from_address(

data_soa_address

),

)

for logit_processor in logits_processor:

recarray.logit[:] = logit_processor(self._input_ids, recarray.logit)

sampler.add_custom(apply_func)

sampler.add_penalties(

n_vocab=self._n_vocab,

special_eos_id=self._token_eos,

linefeed_id=self._token_nl,

penalty_last_n=self.last_n_tokens_size,

penalty_repeat=repeat_penalty,

penalty_freq=frequency_penalty,

penalty_present=presence_penalty,

penalize_nl=penalize_nl,

ignore_eos=False,

)

if grammar is not None:

sampler.add_grammar(self._model, grammar)

if temp < 0.0:

sampler.add_softmax()

sampler.add_dist(self._seed)

elif temp == 0.0:

sampler.add_greedy()

else:

if mirostat_mode == 1:

mirostat_m = 100

sampler.add_mirostat(

self._n_vocab,

self._seed,

mirostat_tau,

mirostat_eta,

mirostat_m,

)

elif mirostat_mode == 2:

sampler.add_mirostat_v2(

self._seed,

mirostat_tau,

mirostat_eta,

)

else:

n_probs = 0

min_keep = max(1, n_probs)

sampler.add_top_k(top_k)

sampler.add_tail_free(tfs_z, min_keep)

sampler.add_typical(typical_p, min_keep)

sampler.add_top_p(top_p, min_keep)

sampler.add_min_p(min_p, min_keep)

sampler.add_temp(temp)

sampler.add_dist(self._seed)

return sampler

def sample(

self,

top_k: int = 40,

top_p: float = 0.95,

min_p: float = 0.05,

typical_p: float = 1.0,

temp: float = 0.80,

repeat_penalty: float = 1.0,

frequency_penalty: float = 0.0,

presence_penalty: float = 0.0,

tfs_z: float = 1.0,

mirostat_mode: int = 0,

mirostat_eta: float = 0.1,

mirostat_tau: float = 5.0,

penalize_nl: bool = True,

logits_processor: Optional[LogitsProcessorList] = None,

grammar: Optional[LlamaGrammar] = None,

idx: Optional[int] = None,

):

"""Sample a token from the model.

assert self.n_tokens > 0

tmp_sampler = False

if self._sampler is None:

tmp_sampler = True

self._sampler = self._init_sampler(

top_k=top_k,

top_p=top_p,

min_p=min_p,

typical_p=typical_p,

temp=temp,

repeat_penalty=repeat_penalty,

frequency_penalty=frequency_penalty,

presence_penalty=presence_penalty,

tfs_z=tfs_z,

mirostat_mode=mirostat_mode,

mirostat_tau=mirostat_tau,

mirostat_eta=mirostat_eta,

penalize_nl=penalize_nl,

logits_processor=logits_processor,

grammar=grammar,

)

assert self.ctx is not None

token = self._sampler.sample(self._ctx, -1)

if tmp_sampler:

self._sampler = None

return token

def generate(

self,

tokens: Sequence[int],

top_k: int = 40,

top_p: float = 0.95,

min_p: float = 0.05,

typical_p: float = 1.0,

temp: float = 0.80,

repeat_penalty: float = 1.0,

reset: bool = True,

frequency_penalty: float = 0.0,

presence_penalty: float = 0.0,

tfs_z: float = 1.0,

mirostat_mode: int = 0,

mirostat_tau: float = 5.0,

mirostat_eta: float = 0.1,

penalize_nl: bool = True,

logits_processor: Optional[LogitsProcessorList] = None,

stopping_criteria: Optional[StoppingCriteriaList] = None,

grammar: Optional[LlamaGrammar] = None,

) -> Generator[int, Optional[Sequence[int]], None]:

"""Create a generator of tokens from a prompt.

# Reset mirostat sampling

self._mirostat_mu = ctypes.c_float(2.0 * mirostat_tau)

self._sampler = self._init_sampler(

top_k=top_k,

top_p=top_p,

min_p=min_p,

typical_p=typical_p,

temp=temp,

repeat_penalty=repeat_penalty,

frequency_penalty=frequency_penalty,

presence_penalty=presence_penalty,

tfs_z=tfs_z,

mirostat_mode=mirostat_mode,

mirostat_tau=mirostat_tau,

mirostat_eta=mirostat_eta,

penalize_nl=penalize_nl,

logits_processor=logits_processor,

grammar=grammar,

)

# Check for kv cache prefix match

if reset and self.n_tokens > 0:

longest_prefix = 0

for a, b in zip(self._input_ids, tokens[:-1]):

if a == b:

longest_prefix += 1

else:

break

if longest_prefix > 0:

reset = False

tokens = tokens[longest_prefix:]

self.n_tokens = longest_prefix

if self.verbose:

print(

f"Llama.generate: {longest_prefix} prefix-match hit, "

f"remaining {len(tokens)} prompt tokens to eval",

file=sys.stderr,

)

# Reset the model state

if reset:

self.reset()

# # Reset the grammar

# if grammar is not None:

# grammar.reset()

sample_idx = self.n_tokens + len(tokens) - 1

tokens = list(tokens)

# Eval and sample

while True:

self.eval(tokens)

while sample_idx < self.n_tokens:

token = self.sample(

top_k=top_k,

top_p=top_p,

min_p=min_p,

typical_p=typical_p,

temp=temp,

repeat_penalty=repeat_penalty,

frequency_penalty=frequency_penalty,

presence_penalty=presence_penalty,

tfs_z=tfs_z,

mirostat_mode=mirostat_mode,

mirostat_tau=mirostat_tau,

mirostat_eta=mirostat_eta,

logits_processor=logits_processor,

grammar=grammar,

penalize_nl=penalize_nl,

idx=sample_idx,

)

sample_idx += 1

if stopping_criteria is not None and stopping_criteria(

self._input_ids, self._scores[-1, :]

):

return

tokens_or_none = yield token

tokens.clear()

tokens.append(token)

if tokens_or_none is not None:

tokens.extend(tokens_or_none)

if sample_idx < self.n_tokens and token != self._input_ids[sample_idx]:

self.n_tokens = sample_idx

self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1)

break

if self.draft_model is not None:

self.input_ids[self.n_tokens : self.n_tokens + len(tokens)] = tokens

draft_tokens = self.draft_model(

self.input_ids[: self.n_tokens + len(tokens)]

)

tokens.extend(

draft_tokens.astype(int)[

: self._n_ctx - self.n_tokens - len(tokens)

]

)

def create_embedding(

self, input: Union[str, List[str]], model: Optional[str] = None

) -> CreateEmbeddingResponse:

"""Embed a string.

model_name: str = model if model is not None else self.model_path

input = input if isinstance(input, list) else [input]

# get numeric embeddings

embeds: Union[List[List[float]], List[List[List[float]]]]

total_tokens: int

embeds, total_tokens = self.embed(input, return_count=True) # type: ignore

# convert to CreateEmbeddingResponse

data: List[Embedding] = [

{

"object": "embedding",

"embedding": emb,

"index": idx,

}

for idx, emb in enumerate(embeds)

]

return {

"object": "list",

"data": data,

"model": model_name,

"usage": {

"prompt_tokens": total_tokens,

"total_tokens": total_tokens,

},

}

def embed(

self,

input: Union[str, List[str]],

normalize: bool = False,

truncate: bool = True,