from __future__ import annotations

import os

import ctypes

from typing import (

Dict,

List,

Tuple,

Optional,

Sequence,

)

from dataclasses import dataclass, field

from contextlib import ExitStack

import numpy as np

import numpy.typing as npt

from .llama_types import *

from .llama_grammar import LlamaGrammar

from ._utils import suppress_stdout_stderr

import llama_cpp.llama_cpp as llama_cpp

# Python wrappers over llama.h structs

class LlamaModel:

"""Intermediate Python wrapper for a llama.cpp llama_model.

NOTE: For stability it's recommended you use the Llama class instead."""

def __init__(

self,

*,

path_model: str,

params: llama_cpp.llama_model_params,

verbose: bool = True,

):

self.path_model = path_model

self.params = params

self.verbose = verbose

self._exit_stack = ExitStack()

model = None

if not os.path.exists(path_model):

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

with suppress_stdout_stderr(disable=verbose):

model = llama_cpp.llama_load_model_from_file(

self.path_model.encode("utf-8"), self.params

)

if model is None:

raise ValueError(f"Failed to load model from file: {path_model}")

self.model = model

def free_model():

if self.model is None:

return

llama_cpp.llama_free_model(self.model)

self.model = None

self._exit_stack.callback(free_model)

def close(self):

self._exit_stack.close()

def __del__(self):

self.close()

def vocab_type(self) -> int:

return llama_cpp.llama_vocab_type(self.model)

def n_vocab(self) -> int:

return llama_cpp.llama_n_vocab(self.model)

def n_ctx_train(self) -> int:

return llama_cpp.llama_n_ctx_train(self.model)

def n_embd(self) -> int:

return llama_cpp.llama_n_embd(self.model)

def rope_freq_scale_train(self) -> float:

return llama_cpp.llama_rope_freq_scale_train(self.model)

def desc(self) -> str:

buf = ctypes.create_string_buffer(1024)

llama_cpp.llama_model_desc(self.model, buf, 1024)

return buf.value.decode("utf-8")

def size(self) -> int:

return llama_cpp.llama_model_size(self.model)

def n_params(self) -> int:

return llama_cpp.llama_model_n_params(self.model)

def get_tensor(self, name: str) -> ctypes.c_void_p:

return llama_cpp.llama_get_model_tensor(self.model, name.encode("utf-8"))

# Vocab

def token_get_text(self, token: int) -> str:

return llama_cpp.llama_token_get_text(self.model, token).decode("utf-8")

def token_get_score(self, token: int) -> float:

return llama_cpp.llama_token_get_score(self.model, token)

def token_get_attr(self, token: int) -> int:

return llama_cpp.llama_token_get_attr(self.model, token)

# Special tokens

def token_bos(self) -> int:

return llama_cpp.llama_token_bos(self.model)

def token_eos(self) -> int:

return llama_cpp.llama_token_eos(self.model)

def token_cls(self) -> int:

return llama_cpp.llama_token_cls(self.model)

def token_sep(self) -> int:

return llama_cpp.llama_token_sep(self.model)

def token_nl(self) -> int:

return llama_cpp.llama_token_nl(self.model)

def token_prefix(self) -> int:

return llama_cpp.llama_token_prefix(self.model)

def token_middle(self) -> int:

return llama_cpp.llama_token_middle(self.model)

def token_suffix(self) -> int:

return llama_cpp.llama_token_suffix(self.model)

def token_eot(self) -> int:

return llama_cpp.llama_token_eot(self.model)

def add_bos_token(self) -> bool:

return llama_cpp.llama_add_bos_token(self.model)

def add_eos_token(self) -> bool:

return llama_cpp.llama_add_eos_token(self.model)

# Tokenization

def tokenize(self, text: bytes, add_bos: bool, special: bool):

n_ctx = self.n_ctx_train()

tokens = (llama_cpp.llama_token * n_ctx)()

n_tokens = llama_cpp.llama_tokenize(

self.model, text, len(text), tokens, n_ctx, add_bos, special

)

if n_tokens < 0:

n_tokens = abs(n_tokens)

tokens = (llama_cpp.llama_token * n_tokens)()

n_tokens = llama_cpp.llama_tokenize(

self.model, text, len(text), tokens, n_tokens, add_bos, special

)

if n_tokens < 0:

raise RuntimeError(

f'Failed to tokenize: text="{text}" n_tokens={n_tokens}'

)

return list(tokens[:n_tokens])

def token_to_piece(self, token: int, special: bool = False) -> bytes:

