hymenjj/llama-cpp-python-prebuilt
0
1from __future__ import annotations2 3import multiprocessing4 5from typing import Optional, List, Literal, Union, Dict, cast6from typing_extensions import Self7 8from pydantic import Field, model_validator9from pydantic_settings import BaseSettings10 11import llama_cpp12 13# Disable warning for model and model_alias settings14BaseSettings.model_config["protected_namespaces"] = ()15 16 17class ModelSettings(BaseSettings):18 """Model settings used to load a Llama model."""19 20 model: str = Field(21 description="The path to the model to use for generating completions."22 )23 model_alias: Optional[str] = Field(24 default=None,25 description="The alias of the model to use for generating completions.",26 )27 # Model Params28 n_gpu_layers: int = Field(29 default=0,30 ge=-1,31 description="The number of layers to put on the GPU. The rest will be on the CPU. Set -1 to move all to GPU.",32 )33 split_mode: int = Field(34 default=llama_cpp.LLAMA_SPLIT_MODE_LAYER,35 description="The split mode to use.",36 )37 main_gpu: int = Field(38 default=0,39 ge=0,40 description="Main GPU to use.",41 )42 tensor_split: Optional[List[float]] = Field(43 default=None,44 description="Split layers across multiple GPUs in proportion.",45 )46 vocab_only: bool = Field(47 default=False, description="Whether to only return the vocabulary."48 )49 use_mmap: bool = Field(50 default=llama_cpp.llama_supports_mmap(),51 description="Use mmap.",52 )53 use_mlock: bool = Field(54 default=llama_cpp.llama_supports_mlock(),55 description="Use mlock.",56 )57 kv_overrides: Optional[List[str]] = Field(58 default=None,59 description="List of model kv overrides in the format key=type:value where type is one of (bool, int, float). Valid true values are (true, TRUE, 1), otherwise false.",60 )61 rpc_servers: Optional[str] = Field(62 default=None,63 description="comma seperated list of rpc servers for offloading",64 )65 # Context Params66 seed: int = Field(67 default=llama_cpp.LLAMA_DEFAULT_SEED, description="Random seed. -1 for random."68 )69 n_ctx: int = Field(default=2048, ge=0, description="The context size.")70 n_batch: int = Field(71 default=512, ge=1, description="The batch size to use per eval."72 )73 n_ubatch: int = Field(74 default=512, ge=1, description="The physical batch size used by llama.cpp"75 )76 n_threads: int = Field(77 default=max(multiprocessing.cpu_count() // 2, 1),78 ge=1,79 description="The number of threads to use. Use -1 for max cpu threads",80 )81 n_threads_batch: int = Field(82 default=max(multiprocessing.cpu_count(), 1),83 ge=0,84 description="The number of threads to use when batch processing. Use -1 for max cpu threads",85 )86 rope_scaling_type: int = Field(87 default=llama_cpp.LLAMA_ROPE_SCALING_TYPE_UNSPECIFIED88 )89 rope_freq_base: float = Field(default=0.0, description="RoPE base frequency")90 rope_freq_scale: float = Field(91 default=0.0, description="RoPE frequency scaling factor"92 )93 yarn_ext_factor: float = Field(default=-1.0)94 yarn_attn_factor: float = Field(default=1.0)95 yarn_beta_fast: float = Field(default=32.0)96 yarn_beta_slow: float = Field(default=1.0)97 yarn_orig_ctx: int = Field(default=0)98 mul_mat_q: bool = Field(99 default=True, description="if true, use experimental mul_mat_q kernels"100 )101 logits_all: bool = Field(default=True, description="Whether to return logits.")102 embedding: bool = Field(default=False, description="Whether to use embeddings.")103 offload_kqv: bool = Field(104 default=True, description="Whether to offload kqv to the GPU."105 )106 flash_attn: bool = Field(107 default=False, description="Whether to use flash attention."108 )109 # Sampling Params110 last_n_tokens_size: int = Field(111 default=64,112 ge=0,113 description="Last n tokens to keep for repeat penalty calculation.",114 )115 # LoRA Params116 lora_base: Optional[str] = Field(117 default=None,118 description="Optional path to base model, useful if using a quantized base model and you want to apply LoRA to an f16 model.",119 )120 lora_path: Optional[str] = Field(121 default=None,122 description="Path to a LoRA file to apply to the model.",123 )124 # Backend Params125 numa: Union[bool, int] = Field(126 default=False,127 description="Enable NUMA support.",128 )129 # Chat Format Params130 chat_format: Optional[str] = Field(131 default=None,132 description="Chat format to use.",133 )134 clip_model_path: Optional[str] = Field(135 default=None,136 description="Path to a CLIP model to use for multi-modal chat completion.",137 )138 # Cache Params139 cache: bool = Field(140 default=False,141 description="Use a cache to reduce processing times for evaluated prompts.",142 )143 cache_type: Literal["ram", "disk"] = Field(144 default="ram",145 description="The type of cache to use. Only used if cache is True.",146 )147 cache_size: int = Field(148 default=2 << 30,149 description="The size of the cache in bytes. Only used if cache is True.",150 )151 # Tokenizer Options152 hf_tokenizer_config_path: Optional[str] = Field(153 default=None,154 description="The path to a HuggingFace tokenizer_config.json file.",155 )156 hf_pretrained_model_name_or_path: Optional[str] = Field(157 default=None,158 description="The model name or path to a pretrained HuggingFace tokenizer model. Same as you would pass to AutoTokenizer.from_pretrained().",159 )160 # Loading from HuggingFace Model Hub161 hf_model_repo_id: Optional[str] = Field(162 default=None,163 description="The model repo id to use for the HuggingFace tokenizer model.",164 )165 # Speculative Decoding166 draft_model: Optional[str] = Field(167 default=None,168 description="Method to use for speculative decoding. One of (prompt-lookup-decoding).",169 )170 draft_model_num_pred_tokens: int = Field(171 default=10,172 description="Number of tokens to predict using the draft model.",173 )174 # KV Cache Quantization175 type_k: Optional[int] = Field(176 default=None,177 description="Type of the key cache quantization.",178 )179 type_v: Optional[int] = Field(180 default=None,181 description="Type of the value cache quantization.",182 )183 # Misc184 verbose: bool = Field(185 default=True, description="Whether to print debug information."186 )187 188 @model_validator(189 mode="before"190 ) # pre=True to ensure this runs before any other validation191 def set_dynamic_defaults(self) -> Self:192 # If n_threads or n_threads_batch is -1, set it to multiprocessing.cpu_count()193 cpu_count = multiprocessing.cpu_count()194 values = cast(Dict[str, int], self)195 if values.get("n_threads", 0) == -1:196 values["n_threads"] = cpu_count197 if values.get("n_threads_batch", 0) == -1:198 values["n_threads_batch"] = cpu_count199 return self200 201 202class ServerSettings(BaseSettings):203 """Server settings used to configure the FastAPI and Uvicorn server."""204 205 # Uvicorn Settings206 host: str = Field(default="localhost", description="Listen address")207 port: int = Field(default=8000, description="Listen port")208 ssl_keyfile: Optional[str] = Field(209 default=None, description="SSL key file for HTTPS"210 )211 ssl_certfile: Optional[str] = Field(212 default=None, description="SSL certificate file for HTTPS"213 )214 # FastAPI Settings215 api_key: Optional[str] = Field(216 default=None,217 description="API key for authentication. If set all requests need to be authenticated.",218 )219 interrupt_requests: bool = Field(220 default=True,221 description="Whether to interrupt requests when a new request is received.",222 )223 disable_ping_events: bool = Field(224 default=False,225 description="Disable EventSource pings (may be needed for some clients).",226 )227 root_path: str = Field(228 default="",229 description="The root path for the server. Useful when running behind a reverse proxy.",230 )231 232 233class Settings(ServerSettings, ModelSettings):234 pass235 236 237class ConfigFileSettings(ServerSettings):238 """Configuration file format settings."""239 240 models: List[ModelSettings] = Field(default=[], description="Model configs")241