from __future__ import annotations

import os

import sys

import uuid

import time

import multiprocessing

from typing import (

List,

Optional,

Union,

Generator,

Sequence,

Iterator,

Deque,

Callable,

)

from collections import deque

import ctypes

import numpy as np

import numpy.typing as npt

from .llama_types import *

from .llama_grammar import LlamaGrammar

from .llama_cache import BaseLlamaCache

import llama_cpp.llama_cpp as llama_cpp

import llama_cpp.llama_chat_format as llama_chat_format

from ._internals import _LlamaModel, _LlamaContext, _LlamaBatch, _LlamaTokenDataArray # type: ignore

from ._utils import suppress_stdout_stderr

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_scaling_type: Optional[int] = llama_cpp.LLAMA_ROPE_SCALING_UNSPECIFIED,

rope_freq_base: float = 0.0,

rope_freq_scale: float = 0.0,

yarn_ext_factor: float = -1.0,

yarn_attn_factor: float = 1.0,

yarn_beta_fast: float = 32.0,

yarn_beta_slow: float = 1.0,

yarn_orig_ctx: int = 0,

mul_mat_q: bool = True,

logits_all: bool = False,

embedding: bool = False,

offload_kqv: 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",

chat_handler: Optional[llama_chat_format.LlamaChatCompletionHandler] = None,

# Misc

verbose: bool = True,

# Extra Params

**kwargs, # type: ignore

):

"""Load a llama.cpp model from `model_path`.

Examples:

Basic usage

>>> import llama_cpp

>>> model = llama_cpp.Llama(

... model_path="path/to/model",

... )

>>> print(model("The quick brown fox jumps ", stop=["."])["choices"][0]["text"])

the lazy dog

Loading a chat model

>>> import llama_cpp

>>> model = llama_cpp.Llama(

... model_path="path/to/model",

... chat_format="llama-2",

... )

>>> print(model.create_chat_completion(

... messages=[{

... "role": "user",

... "content": "what is the meaning of life?"

... }]

... ))

Args:

model_path: Path to the model.

n_gpu_layers: Number of layers to offload to GPU (-ngl). If -1, all layers are offloaded.

main_gpu: The GPU that is used for scratch and small tensors.

tensor_split: How split tensors should be distributed across GPUs. If None, the model is not split.

vocab_only: Only load the vocabulary no weights.

use_mmap: Use mmap if possible.

use_mlock: Force the system to keep the model in RAM.

seed: RNG seed, -1 for random

n_ctx: Text context, 0 = from model

n_batch: Prompt processing maximum batch size

n_threads: Number of threads to use for generation

n_threads_batch: Number of threads to use for batch processing

rope_scaling_type: RoPE scaling type, from `enum llama_rope_scaling_type`. ref: https://github.com/ggerganov/llama.cpp/pull/2054

rope_freq_base: RoPE base frequency, 0 = from model

rope_freq_scale: RoPE frequency scaling factor, 0 = from model

yarn_ext_factor: YaRN extrapolation mix factor, negative = from model

yarn_attn_factor: YaRN magnitude scaling factor

yarn_beta_fast: YaRN low correction dim

yarn_beta_slow: YaRN high correction dim

yarn_orig_ctx: YaRN original context size

logits_all: Return logits for all tokens, not just the last token. Must be True for completion to return logprobs.

embedding: Embedding mode only.

offload_kqv: Offload K, Q, V to GPU.

last_n_tokens_size: Maximum number of tokens to keep in the last_n_tokens deque.

lora_base: Optional path to base model, useful if using a quantized base model and you want to apply LoRA to an f16 model.

lora_path: Path to a LoRA file to apply to the model.

numa: 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.

chat_handler: Optional chat handler to use when calling create_chat_completion.

verbose: Print verbose output to stderr.

Raises:

ValueError: If the model path does not exist.

Returns:

A Llama instance.

"""

self.verbose = verbose

self.numa = numa

if not Llama.__backend_initialized:

with suppress_stdout_stderr(disable=self.verbose):

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_scaling_type = (

rope_scaling_type

if rope_scaling_type is not None

else llama_cpp.LLAMA_ROPE_SCALING_UNSPECIFIED

)

self.context_params.rope_freq_base = (

rope_freq_base if rope_freq_base != 0.0 else 0

)

self.context_params.rope_freq_scale = (

rope_freq_scale if rope_freq_scale != 0.0 else 0

)

self.context_params.yarn_ext_factor = (

yarn_ext_factor if yarn_ext_factor != 0.0 else 0

)

