221 lines
8.1 KiB
Python
221 lines
8.1 KiB
Python
"""GGUF metadata + tensor-table reader (stdlib only).
|
|
|
|
Feeds the per-layer context estimator: architecture, layer count, per-layer
|
|
KV head counts (0 = recurrent layer — the hybrid discriminator), head dims,
|
|
sliding-window config, trained context, and exact weight bytes summed from
|
|
the tensor table (validated to within 0.01% of the loader's buffer).
|
|
|
|
Reads the header only (metadata + tensor infos); never touches tensor data,
|
|
so it is fast enough to run at picker time on multi-GB files.
|
|
"""
|
|
|
|
from __future__ import annotations
|
|
|
|
import struct
|
|
from dataclasses import dataclass, field
|
|
from pathlib import Path
|
|
|
|
_GGUF_MAGIC = b"GGUF"
|
|
|
|
# ggml tensor type sizes: type_id -> (block_bytes, block_elems).
|
|
# IQ-family sizes verified against ggml-common.h.
|
|
_GGML_TYPE_SIZES = {
|
|
0: (4, 1), 1: (2, 1), 2: (18, 32), 3: (20, 32), 6: (22, 32), 7: (24, 32),
|
|
8: (34, 32), 9: (36, 32), 10: (84, 256), 11: (110, 256), 12: (144, 256),
|
|
13: (176, 256), 14: (210, 256), 15: (292, 256), 16: (66, 256),
|
|
17: (74, 256), 18: (98, 256), 19: (50, 256), 20: (18, 32),
|
|
21: (110, 256), 22: (82, 256), 23: (136, 256), 24: (1, 1), 25: (2, 1),
|
|
26: (4, 1), 27: (8, 1), 28: (8, 1), 29: (56, 256), 30: (2, 1),
|
|
}
|
|
|
|
# GGUF metadata value types.
|
|
_V_UINT8, _V_INT8, _V_UINT16, _V_INT16 = 0, 1, 2, 3
|
|
_V_UINT32, _V_INT32, _V_FLOAT32, _V_BOOL = 4, 5, 6, 7
|
|
_V_STRING, _V_ARRAY, _V_UINT64, _V_INT64, _V_FLOAT64 = 8, 9, 10, 11, 12
|
|
|
|
_SCALAR_FMT = {
|
|
_V_UINT8: "<B", _V_INT8: "<b", _V_UINT16: "<H", _V_INT16: "<h",
|
|
_V_UINT32: "<I", _V_INT32: "<i", _V_FLOAT32: "<f", _V_BOOL: "<?",
|
|
_V_UINT64: "<Q", _V_INT64: "<q", _V_FLOAT64: "<d",
|
|
}
|
|
|
|
|
|
@dataclass
|
|
class GGUFHeader:
|
|
path: str
|
|
version: int
|
|
metadata: dict = field(default_factory=dict)
|
|
n_tensors: int = 0
|
|
tensor_bytes: int = 0 # exact sum over the tensor table
|
|
embd_table_bytes: int = 0 # token_embd.weight (duplicated host-side
|
|
# when fully offloaded)
|
|
|
|
# ── typed accessors ──────────────────────────────────────
|
|
|
|
@property
|
|
def architecture(self) -> str:
|
|
return str(self.metadata.get("general.architecture", ""))
|
|
|
|
def _arch_key(self, suffix: str):
|
|
return self.metadata.get(f"{self.architecture}.{suffix}")
|
|
|
|
@property
|
|
def n_layer(self) -> int:
|
|
return int(self._arch_key("block_count") or 0)
|
|
|
|
@property
|
|
def n_vocab(self) -> int:
|
|
"""Vocabulary size: prices the GPU logits buffers (they scale
|
|
ubatch x vocab). vocab_size metadata when present, else the
|
|
tokenizer list length."""
|
|
v = self._arch_key("vocab_size")
|
|
if v:
|
|
return int(v)
|
|
toks = self.metadata.get("tokenizer.ggml.tokens")
|
|
return len(toks) if isinstance(toks, list) else 0
|
|
|
|
@property
|
|
def n_ctx_train(self) -> int:
|
|
return int(self._arch_key("context_length") or 0)
|
|
|
|
@property
|
|
def sampling_defaults(self) -> dict:
|
|
"""Upstream's recommended sampling, when the file carries it.
|
|
|
|
Model publishers bake general.sampling.* keys into the GGUF
|
|
(llama-server reads them as that model's default generation
|
|
settings), so the file itself is the source of truth for how its
|
|
publisher wants it run — it arrives with the download and updates
|
|
with every re-upload, no catalog required. Returned as preset INI
|
|
keys; empty when the file carries none.
|
|
"""
|
|
ini_key = {"temp": "temp", "temperature": "temp", "top_p": "top-p",
|
|
"top_k": "top-k", "min_p": "min-p",
|
|
"repeat_penalty": "repeat-penalty",
|
|
"presence_penalty": "presence-penalty"}
|
|
out = {}
|
|
for key, value in self.metadata.items():
|
|
if not key.startswith("general.sampling."):
|
|
continue
|
|
name = ini_key.get(key.rsplit(".", 1)[-1])
|
|
if name is not None and isinstance(value, (int, float)):
|
|
num = round(float(value), 4)
|
|
out[name] = str(int(num)) if num == int(num) else str(num)
|
|
return out
|
|
|
|
@property
|
|
def n_embd(self) -> int:
|
|
return int(self._arch_key("embedding_length") or 0)
|
|
|
|
@property
|
|
def n_head(self) -> int:
|
|
v = self._arch_key("attention.head_count")
|
|
if isinstance(v, list):
|
|
return int(max(v))
|
|
return int(v or 0)
|
|
|
|
@property
|
|
def full_attention_interval(self) -> int:
|
|
"""GDN-hybrid discriminator (qwen35 family): every Nth layer is full
|
|
attention, the rest are linear/recurrent. 0 = not present."""
|
|
return int(self._arch_key("full_attention_interval") or 0)
|
|
|
|
def head_counts_kv(self) -> list[int]:
|
|
"""Per-layer KV head counts; 0 marks a recurrent/linear layer (the
|
|
n_head_kv == 0 discriminator).
|
|
|
|
Three GGUF shapes, each verified against real files:
|
|
- per-layer array (nemotron_h_moe): use as-is;
|
|
- scalar + full_attention_interval (qwen35): the scalar applies to
|
|
every INTERVAL-th layer (1-indexed: layers where (i+1) % N == 0),
|
|
zero elsewhere — pricing all layers as attention was a 4x
|
|
overestimate on Qwen3.6-27B;
|
|
- plain scalar (dense): broadcast to every layer.
|
|
"""
|
|
v = self._arch_key("attention.head_count_kv")
|
|
if isinstance(v, list):
|
|
return [int(x) for x in v]
|
|
scalar = int(v or 0)
|
|
interval = self.full_attention_interval
|
|
if interval > 1:
|
|
return [scalar if (i + 1) % interval == 0 else 0
|
|
for i in range(self.n_layer)]
|
|
return [scalar] * self.n_layer
|
|
|
|
@property
|
|
def head_dim_k(self) -> int:
|
|
v = self._arch_key("attention.key_length")
|
|
if v:
|
|
return int(v)
|
|
return self.n_embd // self.n_head if self.n_head else 0
|
|
|
|
@property
|
|
def head_dim_v(self) -> int:
|
|
v = self._arch_key("attention.value_length")
|
|
if v:
|
|
return int(v)
|
|
return self.head_dim_k
|
|
|
|
@property
|
|
def sliding_window(self) -> int:
|
|
return int(self._arch_key("attention.sliding_window") or 0)
|
|
|
|
@property
|
|
def expert_count(self) -> int:
|
|
return int(self._arch_key("expert_count") or 0)
|
|
|
|
|
|
def read_gguf_header(path: str | Path) -> GGUFHeader:
|
|
path = Path(path)
|
|
|
|
def read_str(f) -> str:
|
|
(n,) = struct.unpack("<Q", f.read(8))
|
|
return f.read(n).decode("utf-8", errors="replace")
|
|
|
|
def read_value(f, vtype: int):
|
|
if vtype == _V_STRING:
|
|
return read_str(f)
|
|
if vtype == _V_ARRAY:
|
|
(etype,) = struct.unpack("<I", f.read(4))
|
|
(n,) = struct.unpack("<Q", f.read(8))
|
|
return [read_value(f, etype) for _ in range(n)]
|
|
fmt = _SCALAR_FMT[vtype]
|
|
(value,) = struct.unpack(fmt, f.read(struct.calcsize(fmt)))
|
|
return value
|
|
|
|
with open(path, "rb") as f:
|
|
if f.read(4) != _GGUF_MAGIC:
|
|
raise ValueError(f"not a GGUF file: {path}")
|
|
(version,) = struct.unpack("<I", f.read(4))
|
|
n_tensors, n_kv = struct.unpack("<QQ", f.read(16))
|
|
|
|
metadata: dict = {}
|
|
for _ in range(n_kv):
|
|
key = read_str(f)
|
|
(vtype,) = struct.unpack("<I", f.read(4))
|
|
metadata[key] = read_value(f, vtype)
|
|
|
|
tensor_bytes = 0
|
|
embd_bytes = 0
|
|
for _ in range(n_tensors):
|
|
name = read_str(f)
|
|
(n_dims,) = struct.unpack("<I", f.read(4))
|
|
dims = struct.unpack(f"<{n_dims}Q", f.read(8 * n_dims))
|
|
(ttype,) = struct.unpack("<I", f.read(4))
|
|
f.read(8) # offset
|
|
size = _GGML_TYPE_SIZES.get(ttype)
|
|
if size is None:
|
|
raise ValueError(f"unknown ggml tensor type {ttype} in {path}")
|
|
block_bytes, block_elems = size
|
|
elems = 1
|
|
for d in dims:
|
|
elems *= d
|
|
nbytes = (elems // block_elems) * block_bytes
|
|
tensor_bytes += nbytes
|
|
if name == "token_embd.weight":
|
|
embd_bytes = nbytes
|
|
|
|
return GGUFHeader(path=str(path), version=version, metadata=metadata,
|
|
n_tensors=n_tensors, tensor_bytes=tensor_bytes,
|
|
embd_table_bytes=embd_bytes)
|