llama.cpp/llama_cpp/llama.py

161 lines
4.8 KiB
Python
Raw Normal View History

2023-03-23 09:33:06 +00:00
import uuid
import time
import multiprocessing
from typing import List, Optional
from . import llama_cpp
2023-03-24 18:35:41 +00:00
2023-03-23 09:33:06 +00:00
class Llama:
def __init__(
self,
model_path: str,
n_ctx: int = 512,
n_parts: int = -1,
seed: int = 1337,
f16_kv: bool = False,
logits_all: bool = False,
vocab_only: bool = False,
n_threads: Optional[int] = None,
):
self.model_path = model_path
self.last_n = 64
self.max_chunk_size = 32
2023-03-23 09:33:06 +00:00
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.n_threads = n_threads or multiprocessing.cpu_count()
self.tokens = (llama_cpp.llama_token * self.params.n_ctx)()
self.ctx = llama_cpp.llama_init_from_file(
self.model_path.encode("utf-8"), self.params
)
def __call__(
self,
prompt: str,
suffix: Optional[str] = None,
max_tokens: int = 16,
temperature: float = 0.8,
top_p: float = 0.95,
2023-03-23 19:51:05 +00:00
logprobs: Optional[int] = None,
2023-03-23 09:33:06 +00:00
echo: bool = False,
stop: List[str] = [],
repeat_penalty: float = 1.1,
top_k: int = 40,
):
text = b""
2023-03-23 09:33:06 +00:00
finish_reason = "length"
completion_tokens = 0
if stop is not None:
stop = [s.encode("utf-8") for s in stop]
2023-03-23 09:33:06 +00:00
prompt_tokens = llama_cpp.llama_tokenize(
2023-03-24 18:59:29 +00:00
self.ctx,
prompt.encode("utf-8"),
self.tokens,
llama_cpp.llama_n_ctx(self.ctx),
True,
2023-03-23 09:33:06 +00:00
)
if prompt_tokens + max_tokens > self.params.n_ctx:
2023-03-23 09:33:06 +00:00
raise ValueError(
f"Requested tokens exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}"
2023-03-23 09:33:06 +00:00
)
# Process prompt in chunks to avoid running out of memory
for i in range(0, prompt_tokens, self.max_chunk_size):
chunk = self.tokens[i : min(prompt_tokens, i + self.max_chunk_size)]
rc = llama_cpp.llama_eval(
self.ctx,
(llama_cpp.llama_token * len(chunk))(*chunk),
len(chunk),
max(0, i - 1),
self.n_threads,
2023-03-23 09:33:06 +00:00
)
if rc != 0:
raise RuntimeError(f"Failed to evaluate prompt: {rc}")
2023-03-23 09:33:06 +00:00
for i in range(max_tokens):
tokens_seen = prompt_tokens + completion_tokens
2023-03-24 18:35:41 +00:00
last_n_tokens = [0] * max(0, self.last_n - tokens_seen) + [
self.tokens[j]
for j in range(max(tokens_seen - self.last_n, 0), tokens_seen)
]
2023-03-23 09:33:06 +00:00
token = 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),
2023-03-23 09:33:06 +00:00
top_k=top_k,
top_p=top_p,
temp=temperature,
repeat_penalty=repeat_penalty,
)
if token == llama_cpp.llama_token_eos():
finish_reason = "stop"
break
text += llama_cpp.llama_token_to_str(self.ctx, token)
2023-03-23 09:33:06 +00:00
self.tokens[prompt_tokens + i] = token
completion_tokens += 1
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
llama_cpp.llama_eval(
self.ctx,
(llama_cpp.llama_token * 1)(self.tokens[prompt_tokens + i]),
2023-03-23 09:33:06 +00:00
1,
prompt_tokens + completion_tokens,
self.n_threads,
)
text = text.decode("utf-8")
2023-03-23 09:33:06 +00:00
if echo:
text = prompt + text
if suffix is not None:
text = text + suffix
2023-03-23 19:51:05 +00:00
if logprobs is not None:
logprobs = llama_cpp.llama_get_logits(
self.ctx,
)[:logprobs]
2023-03-23 09:33:06 +00:00
return {
2023-03-24 18:35:41 +00:00
"id": f"cmpl-{str(uuid.uuid4())}", # Likely to change
2023-03-23 09:33:06 +00:00
"object": "text_completion",
"created": int(time.time()),
2023-03-24 08:04:29 +00:00
"model": self.model_path,
2023-03-23 09:33:06 +00:00
"choices": [
{
"text": text,
"index": 0,
2023-03-23 19:51:05 +00:00
"logprobs": logprobs,
2023-03-23 09:33:06 +00:00
"finish_reason": finish_reason,
}
],
"usage": {
"prompt_tokens": prompt_tokens,
"completion_tokens": completion_tokens,
"total_tokens": prompt_tokens + completion_tokens,
},
}
def __del__(self):
llama_cpp.llama_free(self.ctx)