gguf_writer.py 47 KB

123456789101112131415161718192021222324252627282930313233343536373839404142434445464748495051525354555657585960616263646566676869707172737475767778798081828384858687888990919293949596979899100101102103104105106107108109110111112113114115116117118119120121122123124125126127128129130131132133134135136137138139140141142143144145146147148149150151152153154155156157158159160161162163164165166167168169170171172173174175176177178179180181182183184185186187188189190191192193194195196197198199200201202203204205206207208209210211212213214215216217218219220221222223224225226227228229230231232233234235236237238239240241242243244245246247248249250251252253254255256257258259260261262263264265266267268269270271272273274275276277278279280281282283284285286287288289290291292293294295296297298299300301302303304305306307308309310311312313314315316317318319320321322323324325326327328329330331332333334335336337338339340341342343344345346347348349350351352353354355356357358359360361362363364365366367368369370371372373374375376377378379380381382383384385386387388389390391392393394395396397398399400401402403404405406407408409410411412413414415416417418419420421422423424425426427428429430431432433434435436437438439440441442443444445446447448449450451452453454455456457458459460461462463464465466467468469470471472473474475476477478479480481482483484485486487488489490491492493494495496497498499500501502503504505506507508509510511512513514515516517518519520521522523524525526527528529530531532533534535536537538539540541542543544545546547548549550551552553554555556557558559560561562563564565566567568569570571572573574575576577578579580581582583584585586587588589590591592593594595596597598599600601602603604605606607608609610611612613614615616617618619620621622623624625626627628629630631632633634635636637638639640641642643644645646647648649650651652653654655656657658659660661662663664665666667668669670671672673674675676677678679680681682683684685686687688689690691692693694695696697698699700701702703704705706707708709710711712713714715716717718719720721722723724725726727728729730731732733734735736737738739740741742743744745746747748749750751752753754755756757758759760761762763764765766767768769770771772773774775776777778779780781782783784785786787788789790791792793794795796797798799800801802803804805806807808809810811812813814815816817818819820821822823824825826827828829830831832833834835836837838839840841842843844845846847848849850851852853854855856857858859860861862863864865866867868869870871872873874875876877878879880881882883884885886887888889890891892893894895896897898899900901902903904905906907908909910911912913914915916917918919920921922923924925926927928929930931932933934935936937938939940941942943944945946947948949950951952953954955956957958959960961962963964965966967968969970971972973974975976977978979980981982983984985986987988989990991992993994995996997998999100010011002100310041005100610071008100910101011101210131014101510161017101810191020102110221023102410251026102710281029103010311032103310341035103610371038103910401041104210431044104510461047104810491050105110521053105410551056105710581059106010611062106310641065106610671068106910701071107210731074107510761077107810791080108110821083108410851086108710881089109010911092109310941095109610971098109911001101110211031104110511061107110811091110111111121113111411151116111711181119112011211122112311241125112611271128112911301131113211331134113511361137113811391140114111421143114411451146114711481149115011511152115311541155115611571158115911601161116211631164116511661167116811691170
  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 math import prod
  10. from pathlib import Path
  11. from io import BufferedWriter
  12. from typing import IO, Any, Sequence, Mapping
  13. from string import ascii_letters, digits
  14. import numpy as np
  15. from .constants import (
  16. GGUF_DEFAULT_ALIGNMENT,
  17. GGUF_MAGIC,
  18. GGUF_VERSION,
  19. GGMLQuantizationType,
  20. GGUFEndian,
  21. GGUFValueType,
  22. Keys,
  23. RopeScalingType,
  24. PoolingType,
  25. TokenType,
  26. ExpertGatingFuncType,
  27. )
  28. from .quants import quant_shape_from_byte_shape
  29. logger = logging.getLogger(__name__)
  30. SHARD_NAME_FORMAT = "{:s}-{:05d}-of-{:05d}.gguf"
  31. @dataclass
  32. class TensorInfo:
  33. shape: Sequence[int]
  34. dtype: GGMLQuantizationType
  35. nbytes: int
  36. tensor: np.ndarray[Any, Any] | None = None
  37. @dataclass
  38. class GGUFValue:
  39. value: Any
  40. type: GGUFValueType
  41. sub_type: GGUFValueType | None = None
  42. class WriterState(Enum):
  43. NO_FILE = auto()
  44. EMPTY = auto()
  45. HEADER = auto()
  46. KV_DATA = auto()
  47. TI_DATA = auto()
  48. WEIGHTS = auto()
  49. class GGUFWriter:
  50. fout: list[BufferedWriter] | None
  51. path: Path | None
  52. temp_file: tempfile.SpooledTemporaryFile[bytes] | None
  53. tensors: list[dict[str, TensorInfo]]
  54. kv_data: list[dict[str, GGUFValue]]
  55. state: WriterState
  56. _simple_value_packing = {
  57. GGUFValueType.UINT8: "B",
  58. GGUFValueType.INT8: "b",
  59. GGUFValueType.UINT16: "H",
  60. GGUFValueType.INT16: "h",
  61. GGUFValueType.UINT32: "I",
  62. GGUFValueType.INT32: "i",
  63. GGUFValueType.FLOAT32: "f",
  64. GGUFValueType.UINT64: "Q",
  65. GGUFValueType.INT64: "q",
  66. GGUFValueType.FLOAT64: "d",
  67. GGUFValueType.BOOL: "?",
  68. }
  69. def __init__(
  70. self, path: os.PathLike[str] | str | None, arch: str, use_temp_file: bool = False, endianess: GGUFEndian = GGUFEndian.LITTLE,
  71. split_max_tensors: int = 0, split_max_size: int = 0, dry_run: bool = False, small_first_shard: bool = False
  72. ):
  73. self.fout = None
  74. self.path = Path(path) if path else None
  75. self.arch = arch
  76. self.endianess = endianess
  77. self.data_alignment = GGUF_DEFAULT_ALIGNMENT
  78. self.use_temp_file = use_temp_file
  79. self.temp_file = None
  80. self.tensors = [{}]
  81. self.kv_data = [{}]
  82. self.split_max_tensors = split_max_tensors
  83. self.split_max_size = split_max_size
  84. self.dry_run = dry_run
  85. self.small_first_shard = small_first_shard
  86. logger.info("gguf: This GGUF file is for {0} Endian only".format(
  87. "Big" if self.endianess == GGUFEndian.BIG else "Little",
  88. ))
  89. self.state = WriterState.NO_FILE
  90. if self.small_first_shard:
  91. self.tensors.append({})
  92. self.add_architecture()
  93. def get_total_parameter_count(self) -> tuple[int, int, int, int]:
  94. total_params = 0
  95. shared_params = 0
  96. expert_params = 0
  97. expert_sum = 0
  98. n_expert_tensors = 0
  99. last_lora_a: tuple[str, TensorInfo] | None = None
  100. for tensors in self.tensors:
  101. for name, info in tensors.items():
  102. shape = info.shape
  103. if name.endswith(".lora_a"):
  104. last_lora_a = (name, info)
  105. continue
  106. elif name.endswith(".lora_b"):
  107. if last_lora_a is None or last_lora_a[0] != name[:-1] + "a":
  108. # Bail when the LoRA pair can't be found trivially
  109. logger.warning("can't measure LoRA size correctly, tensor order is unusual")
  110. return 0, 0, 0, 0
  111. else:
  112. shape = (*shape[:-1], last_lora_a[1].shape[-1])
  113. size = prod(shape)
  114. if "_exps." in name:
  115. expert_count = shape[-2 if ".bias" in name else -3]
  116. expert_params += (size // expert_count)
  117. expert_sum += expert_count
  118. n_expert_tensors += 1
  119. else:
  120. shared_params += size
  121. total_params += size
  122. # Hopefully this should work even for variable-expert-count models
  123. expert_count = (expert_sum // n_expert_tensors) if n_expert_tensors > 0 else 0
  124. # Negate the total to signal it's likely not exact
  125. if last_lora_a is not None:
  126. total_params = -total_params
  127. # NOTE: keep the output in the same order as accepted by 'size_label' in gguf-py/gguf/utility.py
  128. return total_params, shared_params, expert_params, expert_count
  129. def format_shard_names(self, path: Path) -> list[Path]:
  130. if len(self.tensors) == 1:
  131. return [path]
  132. return [path.with_name(SHARD_NAME_FORMAT.format(path.stem, i + 1, len(self.tensors))) for i in range(len(self.tensors))]
  133. def open_output_file(self, path: Path | None = None) -> None:
  134. if self.state is WriterState.EMPTY and self.fout is not None and (path is None or path == self.path):
  135. # allow calling this multiple times as long as the path is the same
  136. return
  137. if self.state is not WriterState.NO_FILE:
  138. raise ValueError(f'Expected output file to be not yet opened, got {self.state}')
  139. if path is not None:
  140. self.path = path
  141. if self.path is not None:
  142. filenames = self.print_plan()
  143. self.fout = [open(filename, "wb") for filename in filenames]
  144. self.state = WriterState.EMPTY
  145. def print_plan(self) -> list[Path]:
  146. logger.info("Writing the following files:")
  147. assert self.path is not None
  148. filenames = self.format_shard_names(self.path)
  149. assert len(filenames) == len(self.tensors)
  150. for name, tensors in zip(filenames, self.tensors):
  151. logger.info(f"{name}: n_tensors = {len(tensors)}, total_size = {GGUFWriter.format_n_bytes_to_str(sum(ti.nbytes for ti in tensors.values()))}")
  152. if self.dry_run:
  153. logger.info("Dry run, not writing files")
  154. for name in filenames:
  155. print(name) # noqa: NP100
  156. exit()
  157. return filenames
  158. def add_shard_kv_data(self) -> None:
  159. if len(self.tensors) == 1:
  160. return
  161. total_tensors = sum(len(t) for t in self.tensors)
  162. assert self.fout is not None
  163. total_splits = len(self.fout)
  164. self.kv_data.extend({} for _ in range(len(self.kv_data), total_splits))
  165. for i, kv_data in enumerate(self.kv_data):
  166. kv_data[Keys.Split.LLM_KV_SPLIT_NO] = GGUFValue(i, GGUFValueType.UINT16)
  167. kv_data[Keys.Split.LLM_KV_SPLIT_COUNT] = GGUFValue(total_splits, GGUFValueType.UINT16)
  168. kv_data[Keys.Split.LLM_KV_SPLIT_TENSORS_COUNT] = GGUFValue(total_tensors, GGUFValueType.INT32)
  169. def write_header_to_file(self, path: Path | None = None) -> None:
  170. if len(self.tensors) == 1 and (self.split_max_tensors != 0 or self.split_max_size != 0):
  171. logger.warning("Model fails split requirements, not splitting")
  172. self.open_output_file(path)
  173. if self.state is not WriterState.EMPTY:
  174. raise ValueError(f'Expected output file to be empty, got {self.state}')
  175. assert self.fout is not None
  176. assert len(self.fout) == len(self.tensors)
  177. assert len(self.kv_data) == 1
  178. self.add_shard_kv_data()
  179. for fout, tensors, kv_data in zip(self.fout, self.tensors, self.kv_data):
  180. fout.write(self._pack("<I", GGUF_MAGIC, skip_pack_prefix = True))
  181. fout.write(self._pack("I", GGUF_VERSION))
  182. fout.write(self._pack("Q", len(tensors)))
  183. fout.write(self._pack("Q", len(kv_data)))
  184. fout.flush()
  185. self.state = WriterState.HEADER
  186. def write_kv_data_to_file(self) -> None:
  187. if self.state is not WriterState.HEADER:
  188. raise ValueError(f'Expected output file to contain the header, got {self.state}')
  189. assert self.fout is not None
  190. for fout, kv_data in zip(self.fout, self.kv_data):
  191. kv_bytes = bytearray()
  192. for key, val in kv_data.items():
  193. kv_bytes += self._pack_val(key, GGUFValueType.STRING, add_vtype=False)
  194. kv_bytes += self._pack_val(val.value, val.type, add_vtype=True, sub_type=val.sub_type)
  195. fout.write(kv_bytes)
  196. self.flush()
  197. self.state = WriterState.KV_DATA
  198. def write_ti_data_to_file(self) -> None:
  199. if self.state is not WriterState.KV_DATA:
  200. raise ValueError(f'Expected output file to contain KV data, got {self.state}')
  201. assert self.fout is not None
  202. for fout, tensors in zip(self.fout, self.tensors):
  203. ti_data = bytearray()
  204. offset_tensor = 0
  205. for name, ti in tensors.items():
  206. ti_data += self._pack_val(name, GGUFValueType.STRING, add_vtype=False)
  207. n_dims = len(ti.shape)
  208. ti_data += self._pack("I", n_dims)
  209. for j in range(n_dims):
  210. ti_data += self._pack("Q", ti.shape[n_dims - 1 - j])
  211. ti_data += self._pack("I", ti.dtype)
  212. ti_data += self._pack("Q", offset_tensor)
  213. offset_tensor += GGUFWriter.ggml_pad(ti.nbytes, self.data_alignment)
  214. fout.write(ti_data)
  215. fout.flush()
  216. self.state = WriterState.TI_DATA
  217. def add_key_value(self, key: str, val: Any, vtype: GGUFValueType, sub_type: GGUFValueType | None = None) -> None:
  218. if any(key in kv_data for kv_data in self.kv_data):
  219. logger.warning(f'Duplicated key name {key!r}, overwriting it with new value {val!r} of type {vtype.name}')
  220. self.kv_data[0][key] = GGUFValue(value=val, type=vtype, sub_type=sub_type)
  221. def add_uint8(self, key: str, val: int) -> None:
  222. self.add_key_value(key,val, GGUFValueType.UINT8)
  223. def add_int8(self, key: str, val: int) -> None:
  224. self.add_key_value(key, val, GGUFValueType.INT8)
  225. def add_uint16(self, key: str, val: int) -> None:
  226. self.add_key_value(key, val, GGUFValueType.UINT16)
  227. def add_int16(self, key: str, val: int) -> None:
  228. self.add_key_value(key, val, GGUFValueType.INT16)
  229. def add_uint32(self, key: str, val: int) -> None:
  230. self.add_key_value(key, val, GGUFValueType.UINT32)
  231. def add_int32(self, key: str, val: int) -> None:
  232. self.add_key_value(key, val, GGUFValueType.INT32)
  233. def add_float32(self, key: str, val: float) -> None:
  234. self.add_key_value(key, val, GGUFValueType.FLOAT32)
  235. def add_uint64(self, key: str, val: int) -> None:
  236. self.add_key_value(key, val, GGUFValueType.UINT64)
  237. def add_int64(self, key: str, val: int) -> None:
  238. self.add_key_value(key, val, GGUFValueType.INT64)
  239. def add_float64(self, key: str, val: float) -> None:
  240. self.add_key_value(key, val, GGUFValueType.FLOAT64)
  241. def add_bool(self, key: str, val: bool) -> None:
  242. self.add_key_value(key, val, GGUFValueType.BOOL)
  243. def add_string(self, key: str, val: str) -> None:
  244. if not val:
  245. return
  246. self.add_key_value(key, val, GGUFValueType.STRING)
  247. def add_array(self, key: str, val: Sequence[Any]) -> None:
  248. if len(val) == 0:
  249. return
  250. self.add_key_value(key, val, GGUFValueType.ARRAY)
  251. @staticmethod
  252. def ggml_pad(x: int, n: int) -> int:
  253. return ((x + n - 1) // n) * n
  254. def add_tensor_info(
  255. self, name: str, tensor_shape: Sequence[int], tensor_dtype: np.dtype,
  256. tensor_nbytes: int, raw_dtype: GGMLQuantizationType | None = None,
  257. ) -> None:
  258. if self.state is not WriterState.NO_FILE:
  259. raise ValueError(f'Expected output file to be not yet opened, got {self.state}')
  260. if any(name in tensors for tensors in self.tensors):
  261. raise ValueError(f'Duplicated tensor name {name!r}')
  262. if raw_dtype is None:
  263. if tensor_dtype == np.float16:
  264. dtype = GGMLQuantizationType.F16
  265. elif tensor_dtype == np.float32:
  266. dtype = GGMLQuantizationType.F32
  267. elif tensor_dtype == np.float64:
  268. dtype = GGMLQuantizationType.F64
  269. elif tensor_dtype == np.int8:
  270. dtype = GGMLQuantizationType.I8
  271. elif tensor_dtype == np.int16:
  272. dtype = GGMLQuantizationType.I16
  273. elif tensor_dtype == np.int32:
  274. dtype = GGMLQuantizationType.I32
  275. elif tensor_dtype == np.int64:
  276. dtype = GGMLQuantizationType.I64
  277. else:
  278. raise ValueError("Only F16, F32, F64, I8, I16, I32, I64 tensors are supported for now")
  279. else:
  280. dtype = raw_dtype
  281. if tensor_dtype == np.uint8:
  282. tensor_shape = quant_shape_from_byte_shape(tensor_shape, raw_dtype)
  283. # make sure there is at least one tensor before splitting
  284. if len(self.tensors[-1]) > 0:
  285. if ( # split when over tensor limit
  286. self.split_max_tensors != 0
  287. and len(self.tensors[-1]) >= self.split_max_tensors
  288. ) or ( # split when over size limit
  289. self.split_max_size != 0
  290. and sum(ti.nbytes for ti in self.tensors[-1].values()) + tensor_nbytes > self.split_max_size
  291. ):
  292. self.tensors.append({})
  293. self.tensors[-1][name] = TensorInfo(shape=tensor_shape, dtype=dtype, nbytes=tensor_nbytes)
  294. def add_tensor(
  295. self, name: str, tensor: np.ndarray[Any, Any], raw_shape: Sequence[int] | None = None,
  296. raw_dtype: GGMLQuantizationType | None = None,
  297. ) -> None:
  298. if self.endianess == GGUFEndian.BIG:
  299. tensor.byteswap(inplace=True)
  300. if self.use_temp_file and self.temp_file is None:
  301. fp = tempfile.SpooledTemporaryFile(mode="w+b", max_size=256 * 1024 * 1024)
  302. fp.seek(0)
  303. self.temp_file = fp
  304. shape: Sequence[int] = raw_shape if raw_shape is not None else tensor.shape
  305. self.add_tensor_info(name, shape, tensor.dtype, tensor.nbytes, raw_dtype=raw_dtype)
  306. if self.temp_file is None:
  307. self.tensors[-1][name].tensor = tensor
  308. return
  309. tensor.tofile(self.temp_file)
  310. self.write_padding(self.temp_file, tensor.nbytes)
  311. def write_padding(self, fp: IO[bytes], n: int, align: int | None = None) -> None:
  312. pad = GGUFWriter.ggml_pad(n, align if align is not None else self.data_alignment) - n
  313. if pad != 0:
  314. fp.write(bytes([0] * pad))
  315. def write_tensor_data(self, tensor: np.ndarray[Any, Any]) -> None:
  316. if self.state is not WriterState.TI_DATA and self.state is not WriterState.WEIGHTS:
  317. raise ValueError(f'Expected output file to contain tensor info or weights, got {self.state}')
  318. assert self.fout is not None
  319. if self.endianess == GGUFEndian.BIG:
  320. tensor.byteswap(inplace=True)
  321. file_id = -1
  322. for i, tensors in enumerate(self.tensors):
  323. if len(tensors) > 0:
  324. file_id = i
  325. break
  326. fout = self.fout[file_id]
  327. # pop the first tensor info
  328. # TODO: cleaner way to get the first key
  329. first_tensor_name = [name for name, _ in zip(self.tensors[file_id].keys(), range(1))][0]
  330. ti = self.tensors[file_id].pop(first_tensor_name)
  331. assert ti.nbytes == tensor.nbytes
  332. self.write_padding(fout, fout.tell())
  333. tensor.tofile(fout)
  334. self.write_padding(fout, tensor.nbytes)
  335. self.state = WriterState.WEIGHTS
  336. def write_tensors_to_file(self, *, progress: bool = False) -> None:
  337. self.write_ti_data_to_file()
  338. assert self.fout is not None
  339. for fout in self.fout:
  340. self.write_padding(fout, fout.tell())
  341. if self.temp_file is None:
  342. shard_bar = None
  343. bar = None
  344. if progress:
  345. from tqdm import tqdm
  346. total_bytes = sum(ti.nbytes for t in self.tensors for ti in t.values())
  347. if len(self.fout) > 1:
  348. shard_bar = tqdm(desc=f"Shard (0/{len(self.fout)})", total=None, unit="byte", unit_scale=True)
  349. bar = tqdm(desc="Writing", total=total_bytes, unit="byte", unit_scale=True)
  350. for i, (fout, tensors) in enumerate(zip(self.fout, self.tensors)):
  351. if shard_bar is not None:
  352. shard_bar.set_description(f"Shard ({i + 1}/{len(self.fout)})")
  353. total = sum(ti.nbytes for ti in tensors.values())
  354. shard_bar.reset(total=(total if total > 0 else None))
  355. # relying on the fact that Python dicts preserve insertion order (since 3.7)
  356. for ti in tensors.values():
  357. assert ti.tensor is not None # can only iterate once over the tensors
  358. assert ti.tensor.nbytes == ti.nbytes
  359. ti.tensor.tofile(fout)
  360. if shard_bar is not None:
  361. shard_bar.update(ti.nbytes)
  362. if bar is not None:
  363. bar.update(ti.nbytes)
  364. self.write_padding(fout, ti.nbytes)
  365. ti.tensor = None
  366. else:
  367. self.temp_file.seek(0)
  368. shutil.copyfileobj(self.temp_file, self.fout[0 if not self.small_first_shard else 1])
  369. self.flush()
  370. self.temp_file.close()
  371. self.state = WriterState.WEIGHTS
  372. def flush(self) -> None:
  373. assert self.fout is not None
  374. for fout in self.fout:
  375. fout.flush()
  376. def close(self) -> None:
  377. if self.fout is not None:
  378. for fout in self.fout:
  379. fout.close()
  380. self.fout = None
  381. def add_type(self, type_name: str) -> None:
  382. self.add_string(Keys.General.TYPE, type_name)
  383. def add_architecture(self) -> None:
  384. self.add_string(Keys.General.ARCHITECTURE, self.arch)
  385. def add_quantization_version(self, quantization_version: int) -> None:
  386. self.add_uint32(Keys.General.QUANTIZATION_VERSION, quantization_version)
  387. def add_custom_alignment(self, alignment: int) -> None:
  388. self.data_alignment = alignment
  389. self.add_uint32(Keys.General.ALIGNMENT, alignment)
  390. def add_file_type(self, ftype: int) -> None:
  391. self.add_uint32(Keys.General.FILE_TYPE, ftype)
  392. def add_name(self, name: str) -> None:
  393. self.add_string(Keys.General.NAME, name)
  394. def add_author(self, author: str) -> None:
  395. self.add_string(Keys.General.AUTHOR, author)
  396. def add_version(self, version: str) -> None:
  397. self.add_string(Keys.General.VERSION, version)
  398. def add_organization(self, organization: str) -> None:
  399. self.add_string(Keys.General.ORGANIZATION, organization)
  400. def add_finetune(self, finetune: str) -> None:
  401. self.add_string(Keys.General.FINETUNE, finetune)
  402. def add_basename(self, basename: str) -> None:
  403. self.add_string(Keys.General.BASENAME, basename)
  404. def add_description(self, description: str) -> None:
  405. self.add_string(Keys.General.DESCRIPTION, description)
  406. def add_quantized_by(self, quantized: str) -> None:
  407. self.add_string(Keys.General.QUANTIZED_BY, quantized)
  408. def add_size_label(self, size_label: str) -> None:
  409. self.add_string(Keys.General.SIZE_LABEL, size_label)
  410. def add_license(self, license: str) -> None:
  411. self.add_string(Keys.General.LICENSE, license)
  412. def add_license_name(self, license: str) -> None:
  413. self.add_string(Keys.General.LICENSE_NAME, license)
  414. def add_license_link(self, license: str) -> None:
  415. self.add_string(Keys.General.LICENSE_LINK, license)
  416. def add_url(self, url: str) -> None:
  417. self.add_string(Keys.General.URL, url)
  418. def add_doi(self, doi: str) -> None:
  419. self.add_string(Keys.General.DOI, doi)
  420. def add_uuid(self, uuid: str) -> None:
  421. self.add_string(Keys.General.UUID, uuid)
  422. def add_repo_url(self, repo_url: str) -> None:
  423. self.add_string(Keys.General.REPO_URL, repo_url)
  424. def add_source_url(self, url: str) -> None:
  425. self.add_string(Keys.General.SOURCE_URL, url)
  426. def add_source_doi(self, doi: str) -> None:
  427. self.add_string(Keys.General.SOURCE_DOI, doi)
  428. def add_source_uuid(self, uuid: str) -> None:
  429. self.add_string(Keys.General.SOURCE_UUID, uuid)
  430. def add_source_repo_url(self, repo_url: str) -> None:
  431. self.add_string(Keys.General.SOURCE_REPO_URL, repo_url)
  432. def add_base_model_count(self, source_count: int) -> None:
  433. self.add_uint32(Keys.General.BASE_MODEL_COUNT, source_count)
  434. def add_base_model_name(self, source_id: int, name: str) -> None:
  435. self.add_string(Keys.General.BASE_MODEL_NAME.format(id=source_id), name)
  436. def add_base_model_author(self, source_id: int, author: str) -> None:
  437. self.add_string(Keys.General.BASE_MODEL_AUTHOR.format(id=source_id), author)
  438. def add_base_model_version(self, source_id: int, version: str) -> None:
  439. self.add_string(Keys.General.BASE_MODEL_VERSION.format(id=source_id), version)
  440. def add_base_model_organization(self, source_id: int, organization: str) -> None:
  441. self.add_string(Keys.General.BASE_MODEL_ORGANIZATION.format(id=source_id), organization)
  442. def add_base_model_description(self, source_id: int, description: str) -> None:
  443. self.add_string(Keys.General.BASE_MODEL_DESCRIPTION.format(id=source_id), description)
  444. def add_base_model_url(self, source_id: int, url: str) -> None:
  445. self.add_string(Keys.General.BASE_MODEL_URL.format(id=source_id), url)
  446. def add_base_model_doi(self, source_id: int, doi: str) -> None:
  447. self.add_string(Keys.General.BASE_MODEL_DOI.format(id=source_id), doi)
  448. def add_base_model_uuid(self, source_id: int, uuid: str) -> None:
  449. self.add_string(Keys.General.BASE_MODEL_UUID.format(id=source_id), uuid)
  450. def add_base_model_repo_url(self, source_id: int, repo_url: str) -> None:
  451. self.add_string(Keys.General.BASE_MODEL_REPO_URL.format(id=source_id), repo_url)
  452. def add_dataset_count(self, source_count: int) -> None:
  453. self.add_uint32(Keys.General.DATASET_COUNT, source_count)
  454. def add_dataset_name(self, source_id: int, name: str) -> None:
  455. self.add_string(Keys.General.DATASET_NAME.format(id=source_id), name)
  456. def add_dataset_author(self, source_id: int, author: str) -> None:
  457. self.add_string(Keys.General.DATASET_AUTHOR.format(id=source_id), author)
  458. def add_dataset_version(self, source_id: int, version: str) -> None:
  459. self.add_string(Keys.General.DATASET_VERSION.format(id=source_id), version)
  460. def add_dataset_organization(self, source_id: int, organization: str) -> None:
  461. self.add_string(Keys.General.DATASET_ORGANIZATION.format(id=source_id), organization)
  462. def add_dataset_description(self, source_id: int, description: str) -> None:
  463. self.add_string(Keys.General.DATASET_DESCRIPTION.format(id=source_id), description)
  464. def add_dataset_url(self, source_id: int, url: str) -> None:
  465. self.add_string(Keys.General.DATASET_URL.format(id=source_id), url)
  466. def add_dataset_doi(self, source_id: int, doi: str) -> None:
  467. self.add_string(Keys.General.DATASET_DOI.format(id=source_id), doi)
  468. def add_dataset_uuid(self, source_id: int, uuid: str) -> None:
  469. self.add_string(Keys.General.DATASET_UUID.format(id=source_id), uuid)
  470. def add_dataset_repo_url(self, source_id: int, repo_url: str) -> None:
  471. self.add_string(Keys.General.DATASET_REPO_URL.format(id=source_id), repo_url)
  472. def add_tags(self, tags: Sequence[str]) -> None:
  473. self.add_array(Keys.General.TAGS, tags)
  474. def add_languages(self, languages: Sequence[str]) -> None:
  475. self.add_array(Keys.General.LANGUAGES, languages)
  476. def add_tensor_data_layout(self, layout: str) -> None:
  477. self.add_string(Keys.LLM.TENSOR_DATA_LAYOUT.format(arch=self.arch), layout)
  478. def add_vocab_size(self, size: int) -> None:
  479. self.add_uint32(Keys.LLM.VOCAB_SIZE.format(arch=self.arch), size)
  480. def add_context_length(self, length: int) -> None:
  481. self.add_uint32(Keys.LLM.CONTEXT_LENGTH.format(arch=self.arch), length)
  482. def add_embedding_length(self, length: int) -> None:
  483. self.add_uint32(Keys.LLM.EMBEDDING_LENGTH.format(arch=self.arch), length)
  484. def add_features_length(self, length: int) -> None:
  485. self.add_uint32(Keys.LLM.FEATURES_LENGTH.format(arch=self.arch), length)
  486. def add_posnet_embedding_length(self, length: int) -> None:
  487. self.add_uint32(Keys.PosNet.EMBEDDING_LENGTH.format(arch=self.arch), length)
  488. def add_posnet_block_count(self, length: int) -> None:
  489. self.add_uint32(Keys.PosNet.BLOCK_COUNT.format(arch=self.arch), length)
  490. def add_convnext_embedding_length(self, length: int) -> None:
  491. self.add_uint32(Keys.ConvNext.EMBEDDING_LENGTH.format(arch=self.arch), length)
  492. def add_convnext_block_count(self, length: int) -> None:
  493. self.add_uint32(Keys.ConvNext.BLOCK_COUNT.format(arch=self.arch), length)
  494. def add_shortconv_l_cache(self, length: int) -> None:
  495. self.add_uint32(Keys.ShortConv.L_CACHE.format(arch=self.arch), length)
  496. def add_block_count(self, length: int) -> None:
  497. self.add_uint32(Keys.LLM.BLOCK_COUNT.format(arch=self.arch), length)
  498. def add_leading_dense_block_count(self, length: int) -> None:
  499. self.add_uint32(Keys.LLM.LEADING_DENSE_BLOCK_COUNT.format(arch=self.arch), length)
  500. def add_feed_forward_length(self, length: int | Sequence[int]) -> None:
  501. if isinstance(length, int):
  502. self.add_uint32(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  503. else:
  504. self.add_array(Keys.LLM.FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  505. def add_expert_feed_forward_length(self, length: int) -> None:
  506. self.add_uint32(Keys.LLM.EXPERT_FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  507. def add_expert_shared_feed_forward_length(self, length: int) -> None:
  508. self.add_uint32(Keys.LLM.EXPERT_SHARED_FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  509. def add_expert_chunk_feed_forward_length(self, length: int) -> None:
  510. self.add_uint32(Keys.LLM.EXPERT_CHUNK_FEED_FORWARD_LENGTH.format(arch=self.arch), length)
  511. def add_parallel_residual(self, use: bool) -> None:
  512. self.add_bool(Keys.LLM.USE_PARALLEL_RESIDUAL.format(arch=self.arch), use)
  513. def add_decoder_start_token_id(self, id: int) -> None:
  514. self.add_uint32(Keys.LLM.DECODER_START_TOKEN_ID.format(arch=self.arch), id)
  515. def add_decoder_block_count(self, value: int) -> None:
  516. self.add_uint32(Keys.LLM.DECODER_BLOCK_COUNT.format(arch=self.arch), value)
  517. def add_embedding_length_per_layer_input(self, value: int) -> None:
  518. self.add_uint32(Keys.LLM.EMBD_LENGTH_PER_LAYER_INP.format(arch=self.arch), value)
  519. def add_altup_active_idx(self, val: int) -> None:
  520. self.add_uint32(Keys.LLM.ALTUP_ACTIVE_IDX.format(arch=self.arch), val)
  521. def add_altup_num_inputs(self, val: int) -> None:
  522. self.add_uint32(Keys.LLM.ALTUP_NUM_INPUTS.format(arch=self.arch), val)
  523. def add_activation_sparsity_scale(self, values: Sequence[float]) -> None:
  524. self.add_array(Keys.LLM.ACTIVATION_SPARSITY_SCALE.format(arch=self.arch), values)
  525. def add_head_count(self, count: int | Sequence[int]) -> None:
  526. if isinstance(count, int):
  527. self.add_uint32(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count)
  528. else:
  529. self.add_array(Keys.Attention.HEAD_COUNT.format(arch=self.arch), count)
  530. def add_head_count_kv(self, count: int | Sequence[int]) -> None:
  531. if isinstance(count, int):
  532. self.add_uint32(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count)
  533. else:
  534. self.add_array(Keys.Attention.HEAD_COUNT_KV.format(arch=self.arch), count)
  535. def add_key_length(self, length: int) -> None:
  536. self.add_uint32(Keys.Attention.KEY_LENGTH.format(arch=self.arch), length)
  537. def add_value_length(self, length: int) -> None:
  538. self.add_uint32(Keys.Attention.VALUE_LENGTH.format(arch=self.arch), length)
  539. def add_key_length_mla(self, length: int) -> None:
  540. self.add_uint32(Keys.Attention.KEY_LENGTH_MLA.format(arch=self.arch), length)
  541. def add_value_length_mla(self, length: int) -> None:
  542. self.add_uint32(Keys.Attention.VALUE_LENGTH_MLA.format(arch=self.arch), length)
  543. def add_max_alibi_bias(self, bias: float) -> None:
  544. self.add_float32(Keys.Attention.MAX_ALIBI_BIAS.format(arch=self.arch), bias)
  545. def add_clamp_kqv(self, value: float) -> None:
  546. self.add_float32(Keys.Attention.CLAMP_KQV.format(arch=self.arch), value)
  547. def add_shared_kv_layers(self, value: int) -> None:
  548. self.add_uint32(Keys.Attention.SHARED_KV_LAYERS.format(arch=self.arch), value)
  549. def add_sliding_window_pattern(self, value: Sequence[bool]) -> None:
  550. self.add_array(Keys.Attention.SLIDING_WINDOW_PATTERN.format(arch=self.arch), value)
  551. def add_dense_features_dims(self, dense:str, in_f:int, out_f:int) -> None:
  552. self.add_uint32(Keys.LLM.DENSE_FEAT_IN_SIZE.format(arch=self.arch, dense=dense), in_f)
  553. self.add_uint32(Keys.LLM.DENSE_FEAT_OUT_SIZE.format(arch=self.arch, dense=dense), out_f)
  554. def add_logit_scale(self, value: float) -> None:
  555. self.add_float32(Keys.LLM.LOGIT_SCALE.format(arch=self.arch), value)
  556. def add_attn_logit_softcapping(self, value: float) -> None:
  557. self.add_float32(Keys.LLM.ATTN_LOGIT_SOFTCAPPING.format(arch=self.arch), value)
  558. def add_router_logit_softcapping(self, value: float) -> None:
  559. self.add_float32(Keys.LLM.ROUTER_LOGIT_SOFTCAPPING.format(arch=self.arch), value)
  560. def add_final_logit_softcapping(self, value: float) -> None:
  561. self.add_float32(Keys.LLM.FINAL_LOGIT_SOFTCAPPING.format(arch=self.arch), value)
  562. def add_expert_count(self, count: int) -> None:
  563. self.add_uint32(Keys.LLM.EXPERT_COUNT.format(arch=self.arch), count)
  564. def add_expert_used_count(self, count: int) -> None:
  565. self.add_uint32(Keys.LLM.EXPERT_USED_COUNT.format(arch=self.arch), count)
  566. def add_expert_shared_count(self, count: int) -> None:
  567. self.add_uint32(Keys.LLM.EXPERT_SHARED_COUNT.format(arch=self.arch), count)
  568. def add_expert_group_count(self, count: int) -> None:
  569. self.add_uint32(Keys.LLM.EXPERT_GROUP_COUNT.format(arch=self.arch), count)
  570. def add_expert_group_used_count(self, count: int) -> None:
  571. self.add_uint32(Keys.LLM.EXPERT_GROUP_USED_COUNT.format(arch=self.arch), count)
  572. def add_expert_weights_scale(self, value: float) -> None:
  573. self.add_float32(Keys.LLM.EXPERT_WEIGHTS_SCALE.format(arch=self.arch), value)
  574. def add_expert_weights_norm(self, value: bool) -> None:
  575. self.add_bool(Keys.LLM.EXPERT_WEIGHTS_NORM.format(arch=self.arch), value)
  576. def add_expert_gating_func(self, value: ExpertGatingFuncType) -> None:
  577. self.add_uint32(Keys.LLM.EXPERT_GATING_FUNC.format(arch=self.arch), value.value)
  578. def add_expert_group_scale(self, value: float) -> None:
  579. self.add_float32(Keys.LLM.EXPERT_GROUP_SCALE.format(arch=self.arch), value)
  580. def add_experts_per_group(self, count: int) -> None:
  581. self.add_uint32(Keys.LLM.EXPERTS_PER_GROUP.format(arch=self.arch), count)
  582. def add_moe_every_n_layers(self, value: int) -> None:
  583. self.add_uint32(Keys.LLM.MOE_EVERY_N_LAYERS.format(arch=self.arch), value)
  584. def add_nextn_predict_layers(self, count: int) -> None:
  585. self.add_uint32(Keys.LLM.NEXTN_PREDICT_LAYERS.format(arch=self.arch), count)
  586. def add_swin_norm(self, value: bool) -> None:
  587. self.add_bool(Keys.LLM.SWIN_NORM.format(arch=self.arch), value)
  588. def add_rescale_every_n_layers(self, count: int) -> None:
  589. self.add_uint32(Keys.LLM.RESCALE_EVERY_N_LAYERS.format(arch=self.arch), count)
  590. def add_time_mix_extra_dim(self, dim: int) -> None:
  591. self.add_uint32(Keys.LLM.TIME_MIX_EXTRA_DIM.format(arch=self.arch), dim)
  592. def add_time_decay_extra_dim(self, dim: int) -> None:
  593. self.add_uint32(Keys.LLM.TIME_DECAY_EXTRA_DIM.format(arch=self.arch), dim)
  594. def add_residual_scale(self, value: float) -> None:
  595. self.add_float32(Keys.LLM.RESIDUAL_SCALE.format(arch=self.arch), value)
  596. def add_embedding_scale(self, value: float) -> None:
  597. self.add_float32(Keys.LLM.EMBEDDING_SCALE.format(arch=self.arch), value)
  598. def add_wkv_head_size(self, size: int) -> None:
  599. self.add_uint32(Keys.WKV.HEAD_SIZE.format(arch=self.arch), size)
  600. def add_token_shift_count(self, count: int) -> None:
  601. self.add_uint32(Keys.LLM.TOKEN_SHIFT_COUNT.format(arch=self.arch), count)
  602. def add_interleave_moe_layer_step(self, value: int) -> None:
  603. self.add_uint32(Keys.LLM.INTERLEAVE_MOE_LAYER_STEP.format(arch=self.arch), value)
  604. def add_layer_norm_eps(self, value: float) -> None:
  605. self.add_float32(Keys.Attention.LAYERNORM_EPS.format(arch=self.arch), value)
  606. def add_layer_norm_rms_eps(self, value: float) -> None:
  607. self.add_float32(Keys.Attention.LAYERNORM_RMS_EPS.format(arch=self.arch), value)
  608. def add_group_norm_eps(self, value: float) -> None:
  609. self.add_float32(Keys.Attention.GROUPNORM_EPS.format(arch=self.arch), value)
  610. def add_group_norm_groups(self, value: int) -> None:
  611. self.add_uint32(Keys.Attention.GROUPNORM_GROUPS.format(arch=self.arch), value)
  612. def add_causal_attention(self, value: bool) -> None:
  613. self.add_bool(Keys.Attention.CAUSAL.format(arch=self.arch), value)
  614. def add_q_lora_rank(self, length: int) -> None:
  615. self.add_uint32(Keys.Attention.Q_LORA_RANK.format(arch=self.arch), length)
  616. def add_kv_lora_rank(self, length: int) -> None:
  617. self.add_uint32(Keys.Attention.KV_LORA_RANK.format(arch=self.arch), length)
  618. def add_decay_lora_rank(self, length: int) -> None:
  619. self.add_uint32(Keys.Attention.DECAY_LORA_RANK.format(arch=self.arch), length)
  620. def add_iclr_lora_rank(self, length: int) -> None:
  621. self.add_uint32(Keys.Attention.ICLR_LORA_RANK.format(arch=self.arch), length)
  622. def add_value_residual_mix_lora_rank(self, length: int) -> None:
  623. self.add_uint32(Keys.Attention.VALUE_RESIDUAL_MIX_LORA_RANK.format(arch=self.arch), length)
  624. def add_gate_lora_rank(self, length: int) -> None:
  625. self.add_uint32(Keys.Attention.GATE_LORA_RANK.format(arch=self.arch), length)
  626. def add_relative_attn_buckets_count(self, value: int) -> None:
  627. self.add_uint32(Keys.Attention.REL_BUCKETS_COUNT.format(arch=self.arch), value)
  628. def add_sliding_window(self, value: int) -> None:
  629. self.add_uint32(Keys.Attention.SLIDING_WINDOW.format(arch=self.arch), value)
  630. def add_attention_scale(self, value: float) -> None:
  631. self.add_float32(Keys.Attention.SCALE.format(arch=self.arch), value)
  632. def add_attn_output_scale(self, value: float) -> None:
  633. self.add_float32(Keys.Attention.OUTPUT_SCALE.format(arch=self.arch), value)
  634. def add_attn_temperature_length(self, value: int) -> None:
  635. self.add_uint32(Keys.Attention.TEMPERATURE_LENGTH.format(arch=self.arch), value)
  636. def add_pooling_type(self, value: PoolingType) -> None:
  637. self.add_uint32(Keys.LLM.POOLING_TYPE.format(arch=self.arch), value.value)
  638. def add_rope_dimension_count(self, count: int) -> None:
  639. self.add_uint32(Keys.Rope.DIMENSION_COUNT.format(arch=self.arch), count)
  640. def add_rope_dimension_sections(self, dims: Sequence[int]) -> None:
  641. self.add_array(Keys.Rope.DIMENSION_SECTIONS.format(arch=self.arch), dims)
  642. def add_rope_freq_base(self, value: float) -> None:
  643. self.add_float32(Keys.Rope.FREQ_BASE.format(arch=self.arch), value)
  644. def add_rope_scaling_type(self, value: RopeScalingType) -> None:
  645. self.add_string(Keys.Rope.SCALING_TYPE.format(arch=self.arch), value.value)
  646. def add_rope_scaling_factor(self, value: float) -> None:
  647. self.add_float32(Keys.Rope.SCALING_FACTOR.format(arch=self.arch), value)
  648. def add_rope_scaling_attn_factors(self, value: float) -> None:
  649. self.add_float32(Keys.Rope.SCALING_ATTN_FACTOR.format(arch=self.arch), value)
  650. def add_rope_scaling_orig_ctx_len(self, value: int) -> None:
  651. self.add_uint32(Keys.Rope.SCALING_ORIG_CTX_LEN.format(arch=self.arch), value)
  652. def add_rope_scaling_finetuned(self, value: bool) -> None:
  653. self.add_bool(Keys.Rope.SCALING_FINETUNED.format(arch=self.arch), value)
  654. def add_rope_scaling_yarn_log_mul(self, value: float) -> None:
  655. self.add_float32(Keys.Rope.SCALING_YARN_LOG_MUL.format(arch=self.arch), value)
  656. def add_rope_scaling_yarn_ext_factor(self, value: float) -> None:
  657. self.add_float32(Keys.Rope.SCALING_YARN_EXT_FACTOR.format(arch=self.arch), value)
  658. def add_rope_scaling_yarn_attn_factor(self, value: float) -> None:
  659. self.add_float32(Keys.Rope.SCALING_YARN_ATTN_FACTOR.format(arch=self.arch), value)
  660. def add_rope_scaling_yarn_beta_fast(self, value: float) -> None:
  661. self.add_float32(Keys.Rope.SCALING_YARN_BETA_FAST.format(arch=self.arch), value)
  662. def add_rope_scaling_yarn_beta_slow(self, value: float) -> None:
  663. self.add_float32(Keys.Rope.SCALING_YARN_BETA_SLOW.format(arch=self.arch), value)
  664. def add_ssm_conv_kernel(self, value: int) -> None:
  665. self.add_uint32(Keys.SSM.CONV_KERNEL.format(arch=self.arch), value)
  666. def add_ssm_inner_size(self, value: int) -> None:
  667. self.add_uint32(Keys.SSM.INNER_SIZE.format(arch=self.arch), value)
  668. def add_ssm_state_size(self, value: int) -> None:
  669. self.add_uint32(Keys.SSM.STATE_SIZE.format(arch=self.arch), value)
  670. def add_ssm_time_step_rank(self, value: int) -> None:
  671. self.add_uint32(Keys.SSM.TIME_STEP_RANK.format(arch=self.arch), value)
  672. def add_ssm_group_count(self, value: int) -> None:
  673. self.add_uint32(Keys.SSM.GROUP_COUNT.format(arch=self.arch), value)
  674. def add_ssm_dt_b_c_rms(self, value: bool) -> None:
  675. self.add_bool(Keys.SSM.DT_B_C_RMS.format(arch=self.arch), value)
  676. def add_tokenizer_model(self, model: str) -> None:
  677. self.add_string(Keys.Tokenizer.MODEL, model)
  678. def add_tokenizer_pre(self, pre: str) -> None:
  679. self.add_string(Keys.Tokenizer.PRE, pre)
  680. def add_token_list(self, tokens: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None:
  681. self.add_array(Keys.Tokenizer.LIST, tokens)
  682. def add_token_merges(self, merges: Sequence[str] | Sequence[bytes] | Sequence[bytearray]) -> None:
  683. self.add_array(Keys.Tokenizer.MERGES, merges)
  684. def add_token_types(self, types: Sequence[TokenType] | Sequence[int]) -> None:
  685. self.add_array(Keys.Tokenizer.TOKEN_TYPE, types)
  686. def add_token_type_count(self, value: int) -> None:
  687. self.add_uint32(Keys.Tokenizer.TOKEN_TYPE_COUNT, value)
  688. def add_token_scores(self, scores: Sequence[float]) -> None:
  689. self.add_array(Keys.Tokenizer.SCORES, scores)
  690. def add_bos_token_id(self, id: int) -> None:
  691. self.add_uint32(Keys.Tokenizer.BOS_ID, id)
  692. def add_eos_token_id(self, id: int) -> None:
  693. self.add_uint32(Keys.Tokenizer.EOS_ID, id)
  694. def add_unk_token_id(self, id: int) -> None:
  695. self.add_uint32(Keys.Tokenizer.UNK_ID, id)
  696. def add_sep_token_id(self, id: int) -> None:
  697. self.add_uint32(Keys.Tokenizer.SEP_ID, id)
  698. def add_pad_token_id(self, id: int) -> None:
  699. self.add_uint32(Keys.Tokenizer.PAD_ID, id)
  700. def add_mask_token_id(self, id: int) -> None:
  701. self.add_uint32(Keys.Tokenizer.MASK_ID, id)
  702. def add_add_bos_token(self, value: bool) -> None:
  703. self.add_bool(Keys.Tokenizer.ADD_BOS, value)
  704. def add_add_eos_token(self, value: bool) -> None:
  705. self.add_bool(Keys.Tokenizer.ADD_EOS, value)
  706. def add_add_sep_token(self, value: bool) -> None:
  707. self.add_bool(Keys.Tokenizer.ADD_SEP, value)
  708. def add_add_space_prefix(self, value: bool) -> None:
  709. self.add_bool(Keys.Tokenizer.ADD_PREFIX, value)
  710. def add_remove_extra_whitespaces(self, value: bool) -> None:
  711. self.add_bool(Keys.Tokenizer.REMOVE_EXTRA_WS, value)
  712. def add_precompiled_charsmap(self, charsmap: bytes) -> None:
  713. self.add_array(Keys.Tokenizer.PRECOMPILED_CHARSMAP, charsmap)
  714. def add_chat_template(self, value: str | Sequence[Mapping[str, str]]) -> None:
  715. if not isinstance(value, str):
  716. template_default = None
  717. template_names = set()
  718. for choice in value:
  719. name = choice.get('name', '')
  720. template = choice.get('template')
  721. # Allowing non-alphanumerical characters in template name is probably not a good idea, so filter it
  722. name = ''.join((c if c in ascii_letters + digits else '_' for c in name))
  723. if name and template is not None:
  724. if name == 'default':
  725. template_default = template
  726. else:
  727. template_names.add(name)
  728. self.add_string(Keys.Tokenizer.CHAT_TEMPLATE_N.format(name=name), template)
  729. if template_names:
  730. self.add_array(Keys.Tokenizer.CHAT_TEMPLATES, list(template_names))
  731. if template_default is None:
  732. return
  733. value = template_default
  734. self.add_string(Keys.Tokenizer.CHAT_TEMPLATE, value)
  735. def add_eot_token_id(self, id: int) -> None:
  736. self.add_uint32(Keys.Tokenizer.EOT_ID, id)
  737. def add_eom_token_id(self, id: int) -> None:
  738. self.add_uint32(Keys.Tokenizer.EOM_ID, id)
  739. def add_classifier_output_labels(self, labels: Sequence[str]) -> None:
  740. self.add_array(Keys.Classifier.OUTPUT_LABELS.format(arch=self.arch), labels)
  741. # for vision models
  742. def add_clip_has_vision_encoder(self, value: bool) -> None:
  743. self.add_bool(Keys.Clip.HAS_VISION_ENCODER, value)
  744. def add_clip_has_audio_encoder(self, value: bool) -> None:
  745. self.add_bool(Keys.Clip.HAS_AUDIO_ENCODER, value)
  746. def add_clip_projector_type(self, value: str) -> None:
  747. self.add_string(Keys.Clip.PROJECTOR_TYPE, value)
  748. def add_vision_projection_dim(self, value: int) -> None:
  749. self.add_uint32(Keys.ClipVision.PROJECTION_DIM, value)
  750. def add_vision_patch_size(self, value: int) -> None:
  751. self.add_uint32(Keys.ClipVision.PATCH_SIZE, value)
  752. def add_vision_embedding_length(self, value: int) -> None:
  753. self.add_uint32(Keys.ClipVision.EMBEDDING_LENGTH, value)
  754. def add_vision_feed_forward_length(self, value: int) -> None:
  755. self.add_uint32(Keys.ClipVision.FEED_FORWARD_LENGTH, value)
  756. def add_vision_block_count(self, value: int) -> None:
  757. self.add_uint32(Keys.ClipVision.BLOCK_COUNT, value)
  758. def add_vision_head_count(self, value: int) -> None:
  759. self.add_uint32(Keys.ClipVision.Attention.HEAD_COUNT, value)
  760. def add_vision_attention_layernorm_eps(self, value: float) -> None:
  761. self.add_float32(Keys.ClipVision.Attention.LAYERNORM_EPS, value)
  762. def add_vision_image_size(self, value: int) -> None:
  763. self.add_uint32(Keys.ClipVision.IMAGE_SIZE, value)
  764. def add_vision_preproc_image_size(self, value: int) -> None:
  765. self.add_uint32(Keys.ClipVision.PREPROC_IMAGE_SIZE, value)
  766. def add_vision_image_mean(self, values: Sequence[float]) -> None:
  767. self.add_array(Keys.ClipVision.IMAGE_MEAN, values)
  768. def add_vision_image_std(self, values: Sequence[float]) -> None:
  769. self.add_array(Keys.ClipVision.IMAGE_STD, values)
  770. def add_vision_spatial_merge_size(self, value: int) -> None:
  771. self.add_uint32(Keys.ClipVision.SPATIAL_MERGE_SIZE, value)
  772. def add_vision_use_gelu(self, value: bool) -> None:
  773. self.add_bool(Keys.ClipVision.USE_GELU, value)
  774. def add_vision_use_silu(self, value: bool) -> None:
  775. self.add_bool(Keys.ClipVision.USE_SILU, value)
  776. def add_vision_projector_scale_factor(self, value: int) -> None:
  777. self.add_uint32(Keys.ClipVision.Projector.SCALE_FACTOR, value)
  778. def add_vision_n_wa_pattern(self, value: int) -> None:
  779. self.add_uint32(Keys.ClipVision.N_WA_PATTERN, value)
  780. # audio models
  781. def add_audio_projection_dim(self, value: int) -> None:
  782. self.add_uint32(Keys.ClipAudio.PROJECTION_DIM, value)
  783. def add_audio_embedding_length(self, value: int) -> None:
  784. self.add_uint32(Keys.ClipAudio.EMBEDDING_LENGTH, value)
  785. def add_audio_feed_forward_length(self, value: int) -> None:
  786. self.add_uint32(Keys.ClipAudio.FEED_FORWARD_LENGTH, value)
  787. def add_audio_block_count(self, value: int) -> None:
  788. self.add_uint32(Keys.ClipAudio.BLOCK_COUNT, value)
  789. def add_audio_head_count(self, value: int) -> None:
  790. self.add_uint32(Keys.ClipAudio.Attention.HEAD_COUNT, value)
  791. def add_audio_attention_layernorm_eps(self, value: float) -> None:
  792. self.add_float32(Keys.ClipAudio.Attention.LAYERNORM_EPS, value)
  793. def add_audio_num_mel_bins(self, value: int) -> None:
  794. self.add_uint32(Keys.ClipAudio.NUM_MEL_BINS, value)
  795. def add_audio_stack_factor(self, value: int) -> None:
  796. self.add_uint32(Keys.ClipAudio.Projector.STACK_FACTOR, value)
  797. def add_xielu_alpha_p(self, values: Sequence[float]):
  798. self.add_array(Keys.xIELU.ALPHA_P, values)
  799. def add_xielu_alpha_n(self, values: Sequence[float]):
  800. self.add_array(Keys.xIELU.ALPHA_N, values)
  801. def add_xielu_beta(self, values: Sequence[float]):
  802. self.add_array(Keys.xIELU.BETA, values)
  803. def add_xielu_eps(self, values: Sequence[float]):
  804. self.add_array(Keys.xIELU.EPS, values)
  805. # diffusion models
  806. def add_diffusion_shift_logits(self, value: bool) -> None:
  807. self.add_bool(Keys.Diffusion.SHIFT_LOGITS, value)
  808. def _pack(self, fmt: str, value: Any, skip_pack_prefix: bool = False) -> bytes:
  809. pack_prefix = ''
  810. if not skip_pack_prefix:
  811. pack_prefix = '<' if self.endianess == GGUFEndian.LITTLE else '>'
  812. return struct.pack(f'{pack_prefix}{fmt}', value)
  813. def _pack_val(self, val: Any, vtype: GGUFValueType, add_vtype: bool, sub_type: GGUFValueType | None = None) -> bytes:
  814. kv_data = bytearray()
  815. if add_vtype:
  816. kv_data += self._pack("I", vtype)
  817. pack_fmt = self._simple_value_packing.get(vtype)
  818. if pack_fmt is not None:
  819. kv_data += self._pack(pack_fmt, val, skip_pack_prefix = vtype == GGUFValueType.BOOL)
  820. elif vtype == GGUFValueType.STRING:
  821. encoded_val = val.encode("utf-8") if isinstance(val, str) else val
  822. kv_data += self._pack("Q", len(encoded_val))
  823. kv_data += encoded_val
  824. elif vtype == GGUFValueType.ARRAY:
  825. if not isinstance(val, Sequence):
  826. raise ValueError("Invalid GGUF metadata array, expecting sequence")
  827. if len(val) == 0:
  828. raise ValueError("Invalid GGUF metadata array. Empty array")
  829. if sub_type is not None:
  830. ltype = sub_type
  831. elif isinstance(val, bytes):
  832. ltype = GGUFValueType.UINT8
  833. else:
  834. ltype = GGUFValueType.get_type(val[0])
  835. if not all(GGUFValueType.get_type(i) is ltype for i in val[1:]):
  836. raise ValueError("All items in a GGUF array should be of the same type")
  837. kv_data += self._pack("I", ltype)
  838. kv_data += self._pack("Q", len(val))
  839. for item in val:
  840. kv_data += self._pack_val(item, ltype, add_vtype=False)
  841. else:
  842. raise ValueError("Invalid GGUF metadata value type or value")
  843. return kv_data
  844. @staticmethod
  845. def format_n_bytes_to_str(num: int) -> str:
  846. if num == 0:
  847. return "negligible - metadata only"
  848. fnum = float(num)
  849. for unit in ("", "K", "M", "G"):
  850. if abs(fnum) < 1000.0:
  851. return f"{fnum:3.1f}{unit}"
  852. fnum /= 1000.0
  853. return f"{fnum:.1f}T - over 1TB, split recommended"