from llama_cpp import *
from ctypes import POINTER, c_size_t
from llama_cpp._internals import (
_LlamaModel, # type: ignore
_LlamaContext, # type: ignore
_LlamaBatch, # type: ignore
_LlamaTokenDataArray, # type: ignore
)
from KMP_list import kmp_search, compute_lps_array
from Turbo_Colormap import map_value_to_color, NOCOLOR, LEGEND, BACK_WHITE
class LLMGenerate:
def __init__(
self,
model,
n_keep,
n_discard: int = 256,
im_start=None,
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.1,
repeat_last_n: int = 64,
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
):
def _eval_t(tokens):
return model.eval_t(
tokens=tokens,
n_keep=n_keep,
n_discard=n_discard,
im_start=im_start
)
def _sample_t(logits_processor):
return model.sample_t(
top_k=top_k,
top_p=top_p,
min_p=min_p,
typical_p=typical_p,
temp=temp,
repeat_penalty=repeat_penalty,
repeat_last_n=repeat_last_n,
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
)
self._eval_t = _eval_t
self._sample_t = _sample_t
self.str_detokenize = model.str_detokenize
self.venv_pop_token = model.venv_pop_token
# ========== 保存输出 ==========
self.t_bot = []
self.completion_tokens = []
self.history = ''
self.token = None
def eval_t(self, tokens):
# ========== 避免不完整的utf-8编码 ==========
self.completion_tokens.extend(tokens)
all_text = self.str_detokenize(self.completion_tokens)
if all_text:
self.t_bot.extend(self.completion_tokens)
self.history += all_text
self.completion_tokens = []
return self._eval_t(tokens)
def sample_t(self, logits_processor):
self.token = self._sample_t(logits_processor)
return self.token
def detokenize_sample_t(self):
self.completion_tokens.append(self.token)
all_text = self.str_detokenize(self.completion_tokens)
if not all_text:
return False
self.t_bot.extend(self.completion_tokens)
self.history += all_text
self.completion_tokens = []
return True
def eval_sample_t(self):
return self._eval_t([self.token])
def endswith_t(self, token_list):
return self.token in token_list
def endswith_s(self, start_func, str_list, com_func=str.rstrip):
if self.completion_tokens: # 不完整
return False
history = self.history
t_bot = self.t_bot
if start_func(history):
history = com_func(history)
for x in str_list:
if history.endswith(x):
n = len(t_bot)
for i in range(1, n): # 找出需要弃置的tokens长度
tmp = self.str_detokenize(t_bot[n - i:])
tmp = com_func(tmp)
if tmp.endswith(x):
if i > 1: # 最后一个token并未进入kv_cache
self.venv_pop_token(i - 1)
if history.endswith(tmp):
self.history = history[:-len(tmp)] # 移除末尾的tmp
return True
return False
kv_cache_type = {
'f32': 0,
'f16': 1,
'q8_0': 8,
'q4_0': 2,
'q4_1': 3,
'iq4_nl': 20,
'q5_0': 6,
'q5_1': 7
}
class StreamingLLM(Llama):
__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,
vocab_only: bool = False,
use_mmap: bool = True,
use_mlock: bool = False,
kv_overrides: Optional[Dict[str, Union[bool, int, float]]] = None,
# Context Params
seed: int = llama_cpp.LLAMA_DEFAULT_SEED,
n_ctx: int = 512,
n_batch: 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,
# 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,
# Misc
verbose: bool = True,
# Extra Params
type_k: str = 'f16',
type_v: str = 'f16',
**kwargs, # type: ignore
):
"""Load a llama.cpp model from `model_path`.
Examples:
Basic usage
>>> import llama_cpp
>>> model = llama_cpp.Llama(
... model_path="path/to/model",
... )
>>> print(model("The quick brown fox jumps ", stop=["."])["choices"][0]["text"])
the lazy dog
Loading a chat model
>>> import llama_cpp
>>> model = llama_cpp.Llama(
... model_path="path/to/model",
... chat_format="llama-2",
... )
>>> print(model.create_chat_completion(
... messages=[{
... "role": "user",
... "content": "what is the meaning of life?"
... }]
... ))
Args:
model_path: Path to the model.
n_gpu_layers: Number of layers to offload to GPU (-ngl). If -1, all layers are offloaded.
split_mode: How to split the model across GPUs. See llama_cpp.LLAMA_SPLIT_* for options.
main_gpu: 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: How split tensors should be distributed across GPUs. If None, the model is not split.
vocab_only: Only load the vocabulary no weights.
use_mmap: Use mmap if possible.
use_mlock: Force the system to keep the model in RAM.
kv_overrides: Key-value overrides for the model.
seed: RNG seed, -1 for random
n_ctx: Text context, 0 = from model
n_batch: Prompt processing maximum batch size
n_threads: Number of threads to use for generation
n_threads_batch: Number of threads to use for batch processing
rope_scaling_type: RoPE scaling type, from `enum llama_rope_scaling_type`. ref: https://github.com/ggerganov/llama.cpp/pull/2054
pooling_type: Pooling type, from `enum llama_pooling_type`.
rope_freq_base: RoPE base frequency, 0 = from model
rope_freq_scale: RoPE frequency scaling factor, 0 = from model
yarn_ext_factor: YaRN extrapolation mix factor, negative = from model
yarn_attn_factor: YaRN magnitude scaling factor
yarn_beta_fast: YaRN low correction dim
yarn_beta_slow: YaRN high correction dim
yarn_orig_ctx: YaRN original context size
logits_all: Return logits for all tokens, not just the last token. Must be True for completion to return logprobs.
embedding: Embedding mode only.
offload_kqv: Offload K, Q, V to GPU.
last_n_tokens_size: Maximum number of tokens to keep in the last_n_tokens deque.
lora_base: Optional path to base model, useful if using a quantized base model and you want to apply LoRA to an f16 model.
lora_path: Path to a LoRA file to apply to the model.
numa: numa policy
chat_format: String specifying the chat format to use when calling create_chat_completion.
chat_handler: Optional chat handler to use when calling create_chat_completion.
draft_model: Optional draft model to use for speculative decoding.
tokenizer: Optional tokenizer to override the default tokenizer from llama.cpp.
verbose: Print verbose output to stderr.
Raises:
ValueError: If the model path does not exist.
Returns:
A Llama instance.
"""
self.verbose = verbose
set_verbose(verbose)
if not StreamingLLM.__backend_initialized:
with suppress_stdout_stderr(disable=verbose):
llama_cpp.llama_backend_init()
StreamingLLM.__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
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.bool_value = v
elif isinstance(v, int):
self._kv_overrides_array[i].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_INT
self._kv_overrides_array[i].value.int_value = v
elif isinstance(v, float):
self._kv_overrides_array[i].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_FLOAT
self._kv_overrides_array[i].value.float_value = v
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 max(
multiprocessing.cpu_count() // 2, 1
)
# Context Params
self.context_params = llama_cpp.llama_context_default_params()
self.context_params.seed = seed
self.context_params.n_ctx = n_ctx
self.context_params.n_batch = self.n_batch
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
# KV cache quantization
print(self.context_params.type_k, self.context_params.type_v)
self.context_params.type_k = kv_cache_type[type_k]
self.context_params.type_v = kv_cache_type[type_v]
self.context_params.offload_kqv = offload_kqv
# 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
if not os.path.exists(model_path):
raise ValueError(f"Model path does not exist: {model_path}")
self._model = _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._ctx = _LlamaContext(
model=self._model,
params=self.context_params,
verbose=self.verbose,
)
self._batch = _LlamaBatch(
n_tokens=self.n_batch,
embd=0,
n_seq_max=self.context_params.n_ctx,
verbose=self.verbose,
)
if self.lora_path:
if self._model.apply_lora_from_file(
self.lora_path,
self.lora_scale,
self.lora_base,
self.n_threads,
):
raise RuntimeError(
f"Failed to apply LoRA from lora path: {self.lora_path} to base path: {self.lora_base}"
)
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.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 = _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, 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)
if (
self.chat_format is None
and self.chat_handler is None
and "tokenizer.chat_template" in self.metadata
):
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:
template = self.metadata["tokenizer.chat_template"]
try:
eos_token_id = int(self.metadata["tokenizer.ggml.eos_token_id"])
except:
eos_token_id = self.token_eos()
try:
bos_token_id = int(self.metadata["tokenizer.ggml.bos_token_id"])
except:
bos_token_id = self.token_bos()
eos_token = self._model.token_get_text(eos_token_id)
bos_token = self._model.token_get_text(bos_token_id)
if self.verbose:
print(f"Using gguf chat template: {template}", 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_handler = llama_chat_format.Jinja2ChatFormatter(
template=template, eos_token=eos_token, bos_token=bos_token
).to_chat_handler()
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: {chat_format}", file=sys.stderr)
self._venv_init()
def str_detokenize(self, tokens) -> str:
return self.detokenize(tokens).decode('utf-8', errors='ignore')
def kv_cache_seq_trim(self):
self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1)
def _venv_init(self):
self.venv = [0]
self.venv_idx_map = []
def venv_create(self, name: str):
self.venv.append(0)
self.venv_idx_map.append(name)
return name
def venv_disband(self, name_set):
if len(self.venv) <= 1:
return False
name_set = {x for x in name_set if x in self.venv_idx_map}
if not name_set:
return False
while self.venv_idx_map:
if self.venv_idx_map[0] in name_set:
self.venv_idx_map.pop(0) # 删除
tmp = self.venv.pop(1) # 对应的 venv 移入上一层
self.venv[0] += tmp
else:
break
return True
def venv_revision(self, name: str):
if len(self.venv) <= 1:
return False
if name not in self.venv_idx_map:
return False
_s = 0
while self.venv_idx_map:
if self.venv_idx_map[-1] == name:
break
self.venv_idx_map.pop() # 删除
_s += self.venv.pop()
if _s:
self.n_tokens -= min(_s, self.n_tokens)
self.kv_cache_seq_trim()
return True
def venv_remove(self, name: str, keep_last=0):
if len(self.venv) <= 1:
return False
if name not in self.venv_idx_map:
return False
venv_idx = self.venv_idx_map.index(name) + 1
count_name = self.venv_idx_map.count(name) if keep_last else 0
while self.venv_idx_map:
if keep_last and count_name <= keep_last:
break # 保留最后n个
self.venv_idx_map.pop(venv_idx - 1) # 删除
if venv_idx == len(self.venv) - 1:
# 最后一层
self.n_tokens -= min(self.venv.pop(), self.n_tokens)
self.kv_cache_seq_trim()
break
else:
# 非最后一层
n_keep = self.n_tokens - sum(self.venv[i] for i in range(venv_idx, len(self.venv)))
n_discard = self.venv.pop(venv_idx)
self.kv_cache_seq_ltrim(n_keep, n_discard)
try:
venv_idx = self.venv_idx_map.index(name, venv_idx - 1) + 1
except ValueError: # 没有了
break
count_name -= 1 # 计数减一
return True
def venv_pop_token(self, n=1):
self.n_tokens -= n
self.venv[-1] -= n
self.kv_cache_seq_trim()
@property
def venv_info(self):
return str((self.n_tokens, self.venv, self.venv_idx_map))
def venv_viz(self):
completion_tokens = []
history = LEGEND + '\n'
text_color = NOCOLOR
for i in range(self.venv[-1]):
idx = self.n_tokens - self.venv[-1] + i
token = self._input_ids[idx]
if not completion_tokens: # 不完整则是第一个token
# ========== 获取对应token的概率 ==========
score = self.scores[idx-1: idx, :].ravel() # 第i个token的分数是前i-1个token预测的，所以减一
score = np.exp(score) # 空白则全1，但无所谓了
sum_score = np.sum(score)
probabilities = score[token] / sum_score
if probabilities < 0.001:
text_color = NOCOLOR
else:
if text_color is NOCOLOR:
text_color = BACK_WHITE + map_value_to_color(probabilities)
else:
text_color = map_value_to_color(probabilities)
history += text_color
# ========== 避免不完整的utf-8编码 ==========
completion_tokens.append(token)
all_text = self.str_detokenize(completion_tokens)
if not all_text:
continue
completion_tokens = [] # 完整则清空缓存
history += repr(all_text)[1:-1]
return history + NOCOLOR
def kv_cache_seq_ltrim(self, n_keep, n_discard=256, n_past=-1, im_start=None):
if n_keep < 0:
return
if n_past < 0:
n_past = self.n_tokens
if im_start is not None: # [<|im_start|>, name, nl]
lps = compute_lps_array(im_start)
_idx = kmp_search(self.input_ids, im_start, n_keep + n_discard, n_past, lps)
if _idx >= n_keep: # 其实是大于等于 n_keep + n_discard
n_discard = _idx - n_keep # 截断到最近的 im_start 序列结构
else:
_idx = kmp_search(self.input_ids, im_start, n_keep, n_past, lps)
if _idx >= n_keep:
n_keep = _idx + len(im_start) # 至少保留一个 im_start 序列结构
print(im_start, n_keep, n_discard, _idx)
self._ctx.kv_cache_seq_rm(-1, n_keep, n_keep + n_discard)
self._ctx.kv_cache_seq_shift(0, n_keep + n_discard, n_past, -n_discard)
self.input_ids[n_keep:n_past - n_discard] = self.input_ids[n_keep + n_discard:n_past]
self.n_tokens = n_past - n_discard
def eval_t(self, tokens, n_keep=4, n_discard=256, im_start=None):
if self._n_ctx < self.n_tokens + len(tokens):
tmp_n_discard = max(n_discard, self.n_tokens + len(tokens) - self._n_ctx)
self.kv_cache_seq_ltrim(n_keep, tmp_n_discard, im_start=im_start)
for i in range(0, len(tokens), self.n_batch):
batch = tokens[i: 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
rows = n_tokens
cols = self._n_vocab
offset = (
0 if self.context_params.logits_all else n_tokens - 1
) # NOTE: Only save the last token logits if logits_all is False
self.scores[n_past + offset: n_past + n_tokens, :].reshape(-1)[
:
] = self._ctx.get_logits()[offset * cols: rows * cols]
# Update n_tokens
self.n_tokens += n_tokens
self.venv[-1] += n_tokens
return self.n_tokens
def sample_t(
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.1,
repeat_last_n: int = 64,
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=None,
grammar: Optional[LlamaGrammar] = None,
):
last_n_tokens_data = [llama_cpp.llama_token(0)] * max(
0, repeat_last_n - self.n_tokens
) + self._input_ids[-repeat_last_n:].tolist()
last_n_tokens_size = len(last_n_tokens_data)
n_vocab = self._n_vocab
n_ctx = self._n_ctx
top_k = n_vocab if top_k <= 0 else top_k
last_n_tokens_size = n_ctx if last_n_tokens_size < 0 else last_n_tokens_size
last_n_tokens_data_c = (llama_cpp.llama_token * last_n_tokens_size)(
*last_n_tokens_data
)
logits: npt.NDArray[np.single] = self.scores[self.n_tokens - 1: self.n_tokens, :].ravel()
if logits_processor is not None:
logits[:] = logits_processor(self._input_ids, logits)
self._candidates.copy_logits(logits)
self._ctx.sample_repetition_penalties(
candidates=self._candidates,
last_tokens_data=last_n_tokens_data_c,
penalty_last_n=last_n_tokens_size,
penalty_repeat=repeat_penalty,
penalty_freq=frequency_penalty,
penalty_present=presence_penalty,
)
if not penalize_nl:
nl_logit = logits[self._token_nl]
self._candidates.candidates.data[self._token_nl].logit = llama_cpp.c_float(
nl_logit
)
if grammar is not None:
self._ctx.sample_grammar(
candidates=self._candidates,
grammar=grammar,
)
if temp < 0.0:
self._ctx.sample_softmax(candidates=self._candidates)
id_ = self._candidates.candidates.data[0].id
elif temp == 0.0:
id_ = self._ctx.sample_token_greedy(candidates=self._candidates)
elif mirostat_mode == 1:
self._ctx.sample_temp(candidates=self._candidates, temp=temp)
id_ = self._ctx.sample_token_mirostat(
candidates=self._candidates,
tau=mirostat_tau,
eta=mirostat_eta,
mu=2.0 * mirostat_tau,
m=100,
)
elif mirostat_mode == 2:
self._ctx.sample_temp(candidates=self._candidates, temp=temp)
id_ = self._ctx.sample_token_mirostat_v2(
candidates=self._candidates,
tau=mirostat_tau,
eta=mirostat_eta,
mu=2.0 * mirostat_tau,
)
else:
self._ctx.sample_top_k(candidates=self._candidates, k=top_k, min_keep=1)
self._ctx.sample_tail_free(candidates=self._candidates, z=tfs_z, min_keep=1)
self._ctx.sample_typical(
candidates=self._candidates, p=typical_p, min_keep=1
)
self._ctx.sample_top_p(candidates=self._candidates, p=top_p, min_keep=1)
self._ctx.sample_min_p(candidates=self._candidates, p=min_p, min_keep=1)
self._ctx.sample_temp(candidates=self._candidates, temp=temp)
id_ = self._ctx.sample_token(candidates=self._candidates)
if grammar is not None:
self._ctx.grammar_accept_token(grammar=grammar, token=id_)
return id_
def generate_t(
self,
tokens: Sequence[int],
n_keep,
n_discard: int = 256,
im_start=None,
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.1,
repeat_last_n: int = 64,
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,
logits_processor: Optional[LogitsProcessorList] = None,
stopping_criteria: Optional[StoppingCriteriaList] = None,
grammar: Optional[LlamaGrammar] = None,
) -> Generator[int, Optional[Sequence[int]], None]:
typical_p = float(typical_p)
frequency_penalty = float(frequency_penalty)
presence_penalty = float(presence_penalty)
tfs_z = float(tfs_z)
mirostat_tau = float(mirostat_tau)
while True:
self.eval_t(tokens, n_keep, n_discard, im_start=im_start)
token = self.sample_t(
top_k=top_k,
top_p=top_p,
min_p=min_p,
typical_p=typical_p,
temp=temp,
repeat_penalty=repeat_penalty,
repeat_last_n=repeat_last_n,
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,
)
if stopping_criteria is not None and stopping_criteria(
self._input_ids, self._scores[-1, :]
):
return
tokens = yield token
if tokens is None:
tokens = [token]
def load_session(self, filepath: str):
n_tokens = POINTER(c_size_t)(c_size_t(0))
tokens = (llama_cpp.llama_token * self.n_ctx())()
retn = llama_cpp.llama_load_session_file(self._ctx.ctx,
filepath.encode('utf-8'),
tokens,
self.n_ctx(),
n_tokens)
self.n_tokens = n_tokens.contents.value
self.input_ids[:self.n_tokens] = tokens[:self.n_tokens]
self._venv_init()
return retn
def save_session(self, filepath: str):
tokens = self._input_ids.tolist()
tokens = (llama_cpp.llama_token * len(tokens))(*tokens)
return llama_cpp.llama_save_session_file(self._ctx.ctx, filepath.encode('utf-8'), tokens, self.n_tokens)