buf = ctypes.create_string_buffer(32)

llama_cpp.llama_token_to_piece(self.model, token, buf, 32, 0, special)

return bytes(buf)

def detokenize(self, tokens: List[int], special: bool = False) -> bytes:

output = b""

size = 32

buffer = (ctypes.c_char * size)()

for token in tokens:

n = llama_cpp.llama_token_to_piece(

self.model, llama_cpp.llama_token(token), buffer, size, 0, special

)

assert n <= size

output += bytes(buffer[:n])

# NOTE: Llama1 models automatically added a space at the start of the prompt

# this line removes a leading space if the first token is a beginning of sentence token

return (

output[1:]

if len(tokens) > 0 and tokens[0] == self.token_bos() and output[0:1] == b" "

else output

)

# Extra

def metadata(self) -> Dict[str, str]:

metadata: Dict[str, str] = {}

buffer_size = 1024

buffer = ctypes.create_string_buffer(buffer_size)

# zero the buffer

buffer.value = b"\0" * buffer_size

# iterate over model keys

for i in range(llama_cpp.llama_model_meta_count(self.model)):

nbytes = llama_cpp.llama_model_meta_key_by_index(

self.model, i, buffer, buffer_size

)

if nbytes > buffer_size:

buffer_size = nbytes + 1

buffer = ctypes.create_string_buffer(buffer_size)

nbytes = llama_cpp.llama_model_meta_key_by_index(

self.model, i, buffer, buffer_size

)

key = buffer.value.decode("utf-8")

nbytes = llama_cpp.llama_model_meta_val_str_by_index(

self.model, i, buffer, buffer_size

)

if nbytes > buffer_size:

buffer_size = nbytes + 1

buffer = ctypes.create_string_buffer(buffer_size)

nbytes = llama_cpp.llama_model_meta_val_str_by_index(

self.model, i, buffer, buffer_size

)

value = buffer.value.decode("utf-8")

metadata[key] = value

return metadata

@staticmethod

def default_params():

"""Get the default llama_model_params."""

return llama_cpp.llama_model_default_params()

class LlamaContext:

"""Intermediate Python wrapper for a llama.cpp llama_context.

NOTE: For stability it's recommended you use the Llama class instead."""

def __init__(

self,

*,

model: LlamaModel,

params: llama_cpp.llama_context_params,

verbose: bool = True,

):

self.model = model

self.params = params

self.verbose = verbose

self._exit_stack = ExitStack()

ctx = llama_cpp.llama_new_context_with_model(self.model.model, self.params)

if ctx is None:

raise ValueError("Failed to create llama_context")

self.ctx = ctx

def free_ctx():

if self.ctx is None:

return

llama_cpp.llama_free(self.ctx)

self.ctx = None

self._exit_stack.callback(free_ctx)

def close(self):

self._exit_stack.close()

def __del__(self):

self.close()

def n_ctx(self) -> int:

return llama_cpp.llama_n_ctx(self.ctx)

def pooling_type(self) -> int:

return llama_cpp.llama_pooling_type(self.ctx)

def kv_cache_clear(self):

llama_cpp.llama_kv_cache_clear(self.ctx)

def kv_cache_seq_rm(self, seq_id: int, p0: int, p1: int):

llama_cpp.llama_kv_cache_seq_rm(self.ctx, seq_id, p0, p1)

def kv_cache_seq_cp(self, seq_id_src: int, seq_id_dst: int, p0: int, p1: int):

llama_cpp.llama_kv_cache_seq_cp(self.ctx, seq_id_src, seq_id_dst, p0, p1)

def kv_cache_seq_keep(self, seq_id: int):

llama_cpp.llama_kv_cache_seq_keep(self.ctx, seq_id)

def kv_cache_seq_shift(self, seq_id: int, p0: int, p1: int, shift: int):

llama_cpp.llama_kv_cache_seq_add(self.ctx, seq_id, p0, p1, shift)

def get_state_size(self) -> int:

return llama_cpp.llama_get_state_size(self.ctx)

# TODO: copy_state_data

# TODO: set_state_data

# TODO: llama_load_session_file

# TODO: llama_save_session_file

def decode(self, batch: LlamaBatch):

return_code = llama_cpp.llama_decode(

self.ctx,

batch.batch,

)

if return_code != 0:

raise RuntimeError(f"llama_decode returned {return_code}")

def set_n_threads(self, n_threads: int, n_threads_batch: int):