self.context_params.yarn_attn_factor = (

yarn_attn_factor if yarn_attn_factor != 0.0 else 0

)

self.context_params.yarn_beta_fast = (

yarn_beta_fast if yarn_beta_fast != 0.0 else 0

)

self.context_params.yarn_beta_slow = (

yarn_beta_slow if yarn_beta_slow != 0.0 else 0

)

self.context_params.yarn_orig_ctx = yarn_orig_ctx if yarn_orig_ctx != 0 else 0

self.context_params.mul_mat_q = mul_mat_q

self.context_params.logits_all = logits_all

self.context_params.embedding = embedding

self.context_params.offload_kqv = offload_kqv

# Sampling Params

self.last_n_tokens_size = last_n_tokens_size

self.cache: Optional[BaseLlamaCache] = None

self.lora_base = lora_base

self.lora_scale = lora_scale

self.lora_path = lora_path

if not os.path.exists(model_path):

raise ValueError(f"Model path does not exist: {model_path}")

self._model = _LlamaModel(

path_model=self.model_path, params=self.model_params, verbose=self.verbose

)

# Set the default value for the context and correct the batch

if n_ctx == 0:

n_ctx = self._model.n_ctx_train()

self.n_batch = min(n_ctx, n_batch)

self.context_params.n_ctx = self._model.n_ctx_train()

self.context_params.n_batch = self.n_batch

self._ctx = _LlamaContext(

model=self._model,

params=self.context_params,

verbose=self.verbose,

)

self._batch = _LlamaBatch(

n_tokens=self.n_batch,

embd=0,

n_seq_max=self.context_params.n_ctx,

verbose=self.verbose,

)

if self.lora_path:

if self._model.apply_lora_from_file(

self.lora_path,

self.lora_scale,

self.lora_base,

self.n_threads,

):

raise RuntimeError(

f"Failed to apply LoRA from lora path: {self.lora_path} to base path: {self.lora_base}"

)

if self.verbose:

print(llama_cpp.llama_print_system_info().decode("utf-8"), file=sys.stderr)

self.chat_format = chat_format

self.chat_handler = chat_handler

self._n_vocab = self.n_vocab()

self._n_ctx = self.n_ctx()

self._token_nl = self.token_nl()

self._token_eos = self.token_eos()

self._candidates = _LlamaTokenDataArray(n_vocab=self._n_vocab)

self.n_tokens = 0

self.input_ids: npt.NDArray[np.intc] = np.ndarray((n_ctx,), dtype=np.intc)

self.scores: npt.NDArray[np.single] = np.ndarray(

(n_ctx, self._n_vocab), dtype=np.single

)

@property

def ctx(self) -> llama_cpp.llama_context_p:

assert self._ctx.ctx is not None

return self._ctx.ctx

@property

def model(self) -> llama_cpp.llama_model_p:

assert self._model.model is not None

return self._model.model

@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.

"""

return self._model.tokenize(text, add_bos, special)

def detokenize(self, tokens: List[int]) -> bytes:

"""Detokenize a list of tokens.

Args:

tokens: The list of tokens to detokenize.

Returns:

The detokenized string.

"""

return self._model.detokenize(tokens)

def set_cache(self, cache: Optional[BaseLlamaCache]):

"""Set the cache.

Args:

cache: The cache to set.

"""

self.cache = cache

def set_seed(self, seed: int):

"""Set the random seed.

Args:

seed: The random seed.

