llama.cpp/llama_cpp/llama.py

1554 lines
57 KiB
Python
Raw Normal View History

2023-03-24 19:47:17 +00:00
import os
2023-04-04 17:09:24 +00:00
import sys
2023-03-23 09:33:06 +00:00
import uuid
import time
import math
2023-03-23 09:33:06 +00:00
import multiprocessing
from abc import ABC
from typing import (
List,
Optional,
Union,
Generator,
Sequence,
Iterator,
Deque,
Tuple,
Callable,
)
from collections import deque, OrderedDict
2023-03-23 09:33:06 +00:00
import diskcache
2023-03-23 09:33:06 +00:00
from . import llama_cpp
2023-04-01 17:01:27 +00:00
from .llama_types import *
2023-03-23 09:33:06 +00:00
2023-05-26 20:12:45 +00:00
import numpy as np
import numpy.typing as npt
2023-03-24 18:35:41 +00:00
class LlamaCache(ABC):
"""Base cache class for a llama.cpp model."""
2023-04-15 16:03:09 +00:00
def __init__(self, capacity_bytes: int = (2 << 30)):
pass
@property
def cache_size(self):
return 0
def _find_longest_prefix_key(
self,
key: Tuple[int, ...],
) -> Optional[Tuple[int, ...]]:
pass
def __getitem__(self, key: Sequence[int]) -> "LlamaState":
pass
def __contains__(self, key: Sequence[int]) -> bool:
pass
def __setitem__(self, key: Sequence[int], value: "LlamaState"):
pass
class LlamaRAMCache(LlamaCache):
"""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()
2023-04-15 16:03:09 +00:00
@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
2023-04-15 16:03:09 +00:00
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:
self.cache_state.popitem(last=False)
2023-04-15 16:03:09 +00:00
class LlamaDiskCache(LlamaCache):
"""Cache for a llama.cpp model using disk."""
2023-04-15 16:03:09 +00:00
def __init__(self, cache_dir="./llama_cache", capacity_bytes: int = (2 << 30)):
super().__init__(capacity_bytes)
self.cache = diskcache.Cache(cache_dir)
2023-04-15 16:03:09 +00:00
@property
def cache_size(self):
return self.cache.volume()
def _find_longest_prefix_key(
self,
key: Tuple[int, ...],
) -> Optional[Tuple[int, ...]]:
min_len = 0
min_key = None
for k in self.cache.iterkeys():
prefix_len = Llama.longest_token_prefix(k, key)
if prefix_len > min_len:
min_len = prefix_len
min_key = k
return min_key
2023-04-15 16:03:09 +00:00
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.pop(_key)
self.cache.push(_key)
return value
def __setitem__(self, key: Sequence[int], value: "LlamaState"):
key = tuple(key)
if key in self.cache:
del self.cache[key]
self.cache[key] = value
while self.cache_size > self.capacity_bytes:
key_to_remove = next(iter(self.cache))
del self.cache[key_to_remove]
2023-04-15 16:03:09 +00:00
class LlamaState:
def __init__(
self,
eval_tokens: Deque[int],
eval_logits: Deque[List[float]],
input_ids: npt.NDArray[np.intc],
scores: npt.NDArray[np.single],
llama_state, # type: llama_cpp.Array[llama_cpp.c_uint8]
llama_state_size: int,
):
self.eval_tokens = eval_tokens
self.eval_logits = eval_logits
2023-05-26 20:12:45 +00:00
self.input_ids = input_ids
self.scores = scores
self.llama_state = llama_state
2023-05-03 13:33:50 +00:00
self.llama_state_size = llama_state_size
2023-04-15 16:03:09 +00:00
LogitsProcessor = Callable[[List[int], List[float]], List[float]]
class LogitsProcessorList(List[LogitsProcessor]):
def __call__(self, input_ids: List[int], scores: List[float]) -> List[float]:
for processor in self:
scores = processor(input_ids, scores)
return scores
StoppingCriteria = Callable[[List[int], List[float]], bool]
class StoppingCriteriaList(List[StoppingCriteria]):
def __call__(self, input_ids: List[int], logits: List[float]) -> bool:
return any([stopping_criteria(input_ids, logits) for stopping_criteria in self])
2023-03-23 09:33:06 +00:00
class Llama:
2023-03-24 22:57:59 +00:00
"""High-level Python wrapper for a llama.cpp model."""
2023-03-23 09:33:06 +00:00
def __init__(
self,
model_path: str,
# NOTE: These parameters are likely to change in the future.
n_ctx: int = 512,
n_parts: int = -1,
n_gpu_layers: int = 0,
seed: int = 1337,
f16_kv: bool = True,
logits_all: bool = False,
vocab_only: bool = False,
use_mmap: bool = True,
use_mlock: bool = False,
embedding: bool = False,
n_threads: Optional[int] = None,
n_batch: int = 512,
last_n_tokens_size: int = 64,
lora_base: Optional[str] = None,
lora_path: Optional[str] = None,
verbose: bool = True,
2023-04-01 17:01:27 +00:00
):
2023-03-24 22:57:59 +00:00
"""Load a llama.cpp model from `model_path`.
Args:
2023-03-25 16:33:18 +00:00
model_path: Path to the model.
n_ctx: Maximum context size.
2023-03-24 22:57:59 +00:00
n_parts: Number of parts to split the model into. If -1, the number of parts is automatically determined.
2023-03-25 16:33:18 +00:00
seed: Random seed. 0 for random.
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.
2023-04-10 06:11:35 +00:00
use_mmap: Use mmap if possible.
2023-03-25 20:26:23 +00:00
use_mlock: Force the system to keep the model in RAM.
embedding: Embedding mode only.
2023-03-24 22:57:59 +00:00
n_threads: Number of threads to use. If None, the number of threads is automatically determined.
2023-04-01 17:01:27 +00:00
n_batch: Maximum number of prompt tokens to batch together when calling llama_eval.
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.
2023-04-18 05:43:44 +00:00
lora_path: Path to a LoRA file to apply to the model.
2023-04-04 17:09:24 +00:00
verbose: Print verbose output to stderr.
2023-03-24 22:57:59 +00:00
Raises:
ValueError: If the model path does not exist.
Returns:
A Llama instance.
"""
2023-04-04 17:09:24 +00:00
self.verbose = verbose
2023-03-23 09:33:06 +00:00
self.model_path = model_path
self.params = llama_cpp.llama_context_default_params()
self.params.n_ctx = n_ctx
2023-05-14 04:04:22 +00:00
self.params.n_gpu_layers = n_gpu_layers
2023-03-23 09:33:06 +00:00
self.params.seed = seed
self.params.f16_kv = f16_kv
self.params.logits_all = logits_all
self.params.vocab_only = vocab_only
self.params.use_mmap = use_mmap if lora_path is None else False
2023-03-25 20:26:23 +00:00
self.params.use_mlock = use_mlock
self.params.embedding = embedding
2023-03-23 09:33:06 +00:00
2023-04-01 17:01:27 +00:00
self.last_n_tokens_size = last_n_tokens_size
self.n_batch = min(n_ctx, n_batch)
self.eval_tokens: Deque[int] = deque(maxlen=n_ctx)
2023-05-07 17:19:56 +00:00
self.eval_logits: Deque[List[float]] = deque(maxlen=n_ctx if logits_all else 1)
2023-03-23 09:33:06 +00:00
self.cache: Optional[LlamaCache] = None
2023-04-15 16:03:09 +00:00
2023-04-08 23:54:04 +00:00
self.n_threads = n_threads or max(multiprocessing.cpu_count() // 2, 1)
2023-03-23 09:33:06 +00:00
self.lora_base = lora_base
self.lora_path = lora_path
2023-05-21 21:47:21 +00:00
### DEPRECATED ###
self.n_parts = n_parts
### DEPRECATED ###
2023-03-24 19:47:17 +00:00
if not os.path.exists(model_path):
raise ValueError(f"Model path does not exist: {model_path}")
2023-03-23 09:33:06 +00:00
self.ctx = llama_cpp.llama_init_from_file(
self.model_path.encode("utf-8"), self.params
)
assert self.ctx is not None
if self.lora_path:
2023-04-18 05:43:44 +00:00
if llama_cpp.llama_apply_lora_from_file(
self.ctx,
llama_cpp.c_char_p(self.lora_path.encode("utf-8")),
llama_cpp.c_char_p(self.lora_base.encode("utf-8"))
if self.lora_base is not None
else llama_cpp.c_char_p(0),
llama_cpp.c_int(self.n_threads),
2023-04-18 05:43:44 +00:00
):
raise RuntimeError(
f"Failed to apply LoRA from lora path: {self.lora_path} to base path: {self.lora_base}"
)
2023-04-18 05:43:44 +00:00
2023-04-04 17:09:24 +00:00
if self.verbose:
2023-04-28 10:50:30 +00:00
print(llama_cpp.llama_print_system_info().decode("utf-8"), file=sys.stderr)
2023-05-23 21:56:21 +00:00
self._n_vocab = self.n_vocab()
self._n_ctx = self.n_ctx()
size = llama_cpp.c_size_t(self._n_vocab)
2023-05-26 20:12:45 +00:00
sorted = llama_cpp.c_bool(False)
self._candidates_data = np.array(
2023-05-27 02:04:31 +00:00
[],
dtype=np.dtype(
[("id", np.intc), ("logit", np.single), ("p", np.single)], align=True
),
2023-05-26 20:12:45 +00:00
)
2023-06-06 15:37:57 +00:00
self._candidates_data.resize(3, self._n_vocab, refcheck=False)
candidates = llama_cpp.llama_token_data_array(
2023-05-26 20:12:45 +00:00
data=self._candidates_data.ctypes.data_as(llama_cpp.llama_token_data_p),
size=size,
sorted=sorted,
)
self._candidates = candidates
self._token_nl = Llama.token_nl()
self._token_eos = Llama.token_eos()
2023-04-04 17:09:24 +00:00
2023-05-26 20:12:45 +00:00
self._input_ids = np.array([], dtype=np.intc)
self._scores = np.ndarray((0, self._n_vocab), dtype=np.single)
def tokenize(self, text: bytes, add_bos: bool = True) -> List[int]:
"""Tokenize a string.
Args:
text: The utf-8 encoded string to tokenize.
2023-04-01 17:01:27 +00:00
Raises:
RuntimeError: If the tokenization failed.
Returns:
A list of tokens.
"""
2023-04-01 17:01:27 +00:00
assert self.ctx is not None
2023-05-23 21:56:21 +00:00
n_ctx = self._n_ctx
tokens = (llama_cpp.llama_token * n_ctx)()
n_tokens = llama_cpp.llama_tokenize(
self.ctx,
text,
tokens,
llama_cpp.c_int(n_ctx),
llama_cpp.c_bool(add_bos),
)
2023-05-23 21:56:21 +00:00
if n_tokens < 0:
n_tokens = abs(n_tokens)
2023-05-23 21:56:21 +00:00
tokens = (llama_cpp.llama_token * n_tokens)()
n_tokens = llama_cpp.llama_tokenize(
self.ctx,
text,
tokens,
llama_cpp.c_int(n_tokens),
llama_cpp.c_bool(add_bos),
)
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.
"""
2023-04-01 17:01:27 +00:00
assert self.ctx is not None
output = b""
for token in tokens:
output += llama_cpp.llama_token_to_str(
self.ctx, llama_cpp.llama_token(token)
)
return output
2023-04-15 16:03:09 +00:00
def set_cache(self, cache: Optional[LlamaCache]):
"""Set the cache.
Args:
cache: The cache to set.
"""
self.cache = cache
2023-04-15 16:03:09 +00:00
def reset(self):
"""Reset the model state."""
self.eval_tokens.clear()
self.eval_logits.clear()
self._input_ids = np.array([], dtype=np.intc)
self._scores = np.ndarray((0, self._n_vocab), dtype=np.single)
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
2023-05-23 21:56:21 +00:00
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=llama_cpp.c_int(n_tokens),
n_past=llama_cpp.c_int(n_past),
n_threads=llama_cpp.c_int(self.n_threads),
)
2023-05-23 21:56:21 +00:00
if return_code != 0:
raise RuntimeError(f"llama_eval returned {return_code}")
2023-05-01 18:47:55 +00:00
# Save tokens
self.eval_tokens.extend(batch)
2023-05-26 20:12:45 +00:00
self._input_ids: npt.NDArray[np.intc] = np.concatenate(
(self._input_ids, np.array(batch, dtype=np.intc)), axis=0
)
2023-05-01 18:47:55 +00:00
# Save logits
rows = n_tokens if self.params.logits_all else 1
2023-05-23 21:56:21 +00:00
n_vocab = self._n_vocab
cols = n_vocab
2023-05-01 18:47:55 +00:00
logits_view = llama_cpp.llama_get_logits(self.ctx)
logits = [logits_view[i * cols: (i + 1) * cols] for i in range(rows)]
2023-05-01 18:47:55 +00:00
self.eval_logits.extend(logits)
2023-05-26 20:12:45 +00:00
self._scores: npt.NDArray[np.single] = np.concatenate(
(self._scores, np.array(logits, dtype=np.single)), axis=0
)
2023-05-01 18:47:55 +00:00
def _sample(
self,
last_n_tokens_data, # type: llama_cpp.Array[llama_cpp.llama_token]
last_n_tokens_size: llama_cpp.c_int,
top_k: llama_cpp.c_int,
top_p: llama_cpp.c_float,
temp: llama_cpp.c_float,
tfs_z: llama_cpp.c_float,
repeat_penalty: llama_cpp.c_float,
frequency_penalty: llama_cpp.c_float,
presence_penalty: llama_cpp.c_float,
mirostat_mode: llama_cpp.c_int,
mirostat_tau: llama_cpp.c_float,
mirostat_eta: llama_cpp.c_float,
penalize_nl: bool = True,
logits_processor: Optional[LogitsProcessorList] = None,
2023-05-01 18:47:55 +00:00
):
assert self.ctx is not None
assert len(self.eval_logits) > 0
assert self._scores.shape[0] > 0
2023-05-23 21:56:21 +00:00
n_vocab = self._n_vocab
n_ctx = self._n_ctx
2023-05-17 05:41:42 +00:00
top_k = llama_cpp.c_int(n_vocab) if top_k.value <= 0 else top_k
2023-05-17 06:00:39 +00:00
last_n_tokens_size = (
llama_cpp.c_int(n_ctx)
if last_n_tokens_size.value < 0
else last_n_tokens_size
)
2023-05-26 20:12:45 +00:00
logits: npt.NDArray[np.single] = self._scores[-1, :]
if logits_processor is not None:
2023-05-26 20:12:45 +00:00
logits = np.array(
logits_processor(self._input_ids.tolist(), logits.tolist()),
2023-05-26 20:12:45 +00:00
dtype=np.single,
)
self._scores[-1, :] = logits
self.eval_logits[-1] = logits.tolist()
nl_logit = logits[self._token_nl]
candidates = self._candidates
2023-05-26 20:12:45 +00:00
candidates_data = self._candidates_data
candidates_data["id"] = np.arange(n_vocab, dtype=np.intc) # type: ignore
candidates_data["logit"] = logits
candidates_data["p"] = np.zeros(n_vocab, dtype=np.single)
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)
2023-05-01 18:47:55 +00:00
llama_cpp.llama_sample_repetition_penalty(
ctx=self.ctx,
last_tokens_data=last_n_tokens_data,
last_tokens_size=last_n_tokens_size,
2023-05-08 00:01:34 +00:00
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
2023-05-01 18:47:55 +00:00
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,
)
2023-05-17 05:53:26 +00:00
if not penalize_nl:
candidates.data[self._token_nl].logit = llama_cpp.c_float(nl_logit)
if temp.value == 0.0:
return llama_cpp.llama_sample_token_greedy(
ctx=self.ctx,
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
)
elif mirostat_mode.value == 1:
mirostat_mu = llama_cpp.c_float(2.0 * mirostat_tau.value)
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,
)
return 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.value == 2:
mirostat_mu = llama_cpp.c_float(2.0 * mirostat_tau.value)
llama_cpp.llama_sample_temperature(
ctx=self.ctx,
candidates=llama_cpp.ctypes.pointer(candidates),
temp=temp,
)
return 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
2023-05-01 18:47:55 +00:00
)
else:
llama_cpp.llama_sample_top_k(
ctx=self.ctx,
2023-05-08 00:01:34 +00:00
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
2023-05-01 18:47:55 +00:00
k=top_k,
2023-05-07 04:12:47 +00:00
min_keep=llama_cpp.c_size_t(1),
2023-05-01 18:47:55 +00:00
)
llama_cpp.llama_sample_tail_free(
ctx=self.ctx,
2023-05-08 00:01:34 +00:00
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
z=tfs_z,
2023-05-07 04:12:47 +00:00
min_keep=llama_cpp.c_size_t(1),
2023-05-01 18:47:55 +00:00
)
llama_cpp.llama_sample_typical(
ctx=self.ctx,
2023-05-08 00:01:34 +00:00
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
2023-05-02 01:51:16 +00:00
p=llama_cpp.c_float(1.0),
2023-05-07 04:12:47 +00:00
min_keep=llama_cpp.c_size_t(1),
2023-05-01 18:47:55 +00:00
)
llama_cpp.llama_sample_top_p(
ctx=self.ctx,
2023-05-08 00:01:34 +00:00
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
2023-05-01 18:47:55 +00:00
p=top_p,
2023-05-07 04:12:47 +00:00
min_keep=llama_cpp.c_size_t(1),
2023-05-01 18:47:55 +00:00
)
llama_cpp.llama_sample_temperature(
ctx=self.ctx,
2023-05-08 00:01:34 +00:00
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
2023-05-01 18:47:55 +00:00
temp=temp,
)
return llama_cpp.llama_sample_token(
ctx=self.ctx,
2023-05-08 00:01:34 +00:00
candidates=llama_cpp.ctypes.byref(candidates), # type: ignore
2023-05-01 18:47:55 +00:00
)
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,
):
"""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=llama_cpp.c_int(self.last_n_tokens_size),
top_k=llama_cpp.c_int(top_k),
top_p=llama_cpp.c_float(top_p),
temp=llama_cpp.c_float(temp),
tfs_z=llama_cpp.c_float(tfs_z),
repeat_penalty=llama_cpp.c_float(repeat_penalty),
frequency_penalty=llama_cpp.c_float(frequency_penalty),
presence_penalty=llama_cpp.c_float(presence_penalty),
mirostat_mode=llama_cpp.c_int(mirostat_mode),
mirostat_tau=llama_cpp.c_float(mirostat_tau),
mirostat_eta=llama_cpp.c_float(mirostat_eta),
2023-05-17 05:53:26 +00:00
penalize_nl=penalize_nl,
logits_processor=logits_processor,
)
2023-04-01 17:01:27 +00:00
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,
) -> Generator[int, Optional[Sequence[int]], None]:
"""Create a generator of tokens from a prompt.
2023-04-01 21:36:30 +00:00
2023-04-01 21:39:35 +00:00
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]))
2023-04-01 21:36:30 +00:00
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.
2023-04-01 21:36:30 +00:00
Yields:
The generated tokens.
"""
2023-04-01 17:01:27 +00:00
assert self.ctx is not None
if reset and len(self._input_ids) > 0:
2023-05-05 01:58:27 +00:00
longest_prefix = 0
for a, b in zip(self._input_ids, tokens[:-1]):
2023-05-05 01:58:27 +00:00
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._input_ids = self._input_ids[:longest_prefix]
self._scores = self._scores[:longest_prefix, :]
2023-05-05 01:58:27 +00:00
for _ in range(len(self.eval_tokens) - longest_prefix):
self.eval_tokens.pop()
try:
self.eval_logits.pop()
except IndexError:
pass
if reset:
self.reset()
2023-05-05 01:58:27 +00:00
2023-04-01 17:01:27 +00:00
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,
2023-05-12 01:56:19 +00:00
tfs_z=tfs_z,
mirostat_mode=mirostat_mode,
mirostat_tau=mirostat_tau,
mirostat_eta=mirostat_eta,
logits_processor=logits_processor,
2023-04-01 17:01:27 +00:00
)
if stopping_criteria is not None and stopping_criteria(
self._input_ids.tolist(), self._scores[-1, :].tolist()
):
return
2023-04-01 17:01:27 +00:00
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
) -> Embedding:
"""Embed a string.
Args:
2023-04-01 17:01:27 +00:00
input: The utf-8 encoded string to embed.
Returns:
2023-04-01 17:01:27 +00:00
An embedding object.
"""
2023-04-01 17:01:27 +00:00
assert self.ctx is not None
2023-05-16 22:07:25 +00:00
model_name: str = model if model is not None else self.model_path
2023-04-04 17:09:24 +00:00
if self.params.embedding == False:
raise RuntimeError(
"Llama model must be created with embedding=True to call this method"
)
2023-04-04 17:09:24 +00:00
if self.verbose:
llama_cpp.llama_reset_timings(self.ctx)
if isinstance(input, str):
inputs = [input]
else:
inputs = input
2023-04-04 17:09:24 +00:00
2023-05-22 01:30:03 +00:00
data: List[EmbeddingData] = []
total_tokens = 0
2023-05-22 01:30:03 +00:00
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.ctx)
]
2023-04-04 17:09:24 +00:00
data.append(
2023-04-01 17:01:27 +00:00
{
"object": "embedding",
"embedding": embedding,
2023-05-22 01:30:03 +00:00
"index": index,
2023-04-01 17:01:27 +00:00
}
)
2023-05-22 01:30:03 +00:00
if self.verbose:
llama_cpp.llama_print_timings(self.ctx)
return {
"object": "list",
"data": data,
2023-05-22 01:30:03 +00:00
"model": model_name,
2023-04-01 17:01:27 +00:00
"usage": {
"prompt_tokens": total_tokens,
"total_tokens": total_tokens,
2023-04-01 17:01:27 +00:00
},
}
2023-03-28 06:42:22 +00:00
2023-04-03 22:46:19 +00:00
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"]))
2023-04-01 17:01:27 +00:00
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,
2023-04-12 18:06:22 +00:00
) -> Union[Iterator[Completion], Iterator[CompletionChunk]]:
2023-04-01 17:01:27 +00:00
assert self.ctx is not None
2023-05-24 20:02:06 +00:00
2023-04-15 15:39:21 +00:00
completion_id: str = f"cmpl-{str(uuid.uuid4())}"
created: int = int(time.time())
completion_tokens: List[int] = []
2023-04-01 17:01:27 +00:00
# Add blank space to start of prompt to match OG llama tokenizer
prompt_tokens: List[int] = self.tokenize(b" " + prompt.encode("utf-8"))
2023-04-15 15:39:21 +00:00
text: bytes = b""
2023-05-18 15:35:59 +00:00
returned_tokens: int = 0
stop = (
stop if isinstance(stop, list) else [stop] if isinstance(stop, str) else []
)
2023-05-16 22:07:25 +00:00
model_name: str = model if model is not None else self.model_path
2023-03-23 09:33:06 +00:00
2023-04-04 17:09:24 +00:00
if self.verbose:
llama_cpp.llama_reset_timings(self.ctx)
2023-05-23 21:56:21 +00:00
if len(prompt_tokens) + max_tokens > self._n_ctx:
2023-05-26 07:00:35 +00:00
raise ValueError(f"Requested tokens exceed context window of {self._n_ctx}")
2023-03-23 09:33:06 +00:00
2023-04-01 17:01:27 +00:00
if stop != []:
stop_sequences = [s.encode("utf-8") for s in stop]
2023-04-01 17:01:27 +00:00
else:
stop_sequences = []
if logprobs is not None and self.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)
2023-04-15 16:03:09 +00:00
finish_reason = "length"
2023-04-28 10:50:30 +00:00
multibyte_fix = 0
2023-04-14 13:59:33 +00:00
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,
2023-04-14 13:59:33 +00:00
):
if token == self._token_eos:
text = self.detokenize(completion_tokens)
2023-03-23 09:33:06 +00:00
finish_reason = "stop"
break
completion_tokens.append(token)
2023-03-23 09:33:06 +00:00
all_text = self.detokenize(completion_tokens)
# Contains multi-byte UTF8
2023-05-01 18:47:55 +00:00
for k, char in enumerate(all_text[-3:]):
k = 3 - k
2023-05-01 18:47:55 +00:00
for num, pattern in [(2, 192), (3, 224), (4, 240)]:
# Bitwise AND check
2023-05-01 18:47:55 +00:00
if num > k and pattern & char == pattern:
multibyte_fix = num - k
2023-04-28 10:50:30 +00:00
# Stop incomplete bytes from passing
2023-05-01 18:47:55 +00:00
if multibyte_fix > 0:
2023-04-28 10:50:30 +00:00
multibyte_fix -= 1
continue
any_stop = [s for s in stop_sequences if s in all_text]
2023-03-23 09:33:06 +00:00
if len(any_stop) > 0:
first_stop = any_stop[0]
text = all_text[: all_text.index(first_stop)]
2023-03-23 09:33:06 +00:00
finish_reason = "stop"
break
if stream:
2023-05-27 00:23:49 +00:00
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:
2023-05-27 00:23:49 +00:00
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
2023-05-18 15:35:59 +00:00
token_end_position = 0
2023-05-18 15:35:59 +00:00
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 - 1
):
2023-05-18 15:35:59 +00:00
break
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
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": [sorted_logprobs[int(token)][0]],
"top_logprobs": [top_logprob],
}
2023-05-18 15:35:59 +00:00
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,
2023-05-18 15:35:59 +00:00
"finish_reason": None,
}
],
}
if len(completion_tokens) >= max_tokens:
text = self.detokenize(completion_tokens)
finish_reason = "length"
break
2023-03-23 09:33:06 +00:00
if stopping_criteria is not None and stopping_criteria(
self._input_ids.tolist(), self._scores[-1, :].tolist()
):
text = self.detokenize(completion_tokens)
finish_reason = "stop"
if self.verbose:
llama_cpp.llama_print_timings(self.ctx)
if stream:
2023-05-18 15:35:59 +00:00
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
2023-05-18 15:35:59 +00:00
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": [sorted_logprobs[int(token)][0]],
"top_logprobs": [top_logprob],
}
if token_end_position >= end:
2023-05-18 15:35:59 +00:00
last_text = self.detokenize([token])
if token_end_position == end - 1:
2023-05-18 15:35:59 +00:00
break
returned_tokens += 1
2023-05-18 15:35:59 +00:00
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"),
2023-05-18 15:35:59 +00:00
"index": 0,
"logprobs": logprobs_or_none,
2023-05-18 15:35:59 +00:00
"finish_reason": finish_reason,
}
],
}
2023-05-18 15:35:59 +00:00
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,
2023-05-18 15:35:59 +00:00
"finish_reason": finish_reason
if returned_tokens == len(completion_tokens)
else None,
}
],
}
if self.cache:
if self.verbose:
print("Llama._create_completion: cache save", file=sys.stderr)
self.cache[prompt_tokens + completion_tokens] = self.save_state()
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()
2023-04-26 12:37:06 +00:00
text_str = text.decode("utf-8", errors="ignore")
2023-03-23 09:33:06 +00:00
if echo:
2023-04-15 16:03:09 +00:00
text_str = prompt + text_str
2023-03-23 09:33:06 +00:00
if suffix is not None:
2023-04-15 16:03:09 +00:00
text_str = text_str + suffix
2023-03-23 09:33:06 +00:00
logprobs_or_none: Optional[CompletionLogprobs] = None
2023-03-23 19:51:05 +00:00
if logprobs is not None:
text_offset = 0 if echo else len(prompt)
token_offset = 0 if echo else len(prompt_tokens[1:])
2023-04-14 13:59:33 +00:00
text_offsets: List[int] = []
token_logprobs: List[Optional[float]] = []
2023-04-14 13:59:33 +00:00
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
2023-04-14 13:59:33 +00:00
all_token_strs = [
2023-05-01 18:47:55 +00:00
self.detokenize([token]).decode("utf-8", errors="ignore")
for token in all_tokens
2023-04-14 13:59:33 +00:00
]
2023-05-05 01:58:36 +00:00
all_logprobs = [
Llama.logits_to_logprobs(row.tolist()) for row in self._scores
][token_offset:]
2023-04-14 13:59:33 +00:00
for token, token_str, logprobs_token in zip(
all_tokens, all_token_strs, all_logprobs
2023-04-14 13:59:33 +00:00
):
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(sorted_logprobs[int(token)][0])
top_logprob: Optional[Dict[str, float]] = {
self.detokenize([i]).decode("utf-8", errors="ignore"): logprob
2023-04-14 13:59:33 +00:00
for logprob, i in sorted_logprobs[:logprobs]
}
top_logprob.update({token_str: logprobs_token[int(token)]})
2023-04-14 13:59:33 +00:00
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,
}
2023-03-23 19:51:05 +00:00
yield {
2023-03-28 06:42:22 +00:00
"id": completion_id,
2023-03-23 09:33:06 +00:00
"object": "text_completion",
2023-03-28 06:42:22 +00:00
"created": created,
2023-05-16 22:07:25 +00:00
"model": model_name,
2023-03-23 09:33:06 +00:00
"choices": [
{
2023-04-15 16:03:09 +00:00
"text": text_str,
2023-03-23 09:33:06 +00:00
"index": 0,
"logprobs": logprobs_or_none,
2023-03-23 09:33:06 +00:00
"finish_reason": finish_reason,
}
],
"usage": {
"prompt_tokens": len(prompt_tokens),
"completion_tokens": len(completion_tokens),
"total_tokens": len(prompt_tokens) + len(completion_tokens),
2023-03-23 09:33:06 +00:00
},
}
2023-04-01 17:01:27 +00:00
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,
) -> Union[Completion, Iterator[CompletionChunk]]:
2023-04-01 17:01:27 +00:00
"""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.
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,
2023-04-01 17:01:27 +00:00
repeat_penalty=repeat_penalty,
top_k=top_k,
stream=stream,
2023-05-12 01:56:19 +00:00
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,
2023-04-01 17:01:27 +00:00
)
if stream:
chunks: Iterator[CompletionChunk] = completion_or_chunks
2023-04-01 17:01:27 +00:00
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,
) -> 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.
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.
"""
2023-04-01 17:01:27 +00:00
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,
2023-05-12 01:56:19 +00:00
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,
)
2023-04-04 00:12:44 +00:00
def _convert_text_completion_to_chat(
self, completion: Completion
2023-04-04 00:12:44 +00:00
) -> 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],
2023-04-04 00:12:44 +00:00
) -> 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"],
},
"finish_reason": chunk["choices"][0]["finish_reason"],
}
],
}
def create_chat_completion(
self,
messages: List[ChatCompletionMessage],
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,
2023-04-04 00:12:44 +00:00
) -> 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.
repeat_penalty: The penalty to apply to repeated tokens.
Returns:
Generated chat completion or a stream of chat completion chunks.
"""
stop = (
stop if isinstance(stop, list) else [stop] if isinstance(stop, str) else []
)
2023-04-15 15:58:19 +00:00
chat_history = "".join(
f'### {"Human" if message["role"] == "user" else "Assistant"}:{message["content"]}'
2023-04-04 00:12:44 +00:00
for message in messages
)
2023-04-15 15:58:19 +00:00
PROMPT = chat_history + "### Assistant:"
2023-04-15 16:02:48 +00:00
PROMPT_STOP = ["### Assistant:", "### Human:"]
2023-04-04 00:12:44 +00:00
completion_or_chunks = self(
prompt=PROMPT,
stop=PROMPT_STOP + stop,
temperature=temperature,
top_p=top_p,
top_k=top_k,
stream=stream,
max_tokens=max_tokens,
repeat_penalty=repeat_penalty,
presence_penalty=presence_penalty,
frequency_penalty=frequency_penalty,
2023-05-12 01:56:19 +00:00
tfs_z=tfs_z,
mirostat_mode=mirostat_mode,
mirostat_tau=mirostat_tau,
mirostat_eta=mirostat_eta,
model=model,
2023-04-04 00:12:44 +00:00
)
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)
2023-03-23 09:33:06 +00:00
def __del__(self):
2023-04-01 17:01:27 +00:00
if self.ctx is not None:
llama_cpp.llama_free(self.ctx)
self.ctx = None
def __getstate__(self):
return dict(
verbose=self.verbose,
model_path=self.model_path,
n_ctx=self.params.n_ctx,
2023-05-14 04:04:22 +00:00
n_gpu_layers=self.params.n_gpu_layers,
seed=self.params.seed,
f16_kv=self.params.f16_kv,
logits_all=self.params.logits_all,
vocab_only=self.params.vocab_only,
2023-04-10 06:11:35 +00:00
use_mmap=self.params.use_mmap,
use_mlock=self.params.use_mlock,
embedding=self.params.embedding,
last_n_tokens_size=self.last_n_tokens_size,
n_batch=self.n_batch,
n_threads=self.n_threads,
lora_base=self.lora_base,
2023-04-18 05:43:44 +00:00
lora_path=self.lora_path,
2023-05-21 21:47:21 +00:00
### DEPRECATED ###
n_parts=self.n_parts,
### DEPRECATED ###
)
def __setstate__(self, state):
self.__init__(
model_path=state["model_path"],
n_ctx=state["n_ctx"],
n_parts=state["n_parts"],
2023-05-14 04:04:22 +00:00
n_gpu_layers=state["n_gpu_layers"],
seed=state["seed"],
f16_kv=state["f16_kv"],
logits_all=state["logits_all"],
vocab_only=state["vocab_only"],
2023-04-10 06:11:35 +00:00
use_mmap=state["use_mmap"],
use_mlock=state["use_mlock"],
embedding=state["embedding"],
n_threads=state["n_threads"],
n_batch=state["n_batch"],
last_n_tokens_size=state["last_n_tokens_size"],
lora_base=state["lora_base"],
2023-04-18 05:43:44 +00:00
lora_path=state["lora_path"],
verbose=state["verbose"],
)
def save_state(self) -> LlamaState:
assert self.ctx is not None
state_size = llama_cpp.llama_get_state_size(self.ctx)
llama_state = (llama_cpp.c_uint8 * int(state_size))()
2023-05-03 13:33:50 +00:00
n_bytes = llama_cpp.llama_copy_state_data(self.ctx, llama_state)
if int(n_bytes) > int(state_size):
raise RuntimeError("Failed to copy llama state data")
2023-05-03 13:33:50 +00:00
llama_state_compact = (llama_cpp.c_uint8 * int(n_bytes))()
llama_cpp.ctypes.memmove(llama_state_compact, llama_state, int(n_bytes))
2023-05-03 14:28:10 +00:00
if self.verbose:
2023-05-05 01:58:36 +00:00
print(
f"Llama.save_state: saving {n_bytes} bytes of llama state",
file=sys.stderr,
)
return LlamaState(
eval_tokens=self.eval_tokens.copy(),
eval_logits=self.eval_logits.copy(),
2023-05-26 20:12:45 +00:00
scores=self._scores.copy(),
input_ids=self._input_ids.copy(),
2023-05-03 13:33:50 +00:00
llama_state=llama_state_compact,
llama_state_size=n_bytes,
)
def load_state(self, state: LlamaState) -> None:
assert self.ctx is not None
self.eval_tokens = state.eval_tokens.copy()
self.eval_logits = state.eval_logits.copy()
2023-05-26 20:12:45 +00:00
self._scores = state.scores.copy()
self._input_ids = state.input_ids.copy()
2023-05-03 13:33:50 +00:00
state_size = state.llama_state_size
if llama_cpp.llama_set_state_data(self.ctx, state.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.ctx is not None
return llama_cpp.llama_n_embd(self.ctx)
def n_vocab(self) -> int:
"""Return the vocabulary size."""
assert self.ctx is not None
return llama_cpp.llama_n_vocab(self.ctx)
2023-05-25 18:11:33 +00:00
def tokenizer(self) -> "LlamaTokenizer":
"""Return the tokenizer for this model."""
assert self.ctx is not None
return LlamaTokenizer(self)
@staticmethod
def token_eos() -> int:
"""Return the end-of-sequence token."""
return llama_cpp.llama_token_eos()
@staticmethod
def token_bos() -> int:
"""Return the beginning-of-sequence token."""
return llama_cpp.llama_token_bos()
2023-05-17 05:53:26 +00:00
@staticmethod
def token_nl() -> int:
2023-05-17 05:53:26 +00:00
"""Return the newline token."""
return llama_cpp.llama_token_nl()
@staticmethod
2023-05-04 16:18:40 +00:00
def logits_to_logprobs(logits: List[float]) -> List[float]:
2023-05-01 21:45:08 +00:00
exps = [math.exp(float(x)) for x in logits]
sum_exps = sum(exps)
2023-05-04 16:18:40 +00:00
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
2023-05-25 18:11:33 +00:00
class LlamaTokenizer:
def __init__(self, llama: Llama):
self.llama = llama
2023-05-26 07:00:51 +00:00
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
)
2023-05-25 18:11:33 +00:00
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))