llama_cpp.llama_set_n_threads(self.ctx, n_threads, n_threads_batch)

def get_logits(self):

return llama_cpp.llama_get_logits(self.ctx)

def get_logits_ith(self, i: int):

return llama_cpp.llama_get_logits_ith(self.ctx, i)

def get_embeddings(self):

return llama_cpp.llama_get_embeddings(self.ctx)

# Sampling functions

def set_rng_seed(self, seed: int):

# TODO: Fix

llama_cpp.llama_set_rng_seed(self.ctx, seed)

def sample_repetition_penalties(

self,

candidates: "_LlamaTokenDataArray",

last_tokens_data: "llama_cpp.Array[llama_cpp.llama_token]",

penalty_last_n: int,

penalty_repeat: float,

penalty_freq: float,

penalty_present: float,

):

llama_cpp.llama_sample_repetition_penalties(

self.ctx,

llama_cpp.byref(candidates.candidates),

last_tokens_data,

penalty_last_n,

penalty_repeat,

penalty_freq,

penalty_present,

)

def sample_softmax(self, candidates: "_LlamaTokenDataArray"):

llama_cpp.llama_sample_softmax(

self.ctx,

llama_cpp.byref(candidates.candidates),

)

def sample_top_k(self, candidates: "_LlamaTokenDataArray", k: int, min_keep: int):

llama_cpp.llama_sample_top_k(

self.ctx, llama_cpp.byref(candidates.candidates), k, min_keep

)

def sample_top_p(self, candidates: "_LlamaTokenDataArray", p: float, min_keep: int):

llama_cpp.llama_sample_top_p(

self.ctx, llama_cpp.byref(candidates.candidates), p, min_keep

)

def sample_min_p(self, candidates: "_LlamaTokenDataArray", p: float, min_keep: int):

llama_cpp.llama_sample_min_p(

self.ctx, llama_cpp.byref(candidates.candidates), p, min_keep

)

def sample_typical(

self, candidates: "_LlamaTokenDataArray", p: float, min_keep: int

):

llama_cpp.llama_sample_typical(

self.ctx, llama_cpp.byref(candidates.candidates), p, min_keep

)

def sample_temp(self, candidates: "_LlamaTokenDataArray", temp: float):

llama_cpp.llama_sample_temp(

self.ctx, llama_cpp.byref(candidates.candidates), temp

)

def sample_grammar(self, candidates: "_LlamaTokenDataArray", grammar: LlamaGrammar):

llama_cpp.llama_sample_grammar(

self.ctx,

llama_cpp.byref(candidates.candidates),

grammar.grammar,

)

def sample_token_mirostat(

self,

candidates: "_LlamaTokenDataArray",

tau: float,

eta: float,

m: int,

mu: llama_cpp.CtypesPointerOrRef[ctypes.c_float],

) -> int:

return llama_cpp.llama_sample_token_mirostat(

self.ctx,

llama_cpp.byref(candidates.candidates),

tau,

eta,

m,

mu,

)

def sample_token_mirostat_v2(

self,

candidates: "_LlamaTokenDataArray",

tau: float,

eta: float,

mu: llama_cpp.CtypesPointerOrRef[ctypes.c_float],

) -> int:

return llama_cpp.llama_sample_token_mirostat_v2(

self.ctx,

llama_cpp.byref(candidates.candidates),

tau,

eta,

mu,

)

def sample_token_greedy(self, candidates: "_LlamaTokenDataArray") -> int:

return llama_cpp.llama_sample_token_greedy(

self.ctx,

llama_cpp.byref(candidates.candidates),

)

def sample_token(self, candidates: "_LlamaTokenDataArray") -> int:

return llama_cpp.llama_sample_token(

self.ctx,

llama_cpp.byref(candidates.candidates),

)

# Grammar

def grammar_accept_token(self, grammar: LlamaGrammar, token: int):

llama_cpp.llama_grammar_accept_token(grammar.grammar, self.ctx, token)

def reset_timings(self):

llama_cpp.llama_perf_context_reset(self.ctx)

def print_timings(self):

llama_cpp.llama_perf_context_print(self.ctx)

# Utility functions

@staticmethod

def default_params():

"""Get the default llama_context_params."""

return llama_cpp.llama_context_default_params()

class LlamaBatch:

def __init__(

self, *, n_tokens: int, embd: int, n_seq_max: int, verbose: bool = True

):

self._n_tokens = n_tokens

self.embd = embd

self.n_seq_max = n_seq_max

self.verbose = verbose