"""

assert self._ctx.ctx is not None

llama_cpp.llama_set_rng_seed(self._ctx.ctx, seed)

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.ctx is not None

assert self._batch.batch is not None

self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1)

for i in range(0, len(tokens), self.n_batch):

batch = tokens[i : min(len(tokens), i + self.n_batch)]

n_past = self.n_tokens

n_tokens = len(batch)

self._batch.set_batch(

batch=batch, n_past=n_past, logits_all=self.context_params.logits_all

)

self._ctx.decode(self._batch)

# Save tokens

self.input_ids[n_past : n_past + n_tokens] = batch

# Save logits

rows = n_tokens

cols = self._n_vocab

offset = (

0 if self.context_params.logits_all else n_tokens - 1

) # NOTE: Only save the last token logits if logits_all is False

self.scores[n_past + offset : n_past + n_tokens, :].reshape(-1)[

:

] = self._ctx.get_logits()[offset * cols : rows * cols]

# Update n_tokens

self.n_tokens += n_tokens

def sample(

self,

top_k: int = 40,

top_p: float = 0.95,

min_p: float = 0.05,

typical_p: float = 1.0,

temp: float = 0.80,

repeat_penalty: float = 1.1,

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

assert self.n_tokens > 0

last_n_tokens_data = [llama_cpp.llama_token(0)] * max(

0, self.last_n_tokens_size - self.n_tokens

) + self._input_ids[-self.last_n_tokens_size :].tolist()

last_n_tokens_size = len(last_n_tokens_data)

n_vocab = self._n_vocab

n_ctx = self._n_ctx

top_k = n_vocab if top_k <= 0 else top_k

last_n_tokens_size = n_ctx if last_n_tokens_size < 0 else last_n_tokens_size

last_n_tokens_data_c = (llama_cpp.llama_token * last_n_tokens_size)(

*last_n_tokens_data

)

logits: npt.NDArray[np.single] = self._scores[-1, :]

if logits_processor is not None:

logits[:] = logits_processor(self._input_ids, logits)

nl_logit = logits[self._token_nl]

self._candidates.copy_logits(logits)

self._ctx.sample_repetition_penalties(

candidates=self._candidates,

last_tokens_data=last_n_tokens_data_c,

penalty_last_n=last_n_tokens_size,

penalty_repeat=repeat_penalty,

penalty_freq=frequency_penalty,

penalty_present=presence_penalty,

)

if not penalize_nl:

self._candidates.candidates.data[self._token_nl].logit = llama_cpp.c_float(

nl_logit

)

if grammar is not None:

self._ctx.sample_grammar(

candidates=self._candidates,

grammar=grammar,

)

if temp < 0.0:

self._ctx.sample_softmax(candidates=self._candidates)

id = self._candidates.candidates.data[0].id

elif temp == 0.0:

id = self._ctx.sample_token_greedy(candidates=self._candidates)

elif mirostat_mode == 1:

mu = 2.0 * mirostat_tau

self._ctx.sample_temp(candidates=self._candidates, temp=temp)

id = self._ctx.sample_token_mirostat(

candidates=self._candidates,

tau=mirostat_tau,

eta=mirostat_eta,

mu=ctypes.pointer(ctypes.c_float(mu)),

m=100,

)

elif mirostat_mode == 2:

mu = 2.0 * mirostat_tau

self._ctx.sample_temp(candidates=self._candidates, temp=temp)

id = self._ctx.sample_token_mirostat_v2(

candidates=self._candidates,

tau=mirostat_tau,

eta=mirostat_eta,

mu=ctypes.pointer(ctypes.c_float(mu)),

)

else:

self._ctx.sample_top_k(candidates=self._candidates, k=top_k, min_keep=1)

self._ctx.sample_tail_free(candidates=self._candidates, z=tfs_z, min_keep=1)

self._ctx.sample_typical(

candidates=self._candidates, p=typical_p, min_keep=1

)

self._ctx.sample_top_p(candidates=self._candidates, p=top_p, min_keep=1)

self._ctx.sample_min_p(candidates=self._candidates, p=min_p, min_keep=1)

self._ctx.sample_temp(candidates=self._candidates, temp=temp)

id = self._ctx.sample_token(candidates=self._candidates)

if grammar is not None:

self._ctx.grammar_accept_token(grammar=grammar, token=id)

return id

def generate(

self,

tokens: Sequence[int],

top_k: int = 40,

top_p: float = 0.95,

min_p: float = 0.05,

typical_p: float = 1.0,

temp: float = 0.80,

repeat_penalty: float = 1.1,

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.

"""

if reset and self.n_tokens > 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,

min_p=min_p,

typical_p=typical_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.ctx is not None

assert self._model.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.ctx)

if isinstance(input, str):

inputs = [input]

else:

inputs = input

data: List[Embedding] = []

total_tokens = 0

for index, input in enumerate(inputs):

tokens = self.tokenize(input.encode("utf-8"), special=True)

self.reset()

self.eval(tokens)

n_tokens = len(tokens)

total_tokens += n_tokens

embedding = llama_cpp.llama_get_embeddings(self._ctx.ctx)[

: llama_cpp.llama_n_embd(self._model.model)

]

data.append(

{

"object": "embedding",

"embedding": embedding,

"index": index,

}

)

if self.verbose:

llama_cpp.llama_print_timings(self._ctx.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: Union[str, List[int]],

suffix: Optional[str] = None,

max_tokens: Optional[int] = 16,

temperature: float = 0.8,

top_p: float = 0.95,

min_p: float = 0.05,

typical_p: float = 1.0,

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,

seed: Optional[int] = None,

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,

logit_bias: Optional[Dict[str, float]] = None,

) -> Union[

Iterator[CreateCompletionResponse], Iterator[CreateCompletionStreamResponse]

]:

