329 lines
11 KiB
Python
329 lines
11 KiB
Python
import os
|
|
import uuid
|
|
import time
|
|
import multiprocessing
|
|
from typing import List, Optional
|
|
from collections import deque
|
|
|
|
from . import llama_cpp
|
|
|
|
|
|
class Llama:
|
|
"""High-level Python wrapper for a llama.cpp model."""
|
|
|
|
def __init__(
|
|
self,
|
|
model_path: str,
|
|
# NOTE: The following parameters are likely to change in the future.
|
|
n_ctx: int = 512,
|
|
n_parts: int = -1,
|
|
seed: int = 1337,
|
|
f16_kv: bool = False,
|
|
logits_all: bool = False,
|
|
vocab_only: bool = False,
|
|
use_mlock: bool = False,
|
|
embedding: bool = False,
|
|
n_threads: Optional[int] = None,
|
|
) -> "Llama":
|
|
"""Load a llama.cpp model from `model_path`.
|
|
|
|
Args:
|
|
model_path: Path to the model.
|
|
n_ctx: Maximum context size.
|
|
n_parts: Number of parts to split the model into. If -1, the number of parts is automatically determined.
|
|
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.
|
|
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.
|
|
|
|
Raises:
|
|
ValueError: If the model path does not exist.
|
|
|
|
Returns:
|
|
A Llama instance.
|
|
"""
|
|
self.model_path = model_path
|
|
|
|
self.params = llama_cpp.llama_context_default_params()
|
|
self.params.n_ctx = n_ctx
|
|
self.params.n_parts = n_parts
|
|
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_mlock = use_mlock
|
|
self.params.embedding = embedding
|
|
|
|
self.last_n = 64
|
|
self.max_chunk_size = n_ctx
|
|
|
|
self.n_threads = n_threads or multiprocessing.cpu_count()
|
|
|
|
if not os.path.exists(model_path):
|
|
raise ValueError(f"Model path does not exist: {model_path}")
|
|
|
|
self.ctx = llama_cpp.llama_init_from_file(
|
|
self.model_path.encode("utf-8"), self.params
|
|
)
|
|
|
|
def tokenize(self, text: bytes) -> List[int]:
|
|
"""Tokenize a string.
|
|
|
|
Args:
|
|
text: The utf-8 encoded string to tokenize.
|
|
|
|
Returns:
|
|
A list of tokens.
|
|
"""
|
|
n_ctx = llama_cpp.llama_n_ctx(self.ctx)
|
|
tokens = (llama_cpp.llama_token * n_ctx)()
|
|
n_tokens = llama_cpp.llama_tokenize(
|
|
self.ctx,
|
|
text,
|
|
tokens,
|
|
n_ctx,
|
|
True,
|
|
)
|
|
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.
|
|
"""
|
|
output = b""
|
|
for token in tokens:
|
|
output += llama_cpp.llama_token_to_str(self.ctx, token)
|
|
return output
|
|
|
|
def _eval(self, tokens: List[int], n_past):
|
|
rc = llama_cpp.llama_eval(
|
|
self.ctx,
|
|
(llama_cpp.llama_token * len(tokens))(*tokens),
|
|
len(tokens),
|
|
n_past,
|
|
self.n_threads,
|
|
)
|
|
if rc != 0:
|
|
raise RuntimeError(f"Failed to evaluate: {rc}")
|
|
|
|
def _sample(self, last_n_tokens, top_p, top_k, temp, repeat_penalty):
|
|
return llama_cpp.llama_sample_top_p_top_k(
|
|
self.ctx,
|
|
(llama_cpp.llama_token * len(last_n_tokens))(*last_n_tokens),
|
|
len(last_n_tokens),
|
|
top_k=top_k,
|
|
top_p=top_p,
|
|
temp=temp,
|
|
repeat_penalty=repeat_penalty,
|
|
)
|
|
|
|
def _generate(self, past_tokens, max_tokens, top_p, top_k, temp, repeat_penalty):
|
|
last_n_tokens = deque([0] * self.last_n, maxlen=self.last_n)
|
|
last_n_tokens.extend(past_tokens)
|
|
for i in range(max_tokens):
|
|
token = self._sample(
|
|
last_n_tokens,
|
|
top_p=top_p,
|
|
top_k=top_k,
|
|
temp=temp,
|
|
repeat_penalty=repeat_penalty,
|
|
)
|
|
yield token
|
|
self._eval([token], len(past_tokens) + i)
|
|
|
|
def _call(
|
|
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: List[str] = [],
|
|
repeat_penalty: float = 1.1,
|
|
top_k: int = 40,
|
|
stream: bool = False,
|
|
):
|
|
completion_id = f"cmpl-{str(uuid.uuid4())}"
|
|
created = int(time.time())
|
|
completion_tokens = []
|
|
prompt_tokens = self.tokenize(prompt.encode("utf-8"))
|
|
|
|
if len(prompt_tokens) + max_tokens > llama_cpp.llama_n_ctx(self.ctx):
|
|
raise ValueError(
|
|
f"Requested tokens exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}"
|
|
)
|
|
|
|
# Process prompt in chunks to avoid running out of memory
|
|
for i in range(0, len(prompt_tokens), self.max_chunk_size):
|
|
chunk = prompt_tokens[i : min(len(prompt_tokens), i + self.max_chunk_size)]
|
|
self._eval(chunk, n_past=i)
|
|
|
|
if stop is not None:
|
|
stop = [s.encode("utf-8") for s in stop]
|
|
|
|
finish_reason = None
|
|
for token in self._generate(
|
|
prompt_tokens, max_tokens, top_p, top_k, temperature, repeat_penalty
|
|
):
|
|
if token == llama_cpp.llama_token_eos():
|
|
finish_reason = "stop"
|
|
break
|
|
completion_tokens.append(token)
|
|
|
|
text = self.detokenize(completion_tokens)
|
|
any_stop = [s for s in stop if s in text]
|
|
if len(any_stop) > 0:
|
|
first_stop = any_stop[0]
|
|
text = text[: text.index(first_stop)]
|
|
finish_reason = "stop"
|
|
break
|
|
|
|
if stream:
|
|
start = len(self.detokenize(completion_tokens[:-1]))
|
|
longest = 0
|
|
for s in stop:
|
|
for i in range(len(s), 0, -1):
|
|
if s[-i:] == text[-i:]:
|
|
if i > longest:
|
|
longest = i
|
|
break
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": self.model_path,
|
|
"choices": [
|
|
{
|
|
"text": text[start : len(text) - longest].decode("utf-8"),
|
|
"index": 0,
|
|
"logprobs": None,
|
|
"finish_reason": None,
|
|
}
|
|
],
|
|
}
|
|
|
|
if finish_reason is None:
|
|
finish_reason = "length"
|
|
|
|
if stream:
|
|
if finish_reason == "stop":
|
|
start = len(self.detokenize(completion_tokens[:-1]))
|
|
text = text[start:].decode("utf-8")
|
|
else:
|
|
text = ""
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": self.model_path,
|
|
"choices": [
|
|
{
|
|
"text": text,
|
|
"index": 0,
|
|
"logprobs": None,
|
|
"finish_reason": finish_reason,
|
|
}
|
|
],
|
|
}
|
|
return
|
|
|
|
text = text.decode("utf-8")
|
|
|
|
if echo:
|
|
text = prompt + text
|
|
|
|
if suffix is not None:
|
|
text = text + suffix
|
|
|
|
if logprobs is not None:
|
|
logprobs = llama_cpp.llama_get_logits(
|
|
self.ctx,
|
|
)[:logprobs]
|
|
|
|
yield {
|
|
"id": completion_id,
|
|
"object": "text_completion",
|
|
"created": created,
|
|
"model": self.model_path,
|
|
"choices": [
|
|
{
|
|
"text": text,
|
|
"index": 0,
|
|
"logprobs": logprobs,
|
|
"finish_reason": finish_reason,
|
|
}
|
|
],
|
|
"usage": {
|
|
"prompt_tokens": len(prompt_tokens),
|
|
"completion_tokens": len(completion_tokens),
|
|
"total_tokens": len(prompt_tokens) + len(completion_tokens),
|
|
},
|
|
}
|
|
|
|
def __call__(
|
|
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: List[str] = [],
|
|
repeat_penalty: float = 1.1,
|
|
top_k: int = 40,
|
|
stream: bool = False,
|
|
):
|
|
"""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.
|
|
"""
|
|
call = self._call(
|
|
prompt=prompt,
|
|
suffix=suffix,
|
|
max_tokens=max_tokens,
|
|
temperature=temperature,
|
|
top_p=top_p,
|
|
logprobs=logprobs,
|
|
echo=echo,
|
|
stop=stop,
|
|
repeat_penalty=repeat_penalty,
|
|
top_k=top_k,
|
|
stream=stream,
|
|
)
|
|
if stream:
|
|
return call
|
|
return next(call)
|
|
|
|
def __del__(self):
|
|
llama_cpp.llama_free(self.ctx)
|