self._exit_stack = ExitStack()

batch = llama_cpp.llama_batch_init(self._n_tokens, self.embd, self.n_seq_max)

if batch is None:

raise ValueError("Failed to create llama_batch")

self.batch = batch

def free_batch():

if self.batch is None:

return

llama_cpp.llama_batch_free(self.batch)

self.batch = None

self._exit_stack.callback(free_batch)

def close(self):

self._exit_stack.close()

def __del__(self):

self.close()

def n_tokens(self) -> int:

return self.batch.n_tokens

def reset(self):

self.batch.n_tokens = 0

def set_batch(self, batch: Sequence[int], n_past: int, logits_all: bool):

n_tokens = len(batch)

self.batch.n_tokens = n_tokens

for i in range(n_tokens):

self.batch.token[i] = batch[i]

self.batch.pos[i] = n_past + i

self.batch.seq_id[i][0] = 0

self.batch.n_seq_id[i] = 1

self.batch.logits[i] = logits_all

self.batch.logits[n_tokens - 1] = True

def add_sequence(self, batch: Sequence[int], seq_id: int, logits_all: bool):

n_tokens = len(batch)

n_tokens0 = self.batch.n_tokens

self.batch.n_tokens += n_tokens

for i in range(n_tokens):

j = n_tokens0 + i

self.batch.token[j] = batch[i]

self.batch.pos[j] = i

self.batch.seq_id[j][0] = seq_id

self.batch.n_seq_id[j] = 1

self.batch.logits[j] = logits_all

self.batch.logits[n_tokens - 1] = True

class LlamaTokenDataArray:

def __init__(self, *, n_vocab: int):

self.n_vocab = n_vocab

self.candidates_data = np.recarray(

(self.n_vocab,),

dtype=np.dtype(

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

),

)

self.candidates = llama_cpp.llama_token_data_array(

data=self.candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p),

size=self.n_vocab,

sorted=False,

)

self.default_candidates_data_id = np.arange(self.n_vocab, dtype=np.intc) # type: ignore

self.default_candidates_data_p = np.zeros(self.n_vocab, dtype=np.single)

def copy_logits(self, logits: npt.NDArray[np.single]):

self.candidates_data.id[:] = self.default_candidates_data_id

self.candidates_data.logit[:] = logits

self.candidates_data.p[:] = self.default_candidates_data_p

self.candidates.sorted = False

self.candidates.size = self.n_vocab

# Embedding functions

def normalize_embedding(embedding):

norm = float(np.linalg.norm(embedding))

if norm == 0.0:

return embedding

return [v / norm for v in embedding]

# Python wrappers over common/sampling structs

@dataclass

class LlamaSamplingParams:

n_prev: int = 64

n_probs: int = 0

top_k: int = 40

top_p: float = 0.95

min_p: float = 0.05

tfs_z: float = 1.00

typical_p: float = 1.00

temp: float = 0.80

penalty_last_n: int = 64

penalty_repeat: float = 1.0

penalty_freq: float = 0.00

penalty_present: float = 0.00

mirostat: int = 0

mirostat_tau: float = 5.00

mirostat_eta: float = 0.10

penalize_nl: bool = True

grammar: str = ""

cfg_negative_prompt: str = ""

cfg_scale: float = 1.00

logit_bias: dict[int, float] = field(default_factory=dict)

@dataclass

class LlamaSamplingContext:

params: LlamaSamplingParams = field(default_factory=LlamaSamplingParams)

mirostat_mu: ctypes.c_float = field(default_factory=ctypes.c_float)

grammar: Optional[LlamaGrammar] = None

# NOTE: Missing parsed_grammar

prev: list[int] = field(default_factory=list)

cur: list[llama_cpp.llama_token_data] = field(default_factory=list)

def reset(self):

self.prev = []

self.cur = []

if self.grammar is not None:

self.grammar.reset()

def cp(self):

return LlamaSamplingContext(

params=self.params,

mirostat_mu=self.mirostat_mu,

grammar=self.grammar,

prev=self.prev.copy(),

cur=self.cur.copy(),

)

def last(self) -> Optional[int]:

if len(self.prev) > 0:

return self.prev[-1]

else:

return None

def prev_str(self, ctx_main: LlamaContext, n: int) -> str:

return ctx_main.model.detokenize(self.prev[-n:]).decode("utf-8")

def sample(

self,

ctx_main: LlamaContext,

idx: int = 0,

logits_array: Optional[npt.NDArray[np.single]] = None,

):

