hymenjj/llama-cpp-python-prebuilt
0
1from __future__ import annotations2 3import os4import sys5import uuid6import time7import json8import ctypes9import typing10import random11import fnmatch12import warnings13import contextlib14import multiprocessing15 16from typing import (17 Any,18 List,19 Literal,20 Optional,21 Union,22 Generator,23 Sequence,24 Iterator,25 Deque,26 Callable,27 Dict,28)29from collections import deque30from pathlib import Path31 32 33from .llama_types import *34from .llama_grammar import LlamaGrammar35from .llama_cache import (36 BaseLlamaCache,37 LlamaCache, # type: ignore38 LlamaDiskCache, # type: ignore39 LlamaRAMCache, # type: ignore40)41from .llama_tokenizer import BaseLlamaTokenizer, LlamaTokenizer42import llama_cpp.llama_cpp as llama_cpp43import llama_cpp.llama_chat_format as llama_chat_format44 45from llama_cpp.llama_speculative import LlamaDraftModel46 47import numpy as np48import numpy.typing as npt49 50import llama_cpp._internals as internals51from ._logger import set_verbose52from ._utils import suppress_stdout_stderr53 54 55class Llama:56 """High-level Python wrapper for a llama.cpp model."""57 58 __backend_initialized = False59 60 def __init__(61 self,62 model_path: str,63 *,64 # Model Params65 n_gpu_layers: int = 0,66 split_mode: int = llama_cpp.LLAMA_SPLIT_MODE_LAYER,67 main_gpu: int = 0,68 tensor_split: Optional[List[float]] = None,69 vocab_only: bool = False,70 use_mmap: bool = True,71 use_mlock: bool = False,72 kv_overrides: Optional[Dict[str, Union[bool, int, float, str]]] = None,73 # Context Params74 seed: int = llama_cpp.LLAMA_DEFAULT_SEED,75 n_ctx: int = 512,76 n_batch: int = 512,77 n_ubatch: int = 512,78 n_threads: Optional[int] = None,79 n_threads_batch: Optional[int] = None,80 rope_scaling_type: Optional[81 int82 ] = llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED,83 pooling_type: int = llama_cpp.LLAMA_POOLING_TYPE_UNSPECIFIED,84 rope_freq_base: float = 0.0,85 rope_freq_scale: float = 0.0,86 yarn_ext_factor: float = -1.0,87 yarn_attn_factor: float = 1.0,88 yarn_beta_fast: float = 32.0,89 yarn_beta_slow: float = 1.0,90 yarn_orig_ctx: int = 0,91 logits_all: bool = False,92 embedding: bool = False,93 offload_kqv: bool = True,94 flash_attn: bool = False,95 op_offload: Optional[bool] = None,96 swa_full: Optional[bool] = None,97 # Sampling Params98 no_perf: bool = False,99 last_n_tokens_size: int = 64,100 # LoRA Params101 lora_base: Optional[str] = None,102 lora_scale: float = 1.0,103 lora_path: Optional[str] = None,104 # Backend Params105 numa: Union[bool, int] = False,106 # Chat Format Params107 chat_format: Optional[str] = None,108 chat_handler: Optional[llama_chat_format.LlamaChatCompletionHandler] = None,109 # Speculative Decoding110 draft_model: Optional[LlamaDraftModel] = None,111 # Tokenizer Override112 tokenizer: Optional[BaseLlamaTokenizer] = None,113 # KV cache quantization114 type_k: Optional[int] = None,115 type_v: Optional[int] = None,116 # Misc117 spm_infill: bool = False,118 verbose: bool = True,119 # Extra Params120 **kwargs, # type: ignore121 ):122 """Load a llama.cpp model from `model_path`.123 124 Examples:125 Basic usage126 127 >>> import llama_cpp128 >>> model = llama_cpp.Llama(129 ... model_path="path/to/model",130 ... )131 >>> print(model("The quick brown fox jumps ", stop=["."])["choices"][0]["text"])132 the lazy dog133 134 Loading a chat model135 136 >>> import llama_cpp137 >>> model = llama_cpp.Llama(138 ... model_path="path/to/model",139 ... chat_format="llama-2",140 ... )141 >>> print(model.create_chat_completion(142 ... messages=[{143 ... "role": "user",144 ... "content": "what is the meaning of life?"145 ... }]146 ... ))147 148 Args:149 model_path: Path to the model.150 n_gpu_layers: Number of layers to offload to GPU (-ngl). If -1, all layers are offloaded.151 split_mode: How to split the model across GPUs. See llama_cpp.LLAMA_SPLIT_* for options.152 main_gpu: main_gpu interpretation depends on split_mode: LLAMA_SPLIT_MODE_NONE: the GPU that is used for the entire model. LLAMA_SPLIT_MODE_ROW: the GPU that is used for small tensors and intermediate results. LLAMA_SPLIT_MODE_LAYER: ignored153 tensor_split: How split tensors should be distributed across GPUs. If None, the model is not split.154 vocab_only: Only load the vocabulary no weights.155 use_mmap: Use mmap if possible.156 use_mlock: Force the system to keep the model in RAM.157 kv_overrides: Key-value overrides for the model.158 seed: RNG seed, -1 for random159 n_ctx: Text context, 0 = from model160 n_batch: Prompt processing maximum batch size161 n_ubatch: Physical batch size162 n_threads: Number of threads to use for generation163 n_threads_batch: Number of threads to use for batch processing164 rope_scaling_type: RoPE scaling type, from `enum llama_rope_scaling_type`. ref: https://github.com/ggerganov/llama.cpp/pull/2054165 pooling_type: Pooling type, from `enum llama_pooling_type`.166 rope_freq_base: RoPE base frequency, 0 = from model167 rope_freq_scale: RoPE frequency scaling factor, 0 = from model168 yarn_ext_factor: YaRN extrapolation mix factor, negative = from model169 yarn_attn_factor: YaRN magnitude scaling factor170 yarn_beta_fast: YaRN low correction dim171 yarn_beta_slow: YaRN high correction dim172 yarn_orig_ctx: YaRN original context size173 logits_all: Return logits for all tokens, not just the last token. Must be True for completion to return logprobs.174 embedding: Embedding mode only.175 offload_kqv: Offload K, Q, V to GPU.176 flash_attn: Use flash attention.177 op_offload: offload host tensor operations to device178 swa_full: use full-size SWA cache (https://github.com/ggml-org/llama.cpp/pull/13194#issuecomment-2868343055)179 no_perf: Measure performance timings.180 last_n_tokens_size: Maximum number of tokens to keep in the last_n_tokens deque.181 lora_base: Optional path to base model, useful if using a quantized base model and you want to apply LoRA to an f16 model.182 lora_path: Path to a LoRA file to apply to the model.183 numa: numa policy184 chat_format: String specifying the chat format to use when calling create_chat_completion.185 chat_handler: Optional chat handler to use when calling create_chat_completion.186 draft_model: Optional draft model to use for speculative decoding.187 tokenizer: Optional tokenizer to override the default tokenizer from llama.cpp.188 verbose: Print verbose output to stderr.189 type_k: KV cache data type for K (default: f16)190 type_v: KV cache data type for V (default: f16)191 spm_infill: Use Suffix/Prefix/Middle pattern for infill (instead of Prefix/Suffix/Middle) as some models prefer this.192 193 Raises:194 ValueError: If the model path does not exist.195 196 Returns:197 A Llama instance.198 """199 self.verbose = verbose200 self._stack = contextlib.ExitStack()201 202 set_verbose(verbose)203 204 if not Llama.__backend_initialized:205 with suppress_stdout_stderr(disable=verbose):206 llama_cpp.llama_backend_init()207 Llama.__backend_initialized = True208 209 if isinstance(numa, bool):210 self.numa = (211 llama_cpp.GGML_NUMA_STRATEGY_DISTRIBUTE212 if numa213 else llama_cpp.GGML_NUMA_STRATEGY_DISABLED214 )215 else:216 self.numa = numa217 218 if self.numa != llama_cpp.GGML_NUMA_STRATEGY_DISABLED:219 with suppress_stdout_stderr(disable=verbose):220 llama_cpp.llama_numa_init(self.numa)221 222 self.model_path = model_path223 224 # Model Params225 self.model_params = llama_cpp.llama_model_default_params()226 self.model_params.n_gpu_layers = (227 0x7FFFFFFF if n_gpu_layers == -1 else n_gpu_layers228 ) # 0x7FFFFFFF is INT32 max, will be auto set to all layers229 self.model_params.split_mode = split_mode230 self.model_params.main_gpu = main_gpu231 self.tensor_split = tensor_split232 self._c_tensor_split = None233 if self.tensor_split is not None:234 if len(self.tensor_split) > llama_cpp.LLAMA_MAX_DEVICES:235 raise ValueError(236 f"Attempt to split tensors that exceed maximum supported devices. Current LLAMA_MAX_DEVICES={llama_cpp.LLAMA_MAX_DEVICES}"237 )238 # Type conversion and expand the list to the length of LLAMA_MAX_DEVICES239 FloatArray = ctypes.c_float * llama_cpp.LLAMA_MAX_DEVICES240 self._c_tensor_split = FloatArray(241 *tensor_split # type: ignore242 ) # keep a reference to the array so it is not gc'd243 self.model_params.tensor_split = self._c_tensor_split244 self.model_params.vocab_only = vocab_only245 self.model_params.use_mmap = use_mmap if lora_path is None else False246 self.model_params.use_mlock = use_mlock247 248 # kv_overrides is the original python dict249 self.kv_overrides = kv_overrides250 if kv_overrides is not None:251 # _kv_overrides_array is a ctypes.Array of llama_model_kv_override Structs252 kvo_array_len = len(kv_overrides) + 1 # for sentinel element253 self._kv_overrides_array = (254 llama_cpp.llama_model_kv_override * kvo_array_len255 )()256 257 for i, (k, v) in enumerate(kv_overrides.items()):258 self._kv_overrides_array[i].key = k.encode("utf-8")259 if isinstance(v, bool):260 self._kv_overrides_array[261 i262 ].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_BOOL263 self._kv_overrides_array[i].value.val_bool = v264 elif isinstance(v, int):265 self._kv_overrides_array[266 i267 ].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_INT268 self._kv_overrides_array[i].value.val_i64 = v269 elif isinstance(v, float):270 self._kv_overrides_array[271 i272 ].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_FLOAT273 self._kv_overrides_array[i].value.val_f64 = v274 elif isinstance(v, str): # type: ignore275 v_bytes = v.encode("utf-8")276 if len(v_bytes) > 128: # TODO: Make this a constant277 raise ValueError(f"Value for {k} is too long: {v}")278 v_bytes = v_bytes.ljust(128, b"\0")279 self._kv_overrides_array[280 i281 ].tag = llama_cpp.LLAMA_KV_OVERRIDE_TYPE_STR282 # copy min(v_bytes, 128) to str_value283 address = typing.cast(284 int,285 ctypes.addressof(self._kv_overrides_array[i].value)286 + llama_cpp.llama_model_kv_override_value.val_str.offset,287 )288 buffer_start = ctypes.cast(address, ctypes.POINTER(ctypes.c_char))289 ctypes.memmove(290 buffer_start,291 v_bytes,292 128,293 )294 else:295 raise ValueError(f"Unknown value type for {k}: {v}")296 297 self._kv_overrides_array[298 -1299 ].key = b"\0" # ensure sentinel element is zeroed300 self.model_params.kv_overrides = self._kv_overrides_array301 302 self.n_batch = min(n_ctx, n_batch) # ???303 self.n_threads = n_threads or max(multiprocessing.cpu_count() // 2, 1)304 self.n_threads_batch = n_threads_batch or multiprocessing.cpu_count()305 306 # Used by the sampler307 self._seed = seed or llama_cpp.LLAMA_DEFAULT_SEED308 309 # Context Params310 self.context_params = llama_cpp.llama_context_default_params()311 self.context_params.n_ctx = n_ctx312 self.context_params.n_batch = self.n_batch313 self.context_params.n_ubatch = min(self.n_batch, n_ubatch)314 self.context_params.n_threads = self.n_threads315 self.context_params.n_threads_batch = self.n_threads_batch316 self.context_params.rope_scaling_type = (317 rope_scaling_type318 if rope_scaling_type is not None319 else llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED320 )321 self.context_params.pooling_type = pooling_type322 self.context_params.rope_freq_base = (323 rope_freq_base if rope_freq_base != 0.0 else 0324 )325 self.context_params.rope_freq_scale = (326 rope_freq_scale if rope_freq_scale != 0.0 else 0327 )328 self.context_params.yarn_ext_factor = (329 yarn_ext_factor if yarn_ext_factor != 0.0 else 0330 )331 self.context_params.yarn_attn_factor = (332 yarn_attn_factor if yarn_attn_factor != 0.0 else 0333 )334 self.context_params.yarn_beta_fast = (335 yarn_beta_fast if yarn_beta_fast != 0.0 else 0336 )337 self.context_params.yarn_beta_slow = (338 yarn_beta_slow if yarn_beta_slow != 0.0 else 0339 )340 self.context_params.yarn_orig_ctx = yarn_orig_ctx if yarn_orig_ctx != 0 else 0341 self._logits_all = logits_all if draft_model is None else True342 self.context_params.embeddings = embedding # TODO: Rename to embeddings343 self.context_params.offload_kqv = offload_kqv344 self.context_params.flash_attn = flash_attn345 346 if op_offload is not None:347 self.context_params.op_offload = op_offload348 349 if swa_full is not None:350 self.context_params.swa_full = swa_full351 352 # KV cache quantization353 if type_k is not None:354 self.context_params.type_k = type_k355 if type_v is not None:356 self.context_params.type_v = type_v357 # Sampling Params358 self.context_params.no_perf = no_perf359 self.last_n_tokens_size = last_n_tokens_size360 361 self.cache: Optional[BaseLlamaCache] = None362 363 self.lora_base = lora_base364 self.lora_scale = lora_scale365 self.lora_path = lora_path366 367 self.spm_infill = spm_infill368 369 if not os.path.exists(model_path):370 raise ValueError(f"Model path does not exist: {model_path}")371 372 self._model = self._stack.enter_context(373 contextlib.closing(374 internals.LlamaModel(375 path_model=self.model_path,376 params=self.model_params,377 verbose=self.verbose,378 )379 )380 )381 382 # Override tokenizer383 self.tokenizer_ = tokenizer or LlamaTokenizer(self)384 385 # Set the default value for the context and correct the batch386 if n_ctx == 0:387 n_ctx = self._model.n_ctx_train()388 self.n_batch = min(n_ctx, n_batch)389 self.context_params.n_ctx = self._model.n_ctx_train()390 self.context_params.n_batch = self.n_batch391 self.context_params.n_ubatch = min(self.n_batch, n_ubatch)392 393 self._ctx = self._stack.enter_context(394 contextlib.closing(395 internals.LlamaContext(396 model=self._model,397 params=self.context_params,398 verbose=self.verbose,399 )400 )401 )402 403 self._batch = self._stack.enter_context(404 contextlib.closing(405 internals.LlamaBatch(406 n_tokens=self.n_batch,407 embd=0,408 n_seq_max=self.context_params.n_ctx,409 verbose=self.verbose,410 )411 )412 )413 414 self._lora_adapter: Optional[llama_cpp.llama_adapter_lora_p] = None415 416 if self.lora_path:417 self._lora_adapter = llama_cpp.llama_adapter_lora_init(418 self._model.model,419 self.lora_path.encode("utf-8"),420 )421 if self._lora_adapter is None:422 raise RuntimeError(423 f"Failed to initialize LoRA adapter from lora path: {self.lora_path}"424 )425 426 def free_lora_adapter():427 if self._lora_adapter is None:428 return429 llama_cpp.llama_adapter_lora_free(self._lora_adapter)430 self._lora_adapter = None431 432 self._stack.callback(free_lora_adapter)433 434 if llama_cpp.llama_set_adapter_lora(435 self._ctx.ctx, self._lora_adapter, self.lora_scale436 ):437 raise RuntimeError(438 f"Failed to set LoRA adapter from lora path: {self.lora_path}"439 )440 441 if self.verbose:442 print(llama_cpp.llama_print_system_info().decode("utf-8"), file=sys.stderr)443 444 self.chat_format = chat_format445 self.chat_handler = chat_handler446 self._chat_handlers: Dict[447 str, llama_chat_format.LlamaChatCompletionHandler448 ] = {}449 450 self.draft_model = draft_model451 452 self._n_vocab = self.n_vocab()453 self._n_ctx = self.n_ctx()454 455 self._token_nl = self.token_nl()456 self._token_eos = self.token_eos()457 458 self._candidates = internals.LlamaTokenDataArray(n_vocab=self._n_vocab)459 460 self.n_tokens = 0461 self.input_ids: npt.NDArray[np.intc] = np.ndarray((n_ctx,), dtype=np.intc)462 self.scores: npt.NDArray[np.single] = np.ndarray(463 (n_ctx if logits_all == True else n_batch, self._n_vocab), dtype=np.single464 )465 466 self._mirostat_mu = ctypes.c_float(467 2.0 * 5.0468 ) # TODO: Move this to sampling context469 470 try:471 self.metadata = self._model.metadata()472 except Exception as e:473 self.metadata = {}474 if self.verbose:475 print(f"Failed to load metadata: {e}", file=sys.stderr)476 477 if self.verbose:478 print(f"Model metadata: {self.metadata}", file=sys.stderr)479 480 eos_token_id = self.token_eos()481 bos_token_id = self.token_bos()482 483 eos_token = (484 self._model.token_get_text(eos_token_id) if eos_token_id != -1 else ""485 )486 bos_token = (487 self._model.token_get_text(bos_token_id) if bos_token_id != -1 else ""488 )489 490 # Unfortunately the llama.cpp API does not return metadata arrays, so we can't get template names from tokenizer.chat_templates491 template_choices = dict(492 (name[10:], template)493 for name, template in self.metadata.items()494 if name.startswith("tokenizer.chat_template.")495 )496 497 if "tokenizer.chat_template" in self.metadata:498 template_choices["chat_template.default"] = self.metadata[499 "tokenizer.chat_template"500 ]501 502 if self.verbose and template_choices:503 print(504 f"Available chat formats from metadata: {', '.join(template_choices.keys())}",505 file=sys.stderr,506 )507 508 for name, template in template_choices.items():509 self._chat_handlers[name] = llama_chat_format.Jinja2ChatFormatter(510 template=template,511 eos_token=eos_token,512 bos_token=bos_token,513 stop_token_ids=[eos_token_id],514 ).to_chat_handler()515 516 if (517 self.chat_format is None518 and self.chat_handler is None519 and "chat_template.default" in template_choices520 ):521 chat_format = llama_chat_format.guess_chat_format_from_gguf_metadata(522 self.metadata523 )524 525 if chat_format is not None:526 self.chat_format = chat_format527 if self.verbose:528 print(f"Guessed chat format: {chat_format}", file=sys.stderr)529 else:530 if self.verbose:531 print(532 f"Using gguf chat template: {template_choices['chat_template.default']}",533 file=sys.stderr,534 )535 print(f"Using chat eos_token: {eos_token}", file=sys.stderr)536 print(f"Using chat bos_token: {bos_token}", file=sys.stderr)537 538 self.chat_format = "chat_template.default"539 540 if self.chat_format is None and self.chat_handler is None:541 self.chat_format = "llama-2"542 if self.verbose:543 print(544 f"Using fallback chat format: {self.chat_format}", file=sys.stderr545 )546 547 self._sampler = None548 549 @property550 def ctx(self) -> llama_cpp.llama_context_p:551 return self._ctx.ctx552 553 @property554 def model(self) -> llama_cpp.llama_model_p:555 return self._model.model556 557 @property558 def _input_ids(self) -> npt.NDArray[np.intc]:559 return self.input_ids[: self.n_tokens]560 561 @property562 def _scores(self) -> npt.NDArray[np.single]:563 return self.scores[: self.n_tokens, :]564 565 @property566 def eval_tokens(self) -> Deque[int]:567 return deque(self.input_ids[: self.n_tokens].tolist(), maxlen=self._n_ctx)568 569 @property570 def eval_logits(self) -> Deque[List[float]]:571 return deque(572 self.scores[: self.n_tokens, :].tolist(),573 maxlen=self._n_ctx if self._logits_all else 1,574 )575 576 def tokenize(577 self, text: bytes, add_bos: bool = True, special: bool = False578 ) -> List[int]:579 """Tokenize a string.580 581 Args:582 text: The utf-8 encoded string to tokenize.583 add_bos: Whether to add a beginning of sequence token.584 special: Whether to tokenize special tokens.585 586 Raises:587 RuntimeError: If the tokenization failed.588 589 Returns:590 A list of tokens.591 """592 return self.tokenizer_.tokenize(text, add_bos, special)593 594 def detokenize(595 self,596 tokens: List[int],597 prev_tokens: Optional[List[int]] = None,598 special: bool = False,599 ) -> bytes:600 """Detokenize a list of tokens.601 602 Args:603 tokens: The list of tokens to detokenize.604 prev_tokens: The list of previous tokens. Offset mapping will be performed if provided.605 special: Whether to detokenize special tokens.606 607 Returns:608 The detokenized string.609 """610 return self.tokenizer_.detokenize(611 tokens, prev_tokens=prev_tokens, special=special612 )613 614 def set_cache(self, cache: Optional[BaseLlamaCache]):615 """Set the cache.616 617 Args:618 cache: The cache to set.619 """620 self.cache = cache621 622 def set_seed(self, seed: int):623 """Set the random seed.624 625 Args:626 seed: The random seed.627 """628 self._seed = seed629 630 def reset(self):631 """Reset the model state."""632 self.n_tokens = 0633 634 def eval(self, tokens: Sequence[int]):635 """Evaluate a list of tokens.636 637 Args:638 tokens: The list of tokens to evaluate.639 """640 self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1)641 for i in range(0, len(tokens), self.n_batch):642 batch = tokens[i : min(len(tokens), i + self.n_batch)]643 n_past = self.n_tokens644 n_tokens = len(batch)645 self._batch.set_batch(646 batch=batch, n_past=n_past, logits_all=self._logits_all647 )648 self._ctx.decode(self._batch)649 # Save tokens650 self.input_ids[n_past : n_past + n_tokens] = batch651 # Save logits652 if self._logits_all:653 rows = n_tokens654 cols = self._n_vocab655 logits = np.ctypeslib.as_array(656 self._ctx.get_logits(), shape=(rows * cols,)657 )658 self.scores[n_past : n_past + n_tokens, :].reshape(-1)[::] = logits659 else:660 # rows = 1661 # cols = self._n_vocab662 # logits = np.ctypeslib.as_array(663 # self._ctx.get_logits(), shape=(rows * cols,)664 # )665 # self.scores[n_past + n_tokens - 1, :].reshape(-1)[::] = logits666 # NOTE: Now that sampling is done inside the sampler, logits are only needed for logprobs which requires logits_all667 pass668 # Update n_tokens669 self.n_tokens += n_tokens670 671 def _init_sampler(672 self,673 top_k: int = 40,674 top_p: float = 0.95,675 min_p: float = 0.05,676 typical_p: float = 1.0,677 temp: float = 0.80,678 repeat_penalty: float = 1.0,679 frequency_penalty: float = 0.0,680 presence_penalty: float = 0.0,681 tfs_z: float = 1.0,682 mirostat_mode: int = 0,683 mirostat_eta: float = 0.1,684 mirostat_tau: float = 5.0,685 penalize_nl: bool = True,686 logits_processor: Optional[LogitsProcessorList] = None,687 grammar: Optional[LlamaGrammar] = None,688 ):689 sampler = internals.LlamaSampler()690 691 if logits_processor is not None:692 # Create and add a custom sampler693 def apply_func(token_data_array: llama_cpp.llama_token_data_array_p):694 size = token_data_array.contents.size695 data_soa = token_data_array.contents.data696 data_soa_address = ctypes.addressof(data_soa.contents)697 # NOTE: This is probably broken698 recarray = np.recarray(699 shape=(size,),700 dtype=np.dtype(701 [("id", np.intc), ("logit", np.single), ("p", np.single)],702 align=True,703 ),704 buf=(llama_cpp.llama_token_data * size).from_address(705 data_soa_address706 ),707 )708 for logit_processor in logits_processor:709 recarray.logit[:] = logit_processor(self._input_ids, recarray.logit)710 711 sampler.add_custom(apply_func)712 713 sampler.add_penalties(714 # n_vocab=self._n_vocab,715 # special_eos_id=self._token_eos,716 # linefeed_id=self._token_nl,717 penalty_last_n=self.last_n_tokens_size,718 penalty_repeat=repeat_penalty,719 penalty_freq=frequency_penalty,720 penalty_present=presence_penalty,721 # penalize_nl=penalize_nl,722 # ignore_eos=False,723 )724 725 if grammar is not None:726 sampler.add_grammar(self._model, grammar)727 728 if temp < 0.0:729 sampler.add_softmax()730 sampler.add_dist(self._seed)731 elif temp == 0.0:732 sampler.add_greedy()733 else:734 if mirostat_mode == 1:735 mirostat_m = 100736 sampler.add_mirostat(737 self._n_vocab,738 self._seed,739 mirostat_tau,740 mirostat_eta,741 mirostat_m,742 )743 elif mirostat_mode == 2:744 sampler.add_mirostat_v2(745 self._seed,746 mirostat_tau,747 mirostat_eta,748 )749 else:750 n_probs = 0751 min_keep = max(1, n_probs)752 sampler.add_top_k(top_k)753 sampler.add_typical(typical_p, min_keep)754 sampler.add_top_p(top_p, min_keep)755 sampler.add_min_p(min_p, min_keep)756 sampler.add_temp(temp)757 sampler.add_dist(self._seed)758 return sampler759 760 def sample(761 self,762 top_k: int = 40,763 top_p: float = 0.95,764 min_p: float = 0.05,765 typical_p: float = 1.0,766 temp: float = 0.80,767 repeat_penalty: float = 1.0,768 frequency_penalty: float = 0.0,769 presence_penalty: float = 0.0,770 tfs_z: float = 1.0,771 mirostat_mode: int = 0,772 mirostat_eta: float = 0.1,773 mirostat_tau: float = 5.0,774 penalize_nl: bool = True,775 logits_processor: Optional[LogitsProcessorList] = None,776 grammar: Optional[LlamaGrammar] = None,777 idx: Optional[int] = None,778 ):779 """Sample a token from the model.780 781 Args:782 top_k: The top-k sampling parameter.783 top_p: The top-p sampling parameter.784 temp: The temperature parameter.785 repeat_penalty: The repeat penalty parameter.786 787 Returns:788 The sampled token.789 """790 assert self.n_tokens > 0791 792 tmp_sampler = False793 794 if self._sampler is None:795 tmp_sampler = True796 self._sampler = self._init_sampler(797 top_k=top_k,798 top_p=top_p,799 min_p=min_p,800 typical_p=typical_p,801 temp=temp,802 repeat_penalty=repeat_penalty,803 frequency_penalty=frequency_penalty,804 presence_penalty=presence_penalty,805 tfs_z=tfs_z,806 mirostat_mode=mirostat_mode,807 mirostat_tau=mirostat_tau,808 mirostat_eta=mirostat_eta,809 penalize_nl=penalize_nl,810 logits_processor=logits_processor,811 grammar=grammar,812 )813 814 ridx = idx - self.n_tokens if idx is not None else -1815 816 assert self.ctx is not None817 token = self._sampler.sample(self._ctx, ridx)818 if tmp_sampler:819 self._sampler = None820 return token821 822 def generate(823 self,824 tokens: Sequence[int],825 top_k: int = 40,826 top_p: float = 0.95,827 min_p: float = 0.05,828 typical_p: float = 1.0,829 temp: float = 0.80,830 repeat_penalty: float = 1.0,831 reset: bool = True,832 frequency_penalty: float = 0.0,833 presence_penalty: float = 0.0,834 tfs_z: float = 1.0,835 mirostat_mode: int = 0,836 mirostat_tau: float = 5.0,837 mirostat_eta: float = 0.1,838 penalize_nl: bool = True,839 logits_processor: Optional[LogitsProcessorList] = None,840 stopping_criteria: Optional[StoppingCriteriaList] = None,841 grammar: Optional[LlamaGrammar] = None,842 ) -> Generator[int, Optional[Sequence[int]], None]:843 """Create a generator of tokens from a prompt.844 845 Examples:846 >>> llama = Llama("models/ggml-7b.bin")847 >>> tokens = llama.tokenize(b"Hello, world!")848 >>> for token in llama.generate(tokens, top_k=40, top_p=0.95, temp=1.0, repeat_penalty=1.0):849 ... print(llama.detokenize([token]))850 851 Args:852 tokens: The prompt tokens.853 top_k: The top-k sampling parameter.854 top_p: The top-p sampling parameter.855 temp: The temperature parameter.856 repeat_penalty: The repeat penalty parameter.857 reset: Whether to reset the model state.858 859 Yields:860 The generated tokens.861 """862 # Reset mirostat sampling863 self._mirostat_mu = ctypes.c_float(2.0 * mirostat_tau)864 self._sampler = self._init_sampler(865 top_k=top_k,866 top_p=top_p,867 min_p=min_p,868 typical_p=typical_p,869 temp=temp,870 repeat_penalty=repeat_penalty,871 frequency_penalty=frequency_penalty,872 presence_penalty=presence_penalty,873 tfs_z=tfs_z,874 mirostat_mode=mirostat_mode,875 mirostat_tau=mirostat_tau,876 mirostat_eta=mirostat_eta,877 penalize_nl=penalize_nl,878 logits_processor=logits_processor,879 grammar=grammar,880 )881 882 # Check for kv cache prefix match883 if reset and self.n_tokens > 0:884 longest_prefix = 0885 for a, b in zip(self._input_ids, tokens[:-1]):886 if a == b:887 longest_prefix += 1888 else:889 break890 if longest_prefix > 0:891 reset = False892 tokens = tokens[longest_prefix:]893 self.n_tokens = longest_prefix894 if self.verbose:895 print(896 f"Llama.generate: {longest_prefix} prefix-match hit, "897 f"remaining {len(tokens)} prompt tokens to eval",898 file=sys.stderr,899 )900 901 # Reset the model state902 if reset:903 self.reset()904 905 # # Reset the grammar906 # if grammar is not None:907 # grammar.reset()908 909 sample_idx = self.n_tokens + len(tokens) - 1910 tokens = list(tokens)911 912 # Eval and sample913 while True:914 self.eval(tokens)915 while sample_idx < self.n_tokens:916 token = self.sample(917 top_k=top_k,918 top_p=top_p,919 min_p=min_p,920 typical_p=typical_p,921 temp=temp,922 repeat_penalty=repeat_penalty,923 frequency_penalty=frequency_penalty,924 presence_penalty=presence_penalty,925 tfs_z=tfs_z,926 mirostat_mode=mirostat_mode,927 mirostat_tau=mirostat_tau,928 mirostat_eta=mirostat_eta,929 logits_processor=logits_processor,930 grammar=grammar,931 penalize_nl=penalize_nl,932 idx=sample_idx,933 )934 935 sample_idx += 1936 if stopping_criteria is not None and stopping_criteria(937 self._input_ids[: sample_idx], self._scores[sample_idx - self.n_tokens, :]938 ):939 return940 tokens_or_none = yield token941 tokens.clear()942 tokens.append(token)943 if tokens_or_none is not None:944 tokens.extend(tokens_or_none)945 946 if sample_idx < self.n_tokens and token != self._input_ids[sample_idx]:947 self.n_tokens = sample_idx948 self._ctx.kv_cache_seq_rm(-1, self.n_tokens, -1)949 break950 951 if self.draft_model is not None:952 self.input_ids[self.n_tokens : self.n_tokens + len(tokens)] = tokens953 draft_tokens = self.draft_model(954 self.input_ids[: self.n_tokens + len(tokens)]955 )956 tokens.extend(957 draft_tokens.astype(int)[958 : self._n_ctx - self.n_tokens - len(tokens)959 ]960 )961 962 def create_embedding(963 self, input: Union[str, List[str]], model: Optional[str] = None964 ) -> CreateEmbeddingResponse:965 """Embed a string.966 967 Args:968 input: The utf-8 encoded string to embed.969 970 Returns:971 An embedding object.972 """973 model_name: str = model if model is not None else self.model_path974 975 input = input if isinstance(input, list) else [input]976 977 # get numeric embeddings978 embeds: Union[List[List[float]], List[List[List[float]]]]979 total_tokens: int980 embeds, total_tokens = self.embed(input, return_count=True) # type: ignore981 982 # convert to CreateEmbeddingResponse983 data: List[Embedding] = [984 {985 "object": "embedding",986 "embedding": emb,987 "index": idx,988 }989 for idx, emb in enumerate(embeds)990 ]991 992 return {993 "object": "list",994 "data": data,995 "model": model_name,996 "usage": {997 "prompt_tokens": total_tokens,998 "total_tokens": total_tokens,999 },1000 }1001 1002 def embed(1003 self,1004 input: Union[str, List[str]],1005 normalize: bool = False,1006 truncate: bool = True,1007 return_count: bool = False,1008 ):1009 """Embed a string.1010 1011 Args:1012 input: The utf-8 encoded string to embed.1013 1014 Returns:1015 A list of embeddings1016 """1017 n_embd = self.n_embd()1018 n_batch = self.n_batch1019 1020 # get pooling information1021 pooling_type = self.pooling_type()1022 logits_all = pooling_type == llama_cpp.LLAMA_POOLING_TYPE_NONE1023 1024 if self.context_params.embeddings is False:1025 raise RuntimeError(1026 "Llama model must be created with embedding=True to call this method"1027 )1028 1029 if self.verbose:1030 llama_cpp.llama_perf_context_reset(self._ctx.ctx)1031 1032 if isinstance(input, str):1033 inputs = [input]1034 else:1035 inputs = input1036 1037 # reset batch1038 self._batch.reset()1039 1040 # decode and fetch embeddings1041 data: Union[List[List[float]], List[List[List[float]]]] = []1042 1043 def decode_batch(seq_sizes: List[int]):1044 llama_cpp.llama_kv_self_clear(self._ctx.ctx)1045 self._ctx.decode(self._batch)1046 self._batch.reset()1047 1048 # store embeddings1049 if pooling_type == llama_cpp.LLAMA_POOLING_TYPE_NONE:1050 pos: int = 01051 for i, size in enumerate(seq_sizes):1052 ptr = llama_cpp.llama_get_embeddings(self._ctx.ctx)1053 embedding: List[List[float]] = [1054 ptr[pos + j * n_embd : pos + (j + 1) * n_embd]1055 for j in range(size)1056 ]1057 if normalize:1058 embedding = [1059 internals.normalize_embedding(e) for e in embedding1060 ]1061 data.append(embedding)1062 pos += size1063 else:1064 for i in range(len(seq_sizes)):1065 ptr = llama_cpp.llama_get_embeddings_seq(self._ctx.ctx, i)1066 embedding: List[float] = ptr[:n_embd]1067 if normalize:1068 embedding = internals.normalize_embedding(embedding)1069 data.append(embedding)1070 1071 # init state1072 total_tokens = 01073 s_batch = []1074 t_batch = 01075 p_batch = 01076 1077 # accumulate batches and encode1078 for text in inputs:1079 tokens = self.tokenize(text.encode("utf-8"))1080 if truncate:1081 tokens = tokens[:n_batch]1082 1083 n_tokens = len(tokens)1084 total_tokens += n_tokens1085 1086 # check for overrun1087 if n_tokens > n_batch:1088 raise ValueError(1089 f"Requested tokens ({n_tokens}) exceed batch size of {n_batch}"1090 )1091 1092 # time to eval batch1093 if t_batch + n_tokens > n_batch:1094 decode_batch(s_batch)1095 s_batch = []1096 t_batch = 01097 p_batch = 01098 1099 # add to batch1100 self._batch.add_sequence(tokens, p_batch, logits_all)1101 1102 # update batch stats1103 s_batch.append(n_tokens)1104 t_batch += n_tokens1105 p_batch += 11106 1107 # hanlde last batch1108 decode_batch(s_batch)1109 1110 if self.verbose:1111 llama_cpp.llama_perf_context_print(self._ctx.ctx)1112 1113 output = data[0] if isinstance(input, str) else data1114 1115 llama_cpp.llama_kv_self_clear(self._ctx.ctx)1116 self.reset()1117 1118 if return_count:1119 return output, total_tokens1120 else:1121 return output1122 1123 def _create_completion(1124 self,1125 prompt: Union[str, List[int]],1126 suffix: Optional[str] = None,1127 max_tokens: Optional[int] = 16,1128 temperature: float = 0.8,1129 top_p: float = 0.95,1130 min_p: float = 0.05,1131 typical_p: float = 1.0,1132 logprobs: Optional[int] = None,1133 echo: bool = False,1134 stop: Optional[Union[str, List[str]]] = [],1135 frequency_penalty: float = 0.0,1136 presence_penalty: float = 0.0,1137 repeat_penalty: float = 1.0,1138 top_k: int = 40,1139 stream: bool = False,1140 seed: Optional[int] = None,1141 tfs_z: float = 1.0,1142 mirostat_mode: int = 0,1143 mirostat_tau: float = 5.0,1144 mirostat_eta: float = 0.1,1145 model: Optional[str] = None,1146 stopping_criteria: Optional[StoppingCriteriaList] = None,1147 logits_processor: Optional[LogitsProcessorList] = None,1148 grammar: Optional[LlamaGrammar] = None,1149 logit_bias: Optional[Dict[int, float]] = None,1150 ) -> Union[1151 Iterator[CreateCompletionResponse], Iterator[CreateCompletionStreamResponse]1152 ]:1153 assert suffix is None or suffix.__class__ is str1154 1155 completion_id: str = f"cmpl-{str(uuid.uuid4())}"1156 created: int = int(time.time())1157 bos_token_id: int = self.token_bos()1158 cls_token_id: int = self._model.token_cls()1159 sep_token_id: int = self._model.token_sep()1160 prefix_token_id: int = 0 # self._model.token_prefix() # TODO: Fix1161 middle_token_id: int = 0 # self._model.token_middle() # TODO: Fix1162 suffix_token_id: int = 0 # self._model.token_suffix() # TODO: Fix1163 add_space_prefix: bool = (1164 self.metadata.get("tokenizer.ggml.add_space_prefix", "true") == "true"1165 )1166 bos_tokens: List[int] = [cls_token_id if cls_token_id != -1 else bos_token_id]1167 eos_tokens: List[int] = [1168 sep_token_id if sep_token_id != -1 else self.token_eos()1169 ]1170 1171 if (1172 (isinstance(prompt, list) and suffix is None)1173 or not self._model.add_bos_token()1174 or bos_tokens[:1] == [-1]1175 ):1176 bos_tokens = []1177 1178 if (isinstance(prompt, list) and suffix is None) or (1179 not self._model.add_eos_token() and sep_token_id == -11180 ):1181 eos_tokens = []1182 1183 suffix_space_prefix: int = 01184 # Tokenizer hack to remove leading space1185 if add_space_prefix and suffix_token_id >= 0 and suffix:1186 suffix = "☺" + suffix1187 suffix_space_prefix = 21188 1189 # If prompt is empty, initialize completion with BOS token to avoid1190 # detokenization including a space at the beginning of the completion1191 completion_tokens: List[int] = [] if len(prompt) > 0 else [bos_token_id]1192 # Add blank space to start of prompt to match OG llama tokenizer1193 prefix_tokens: List[int] = (1194 [prefix_token_id] if prefix_token_id >= 0 and suffix is not None else []1195 ) + (1196 (1197 self.tokenize(1198 prompt.encode("utf-8"),1199 add_bos=False,1200 special=(prefix_token_id < 0 or suffix is None),