09a8406c83
When streaming the yield that contains the finish can be skipped. This change ensures that yield isn't skipped.
1840 lines
68 KiB
Python
1840 lines
68 KiB
Python
import os
|
|
import sys
|
|
import uuid
|
|
import time
|
|
import math
|
|
import multiprocessing
|
|
from abc import ABC, abstractmethod
|
|
from typing import (
|
|
List,
|
|
Optional,
|
|
Union,
|
|
Generator,
|
|
Sequence,
|
|
Iterator,
|
|
Deque,
|
|
Tuple,
|
|
Callable,
|
|
)
|
|
from collections import deque, OrderedDict
|
|
|
|
import diskcache
|
|
import ctypes
|
|
|
|
from . import llama_cpp
|
|
from .llama_types import *
|
|
from .llama_grammar import LlamaGrammar
|
|
from . import llama_chat_format
|
|
|
|
import numpy as np
|
|
import numpy.typing as npt
|
|
|
|
from .utils import suppress_stdout_stderr
|
|
|
|
|
|
class BaseLlamaCache(ABC):
|
|
"""Base cache class for a llama.cpp model."""
|
|
|
|
def __init__(self, capacity_bytes: int = (2 << 30)):
|
|
self.capacity_bytes = capacity_bytes
|
|
|
|
@property
|
|
@abstractmethod
|
|
def cache_size(self) -> int:
|
|
raise NotImplementedError
|
|
|
|
def _find_longest_prefix_key(
|
|
self,
|
|
key: Tuple[int, ...],
|
|
) -> Optional[Tuple[int, ...]]:
|
|
pass
|
|
|
|
@abstractmethod
|
|
def __getitem__(self, key: Sequence[int]) -> "LlamaState":
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def __contains__(self, key: Sequence[int]) -> bool:
|
|
raise NotImplementedError
|
|
|
|
@abstractmethod
|
|
def __setitem__(self, key: Sequence[int], value: "LlamaState") -> None:
|
|
raise NotImplementedError
|
|
|
|
|
|
class LlamaRAMCache(BaseLlamaCache):
|
|
"""Cache for a llama.cpp model using RAM."""
|
|
|
|
def __init__(self, capacity_bytes: int = (2 << 30)):
|
|
super().__init__(capacity_bytes)
|
|
self.capacity_bytes = capacity_bytes
|
|
self.cache_state: OrderedDict[Tuple[int, ...], "LlamaState"] = OrderedDict()
|
|
|
|
@property
|
|
def cache_size(self):
|
|
return sum([state.llama_state_size for state in self.cache_state.values()])
|
|
|
|
def _find_longest_prefix_key(
|
|
self,
|
|
key: Tuple[int, ...],
|
|
) -> Optional[Tuple[int, ...]]:
|
|
min_len = 0
|
|
min_key = None
|
|
keys = (
|
|
(k, Llama.longest_token_prefix(k, key)) for k in self.cache_state.keys()
|
|
)
|
|
for k, prefix_len in keys:
|
|
if prefix_len > min_len:
|
|
min_len = prefix_len
|
|
min_key = k
|
|
return min_key
|
|
|
|
def __getitem__(self, key: Sequence[int]) -> "LlamaState":
|
|
key = tuple(key)
|
|
_key = self._find_longest_prefix_key(key)
|
|
if _key is None:
|
|
raise KeyError("Key not found")
|
|
value = self.cache_state[_key]
|
|
self.cache_state.move_to_end(_key)
|
|
return value
|
|
|
|
def __contains__(self, key: Sequence[int]) -> bool:
|
|
return self._find_longest_prefix_key(tuple(key)) is not None
|
|
|
|
def __setitem__(self, key: Sequence[int], value: "LlamaState"):
|
|
key = tuple(key)
|
|
if key in self.cache_state:
|
|
del self.cache_state[key]
|
|
self.cache_state[key] = value
|
|
while self.cache_size > self.capacity_bytes and len(self.cache_state) > 0:
|
|
self.cache_state.popitem(last=False)
|
|
|
|
|
|
# Alias for backwards compatibility
|
|
LlamaCache = LlamaRAMCache
|
|
|
|
|
|
class LlamaDiskCache(BaseLlamaCache):
|
|
"""Cache for a llama.cpp model using disk."""
|
|
|
|
def __init__(
|
|
self, cache_dir: str = ".cache/llama_cache", capacity_bytes: int = (2 << 30)
|
|
):
|
|
super().__init__(capacity_bytes)
|
|
self.cache = diskcache.Cache(cache_dir)
|
|
|
|
@property
|
|
def cache_size(self):
|
|
return int(self.cache.volume()) # type: ignore
|
|
|
|
def _find_longest_prefix_key(
|
|
self,
|
|
key: Tuple[int, ...],
|
|
) -> Optional[Tuple[int, ...]]:
|
|
min_len = 0
|
|
min_key: Optional[Tuple[int, ...]] = None
|
|
for k in self.cache.iterkeys(): # type: ignore
|
|
prefix_len = Llama.longest_token_prefix(k, key)
|
|
if prefix_len > min_len:
|
|
min_len = prefix_len
|
|
min_key = k # type: ignore
|
|
return min_key
|
|
|
|
def __getitem__(self, key: Sequence[int]) -> "LlamaState":
|
|
key = tuple(key)
|
|
_key = self._find_longest_prefix_key(key)
|
|
if _key is None:
|
|
raise KeyError("Key not found")
|
|
value: "LlamaState" = self.cache.pop(_key) # type: ignore
|
|
# NOTE: This puts an integer as key in cache, which breaks,
|
|
# Llama.longest_token_prefix(k, key) above since k is not a tuple of ints/tokens
|
|
# self.cache.push(_key, side="front") # type: ignore
|
|
return value
|
|
|
|
def __contains__(self, key: Sequence[int]) -> bool:
|
|
return self._find_longest_prefix_key(tuple(key)) is not None
|
|
|
|
def __setitem__(self, key: Sequence[int], value: "LlamaState"):
|
|
print("LlamaDiskCache.__setitem__: called", file=sys.stderr)
|
|
key = tuple(key)
|
|
if key in self.cache:
|
|
print("LlamaDiskCache.__setitem__: delete", file=sys.stderr)
|
|
del self.cache[key]
|
|
self.cache[key] = value
|
|
print("LlamaDiskCache.__setitem__: set", file=sys.stderr)
|
|
while self.cache_size > self.capacity_bytes and len(self.cache) > 0:
|
|
key_to_remove = next(iter(self.cache))
|
|
del self.cache[key_to_remove]
|
|
print("LlamaDiskCache.__setitem__: trim", file=sys.stderr)
|
|
|
|
|
|
class LlamaState:
|
|
def __init__(
|
|
self,
|
|
input_ids: npt.NDArray[np.intc],
|
|
scores: npt.NDArray[np.single],
|
|
n_tokens: int,
|
|
llama_state: bytes,
|
|
llama_state_size: int,
|
|
):
|
|
self.input_ids = input_ids
|
|
self.scores = scores
|
|
self.n_tokens = n_tokens
|
|
self.llama_state = llama_state
|
|
self.llama_state_size = llama_state_size
|
|
|
|
|
|
LogitsProcessor = Callable[
|
|
[npt.NDArray[np.intc], npt.NDArray[np.single]], npt.NDArray[np.single]
|
|
]
|
|
|
|
|
|
class LogitsProcessorList(List[LogitsProcessor]):
|
|
def __call__(
|
|
self, input_ids: npt.NDArray[np.intc], scores: npt.NDArray[np.single]
|
|
) -> npt.NDArray[np.single]:
|
|
for processor in self:
|
|
scores = processor(input_ids, scores)
|
|
return scores
|
|
|
|
|
|
StoppingCriteria = Callable[[npt.NDArray[np.intc], npt.NDArray[np.single]], bool]
|
|
|
|
|
|
class StoppingCriteriaList(List[StoppingCriteria]):
|
|
def __call__(
|
|
self, input_ids: npt.NDArray[np.intc], logits: npt.NDArray[np.single]
|
|
) -> bool:
|
|
return any([stopping_criteria(input_ids, logits) for stopping_criteria in self])
|
|
|
|
|
|
class Llama:
|
|
"""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,
|
|
main_gpu: int = 0,
|
|
tensor_split: Optional[List[float]] = None,
|
|
vocab_only: bool = False,
|
|
use_mmap: bool = True,
|
|
use_mlock: bool = False,
|
|
# 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_freq_base: float = 0.0,
|
|
rope_freq_scale: float = 0.0,
|
|
mul_mat_q: bool = True,
|
|
f16_kv: bool = True,
|
|
logits_all: bool = False,
|
|
embedding: 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: bool = False,
|
|
# Chat Format Params
|
|
chat_format: str = "llama-2",
|
|
# Misc
|
|
verbose: bool = True,
|
|
# Extra Params
|
|
**kwargs, # type: ignore
|
|
):
|
|
"""Load a llama.cpp model from `model_path`.
|
|
|
|
Args:
|
|
model_path: Path to the model.
|
|
seed: Random seed. -1 for random.
|
|
n_ctx: Maximum context size.
|
|
n_batch: Maximum number of prompt tokens to batch together when calling llama_eval.
|
|
n_gpu_layers: Number of layers to offload to GPU (-ngl). If -1, all layers are offloaded.
|
|
main_gpu: Main GPU to use.
|
|
tensor_split: Optional list of floats to split the model across multiple GPUs. If None, the model is not split.
|
|
rope_freq_base: Base frequency for rope sampling.
|
|
rope_freq_scale: Scale factor for rope sampling.
|
|
low_vram: Use low VRAM mode.
|
|
mul_mat_q: if true, use experimental mul_mat_q kernels
|
|
f16_kv: Use half-precision for key/value cache.
|
|
logits_all: Return logits for all tokens, not just the last token.
|
|
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.
|
|
embedding: Embedding mode only.
|
|
n_threads: Number of threads to use. If None, the number of threads is automatically determined.
|
|
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: Enable NUMA support. (NOTE: The initial value of this parameter is used for the remainder of the program as this value is set in llama_backend_init)
|
|
chat_format: String specifying the chat format to use when calling create_chat_completion.
|
|
verbose: Print verbose output to stderr.
|
|
kwargs: Unused keyword arguments (for additional backwards compatibility).
|
|
|
|
Raises:
|
|
ValueError: If the model path does not exist.
|
|
|
|
Returns:
|
|
A Llama instance.
|
|
"""
|
|
self.verbose = verbose
|
|
|
|
self.numa = numa
|
|
if not Llama.__backend_initialized:
|
|
if self.verbose:
|
|
llama_cpp.llama_backend_init(self.numa)
|
|
else:
|
|
with suppress_stdout_stderr():
|
|
llama_cpp.llama_backend_init(self.numa)
|
|
Llama.__backend_initialized = True
|
|
|
|
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.main_gpu = main_gpu
|
|
self.tensor_split = tensor_split
|
|
self._p_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
|
|
|
|
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_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.mul_mat_q = mul_mat_q
|
|
self.context_params.f16_kv = f16_kv
|
|
self.context_params.logits_all = logits_all
|
|
self.context_params.embedding = embedding
|
|
|
|
# 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}")
|
|
|
|
if verbose:
|
|
self.model = llama_cpp.llama_load_model_from_file(
|
|
self.model_path.encode("utf-8"), self.model_params
|
|
)
|
|
else:
|
|
with suppress_stdout_stderr():
|
|
self.model = llama_cpp.llama_load_model_from_file(
|
|
self.model_path.encode("utf-8"), self.model_params
|
|
)
|
|
assert self.model is not None
|
|
|
|
if verbose:
|
|
self.ctx = llama_cpp.llama_new_context_with_model(
|
|
self.model, self.context_params
|
|
)
|
|
else:
|
|
with suppress_stdout_stderr():
|
|
self.ctx = llama_cpp.llama_new_context_with_model(
|
|
self.model, self.context_params
|
|
)
|
|
|
|
assert self.ctx is not None
|
|
|
|
if self.lora_path:
|
|
if llama_cpp.llama_model_apply_lora_from_file(
|
|
self.model,
|
|
self.lora_path.encode("utf-8"),
|
|
self.lora_scale,
|
|
self.lora_base.encode("utf-8")
|
|
if self.lora_base is not None
|
|
else llama_cpp.c_char_p(0),
|
|
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._n_vocab = self.n_vocab()
|
|
self._n_ctx = self.n_ctx()
|
|
size = self._n_vocab
|
|
sorted = False
|
|
self._candidates_data = np.array(
|
|
[],
|
|
dtype=np.dtype(
|
|
[("id", np.intc), ("logit", np.single), ("p", np.single)], align=True
|
|
),
|
|
)
|
|
self._candidates_data.resize(3, self._n_vocab, refcheck=False)
|
|
candidates = llama_cpp.llama_token_data_array(
|
|
data=self._candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p),
|
|
size=size,
|
|
sorted=sorted,
|
|
)
|
|
self._candidates = candidates
|
|
self._token_nl = self.token_nl()
|
|
self._token_eos = self.token_eos()
|
|
self._candidates_data_id = np.arange(self._n_vocab, dtype=np.intc) # type: ignore
|
|
self._candidates_data_p = np.zeros(self._n_vocab, dtype=np.single)
|
|
|
|
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
|
|
)
|
|
|
|
@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.
|
|
|
|
Args:
|
|
text: The utf-8 encoded string to tokenize.
|
|
|
|
Raises:
|
|
RuntimeError: If the tokenization failed.
|
|
|
|
Returns:
|
|
A list of tokens.
|
|
"""
|
|
assert self.model is not None
|
|
n_ctx = self._n_ctx
|
|
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 detokenize(self, tokens: List[int]) -> bytes:
|
|
"""Detokenize a list of tokens.
|
|
|
|
Args:
|
|
tokens: The list of tokens to detokenize.
|
|
|
|
Returns:
|
|
The detokenized string.
|
|
"""
|
|
assert self.model is not None
|
|
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
|
|
)
|
|
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() else output
|
|
)
|
|
|
|
def set_cache(self, cache: Optional[BaseLlamaCache]):
|
|
"""Set the cache.
|
|
|
|
Args:
|
|
cache: The cache to set.
|
|
"""
|
|
self.cache = cache
|
|
|
|
def reset(self):
|
|
"""Reset the model state."""
|
|
self.n_tokens = 0
|
|
|
|
def eval(self, tokens: Sequence[int]):
|
|
"""Evaluate a list of tokens.
|
|
|
|
Args:
|
|
tokens: The list of tokens to evaluate.
|
|
"""
|
|
assert self.ctx is not None
|
|
n_ctx = self._n_ctx
|
|
for i in range(0, len(tokens), self.n_batch):
|
|
batch = tokens[i : min(len(tokens), i + self.n_batch)]
|
|
n_past = min(n_ctx - len(batch), len(self._input_ids))
|
|
n_tokens = len(batch)
|
|
return_code = llama_cpp.llama_eval(
|
|
ctx=self.ctx,
|
|
tokens=(llama_cpp.llama_token * len(batch))(*batch),
|
|
n_tokens=n_tokens,
|
|
n_past=n_past,
|
|
)
|
|
if return_code != 0:
|
|
raise RuntimeError(f"llama_eval returned {return_code}")
|
|
# Save tokens
|
|
self.input_ids[self.n_tokens : self.n_tokens + n_tokens] = batch
|
|
# Save logits
|
|
rows = n_tokens if self.context_params.logits_all else 1
|
|
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[self.n_tokens + offset : self.n_tokens + n_tokens, :].reshape(
|
|
-1
|
|
)[:] = llama_cpp.llama_get_logits(self.ctx)[: rows * cols]
|
|
# Update n_tokens
|
|
self.n_tokens += n_tokens
|
|
|
|
def _sample(
|
|
self,
|
|
last_n_tokens_data, # type: llama_cpp.Array[llama_cpp.llama_token]
|
|
last_n_tokens_size: int,
|
|
top_k: int,
|
|
top_p: float,
|
|
temp: float,
|
|
tfs_z: float,
|
|
repeat_penalty: float,
|
|
frequency_penalty: float,
|
|
presence_penalty: float,
|
|
mirostat_mode: float,
|
|
mirostat_tau: float,
|
|
mirostat_eta: float,
|
|
penalize_nl: bool = True,
|
|
logits_processor: Optional[LogitsProcessorList] = None,
|
|
grammar: Optional[LlamaGrammar] = None,
|
|
):
|
|
assert self.ctx is not None
|
|
assert self.n_tokens > 0
|
|
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
|
|
logits: npt.NDArray[np.single] = self._scores[-1, :]
|
|
|
|
if logits_processor is not None:
|
|
logits[:] = logits_processor(self._input_ids, logits)
|
|
|
|
nl_logit = logits[self._token_nl]
|
|
candidates = self._candidates
|
|
candidates_data = self._candidates_data
|
|
candidates_data["id"][:] = self._candidates_data_id # type: ignore
|
|
candidates_data["logit"][:] = logits
|
|
candidates_data["p"][:] = self._candidates_data_p # type: ignore
|
|
candidates.data = candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p)
|
|
candidates.sorted = llama_cpp.c_bool(False)
|
|
candidates.size = llama_cpp.c_size_t(n_vocab)
|
|
llama_cpp.llama_sample_repetition_penalty(
|
|
ctx=self.ctx,
|
|
last_tokens_data=last_n_tokens_data,
|
|
last_tokens_size=last_n_tokens_size,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
penalty=repeat_penalty,
|
|
)
|
|
llama_cpp.llama_sample_frequency_and_presence_penalties(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
last_tokens_data=last_n_tokens_data,
|
|
last_tokens_size=last_n_tokens_size,
|
|
alpha_frequency=frequency_penalty,
|
|
alpha_presence=presence_penalty,
|
|
)
|
|
if not penalize_nl:
|
|
candidates.data[self._token_nl].logit = llama_cpp.c_float(nl_logit)
|
|
|
|
if grammar is not None:
|
|
llama_cpp.llama_sample_grammar(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
grammar=grammar.grammar,
|
|
)
|
|
|
|
if temp == 0.0:
|
|
id = llama_cpp.llama_sample_token_greedy(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
)
|
|
elif mirostat_mode == 1:
|
|
mirostat_mu = llama_cpp.c_float(2.0 * mirostat_tau)
|
|
mirostat_m = llama_cpp.c_int(100)
|
|
llama_cpp.llama_sample_temperature(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
temp=temp,
|
|
)
|
|
id = llama_cpp.llama_sample_token_mirostat(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
tau=mirostat_tau,
|
|
eta=mirostat_eta,
|
|
mu=llama_cpp.ctypes.byref(mirostat_mu), # type: ignore
|
|
m=mirostat_m,
|
|
)
|
|
elif mirostat_mode == 2:
|
|
mirostat_mu = llama_cpp.c_float(2.0 * mirostat_tau)
|
|
llama_cpp.llama_sample_temperature(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
temp=temp,
|
|
)
|
|
id = llama_cpp.llama_sample_token_mirostat_v2(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
tau=mirostat_tau,
|
|
eta=mirostat_eta,
|
|
mu=llama_cpp.ctypes.byref(mirostat_mu), # type: ignore
|
|
)
|
|
else:
|
|
llama_cpp.llama_sample_top_k(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
k=top_k,
|
|
min_keep=llama_cpp.c_size_t(1),
|
|
)
|
|
llama_cpp.llama_sample_tail_free(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
z=tfs_z,
|
|
min_keep=llama_cpp.c_size_t(1),
|
|
)
|
|
llama_cpp.llama_sample_typical(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
p=llama_cpp.c_float(1.0),
|
|
min_keep=llama_cpp.c_size_t(1),
|
|
)
|
|
llama_cpp.llama_sample_top_p(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
p=top_p,
|
|
min_keep=llama_cpp.c_size_t(1),
|
|
)
|
|
llama_cpp.llama_sample_temperature(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
temp=temp,
|
|
)
|
|
id = llama_cpp.llama_sample_token(
|
|
ctx=self.ctx,
|
|
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
|
|
)
|
|
if grammar is not None:
|
|
llama_cpp.llama_grammar_accept_token(
|
|
ctx=self.ctx,
|
|
grammar=grammar.grammar,
|
|
token=llama_cpp.ctypes.c_int(id),
|
|
)
|
|
return id
|
|
|
|
def sample(
|
|
self,
|
|
top_k: int = 40,
|
|
top_p: float = 0.95,
|
|
temp: float = 0.80,
|
|
repeat_penalty: float = 1.1,
|
|
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,
|
|
):
|
|
"""Sample a token from the model.
|
|
|
|
Args:
|
|
top_k: The top-k sampling parameter.
|
|
top_p: The top-p sampling parameter.
|
|
temp: The temperature parameter.
|
|
repeat_penalty: The repeat penalty parameter.
|
|
|
|
Returns:
|
|
The sampled token.
|
|
"""
|
|
assert self.ctx is not None
|
|
last_n_tokens_data = [llama_cpp.llama_token(0)] * max(
|
|
0, self.last_n_tokens_size - len(self._input_ids)
|
|
) + self._input_ids[-self.last_n_tokens_size :].tolist()
|
|
return self._sample(
|
|
last_n_tokens_data=(llama_cpp.llama_token * self.last_n_tokens_size)(
|
|
*last_n_tokens_data
|
|
),
|
|
last_n_tokens_size=self.last_n_tokens_size,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
temp=temp,
|
|
tfs_z=tfs_z,
|
|
repeat_penalty=repeat_penalty,
|
|
frequency_penalty=frequency_penalty,
|
|
presence_penalty=presence_penalty,
|
|
mirostat_mode=mirostat_mode,
|
|
mirostat_tau=mirostat_tau,
|
|
mirostat_eta=mirostat_eta,
|
|
penalize_nl=penalize_nl,
|
|
logits_processor=logits_processor,
|
|
grammar=grammar,
|
|
)
|
|
|
|
def generate(
|
|
self,
|
|
tokens: Sequence[int],
|
|
top_k: int = 40,
|
|
top_p: float = 0.95,
|
|
temp: float = 0.80,
|
|
repeat_penalty: float = 1.1,
|
|
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,
|
|
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.
|
|
|
|
Examples:
|
|
>>> llama = Llama("models/ggml-7b.bin")
|
|
>>> tokens = llama.tokenize(b"Hello, world!")
|
|
>>> for token in llama.generate(tokens, top_k=40, top_p=0.95, temp=1.0, repeat_penalty=1.1):
|
|
... print(llama.detokenize([token]))
|
|
|
|
Args:
|
|
tokens: The prompt tokens.
|
|
top_k: The top-k sampling parameter.
|
|
top_p: The top-p sampling parameter.
|
|
temp: The temperature parameter.
|
|
repeat_penalty: The repeat penalty parameter.
|
|
reset: Whether to reset the model state.
|
|
|
|
Yields:
|
|
The generated tokens.
|
|
"""
|
|
assert self.ctx is not None
|
|
if reset and len(self._input_ids) > 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:
|
|
if self.verbose:
|
|
print("Llama.generate: prefix-match hit", file=sys.stderr)
|
|
reset = False
|
|
tokens = tokens[longest_prefix:]
|
|
self.n_tokens = longest_prefix
|
|
|
|
if reset:
|
|
self.reset()
|
|
|
|
if grammar is not None:
|
|
grammar.reset()
|
|
|
|
while True:
|
|
self.eval(tokens)
|
|
token = self.sample(
|
|
top_k=top_k,
|
|
top_p=top_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,
|
|
)
|
|
if stopping_criteria is not None and stopping_criteria(
|
|
self._input_ids, self._scores[-1, :]
|
|
):
|
|
return
|
|
tokens_or_none = yield token
|
|
tokens = [token]
|
|
if tokens_or_none is not None:
|
|
tokens.extend(tokens_or_none)
|
|
|
|
def create_embedding(
|
|
self, input: Union[str, List[str]], model: Optional[str] = None
|
|
) -> CreateEmbeddingResponse:
|
|
"""Embed a string.
|
|
|
|
Args:
|
|
input: The utf-8 encoded string to embed.
|
|
|
|
Returns:
|
|
An embedding object.
|
|
"""
|
|
assert self.ctx is not None
|
|
assert self.model is not None
|
|
model_name: str = model if model is not None else self.model_path
|
|
|
|
if self.context_params.embedding == False:
|
|
raise RuntimeError(
|
|
"Llama model must be created with embedding=True to call this method"
|
|
)
|
|
|
|
if self.verbose:
|
|
llama_cpp.llama_reset_timings(self.ctx)
|
|
|
|
if isinstance(input, str):
|
|
inputs = [input]
|
|
else:
|
|
inputs = input
|
|
|
|
data: List[EmbeddingData] = []
|
|
total_tokens = 0
|
|
for index, input in enumerate(inputs):
|
|
tokens = self.tokenize(input.encode("utf-8"))
|
|
self.reset()
|
|
self.eval(tokens)
|
|
n_tokens = len(tokens)
|
|
total_tokens += n_tokens
|
|
embedding = llama_cpp.llama_get_embeddings(self.ctx)[
|
|
: llama_cpp.llama_n_embd(self.model)
|
|
]
|
|
|
|
data.append(
|
|
{
|
|
"object": "embedding",
|
|
"embedding": embedding,
|
|
"index": index,
|
|
}
|
|
)
|
|
if self.verbose:
|
|
llama_cpp.llama_print_timings(self.ctx)
|
|
|
|
return {
|
|
"object": "list",
|
|
"data": data,
|
|
"model": model_name,
|
|
"usage": {
|
|
"prompt_tokens": total_tokens,
|
|
"total_tokens": total_tokens,
|
|
},
|
|
}
|
|
|
|
def embed(self, input: str) -> List[float]:
|
|
"""Embed a string.
|
|
|
|
Args:
|
|
input: The utf-8 encoded string to embed.
|
|
|
|
Returns:
|
|
A list of embeddings
|
|
"""
|
|
return list(map(float, self.create_embedding(input)["data"][0]["embedding"]))
|
|
|
|
def _create_completion(
|
|
self,
|
|
prompt: str,
|
|
suffix: Optional[str] = None,
|
|
max_tokens: int = 16,
|
|
temperature: float = 0.8,
|
|
top_p: float = 0.95,
|
|
logprobs: Optional[int] = None,
|
|
echo: bool = False,
|
|
stop: Optional[Union[str, List[str]]] = [],
|
|
frequency_penalty: float = 0.0,
|
|
presence_penalty: float = 0.0,
|
|
repeat_penalty: float = 1.1,
|
|
top_k: int = 40,
|
|
stream: bool = False,
|
|
tfs_z: float = 1.0,
|
|
mirostat_mode: int = 0,
|
|
mirostat_tau: float = 5.0,
|
|
mirostat_eta: float = 0.1,
|
|
model: Optional[str] = None,
|
|
stopping_criteria: Optional[StoppingCriteriaList] = None,
|
|
logits_processor: Optional[LogitsProcessorList] = None,
|
|
grammar: Optional[LlamaGrammar] = None,
|
|
) -> Union[Iterator[Completion], Iterator[CompletionChunk]]:
|
|
assert self.ctx is not None
|
|
|
|
completion_id: str = f"cmpl-{str(uuid.uuid4())}"
|
|
created: int = int(time.time())
|
|
completion_tokens: List[int] = []
|
|
# Add blank space to start of prompt to match OG llama tokenizer
|
|
prompt_tokens: List[int] = (
|
|
self.tokenize(prompt.encode("utf-8"))
|
|
if prompt != ""
|
|
else [self.token_bos()]
|
|
)
|
|
text: bytes = b""
|
|
returned_tokens: int = 0
|
|
stop = (
|
|
stop if isinstance(stop, list) else [stop] if isinstance(stop, str) else []
|
|
)
|
|
model_name: str = model if model is not None else self.model_path
|
|
|
|
if self.verbose:
|
|
llama_cpp.llama_reset_timings(self.ctx)
|
|
|
|
if len(prompt_tokens) >= llama_cpp.llama_n_ctx(self.ctx):
|
|
raise ValueError(
|
|
f"Requested tokens ({len(prompt_tokens)}) exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}"
|
|
)
|
|
|
|
if max_tokens <= 0:
|
|
# Unlimited, depending on n_ctx.
|
|
max_tokens = llama_cpp.llama_n_ctx(self.ctx) - len(prompt_tokens)
|
|
|
|
# Truncate max_tokens if requested tokens would exceed the context window
|
|
max_tokens = (
|
|
max_tokens
|
|
if max_tokens + len(prompt_tokens) < self._n_ctx
|
|
else (self._n_ctx - len(prompt_tokens))
|
|
)
|
|
|
|
if stop != []:
|
|
stop_sequences = [s.encode("utf-8") for s in stop]
|
|
else:
|
|
stop_sequences = []
|
|
|
|
if logprobs is not None and self.context_params.logits_all is False:
|
|
raise ValueError(
|
|
"logprobs is not supported for models created with logits_all=False"
|
|
)
|
|
|
|
if self.cache:
|
|
try:
|
|
cache_item = self.cache[prompt_tokens]
|
|
cache_prefix_len = Llama.longest_token_prefix(
|
|
cache_item.input_ids.tolist(), prompt_tokens
|
|
)
|
|
eval_prefix_len = Llama.longest_token_prefix(
|
|
self._input_ids.tolist(), prompt_tokens
|
|
)
|
|
if cache_prefix_len > eval_prefix_len:
|
|
self.load_state(cache_item)
|
|
if self.verbose:
|
|
print("Llama._create_completion: cache hit", file=sys.stderr)
|
|
except KeyError:
|
|
if self.verbose:
|
|
print("Llama._create_completion: cache miss", file=sys.stderr)
|
|
|
|
finish_reason = "length"
|
|
multibyte_fix = 0
|
|
for token in self.generate(
|
|
prompt_tokens,
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
temp=temperature,
|
|
tfs_z=tfs_z,
|
|
mirostat_mode=mirostat_mode,
|
|
mirostat_tau=mirostat_tau,
|
|
mirostat_eta=mirostat_eta,
|
|
frequency_penalty=frequency_penalty,
|
|
presence_penalty=presence_penalty,
|
|
repeat_penalty=repeat_penalty,
|
|
stopping_criteria=stopping_criteria,
|
|
logits_processor=logits_processor,
|
|
grammar=grammar,
|
|
):
|
|
if token == self._token_eos:
|
|
text = self.detokenize(completion_tokens)
|
|
finish_reason = "stop"
|
|
break
|
|
|
|
completion_tokens.append(token)
|
|
|
|
all_text = self.detokenize(completion_tokens)
|
|
|
|
# Contains multi-byte UTF8
|
|
for k, char in enumerate(all_text[-3:]):
|
|
k = 3 - k
|
|
for num, pattern in [(2, 192), (3, 224), (4, 240)]:
|
|
# Bitwise AND check
|
|
if num > k and pattern & char == pattern:
|
|
multibyte_fix = num - k
|
|
|
|
# Stop incomplete bytes from passing
|
|
if multibyte_fix > 0:
|
|
multibyte_fix -= 1
|
|
continue
|
|
|
|
any_stop = [s for s in stop_sequences if s in all_text]
|
|
if len(any_stop) > 0:
|
|
first_stop = any_stop[0]
|
|
text = all_text[: all_text.index(first_stop)]
|
|
finish_reason = "stop"
|
|
break
|
|
|
|
if stream:
|
|
remaining_tokens = completion_tokens[returned_tokens:]
|
|
remaining_text = self.detokenize(remaining_tokens)
|
|
remaining_length = len(remaining_text)
|
|
|
|
# We want to avoid yielding any characters from
|
|
# the generated text if they are part of a stop
|
|
# sequence.
|
|
first_stop_position = 0
|
|
for s in stop_sequences:
|
|
for i in range(min(len(s), remaining_length), 0, -1):
|
|
if remaining_text.endswith(s[:i]):
|
|
if i > first_stop_position:
|
|
first_stop_position = i
|
|
break
|
|
|
|
token_end_position = 0
|
|
|
|
if logprobs is not None:
|
|
# not sure how to handle this branch when dealing
|
|
# with CJK output, so keep it unchanged
|
|
for token in remaining_tokens:
|
|
token_end_position += len(self.detokenize([token]))
|
|
# Check if stop sequence is in the token
|
|
if token_end_position > (
|
|
remaining_length - first_stop_position
|
|
):
|
|
break
|
|
token_str = self.detokenize([token]).decode(
|
|
"utf-8", errors="ignore"
|
|
)
|
|
text_offset = len(prompt) + len(
|
|
self.detokenize(completion_tokens[:returned_tokens])
|
|
)
|
|
token_offset = len(prompt_tokens) + returned_tokens
|
|
logits = self._scores[token_offset - 1, :].tolist()
|
|
current_logprobs = Llama.logits_to_logprobs(logits)
|
|
sorted_logprobs = list(
|
|
sorted(
|
|
zip(current_logprobs, range(len(current_logprobs))),
|
|
reverse=True,
|
|
)
|
|
)
|
|
top_logprob = {
|
|
self.detokenize([i]).decode(
|
|
"utf-8", errors="ignore"
|
|
): logprob
|
|
for logprob, i in sorted_logprobs[:logprobs]
|
|
}
|
|
top_logprob.update({token_str: current_logprobs[int(token)]})
|
|
logprobs_or_none = {
|
|
"tokens": [
|
|
self.detokenize([token]).decode(
|
|
"utf-8", errors="ignore"
|
|
)
|
|
],
|
|
"text_offset": [text_offset],
|
|
"token_logprobs": [current_logprobs[int(token)]],
|
|
"top_logprobs": [top_logprob],
|
|
}
|
|
returned_tokens += 1
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"text": self.detokenize([token]).decode(
|
|
"utf-8", errors="ignore"
|
|
),
|
|
"index": 0,
|
|
"logprobs": logprobs_or_none,
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
else:
|
|
while len(remaining_tokens) > 0:
|
|
decode_success = False
|
|
for i in range(1, len(remaining_tokens) + 1):
|
|
try:
|
|
bs = self.detokenize(remaining_tokens[:i])
|
|
ts = bs.decode("utf-8")
|
|
decode_success = True
|
|
break
|
|
except UnicodeError:
|
|
pass
|
|
else:
|
|
break
|
|
if not decode_success:
|
|
# all remaining tokens cannot be decoded to a UTF-8 character
|
|
break
|
|
token_end_position += len(bs)
|
|
if token_end_position > (
|
|
remaining_length - first_stop_position
|
|
):
|
|
break
|
|
remaining_tokens = remaining_tokens[i:]
|
|
returned_tokens += i
|
|
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"text": ts,
|
|
"index": 0,
|
|
"logprobs": None,
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
|
|
if len(completion_tokens) >= max_tokens:
|
|
text = self.detokenize(completion_tokens)
|
|
finish_reason = "length"
|
|
break
|
|
|
|
if stopping_criteria is not None and stopping_criteria(
|
|
self._input_ids, self._scores[-1, :]
|
|
):
|
|
text = self.detokenize(completion_tokens)
|
|
finish_reason = "stop"
|
|
|
|
if self.verbose:
|
|
llama_cpp.llama_print_timings(self.ctx)
|
|
|
|
if stream:
|
|
remaining_tokens = completion_tokens[returned_tokens:]
|
|
all_text = self.detokenize(remaining_tokens)
|
|
any_stop = [s for s in stop_sequences if s in all_text]
|
|
if len(any_stop) > 0:
|
|
end = min(all_text.index(stop) for stop in any_stop)
|
|
else:
|
|
end = len(all_text)
|
|
|
|
token_end_position = 0
|
|
for token in remaining_tokens:
|
|
token_end_position += len(self.detokenize([token]))
|
|
|
|
logprobs_or_none: Optional[CompletionLogprobs] = None
|
|
if logprobs is not None:
|
|
token_str = self.detokenize([token]).decode(
|
|
"utf-8", errors="ignore"
|
|
)
|
|
text_offset = len(prompt) + len(
|
|
self.detokenize(completion_tokens[:returned_tokens])
|
|
)
|
|
token_offset = len(prompt_tokens) + returned_tokens - 1
|
|
logits = self._scores[token_offset, :].tolist()
|
|
current_logprobs = Llama.logits_to_logprobs(logits)
|
|
sorted_logprobs = list(
|
|
sorted(
|
|
zip(current_logprobs, range(len(current_logprobs))),
|
|
reverse=True,
|
|
)
|
|
)
|
|
top_logprob = {
|
|
self.detokenize([i]).decode("utf-8", errors="ignore"): logprob
|
|
for logprob, i in sorted_logprobs[:logprobs]
|
|
}
|
|
top_logprob.update({token_str: current_logprobs[int(token)]})
|
|
logprobs_or_none = {
|
|
"tokens": [
|
|
self.detokenize([token]).decode("utf-8", errors="ignore")
|
|
],
|
|
"text_offset": [text_offset],
|
|
"token_logprobs": [current_logprobs[int(token)]],
|
|
"top_logprobs": [top_logprob],
|
|
}
|
|
|
|
if token_end_position >= end:
|
|
last_text = self.detokenize([token])
|
|
if token_end_position == end - 1:
|
|
break
|
|
returned_tokens += 1
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"text": last_text[
|
|
: len(last_text) - (token_end_position - end)
|
|
].decode("utf-8", errors="ignore"),
|
|
"index": 0,
|
|
"logprobs": logprobs_or_none,
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
break
|
|
returned_tokens += 1
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"text": self.detokenize([token]).decode(
|
|
"utf-8", errors="ignore"
|
|
),
|
|
"index": 0,
|
|
"logprobs": logprobs_or_none,
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"text": "",
|
|
"index": 0,
|
|
"logprobs": None,
|
|
"finish_reason": finish_reason,
|
|
}
|
|
],
|
|
}
|
|
if self.cache:
|
|
if self.verbose:
|
|
print("Llama._create_completion: cache save", file=sys.stderr)
|
|
self.cache[prompt_tokens + completion_tokens] = self.save_state()
|
|
print("Llama._create_completion: cache saved", file=sys.stderr)
|
|
return
|
|
|
|
if self.cache:
|
|
if self.verbose:
|
|
print("Llama._create_completion: cache save", file=sys.stderr)
|
|
self.cache[prompt_tokens + completion_tokens] = self.save_state()
|
|
|
|
text_str = text.decode("utf-8", errors="ignore")
|
|
|
|
if echo:
|
|
text_str = prompt + text_str
|
|
|
|
if suffix is not None:
|
|
text_str = text_str + suffix
|
|
|
|
logprobs_or_none: Optional[CompletionLogprobs] = None
|
|
if logprobs is not None:
|
|
text_offset = 0 if echo else len(prompt)
|
|
token_offset = 0 if echo else len(prompt_tokens[1:])
|
|
text_offsets: List[int] = []
|
|
token_logprobs: List[Optional[float]] = []
|
|
tokens: List[str] = []
|
|
top_logprobs: List[Optional[Dict[str, float]]] = []
|
|
|
|
if echo:
|
|
# Remove leading BOS token
|
|
all_tokens = prompt_tokens[1:] + completion_tokens
|
|
else:
|
|
all_tokens = completion_tokens
|
|
|
|
all_token_strs = [
|
|
self.detokenize([token]).decode("utf-8", errors="ignore")
|
|
for token in all_tokens
|
|
]
|
|
all_logprobs = [
|
|
Llama.logits_to_logprobs(row.tolist()) for row in self._scores
|
|
][token_offset:]
|
|
for token, token_str, logprobs_token in zip(
|
|
all_tokens, all_token_strs, all_logprobs
|
|
):
|
|
text_offsets.append(text_offset)
|
|
text_offset += len(token_str)
|
|
tokens.append(token_str)
|
|
sorted_logprobs = list(
|
|
sorted(
|
|
zip(logprobs_token, range(len(logprobs_token))), reverse=True
|
|
)
|
|
)
|
|
token_logprobs.append(logprobs_token[int(token)])
|
|
top_logprob: Optional[Dict[str, float]] = {
|
|
self.detokenize([i]).decode("utf-8", errors="ignore"): logprob
|
|
for logprob, i in sorted_logprobs[:logprobs]
|
|
}
|
|
top_logprob.update({token_str: logprobs_token[int(token)]})
|
|
top_logprobs.append(top_logprob)
|
|
# Weird idosincracy of the OpenAI API where
|
|
# token_logprobs and top_logprobs are null for
|
|
# the first token.
|
|
if echo and len(all_tokens) > 0:
|
|
token_logprobs[0] = None
|
|
top_logprobs[0] = None
|
|
logprobs_or_none = {
|
|
"tokens": tokens,
|
|
"text_offset": text_offsets,
|
|
"token_logprobs": token_logprobs,
|
|
"top_logprobs": top_logprobs,
|
|
}
|
|
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": model_name,
|
|
"choices": [
|
|
{
|
|
"text": text_str,
|
|
"index": 0,
|
|
"logprobs": logprobs_or_none,
|
|
"finish_reason": finish_reason,
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": len(prompt_tokens),
|
|
"completion_tokens": len(completion_tokens),
|
|
"total_tokens": len(prompt_tokens) + len(completion_tokens),
|
|
},
|
|
}
|
|
|
|
def create_completion(
|
|
self,
|
|
prompt: str,
|
|
suffix: Optional[str] = None,
|
|
max_tokens: int = 128,
|
|
temperature: float = 0.8,
|
|
top_p: float = 0.95,
|
|
logprobs: Optional[int] = None,
|
|
echo: bool = False,
|
|
stop: Optional[Union[str, List[str]]] = [],
|
|
frequency_penalty: float = 0.0,
|
|
presence_penalty: float = 0.0,
|
|
repeat_penalty: float = 1.1,
|
|
top_k: int = 40,
|
|
stream: bool = False,
|
|
tfs_z: float = 1.0,
|
|
mirostat_mode: int = 0,
|
|
mirostat_tau: float = 5.0,
|
|
mirostat_eta: float = 0.1,
|
|
model: Optional[str] = None,
|
|
stopping_criteria: Optional[StoppingCriteriaList] = None,
|
|
logits_processor: Optional[LogitsProcessorList] = None,
|
|
grammar: Optional[LlamaGrammar] = None,
|
|
) -> Union[Completion, Iterator[CompletionChunk]]:
|
|
"""Generate text from a prompt.
|
|
|
|
Args:
|
|
prompt: The prompt to generate text from.
|
|
suffix: A suffix to append to the generated text. If None, no suffix is appended.
|
|
max_tokens: The maximum number of tokens to generate. If max_tokens <= 0, the maximum number of tokens to generate is unlimited and depends on n_ctx.
|
|
temperature: The temperature to use for sampling.
|
|
top_p: The top-p value to use for sampling.
|
|
logprobs: The number of logprobs to return. If None, no logprobs are returned.
|
|
echo: Whether to echo the prompt.
|
|
stop: A list of strings to stop generation when encountered.
|
|
repeat_penalty: The penalty to apply to repeated tokens.
|
|
top_k: The top-k value to use for sampling.
|
|
stream: Whether to stream the results.
|
|
|
|
Raises:
|
|
ValueError: If the requested tokens exceed the context window.
|
|
RuntimeError: If the prompt fails to tokenize or the model fails to evaluate the prompt.
|
|
|
|
Returns:
|
|
Response object containing the generated text.
|
|
"""
|
|
completion_or_chunks = self._create_completion(
|
|
prompt=prompt,
|
|
suffix=suffix,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
logprobs=logprobs,
|
|
echo=echo,
|
|
stop=stop,
|
|
frequency_penalty=frequency_penalty,
|
|
presence_penalty=presence_penalty,
|
|
repeat_penalty=repeat_penalty,
|
|
top_k=top_k,
|
|
stream=stream,
|
|
tfs_z=tfs_z,
|
|
mirostat_mode=mirostat_mode,
|
|
mirostat_tau=mirostat_tau,
|
|
mirostat_eta=mirostat_eta,
|
|
model=model,
|
|
stopping_criteria=stopping_criteria,
|
|
logits_processor=logits_processor,
|
|
grammar=grammar,
|
|
)
|
|
if stream:
|
|
chunks: Iterator[CompletionChunk] = completion_or_chunks
|
|
return chunks
|
|
completion: Completion = next(completion_or_chunks) # type: ignore
|
|
return completion
|
|
|
|
def __call__(
|
|
self,
|
|
prompt: str,
|
|
suffix: Optional[str] = None,
|
|
max_tokens: int = 128,
|
|
temperature: float = 0.8,
|
|
top_p: float = 0.95,
|
|
logprobs: Optional[int] = None,
|
|
echo: bool = False,
|
|
stop: Optional[Union[str, List[str]]] = [],
|
|
frequency_penalty: float = 0.0,
|
|
presence_penalty: float = 0.0,
|
|
repeat_penalty: float = 1.1,
|
|
top_k: int = 40,
|
|
stream: bool = False,
|
|
tfs_z: float = 1.0,
|
|
mirostat_mode: int = 0,
|
|
mirostat_tau: float = 5.0,
|
|
mirostat_eta: float = 0.1,
|
|
model: Optional[str] = None,
|
|
stopping_criteria: Optional[StoppingCriteriaList] = None,
|
|
logits_processor: Optional[LogitsProcessorList] = None,
|
|
grammar: Optional[LlamaGrammar] = None,
|
|
) -> Union[Completion, Iterator[CompletionChunk]]:
|
|
"""Generate text from a prompt.
|
|
|
|
Args:
|
|
prompt: The prompt to generate text from.
|
|
suffix: A suffix to append to the generated text. If None, no suffix is appended.
|
|
max_tokens: The maximum number of tokens to generate. If max_tokens <= 0, the maximum number of tokens to generate is unlimited and depends on n_ctx.
|
|
temperature: The temperature to use for sampling.
|
|
top_p: The top-p value to use for sampling.
|
|
logprobs: The number of logprobs to return. If None, no logprobs are returned.
|
|
echo: Whether to echo the prompt.
|
|
stop: A list of strings to stop generation when encountered.
|
|
repeat_penalty: The penalty to apply to repeated tokens.
|
|
top_k: The top-k value to use for sampling.
|
|
stream: Whether to stream the results.
|
|
|
|
Raises:
|
|
ValueError: If the requested tokens exceed the context window.
|
|
RuntimeError: If the prompt fails to tokenize or the model fails to evaluate the prompt.
|
|
|
|
Returns:
|
|
Response object containing the generated text.
|
|
"""
|
|
return self.create_completion(
|
|
prompt=prompt,
|
|
suffix=suffix,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
logprobs=logprobs,
|
|
echo=echo,
|
|
stop=stop,
|
|
frequency_penalty=frequency_penalty,
|
|
presence_penalty=presence_penalty,
|
|
repeat_penalty=repeat_penalty,
|
|
top_k=top_k,
|
|
stream=stream,
|
|
tfs_z=tfs_z,
|
|
mirostat_mode=mirostat_mode,
|
|
mirostat_tau=mirostat_tau,
|
|
mirostat_eta=mirostat_eta,
|
|
model=model,
|
|
stopping_criteria=stopping_criteria,
|
|
logits_processor=logits_processor,
|
|
grammar=grammar,
|
|
)
|
|
|
|
def _convert_text_completion_to_chat(
|
|
self, completion: Completion
|
|
) -> ChatCompletion:
|
|
return {
|
|
"id": "chat" + completion["id"],
|
|
"object": "chat.completion",
|
|
"created": completion["created"],
|
|
"model": completion["model"],
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"message": {
|
|
"role": "assistant",
|
|
"content": completion["choices"][0]["text"],
|
|
},
|
|
"finish_reason": completion["choices"][0]["finish_reason"],
|
|
}
|
|
],
|
|
"usage": completion["usage"],
|
|
}
|
|
|
|
def _convert_text_completion_chunks_to_chat(
|
|
self,
|
|
chunks: Iterator[CompletionChunk],
|
|
) -> Iterator[ChatCompletionChunk]:
|
|
for i, chunk in enumerate(chunks):
|
|
if i == 0:
|
|
yield {
|
|
"id": "chat" + chunk["id"],
|
|
"model": chunk["model"],
|
|
"created": chunk["created"],
|
|
"object": "chat.completion.chunk",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {
|
|
"role": "assistant",
|
|
},
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
yield {
|
|
"id": "chat" + chunk["id"],
|
|
"model": chunk["model"],
|
|
"created": chunk["created"],
|
|
"object": "chat.completion.chunk",
|
|
"choices": [
|
|
{
|
|
"index": 0,
|
|
"delta": {
|
|
"content": chunk["choices"][0]["text"],
|
|
}
|
|
if chunk["choices"][0]["finish_reason"] is None
|
|
else {},
|
|
"finish_reason": chunk["choices"][0]["finish_reason"],
|
|
}
|
|
],
|
|
}
|
|
|
|
def _convert_completion_to_chat(
|
|
self,
|
|
completion_or_chunks: Union[Completion, Iterator[CompletionChunk]],
|
|
stream: bool = False,
|
|
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
|
|
if stream:
|
|
chunks: Iterator[CompletionChunk] = completion_or_chunks # type: ignore
|
|
return self._convert_text_completion_chunks_to_chat(chunks)
|
|
else:
|
|
completion: Completion = completion_or_chunks # type: ignore
|
|
return self._convert_text_completion_to_chat(completion)
|
|
|
|
def create_chat_completion(
|
|
self,
|
|
messages: List[ChatCompletionRequestMessage],
|
|
functions: Optional[List[ChatCompletionFunction]] = None,
|
|
function_call: Optional[Union[str, ChatCompletionFunctionCall]] = None,
|
|
temperature: float = 0.2,
|
|
top_p: float = 0.95,
|
|
top_k: int = 40,
|
|
stream: bool = False,
|
|
stop: Optional[Union[str, List[str]]] = [],
|
|
max_tokens: int = 256,
|
|
presence_penalty: float = 0.0,
|
|
frequency_penalty: float = 0.0,
|
|
repeat_penalty: float = 1.1,
|
|
tfs_z: float = 1.0,
|
|
mirostat_mode: int = 0,
|
|
mirostat_tau: float = 5.0,
|
|
mirostat_eta: float = 0.1,
|
|
model: Optional[str] = None,
|
|
logits_processor: Optional[LogitsProcessorList] = None,
|
|
grammar: Optional[LlamaGrammar] = None,
|
|
) -> Union[ChatCompletion, Iterator[ChatCompletionChunk]]:
|
|
"""Generate a chat completion from a list of messages.
|
|
|
|
Args:
|
|
messages: A list of messages to generate a response for.
|
|
temperature: The temperature to use for sampling.
|
|
top_p: The top-p value to use for sampling.
|
|
top_k: The top-k value to use for sampling.
|
|
stream: Whether to stream the results.
|
|
stop: A list of strings to stop generation when encountered.
|
|
max_tokens: The maximum number of tokens to generate. If max_tokens <= 0, the maximum number of tokens to generate is unlimited and depends on n_ctx.
|
|
repeat_penalty: The penalty to apply to repeated tokens.
|
|
|
|
Returns:
|
|
Generated chat completion or a stream of chat completion chunks.
|
|
"""
|
|
|
|
format = llama_chat_format.get_chat_format(self.chat_format)
|
|
result = format(
|
|
messages=messages,
|
|
)
|
|
prompt = result.prompt
|
|
if result.stop is not None:
|
|
stop = [] if stop is None else [stop] if isinstance(stop, str) else stop
|
|
rstop = result.stop if isinstance(result.stop, list) else [result.stop]
|
|
stop = stop + rstop
|
|
|
|
completion_or_chunks = self.create_completion(
|
|
prompt=prompt,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
stream=stream,
|
|
stop=stop,
|
|
max_tokens=max_tokens,
|
|
presence_penalty=presence_penalty,
|
|
frequency_penalty=frequency_penalty,
|
|
repeat_penalty=repeat_penalty,
|
|
tfs_z=tfs_z,
|
|
mirostat_mode=mirostat_mode,
|
|
mirostat_tau=mirostat_tau,
|
|
mirostat_eta=mirostat_eta,
|
|
model=model,
|
|
logits_processor=logits_processor,
|
|
grammar=grammar,
|
|
)
|
|
return self._convert_completion_to_chat(completion_or_chunks, stream=stream) # type: ignore
|
|
|
|
def _free_model(self):
|
|
if hasattr(self, "model") and self.model is not None:
|
|
llama_cpp.llama_free_model(self.model)
|
|
self.model = None
|
|
if hasattr(self, "ctx") and self.ctx is not None:
|
|
llama_cpp.llama_free(self.ctx)
|
|
self.ctx = None
|
|
|
|
def __del__(self):
|
|
if self.verbose:
|
|
self._free_model()
|
|
else:
|
|
with suppress_stdout_stderr():
|
|
self._free_model()
|
|
|
|
def __getstate__(self):
|
|
return dict(
|
|
model_path=self.model_path,
|
|
# Model Params
|
|
n_gpu_layers=self.model_params.n_gpu_layers,
|
|
main_gpu=self.model_params.main_gpu,
|
|
tensor_split=self.tensor_split,
|
|
vocab_only=self.model_params.vocab_only,
|
|
use_mmap=self.model_params.use_mmap,
|
|
use_mlock=self.model_params.use_mlock,
|
|
# Context Params
|
|
seed=self.context_params.seed,
|
|
n_ctx=self.context_params.n_ctx,
|
|
n_batch=self.n_batch,
|
|
n_threads=self.context_params.n_threads,
|
|
n_threads_batch=self.context_params.n_threads_batch,
|
|
rope_freq_base=self.context_params.rope_freq_base,
|
|
rope_freq_scale=self.context_params.rope_freq_scale,
|
|
mul_mat_q=self.context_params.mul_mat_q,
|
|
f16_kv=self.context_params.f16_kv,
|
|
logits_all=self.context_params.logits_all,
|
|
embedding=self.context_params.embedding,
|
|
# Sampling Params
|
|
last_n_tokens_size=self.last_n_tokens_size,
|
|
# LoRA Params
|
|
lora_base=self.lora_base,
|
|
lora_scale=self.lora_scale,
|
|
lora_path=self.lora_path,
|
|
# Backend Params
|
|
numa=self.numa,
|
|
# Chat Format Params
|
|
chat_format=self.chat_format,
|
|
# Misc
|
|
verbose=self.verbose,
|
|
)
|
|
|
|
def __setstate__(self, state):
|
|
self.__init__(
|
|
model_path=state["model_path"],
|
|
# Model Params
|
|
n_gpu_layers=state["n_gpu_layers"],
|
|
main_gpu=state["main_gpu"],
|
|
tensor_split=state["tensor_split"],
|
|
vocab_only=state["vocab_only"],
|
|
use_mmap=state["use_mmap"],
|
|
use_mlock=state["use_mlock"],
|
|
# Context Params
|
|
seed=state["seed"],
|
|
n_ctx=state["n_ctx"],
|
|
n_batch=state["n_batch"],
|
|
n_threads=state["n_threads"],
|
|
n_threads_batch=state["n_threads_batch"],
|
|
rope_freq_base=state["rope_freq_base"],
|
|
rope_freq_scale=state["rope_freq_scale"],
|
|
mul_mat_q=state["mul_mat_q"],
|
|
f16_kv=state["f16_kv"],
|
|
logits_all=state["logits_all"],
|
|
embedding=state["embedding"],
|
|
# Sampling Params
|
|
last_n_tokens_size=state["last_n_tokens_size"],
|
|
# LoRA Params
|
|
lora_base=state["lora_base"],
|
|
lora_path=state["lora_path"],
|
|
# Backend Params
|
|
numa=state["numa"],
|
|
# Chat Format Params
|
|
chat_format=state["chat_format"],
|
|
# Misc
|
|
verbose=state["verbose"],
|
|
)
|
|
|
|
def save_state(self) -> LlamaState:
|
|
assert self.ctx is not None
|
|
if self.verbose:
|
|
print("Llama.save_state: saving llama state", file=sys.stderr)
|
|
state_size = llama_cpp.llama_get_state_size(self.ctx)
|
|
if self.verbose:
|
|
print(f"Llama.save_state: got state size: {state_size}", file=sys.stderr)
|
|
llama_state = (llama_cpp.c_uint8 * int(state_size))()
|
|
if self.verbose:
|
|
print("Llama.save_state: allocated state", file=sys.stderr)
|
|
n_bytes = llama_cpp.llama_copy_state_data(self.ctx, llama_state)
|
|
if self.verbose:
|
|
print(f"Llama.save_state: copied llama state: {n_bytes}", file=sys.stderr)
|
|
if int(n_bytes) > int(state_size):
|
|
raise RuntimeError("Failed to copy llama state data")
|
|
llama_state_compact = (llama_cpp.c_uint8 * int(n_bytes))()
|
|
llama_cpp.ctypes.memmove(llama_state_compact, llama_state, int(n_bytes))
|
|
if self.verbose:
|
|
print(
|
|
f"Llama.save_state: saving {n_bytes} bytes of llama state",
|
|
file=sys.stderr,
|
|
)
|
|
return LlamaState(
|
|
scores=self.scores.copy(),
|
|
input_ids=self.input_ids.copy(),
|
|
n_tokens=self.n_tokens,
|
|
llama_state=bytes(llama_state_compact),
|
|
llama_state_size=n_bytes,
|
|
)
|
|
|
|
def load_state(self, state: LlamaState) -> None:
|
|
assert self.ctx is not None
|
|
self.scores = state.scores.copy()
|
|
self.input_ids = state.input_ids.copy()
|
|
self.n_tokens = state.n_tokens
|
|
state_size = state.llama_state_size
|
|
LLamaStateArrayType = llama_cpp.c_uint8 * state_size
|
|
llama_state = LLamaStateArrayType.from_buffer_copy(state.llama_state)
|
|
|
|
if llama_cpp.llama_set_state_data(self.ctx, llama_state) != state_size:
|
|
raise RuntimeError("Failed to set llama state data")
|
|
|
|
def n_ctx(self) -> int:
|
|
"""Return the context window size."""
|
|
assert self.ctx is not None
|
|
return llama_cpp.llama_n_ctx(self.ctx)
|
|
|
|
def n_embd(self) -> int:
|
|
"""Return the embedding size."""
|
|
assert self.model is not None
|
|
return llama_cpp.llama_n_embd(self.model)
|
|
|
|
def n_vocab(self) -> int:
|
|
"""Return the vocabulary size."""
|
|
assert self.model is not None
|
|
return llama_cpp.llama_n_vocab(self.model)
|
|
|
|
def tokenizer(self) -> "LlamaTokenizer":
|
|
"""Return the tokenizer for this model."""
|
|
assert self.ctx is not None
|
|
return LlamaTokenizer(self)
|
|
|
|
def token_eos(self) -> int:
|
|
"""Return the end-of-sequence token."""
|
|
assert self.ctx is not None
|
|
return llama_cpp.llama_token_eos(self.ctx)
|
|
|
|
def token_bos(self) -> int:
|
|
"""Return the beginning-of-sequence token."""
|
|
assert self.ctx is not None
|
|
return llama_cpp.llama_token_bos(self.ctx)
|
|
|
|
def token_nl(self) -> int:
|
|
"""Return the newline token."""
|
|
assert self.ctx is not None
|
|
return llama_cpp.llama_token_nl(self.ctx)
|
|
|
|
@staticmethod
|
|
def logits_to_logprobs(logits: List[float]) -> List[float]:
|
|
exps = [math.exp(float(x)) for x in logits]
|
|
sum_exps = sum(exps)
|
|
return [math.log(x / sum_exps) for x in exps]
|
|
|
|
@staticmethod
|
|
def longest_token_prefix(a: Sequence[int], b: Sequence[int]):
|
|
longest_prefix = 0
|
|
for _a, _b in zip(a, b):
|
|
if _a == _b:
|
|
longest_prefix += 1
|
|
else:
|
|
break
|
|
return longest_prefix
|
|
|
|
|
|
class LlamaTokenizer:
|
|
def __init__(self, llama: Llama):
|
|
self.llama = llama
|
|
|
|
def encode(self, text: str, add_bos: bool = True) -> List[int]:
|
|
return self.llama.tokenize(
|
|
text.encode("utf-8", errors="ignore"), add_bos=add_bos
|
|
)
|
|
|
|
def decode(self, tokens: List[int]) -> str:
|
|
return self.llama.detokenize(tokens).decode("utf-8", errors="ignore")
|
|
|
|
@classmethod
|
|
def from_ggml_file(cls, path: str) -> "LlamaTokenizer":
|
|
return cls(Llama(model_path=path, vocab_only=True))
|