n_vocab = ctx_main.model.n_vocab()

id: int = 0

if logits_array is None:

logits = ctx_main.get_logits_ith(idx)

logits_array = np.array(

ctypes.cast(logits, ctypes.POINTER(ctypes.c_float * n_vocab)).contents,

dtype=np.single,

)

# apply logit_bias

for token, logit_bias in self.params.logit_bias.items():

logits_array[token] += logit_bias

token_data_array = LlamaTokenDataArray(

n_vocab=n_vocab

) # TODO: Only create this once

token_data_array.copy_logits(logits_array)

# apply penalties

if len(self.prev) > 0:

nl_token = ctx_main.model.token_nl()

nl_logit = logits_array[nl_token]

last_tokens = self.prev[-self.params.penalty_last_n :]

last_tokens_size = min(len(last_tokens), self.params.penalty_last_n)

if last_tokens_size > 0:

last_tokens_p = (llama_cpp.llama_token * len(last_tokens))(*last_tokens)

ctx_main.sample_repetition_penalties(

token_data_array,

last_tokens_p,

last_tokens_size,

self.params.penalty_repeat,

self.params.penalty_freq,

self.params.penalty_present,

)

if not self.params.penalize_nl:

token_data_array.candidates_data.logit[nl_token] = nl_logit

if self.grammar is not None:

ctx_main.sample_grammar(token_data_array, self.grammar)

if self.params.temp < 0:

ctx_main.sample_softmax(token_data_array)

id = token_data_array.candidates_data.id[0]

elif self.params.temp == 0:

id = ctx_main.sample_token_greedy(token_data_array)

else:

if self.params.mirostat == 1:

mirostat_m = 100

ctx_main.sample_temp(token_data_array, self.params.temp)

id = ctx_main.sample_token_mirostat(

token_data_array,

self.params.mirostat_tau,

self.params.mirostat_eta,

mirostat_m,

ctypes.pointer(self.mirostat_mu),

)

elif self.params.mirostat == 2:

ctx_main.sample_temp(token_data_array, self.params.temp)

id = ctx_main.sample_token_mirostat_v2(

token_data_array,

self.params.mirostat_tau,

self.params.mirostat_eta,

ctypes.pointer(self.mirostat_mu),

)

else:

min_keep = max(1, self.params.n_probs)

ctx_main.sample_top_k(

token_data_array, self.params.top_k, min_keep=min_keep

)

ctx_main.sample_typical(

token_data_array, self.params.typical_p, min_keep=min_keep

)

ctx_main.sample_top_p(

token_data_array, self.params.top_p, min_keep=min_keep

)

ctx_main.sample_min_p(

token_data_array, self.params.min_p, min_keep=min_keep

)

ctx_main.sample_temp(token_data_array, self.params.temp)

id = ctx_main.sample_token(token_data_array)

return id

def accept(self, ctx_main: LlamaContext, id: int, apply_grammar: bool):

if apply_grammar and self.grammar is not None:

ctx_main.grammar_accept_token(self.grammar, id)

self.prev.append(id)

from typing import List, Callable, Optional, Union

import ctypes

import llama_cpp

class CustomSampler:

def __init__(

self, apply_func: typing.Callable[[llama_cpp.llama_token_data_array], None]

):

self.apply_func = apply_func

def apply_wrapper(

sampler: llama_cpp.llama_sampler_p,

cur_p: llama_cpp.llama_token_data_array_p,

):

self.apply_func(cur_p)

def free_wrapper(sampler: llama_cpp.llama_sampler_p):

pass

sampler_i = llama_cpp.llama_sampler_i()

sampler_i.apply = llama_cpp.llama_sampler_i_apply(apply_wrapper)

self._apply_wrapper_ref = apply_wrapper

sampler_i.name = llama_cpp.llama_sampler_i_name(0)

sampler_i.accept = llama_cpp.llama_sampler_i_accept(0)

sampler_i.reset = llama_cpp.llama_sampler_i_reset(0)

sampler_i.clone = llama_cpp.llama_sampler_i_clone(0)

sampler_i.free = llama_cpp.llama_sampler_i_free(0)

self.sampler = llama_cpp.llama_sampler()

self.sampler.iface = ctypes.pointer(sampler_i)

self.sampler.ctx = None

def get_sampler(self) -> llama_cpp.llama_sampler_p:

return ctypes.pointer(self.sampler)

class LlamaSampler:

def __init__(self):

params = llama_cpp.llama_sampler_chain_params()

