gguf_writer.py 29 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751
  1. from __future__ import annotations
  2. import logging
  3. import os
  4. import shutil
  5. import struct
  6. import tempfile
  7. from dataclasses import dataclass
  8. from enum import Enum, auto
  9. from pathlib import Path
  10. from io import BufferedWriter
  11. from typing import IO, Any, Sequence, Mapping
  12. from string import ascii_letters, digits
  13. import numpy as np
  14. from .constants import (
  15. GGUF_DEFAULT_ALIGNMENT,
  16. GGUF_MAGIC,
  17. GGUF_VERSION,
  18. GGMLQuantizationType,
  19. GGUFEndian,
  20. GGUFValueType,
  21. Keys,
  22. RopeScalingType,
  23. PoolingType,
  24. TokenType,
  25. )
  26. from .quants import quant_shape_from_byte_shape
  27. logger = logging.getLogger(__name__)
  28. SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf"
  29. @dataclass
  30. class TensorInfo:
  31. shape: Sequence[int]
  32. dtype: GGMLQuantizationType
  33. nbytes: int
  34. tensor: np.ndarray[Any, Any] | None = None
  35. @dataclass
  36. class GGUFValue:
  37. value: Any
  38. type: GGUFValueType
  39. class WriterState(Enum):
  40. NO_FILE = auto()
  41. EMPTY = auto()
  42. HEADER = auto()
  43. KV_DATA = auto()
  44. TI_DATA = auto()
  45. WEIGHTS = auto()
  46. class GGUFWriter:
  47. fout: list[BufferedWriter] | None
  48. path: Path | None
  49. temp_file: tempfile.SpooledTemporaryFile[bytes] | None
  50. tensors: list[dict[str, TensorInfo]]
  51. kv_data: list[dict[str, GGUFValue]]
  52. state: WriterState
  53. _simple_value_packing = {
  54. GGUFValueType.UINT8: "B",
  55. GGUFValueType.INT8: "b",
  56. GGUFValueType.UINT16: "H",
  57. GGUFValueType.INT16: "h",
  58. GGUFValueType.UINT32: "I",
  59. GGUFValueType.INT32: "i",
  60. GGUFValueType.FLOAT32: "f",
  61. GGUFValueType.UINT64: "Q",
  62. GGUFValueType.INT64: "q",
  63. GGUFValueType.FLOAT64: "d",
  64. GGUFValueType.BOOL: "?",
  65. }
  66. def __init__(
  67. self, path: os.PathLike[str] | str | None, arch: str, use_temp_file: bool = False, endianess: GGUFEndian = GGUFEndian.LITTLE,
  68. split_max_tensors: int = 0, split_max_size: int = 0, dry_run: bool = False, small_first_shard: bool = False
  69. ):
  70. self.fout = None
  71. self.path = Path(path) if path else None
  72. self.arch = arch
  73. self.endianess = endianess
  74. self.data_alignment = GGUF_DEFAULT_ALIGNMENT
  75. self.use_temp_file = use_temp_file
  76. self.temp_file = None
  77. self.tensors = [{}]
  78. self.kv_data = [{}]
  79. self.split_max_tensors = split_max_tensors
  80. self.split_max_size = split_max_size
  81. self.dry_run = dry_run
  82. self.small_first_shard = small_first_shard
  83. logger.info("gguf: This GGUF file is for {0} Endian only".format(
  84. "Big" if self.endianess == GGUFEndian.BIG else "Little",
  85. ))
  86. self.state = WriterState.NO_FILE
  87. if self.small_first_shard:
  88. self.tensors.append({})
  89. self.add_architecture()
  90. def format_shard_names(self, path: Path) -> list[Path]:
  91. if len(self.tensors) == 1:
  92. return [path]
  93. return [path.with_name(SHARD_NAME_FORMAT.format(path.stem, i + 1, len(self.tensors))) for i in range(len(self.tensors))]
  94. def open_output_file(self, path: Path | None = None) -> None:
  95. if self.state is WriterState.EMPTY and self.fout is not None and (path is None or path == self.path):
  96. # allow calling this multiple times as long as the path is the same
  97. return
  98. if self.state is not WriterState.NO_FILE:
  99. raise ValueError(f'Expected output file to be not yet opened, got {self.state}')
  100. if path is not None:
  101. self.path = path
  102. if self.path is not None:
  103. filenames = self.print_plan()
  104. self.fout = [open(filename, "wb") for filename in filenames]
  105. self.state = WriterState.EMPTY
  106. def print_plan(self) -> list[Path]:
  107. logger.info("Writing the following files:")
  108. assert self.path is not None
  109. filenames = self.format_shard_names(self.path)
  110. assert len(filenames) == len(self.tensors)
  111. for name, tensors in zip(filenames, self.tensors):
  112. logger.info(f"{name}: n_tensors = {len(tensors)}, total_size = {GGUFWriter.format_n_bytes_to_str(sum(ti.nbytes for ti in tensors.values()))}")
  113. if self.dry_run:
  114. logger.info("Dry run, not writing files")
  115. exit()
  116. return filenames
  117. def add_shard_kv_data(self) -> None:
  118. if len(self.tensors) == 1:
  119. return
  120. total_tensors = sum(len(t) for t in self.tensors)
  121. assert self.fout is not None
  122. total_splits = len(self.fout)
  123. self.kv_data.extend({} for _ in range(len(self.kv_data), total_splits))
  124. for i, kv_data in enumerate(self.kv_data):
  125. kv_data[Keys.Split.LLM_KV_SPLIT_NO] = GGUFValue(i, GGUFValueType.UINT16)
  126. kv_data[Keys.Split.LLM_KV_SPLIT_COUNT] = GGUFValue(total_splits, GGUFValueType.UINT16)
  127. kv_data[Keys.Split.LLM_KV_SPLIT_TENSORS_COUNT] = GGUFValue(total_tensors, GGUFValueType.INT32)
  128. def write_header_to_file(self, path: Path | None = None) -> None:
  129. if len(self.tensors) == 1 and (self.split_max_tensors != 0 or self.split_max_size != 0):
  130. logger.warning("Model fails split requirements, not splitting")
  131. self.open_output_file(path)
  132. if self.state is not WriterState.EMPTY:
  133. raise ValueError(f'Expected output file to be empty, got {self.state}')
  134. assert self.fout is not None
  135. assert len(self.fout) == len(self.tensors)
  136. assert len(self.kv_data) == 1
  137. self.add_shard_kv_data()
  138. for fout, tensors, kv_data in zip(self.fout, self.tensors, self.kv_data):
  139. fout.write(self._pack("<I", GGUF_MAGIC, skip_pack_prefix = True))
  140. fout.write(self._pack("I", GGUF_VERSION))
  141. fout.write(self._pack("Q", len(tensors)))
  142. fout.write(self._pack("Q", len(kv_data)))
  143. fout.flush()
  144. self.state = WriterState.HEADER
  145. def write_kv_data_to_file(self) -> None:
  146. if self.state is not WriterState.HEADER:
  147. raise ValueError(f'Expected output file to contain the header, got {self.state}')
  148. assert self.fout is not None
  149. for fout, kv_data in zip(self.fout, self.kv_data):
  150. kv_bytes = bytearray()
  151. for key, val in kv_data.items():
  152. kv_bytes += self._pack_val(key, GGUFValueType.STRING, add_vtype=False)
  153. kv_bytes += self._pack_val(val.value, val.type, add_vtype=True)
  154. fout.write(kv_bytes)
  155. self.flush()
  156. self.state = WriterState.KV_DATA
  157. def write_ti_data_to_file(self) -> None:
  158. if self.state is not WriterState.KV_DATA:
  159. raise ValueError(f'Expected output file to contain KV data, got {self.state}')
  160. assert self.fout is not None
  161. for fout, tensors in zip(self.fout, self.tensors):
  162. ti_data = bytearray()
  163. offset_tensor = 0
  164. for name, ti in tensors.items():
  165. ti_data += self._pack_val(name, GGUFValueType.STRING, add_vtype=False)
  166. n_dims = len(ti.shape)
  167. ti_data += self._pack("I", n_dims)
  168. for j in range(n_dims):
  169. ti_data += self._pack("Q", ti.shape[n_dims - 1 - j])
  170. ti_data += self._pack("I", ti.dtype)
  171. ti_data += self._pack("Q", offset_tensor)
  172. offset_tensor += GGUFWriter.ggml_pad(ti.nbytes, self.data_alignment)
  173. fout.write(ti_data)
  174. fout.flush()
  175. self.state = WriterState.TI_DATA
  176. def add_key_value(self, key: str, val: Any, vtype: GGUFValueType) -> None:
  177. if any(key in kv_data for kv_data in self.kv_data):
  178. raise ValueError(f'Duplicated key name {key!r}')
  179. self.kv_data[0][key] = GGUFValue(value=val, type=vtype)
  180. def add_uint8(self, key: str, val: int) -> None:
  181. self.add_key_value(key,val, GGUFValueType.UINT8)
  182. def add_int8(self, key: str, val: int) -> None:
  183. self.add_key_value(key, val, GGUFValueType.INT8)
  184. def add_uint16(self, key: str, val: int) -> None:
  185. self.add_key_value(key, val, GGUFValueType.UINT16)
  186. def add_int16(self, key: str, val: int) -> None:
  187. self.add_key_value(key, val, GGUFValueType.INT16)
  188. def add_uint32(self, key: str, val: int) -> None:
  189. self.add_key_value(key, val, GGUFValueType.UINT32)
  190. def add_int32(self, key: str, val: int) -> None:
  191. self.add_key_value(key, val, GGUFValueType.INT32)
  192. def add_float32(self, key: str, val: float) -> None:
  193. self.add_key_value(key, val, GGUFValueType.FLOAT32)
  194. def add_uint64(self, key: str, val: int) -> None:
  195. self.add_key_value(key, val, GGUFValueType.UINT64)
  196. def add_int64(self, key: str, val: int) -> None:
  197. self.add_key_value(key, val, GGUFValueType.INT64)
  198. def add_float64(self, key: str, val: float) -> None:
  199. self.add_key_value(key, val, GGUFValueType.FLOAT64)
  200. def add_bool(self, key: str, val: bool) -> None:
  201. self.add_key_value(key, val, GGUFValueType.BOOL)
  202. def add_string(self, key: str, val: str) -> None:
  203. if not val:
  204. return
  205. self.add_key_value(key, val, GGUFValueType.STRING)
  206. def add_array(self, key: str, val: Sequence[Any]) -> None:
  207. self.add_key_value(key, val, GGUFValueType.ARRAY)
  208. @staticmethod
  209. def ggml_pad(x: int, n: int) -> int:
  210. return ((x + n - 1) // n) * n
  211. def add_tensor_info(
  212. self, name: str, tensor_shape: Sequence[int], tensor_dtype: np.dtype,
  213. tensor_nbytes: int, raw_dtype: GGMLQuantizationType | None = None,
  214. ) -> None:
  215. if self.state is not WriterState.NO_FILE:
  216. raise ValueError(f'Expected output file to be not yet opened, got {self.state}')
  217. if any(name in tensors for tensors in self.tensors):
  218. raise ValueError(f'Duplicated tensor name {name!r}')
  219. if raw_dtype is None:
  220. if tensor_dtype == np.float16:
  221. dtype = GGMLQuantizationType.F16
  222. elif tensor_dtype == np.float32:
  223. dtype = GGMLQuantizationType.F32
  224. elif tensor_dtype == np.float64:
  225. dtype = GGMLQuantizationType.F64
  226. elif tensor_dtype == np.int8:
  227. dtype = GGMLQuantizationType.I8
  228. elif tensor_dtype == np.int16:
  229. dtype = GGMLQuantizationType.I16
  230. elif tensor_dtype == np.int32:
  231. dtype = GGMLQuantizationType.I32
  232. elif tensor_dtype == np.int64:
  233. dtype = GGMLQuantizationType.I64
  234. else:
  235. raise ValueError("Only F16, F32, F64, I8, I16, I32, I64 tensors are supported for now")
  236. else:
  237. dtype = raw_dtype
  238. if tensor_dtype == np.uint8:
  239. tensor_shape = quant_shape_from_byte_shape(tensor_shape, raw_dtype)
  240. # make sure there is at least one tensor before splitting
  241. if len(self.tensors[-1]) > 0:
  242. if ( # split when over tensor limit
  243. self.split_max_tensors != 0
  244. and len(self.tensors[-1]) >= self.split_max_tensors
  245. ) or ( # split when over size limit
  246. self.split_max_size != 0
  247. and sum(ti.nbytes for ti in self.tensors[-1].values()) + tensor_nbytes > self.split_max_size
  248. ):
  249. self.tensors.append({})
  250. self.tensors[-1][name] = TensorInfo(shape=tensor_shape, dtype=dtype, nbytes=tensor_nbytes)
  251. def add_tensor(
  252. self, name: str, tensor: np.ndarray[Any, Any], raw_shape: Sequence[int] | None = None,
  253. raw_dtype: GGMLQuantizationType | None = None,
  254. ) -> None:
  255. if self.endianess == GGUFEndian.BIG:
  256. tensor.byteswap(inplace=True)
  257. if self.use_temp_file and self.temp_file is None:
  258. fp = tempfile.SpooledTemporaryFile(mode="w+b", max_size=256 * 1024 * 1024)
  259. fp.seek(0)
  260. self.temp_file = fp
  261. shape: Sequence[int] = raw_shape if raw_shape is not None else tensor.shape
  262. self.add_tensor_info(name, shape, tensor.dtype, tensor.nbytes, raw_dtype=raw_dtype)
  263. if self.temp_file is None:
  264. self.tensors[-1][name].tensor = tensor
  265. return
  266. tensor.tofile(self.temp_file)
  267. self.write_padding(self.temp_file, tensor.nbytes)
  268. def write_padding(self, fp: IO[bytes], n: int, align: int | None = None) -> None:
  269. pad = GGUFWriter.ggml_pad(n, align if align is not None else self.data_alignment) - n
  270. if pad != 0:
  271. fp.write(bytes([0] * pad))
  272. def write_tensor_data(self, tensor: np.ndarray[Any, Any]) -> None:
  273. if self.state is not WriterState.TI_DATA and self.state is not WriterState.WEIGHTS:
  274. raise ValueError(f'Expected output file to contain tensor info or weights, got {self.state}')
  275. assert self.fout is not None
  276. if self.endianess == GGUFEndian.BIG:
  277. tensor.byteswap(inplace=True)
  278. file_id = -1
  279. for i, tensors in enumerate(self.tensors):
  280. if len(tensors) > 0:
  281. file_id = i
  282. break
  283. fout = self.fout[file_id]
  284. # pop the first tensor info
  285. # TODO: cleaner way to get the first key
  286. first_tensor_name = [name for name, _ in zip(self.tensors[file_id].keys(), range(1))][0]
  287. ti = self.tensors[file_id].pop(first_tensor_name)
  288. assert ti.nbytes == tensor.nbytes
  289. self.write_padding(fout, fout.tell())
  290. tensor.tofile(fout)
  291. self.write_padding(fout, tensor.nbytes)
  292. self.state = WriterState.WEIGHTS
  293. def write_tensors_to_file(self, *, progress: bool = False) -> None:
  294. self.write_ti_data_to_file()
  295. assert self.fout is not None
  296. for fout in self.fout:
  297. self.write_padding(fout, fout.tell())
  298. if self.temp_file is None:
  299. shard_bar = None
  300. bar = None
  301. if progress:
  302. from tqdm import tqdm
  303. total_bytes = sum(ti.nbytes for t in self.tensors for ti in t.values())
  304. if len(self.fout) > 1:
  305. shard_bar = tqdm(desc=f"Shard (0/{len(self.fout)})", total=None, unit="byte", unit_scale=True)
  306. bar = tqdm(desc="Writing", total=total_bytes, unit="byte", unit_scale=True)
  307. for i, (fout, tensors) in enumerate(zip(self.fout, self.tensors)):
  308. if shard_bar is not None:
  309. shard_bar.set_description(f"Shard ({i + 1}/{len(self.fout)})")
  310. total = sum(ti.nbytes for ti in tensors.values())
  311. shard_bar.reset(total=(total if total > 0 else None))
  312. # relying on the fact that Python dicts preserve insertion order (since 3.7)
  313. for ti in tensors.values():
  314. assert ti.tensor is not None # can only iterate once over the tensors
  315. assert ti.tensor.nbytes == ti.nbytes
  316. ti.tensor.tofile(fout)
  317. if shard_bar is not None:
  318. shard_bar.update(ti.nbytes)
  319. if bar is not None:
  320. bar.update(ti.nbytes)
  321. self.write_padding(fout, ti.nbytes)
  322. ti.tensor = None
  323. else:
  324. self.temp_file.seek(0)
  325. shutil.copyfileobj(self.temp_file, self.fout[0 if not self.small_first_shard else 1])
  326. self.flush()
  327. self.temp_file.close()
  328. self.state = WriterState.WEIGHTS
  329. def flush(self) -> None:
  330. assert self.fout is not None
  331. for fout in self.fout:
  332. fout.flush()
  333. def close(self) -> None:
  334. if self.fout is not None:
  335. for fout in self.fout:
  336. fout.close()
  337. self.fout = None
  338. def add_type(self, type_name: str) -> None:
  339. self.add_string(Keys.General.TYPE, type_name)
  340. def add_architecture(self) -> None:
  341. self.add_string(Keys.General.ARCHITECTURE, self.arch)
  342. def add_author(self, author: str) -> None:
  343. self.add_string(Keys.General.AUTHOR, author)
  344. def add_version(self, version: str) -> None:
  345. self.add_string(Keys.General.VERSION, version)
  346. def add_tensor_data_layout(self, layout: str) -> None:
  347. self.add_string(Keys.LLM.TENSOR_DATA_LAYOUT.format(arch=self.arch), layout)
  348. def add_url(self, url: str) -> None:
  349. self.add_string(Keys.General.URL, url)
  350. def add_description(self, description: str) -> None:
  351. self.add_string(Keys.General.DESCRIPTION, description)
  352. def add_licence(self, licence: str) -> None:
  353. self.add_string(Keys.General.LICENSE, licence)
  354. def add_source_url(self, url: str) -> None:
  355. self.add_string(Keys.General.SOURCE_URL, url)
  356. def add_source_hf_repo(self, repo: str) -> None:
  357. self.add_string(Keys.General.SOURCE_HF_REPO, repo)
  358. def add_file_type(self, ftype: int) -> None:
  359. self.add_uint32(Keys.General.FILE_TYPE, ftype)
  360. def add_name(self, name: str) -> None:
  361. self.add_string(Keys.General.NAME, name)
  362. def add_quantization_version(self, quantization_version: int) -> None:
  363. self.add_uint32(
  364. Keys.General.QUANTIZATION_VERSION, quantization_version)
  365. def add_custom_alignment(self, alignment: int) -> None:
  366. self.data_alignment = alignment
  367. self.add_uint32(Keys.General.ALIGNMENT, alignment)
  368. def add_vocab_size(self, size: int) -> None:
  369. self.add_uint32(Keys.LLM.VOCAB_SIZE.format(arch=self.arch), size)
  370. def add_context_length(self, length: int) -> None:
  371. self.add_uint32(Keys.LLM.CONTEXT_LENGTH.format(arch=self.arch), length)
  372. def add_embedding_length(self, length: int) -> None:
  373. self.add_uint32(Keys.LLM.EMBEDDING_LENGTH.format(arch=self.arch), length)
  374. def add_block_count(self, length: int) -> None:
  375. self.add_uint32(Keys.LLM.BLOCK_COUNT.format(arch=self.arch), length)
  376. def add_leading_dense_block_count(self, length: int) -> None:
  377. self.add_uint32(Keys.LLM.LEADING_DENSE_BLOCK_COUNT.format(arch=self.arch), length)
  378. def add_feed_forward_length(self, length: int | Sequence[int]) -> None:
  379. if isinstance(length, int):
  380. self.add_uint32(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  381. else:
  382. self.add_array(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  383. def add_expert_feed_forward_length(self, length: int) -> None:
  384. self.add_uint32(Keys.LLM.EXPERT_FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  385. def add_expert_shared_feed_forward_length(self, length: int) -> None:
  386. self.add_uint32(Keys.LLM.EXPERT_SHARED_FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  387. def add_parallel_residual(self, use: bool) -> None:
  388. self.add_bool(Keys.LLM.USE_PARALLEL_RESIDUAL.format(arch=self.arch), use)
  389. def add_decoder_start_token_id(self, id: int) -> None:
  390. self.add_uint32(Keys.LLM.DECODER_START_TOKEN_ID.format(arch=self.arch), id)
  391. def add_head_count(self, count: int | Sequence[int]) -> None:
  392. if isinstance(count, int):
  393. self.add_uint32(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count)
  394. else:
  395. self.add_array(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count)
  396. def add_head_count_kv(self, count: int | Sequence[int]) -> None:
  397. if isinstance(count, int):
  398. self.add_uint32(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count)
  399. else:
  400. self.add_array(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count)
  401. def add_key_length(self, length: int) -> None:
  402. self.add_uint32(Keys.Attention.KEY_LENGTH.format(arch=self.arch), length)
  403. def add_value_length(self, length: int) -> None:
  404. self.add_uint32(Keys.Attention.VALUE_LENGTH.format(arch=self.arch), length)
  405. def add_max_alibi_bias(self, bias: float) -> None:
  406. self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias)
  407. def add_clamp_kqv(self, value: float) -> None:
  408. self.add_float32(Keys.Attention.CLAMP_KQV.format(arch=self.arch), value)
  409. def add_logit_scale(self, value: float) -> None:
  410. self.add_float32(Keys.LLM.LOGIT_SCALE.format(arch=self.arch), value)
  411. def add_attn_logit_softcapping(self, value: float) -> None:
  412. self.add_float32(Keys.LLM.ATTN_LOGIT_SOFTCAPPING.format(arch=self.arch), value)
  413. def add_final_logit_softcapping(self, value: float) -> None:
  414. self.add_float32(Keys.LLM.FINAL_LOGIT_SOFTCAPPING.format(arch=self.arch), value)
  415. def add_expert_count(self, count: int) -> None:
  416. self.add_uint32(Keys.LLM.EXPERT_COUNT.format(arch=self.arch), count)
  417. def add_expert_used_count(self, count: int) -> None:
  418. self.add_uint32(Keys.LLM.EXPERT_USED_COUNT.format(arch=self.arch), count)
  419. def add_expert_shared_count(self, count: int) -> None:
  420. self.add_uint32(Keys.LLM.EXPERT_SHARED_COUNT.format(arch=self.arch), count)
  421. def add_expert_weights_scale(self, value: float) -> None:
  422. self.add_float32(Keys.LLM.EXPERT_WEIGHTS_SCALE.format(arch=self.arch), value)
  423. def add_layer_norm_eps(self, value: float) -> None:
  424. self.add_float32(Keys.Attention.LAYERNORM_EPS.format(arch=self.arch), value)
  425. def add_layer_norm_rms_eps(self, value: float) -> None:
  426. self.add_float32(Keys.Attention.LAYERNORM_RMS_EPS.format(arch=self.arch), value)
  427. def add_causal_attention(self, value: bool) -> None:
  428. self.add_bool(Keys.Attention.CAUSAL.format(arch=self.arch), value)
  429. def add_q_lora_rank(self, length: int) -> None:
  430. self.add_uint32(Keys.Attention.Q_LORA_RANK.format(arch=self.arch), length)
  431. def add_kv_lora_rank(self, length: int) -> None:
  432. self.add_uint32(Keys.Attention.KV_LORA_RANK.format(arch=self.arch), length)
  433. def add_relative_attn_buckets_count(self, value: int) -> None:
  434. self.add_uint32(Keys.Attention.REL_BUCKETS_COUNT.format(arch=self.arch), value)
  435. def add_sliding_window(self, value: int) -> None:
  436. self.add_uint32(Keys.Attention.SLIDING_WINDOW.format(arch=self.arch), value)
  437. def add_pooling_type(self, value: PoolingType) -> None:
  438. self.add_uint32(Keys.LLM.POOLING_TYPE.format(arch=self.arch), value.value)
  439. def add_rope_dimension_count(self, count: int) -> None:
  440. self.add_uint32(Keys.Rope.DIMENSION_COUNT.format(arch=self.arch), count)
  441. def add_rope_freq_base(self, value: float) -> None:
  442. self.add_float32(Keys.Rope.FREQ_BASE.format(arch=self.arch), value)
  443. def add_rope_scaling_type(self, value: RopeScalingType) -> None:
  444. self.add_string(Keys.Rope.SCALING_TYPE.format(arch=self.arch), value.value)
  445. def add_rope_scaling_factor(self, value: float) -> None:
  446. self.add_float32(Keys.Rope.SCALING_FACTOR.format(arch=self.arch), value)
  447. def add_rope_scaling_attn_factors(self, value: float) -> None:
  448. self.add_float32(Keys.Rope.SCALING_ATTN_FACTOR.format(arch=self.arch), value)
  449. def add_rope_scaling_orig_ctx_len(self, value: int) -> None:
  450. self.add_uint32(Keys.Rope.SCALING_ORIG_CTX_LEN.format(arch=self.arch), value)
  451. def add_rope_scaling_finetuned(self, value: bool) -> None:
  452. self.add_bool(Keys.Rope.SCALING_FINETUNED.format(arch=self.arch), value)
  453. def add_rope_scaling_yarn_log_mul(self, value: float) -> None:
  454. self.add_float32(Keys.Rope.SCALING_YARN_LOG_MUL.format(arch=self.arch), value)
  455. def add_ssm_conv_kernel(self, value: int) -> None:
  456. self.add_uint32(Keys.SSM.CONV_KERNEL.format(arch=self.arch), value)
  457. def add_ssm_inner_size(self, value: int) -> None:
  458. self.add_uint32(Keys.SSM.INNER_SIZE.format(arch=self.arch), value)
  459. def add_ssm_state_size(self, value: int) -> None:
  460. self.add_uint32(Keys.SSM.STATE_SIZE.format(arch=self.arch), value)
  461. def add_ssm_time_step_rank(self, value: int) -> None:
  462. self.add_uint32(Keys.SSM.TIME_STEP_RANK.format(arch=self.arch), value)
  463. def add_tokenizer_model(self, model: str) -> None:
  464. self.add_string(Keys.Tokenizer.MODEL, model)
  465. def add_tokenizer_pre(self, pre: str) -> None:
  466. self.add_string(Keys.Tokenizer.PRE, pre)
  467. def add_token_list(self, tokens: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None:
  468. self.add_array(Keys.Tokenizer.LIST, tokens)
  469. def add_token_merges(self, merges: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None:
  470. self.add_array(Keys.Tokenizer.MERGES, merges)
  471. def add_token_types(self, types: Sequence[TokenType] | Sequence[int]) -> None:
  472. self.add_array(Keys.Tokenizer.TOKEN_TYPE, types)
  473. def add_token_type_count(self, value: int) -> None:
  474. self.add_uint32(Keys.Tokenizer.TOKEN_TYPE_COUNT, value)
  475. def add_token_scores(self, scores: Sequence[float]) -> None:
  476. self.add_array(Keys.Tokenizer.SCORES, scores)
  477. def add_bos_token_id(self, id: int) -> None:
  478. self.add_uint32(Keys.Tokenizer.BOS_ID, id)
  479. def add_eos_token_id(self, id: int) -> None:
  480. self.add_uint32(Keys.Tokenizer.EOS_ID, id)
  481. def add_unk_token_id(self, id: int) -> None:
  482. self.add_uint32(Keys.Tokenizer.UNK_ID, id)
  483. def add_sep_token_id(self, id: int) -> None:
  484. self.add_uint32(Keys.Tokenizer.SEP_ID, id)
  485. def add_pad_token_id(self, id: int) -> None:
  486. self.add_uint32(Keys.Tokenizer.PAD_ID, id)
  487. def add_cls_token_id(self, id: int) -> None:
  488. self.add_uint32(Keys.Tokenizer.CLS_ID, id)
  489. def add_mask_token_id(self, id: int) -> None:
  490. self.add_uint32(Keys.Tokenizer.MASK_ID, id)
  491. def add_add_bos_token(self, value: bool) -> None:
  492. self.add_bool(Keys.Tokenizer.ADD_BOS, value)
  493. def add_add_eos_token(self, value: bool) -> None:
  494. self.add_bool(Keys.Tokenizer.ADD_EOS, value)
  495. def add_add_space_prefix(self, value: bool) -> None:
  496. self.add_bool(Keys.Tokenizer.ADD_PREFIX, value)
  497. def add_remove_extra_whitespaces(self, value: bool) -> None:
  498. self.add_bool(Keys.Tokenizer.REMOVE_EXTRA_WS, value)
  499. def add_precompiled_charsmap(self, charsmap: Sequence[bytes]) -> None:
  500. self.add_array(Keys.Tokenizer.PRECOMPILED_CHARSMAP, charsmap)
  501. def add_chat_template(self, value: str | Sequence[Mapping[str, str]]) -> None:
  502. if not isinstance(value, str):
  503. template_default = None
  504. template_names = set()
  505. for choice in value:
  506. name = choice.get('name', '')
  507. template = choice.get('template')
  508. # Allowing non-alphanumerical characters in template name is probably not a good idea, so filter it
  509. name = ''.join((c if c in ascii_letters + digits else '_' for c in name))
  510. if name and template is not None:
  511. if name == 'default':
  512. template_default = template
  513. else:
  514. template_names.add(name)
  515. self.add_string(Keys.Tokenizer.CHAT_TEMPLATE_N.format(name=name), template)
  516. if template_names:
  517. self.add_array(Keys.Tokenizer.CHAT_TEMPLATES, list(template_names))
  518. if template_default is None:
  519. return
  520. value = template_default
  521. self.add_string(Keys.Tokenizer.CHAT_TEMPLATE, value)
  522. def add_prefix_token_id(self, id: int) -> None:
  523. self.add_uint32(Keys.Tokenizer.PREFIX_ID, id)
  524. def add_suffix_token_id(self, id: int) -> None:
  525. self.add_uint32(Keys.Tokenizer.SUFFIX_ID, id)
  526. def add_middle_token_id(self, id: int) -> None:
  527. self.add_uint32(Keys.Tokenizer.MIDDLE_ID, id)
  528. def add_eot_token_id(self, id: int) -> None:
  529. self.add_uint32(Keys.Tokenizer.EOT_ID, id)
  530. def _pack(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> bytes:
  531. pack_prefix = ''
  532. if not skip_pack_prefix:
  533. pack_prefix = '<' if self.endianess == GGUFEndian.LITTLE else '>'
  534. return struct.pack(f'{pack_prefix}{fmt}', value)
  535. def _pack_val(self, val: Any, vtype: GGUFValueType, add_vtype: bool) -> bytes:
  536. kv_data = bytearray()
  537. if add_vtype:
  538. kv_data += self._pack("I", vtype)
  539. pack_fmt = self._simple_value_packing.get(vtype)
  540. if pack_fmt is not None:
  541. kv_data += self._pack(pack_fmt, val, skip_pack_prefix = vtype == GGUFValueType.BOOL)
  542. elif vtype == GGUFValueType.STRING:
  543. encoded_val = val.encode("utf-8") if isinstance(val, str) else val
  544. kv_data += self._pack("Q", len(encoded_val))
  545. kv_data += encoded_val
  546. elif vtype == GGUFValueType.ARRAY and isinstance(val, Sequence) and val:
  547. if isinstance(val, bytes):
  548. ltype = GGUFValueType.UINT8
  549. else:
  550. ltype = GGUFValueType.get_type(val[0])
  551. if not all(GGUFValueType.get_type(i) is ltype for i in val[1:]):
  552. raise ValueError("All items in a GGUF array should be of the same type")
  553. kv_data += self._pack("I", ltype)
  554. kv_data += self._pack("Q", len(val))
  555. for item in val:
  556. kv_data += self._pack_val(item, ltype, add_vtype=False)
  557. else:
  558. raise ValueError("Invalid GGUF metadata value type or value")
  559. return kv_data
  560. @staticmethod
  561. def format_n_bytes_to_str(num: int) -> str:
  562. if num == 0:
  563. return "negligible - metadata only"
  564. fnum = float(num)
  565. for unit in ("", "K", "M", "G"):
  566. if abs(fnum) < 1000.0:
  567. return f"{fnum:3.1f}{unit}"
  568. fnum /= 1000.0
  569. return f"{fnum:.1f}T - over 1TB, split recommended"