assert self._ctx is not None

assert suffix is None or suffix.__class__ is str

completion_id: str = f"cmpl-{str(uuid.uuid4())}"

created: int = int(time.time())

# If prompt is empty, initialize completion with BOS token to avoid

# detokenization including a space at the beginning of the completion

completion_tokens: List[int] = [] if len(prompt) > 0 else [self.token_bos()]

# Add blank space to start of prompt to match OG llama tokenizer

prompt_tokens: List[int] = (

(

self.tokenize(prompt.encode("utf-8"), special=True)

if prompt != ""

else [self.token_bos()]

)

if isinstance(prompt, str)

else prompt

)

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

# NOTE: This likely doesn't work correctly for the first token in the prompt

# because of the extra space added to the start of the prompt_tokens

if logit_bias is not None:

logit_bias_map = {int(k): float(v) for k, v in logit_bias.items()}

def logit_bias_processor(

input_ids: npt.NDArray[np.intc],

scores: npt.NDArray[np.single],

) -> npt.NDArray[np.single]:

new_scores = np.copy(

scores

) # Does it make sense to copy the whole array or can we just overwrite the original one?

for input_id, score in logit_bias_map.items():

new_scores[input_id] = score + scores[input_id]

return new_scores

_logit_bias_processor = LogitsProcessorList([logit_bias_processor])

if logits_processor is None:

logits_processor = _logit_bias_processor

else:

logits_processor = logits_processor.extend(_logit_bias_processor)

if self.verbose:

self._ctx.reset_timings()

if len(prompt_tokens) >= self._n_ctx:

raise ValueError(

f"Requested tokens ({len(prompt_tokens)}) exceed context window of {llama_cpp.llama_n_ctx(self.ctx)}"

)

if max_tokens is None or max_tokens <= 0:

# Unlimited, depending on n_ctx.

max_tokens = self._n_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)

if seed is not None:

self._ctx.set_rng_seed(seed)

def _completion_stream_response(text: str, logprobs_or_none: Optional[CompletionLogprobs] = None, finish_reason: Optional[Literal["stop", "length"]] = None) -> CreateCompletionStreamResponse:

return {

"id": completion_id,

"object": "text_completion",

"created": created,

"model": model_name,

"choices": [

{

"text": text,

"index": 0,

"logprobs": logprobs_or_none,

"finish_reason": finish_reason,

}

],

}

def _completion_response(text: str, finish_reason: Literal["stop", "length"], logprobs_or_none: Optional[CompletionLogprobs] = None) -> CreateCompletionResponse:

return {

"id": completion_id,

"object": "text_completion",

"created": created,

"model": model_name,

"choices": [

{

"text": text,

"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 _logprobs_or_none(all_tokens: List[int], all_token_strs: List[str], all_logprobs: List[List[float]], text_offset: int) -> CompletionLogprobs:

tokens: List[str] = []

text_offsets: List[int] = []

token_logprobs: List[Optional[float]] = []

top_logprobs: List[Optional[Dict[str, float]]] = []

for token, token_str, token_logprob in zip(

all_tokens, all_token_strs, all_logprobs

):

if token == self.token_bos():

continue

text_offset += len(token_str)

sorted_logprobs = list(

sorted(

zip(token_logprob, range(len(token_logprob))), 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: token_logprob[int(token)]})

tokens.append(token_str)

text_offsets.append(text_offset)

token_logprobs.append(token_logprob[int(token)])

top_logprobs.append(top_logprob)

return {

"tokens": tokens,

"text_offset": text_offsets,

"token_logprobs": token_logprobs,

"top_logprobs": top_logprobs,

}

finish_reason = "length"

multibyte_fix = 0

for token in self.generate(

prompt_tokens,

top_k=top_k,

top_p=top_p,

min_p=min_p,

typical_p=typical_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:

if token == self.token_bos():

continue

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, :]

token_logprob = Llama.logits_to_logprobs(logits).tolist()

sorted_logprobs = list(

sorted(

zip(token_logprob, range(len(token_logprob))),

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: token_logprob[int(token)]})

logprobs_or_none = {

"tokens": [

self.detokenize([token]).decode(

"utf-8", errors="ignore"

)

],

"text_offset": [text_offset],

"token_logprobs": [token_logprob[int(token)]],

"top_logprobs": [top_logprob],

}

returned_tokens += 1

yield _completion_stream_response(

self.detokenize([token]).decode("utf-8", errors="ignore"),

logprobs_or_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