self.sampler = llama_cpp.llama_sampler_chain_init(params)

self.samplers: List[llama_cpp.llama_sampler_p] = []

self.custom_samplers: List[Tuple[int, CustomSampler]] = []

def add_greedy(self):

sampler = llama_cpp.llama_sampler_init_greedy()

self._add_sampler(sampler)

def add_dist(self, seed: int):

sampler = llama_cpp.llama_sampler_init_dist(seed)

self._add_sampler(sampler)

def add_softmax(self):

sampler = llama_cpp.llama_sampler_init_softmax()

self._add_sampler(sampler)

def add_top_k(self, k: int):

sampler = llama_cpp.llama_sampler_init_top_k(k)

self._add_sampler(sampler)

def add_top_p(self, p: float, min_keep: int):

sampler = llama_cpp.llama_sampler_init_top_p(p, min_keep)

self._add_sampler(sampler)

def add_min_p(self, p: float, min_keep: int):

sampler = llama_cpp.llama_sampler_init_min_p(p, min_keep)

self._add_sampler(sampler)

def add_typical(self, p: float, min_keep: int):

sampler = llama_cpp.llama_sampler_init_typical(p, min_keep)

self._add_sampler(sampler)

def add_temp(self, temp: float):

sampler = llama_cpp.llama_sampler_init_temp(temp)

self._add_sampler(sampler)

def add_temp_ext(self, t: float, delta: float, exponent: float):

sampler = llama_cpp.llama_sampler_init_temp_ext(t, delta, exponent)

self._add_sampler(sampler)

def add_mirostat(self, n_vocab: int, seed: int, tau: float, eta: float, m: int):

sampler = llama_cpp.llama_sampler_init_mirostat(n_vocab, seed, tau, eta, m)

self._add_sampler(sampler)

def add_mirostat_v2(self, seed: int, tau: float, eta: float):

sampler = llama_cpp.llama_sampler_init_mirostat_v2(seed, tau, eta)

self._add_sampler(sampler)

def add_grammar(self, model: LlamaModel, grammar: LlamaGrammar):

sampler = llama_cpp.llama_sampler_init_grammar(

model.model, grammar._grammar.encode("utf-8"), grammar._root.encode("utf-8")

)

self._add_sampler(sampler)

def add_penalties(

self,

n_vocab: int,

special_eos_id: int,

linefeed_id: int,

penalty_last_n: int,

penalty_repeat: float,

penalty_freq: float,

penalty_present: float,

penalize_nl: bool,

ignore_eos: bool,

):

sampler = llama_cpp.llama_sampler_init_penalties(

n_vocab,

special_eos_id,

linefeed_id,

penalty_last_n,

penalty_repeat,

penalty_freq,

penalty_present,

penalize_nl,

ignore_eos,

)

self._add_sampler(sampler)

def init_logit_bias(

self, n_vocab: int, n_logit_bias, logit_bias: llama_cpp.llama_logit_bias_p

):

sampler = llama_cpp.llama_sampler_init_logit_bias(

n_vocab, n_logit_bias, logit_bias

)

self._add_sampler(sampler)

def add_custom(

self, apply_func: Callable[[llama_cpp.llama_token_data_array], None]

):

custom_sampler = CustomSampler(apply_func)

sampler = custom_sampler.get_sampler()

self._add_sampler(sampler)

# NOTE: Must remove custom samplers before free or llama.cpp will try to free them

self.custom_samplers.append(

(llama_cpp.llama_sampler_chain_n(self.sampler) - 1, custom_sampler)

)

def _add_sampler(self, sampler: llama_cpp.llama_sampler_p):

assert self.sampler is not None

llama_cpp.llama_sampler_chain_add(self.sampler, sampler)

self.samplers.append(sampler)

def get_seed(self) -> int:

assert self.sampler is not None

return llama_cpp.llama_sampler_get_seed(self.sampler)

def sample(self, ctx: LlamaContext, idx: int) -> int:

assert self.sampler is not None

return llama_cpp.llama_sampler_sample(self.sampler, ctx.ctx, idx)

def close(self):

if self.sampler:

# NOTE: Must remove custom samplers before free or llama.cpp will try to free them

for i, _ in reversed(self.custom_samplers):

llama_cpp.llama_sampler_chain_remove(self.sampler, i)

llama_cpp.llama_sampler_free(self.sampler)

self.sampler = None

self.samplers.clear()

self.custom_samplers.clear()

def __del__(self):

self.close()