kopitiam_loader/metadata.rs
1//! Format-agnostic model metadata.
2//!
3//! GGUF and SafeTensors disagree about almost everything except that both
4//! files carry *some* key-value bag of model-level facts alongside the
5//! tensors: GGUF has a typed metadata KV store baked into the format;
6//! SafeTensors has an optional `__metadata__` entry whose values happen to
7//! all be strings. [`GgufValue`] is deliberately expressive enough to
8//! represent both without loss, so [`ModelMetadata::raw`] can hold either
9//! format's key-value bag through the same type rather than forcing callers
10//! to match on a format tag before they can even read a key.
11//!
12//! The name `GgufValue` reflects where the taxonomy comes from (GGUF is the
13//! only one of the two formats with typed, non-string values), not that the
14//! type is GGUF-exclusive.
15
16use indexmap::IndexMap;
17
18/// A single metadata value, typed as GGUF's key-value store types it.
19///
20/// Modeling this as a strongly-typed enum — rather than handing callers the
21/// raw bytes and a type tag to `match` on themselves — means "get me
22/// `llama.block_count` as a `u64`" is one call ([`GgufMetadata::get_u64`])
23/// instead of a `match` arm duplicated at every call site.
24#[derive(Debug, Clone, PartialEq)]
25pub enum GgufValue {
26 U8(u8),
27 I8(i8),
28 U16(u16),
29 I16(i16),
30 U32(u32),
31 I32(i32),
32 U64(u64),
33 I64(i64),
34 F32(f32),
35 F64(f64),
36 Bool(bool),
37 String(String),
38 /// A (possibly nested) array of values. GGUF arrays are homogeneous in
39 /// principle, but the format does not forbid nesting `Array` inside
40 /// `Array`, so this has to allow it too.
41 Array(Vec<GgufValue>),
42}
43
44impl GgufValue {
45 /// Widens to `u64` for every unsigned integer variant, `None` otherwise.
46 ///
47 /// The GGUF spec changed most count/length fields from `uint32` to
48 /// `uint64` between v1 and v2 and explicitly recommends readers accept
49 /// both ("Some models may use `uint32` for their values; it is
50 /// recommended that readers support both"). Widening here, once,
51 /// is how this crate honors that recommendation without every call
52 /// site re-deriving it.
53 pub fn as_u64(&self) -> Option<u64> {
54 match *self {
55 Self::U8(v) => Some(v as u64),
56 Self::U16(v) => Some(v as u64),
57 Self::U32(v) => Some(v as u64),
58 Self::U64(v) => Some(v),
59 _ => None,
60 }
61 }
62
63 /// Narrows to `u32`, refusing to silently truncate a `u64` that does not
64 /// fit — a truncated `block_count` is a corrupt-looking model, not a
65 /// smaller one, so this returns `None` rather than lying.
66 pub fn as_u32(&self) -> Option<u32> {
67 match *self {
68 Self::U8(v) => Some(v as u32),
69 Self::U16(v) => Some(v as u32),
70 Self::U32(v) => Some(v),
71 Self::U64(v) => u32::try_from(v).ok(),
72 _ => None,
73 }
74 }
75
76 /// Widens to `i64` for every signed integer variant, `None` otherwise.
77 pub fn as_i64(&self) -> Option<i64> {
78 match *self {
79 Self::I8(v) => Some(v as i64),
80 Self::I16(v) => Some(v as i64),
81 Self::I32(v) => Some(v as i64),
82 Self::I64(v) => Some(v),
83 _ => None,
84 }
85 }
86
87 /// Narrows to `i32`, refusing to silently truncate.
88 pub fn as_i32(&self) -> Option<i32> {
89 match *self {
90 Self::I8(v) => Some(v as i32),
91 Self::I16(v) => Some(v as i32),
92 Self::I32(v) => Some(v),
93 Self::I64(v) => i32::try_from(v).ok(),
94 _ => None,
95 }
96 }
97
98 /// The exact `f32`, if this value was stored as one. Deliberately does
99 /// *not* narrow `F64` to `F32`: unlike the integer widenings above,
100 /// float narrowing is lossy in the common case, not just at the edges,
101 /// so silently doing it would be more surprising than useful.
102 pub fn as_f32(&self) -> Option<f32> {
103 match *self {
104 Self::F32(v) => Some(v),
105 _ => None,
106 }
107 }
108
109 /// Widens to `f64` for either float variant.
110 pub fn as_f64(&self) -> Option<f64> {
111 match *self {
112 Self::F32(v) => Some(v as f64),
113 Self::F64(v) => Some(v),
114 _ => None,
115 }
116 }
117
118 pub fn as_bool(&self) -> Option<bool> {
119 match *self {
120 Self::Bool(v) => Some(v),
121 _ => None,
122 }
123 }
124
125 pub fn as_str(&self) -> Option<&str> {
126 match self {
127 Self::String(v) => Some(v),
128 _ => None,
129 }
130 }
131
132 pub fn as_array(&self) -> Option<&[GgufValue]> {
133 match self {
134 Self::Array(v) => Some(v),
135 _ => None,
136 }
137 }
138}
139
140/// An ordered key-value metadata bag, with typed getters over [`GgufValue`].
141///
142/// Preserves the order keys were encountered in the source file
143/// ([`indexmap::IndexMap`]) purely so that anything that dumps metadata for
144/// a human (a debug CLI command, say) reproduces the file's own ordering
145/// instead of an arbitrary hash order — lookups themselves are still O(1).
146#[derive(Debug, Clone, Default, PartialEq)]
147pub struct GgufMetadata(pub(crate) IndexMap<String, GgufValue>);
148
149impl GgufMetadata {
150 pub(crate) fn new() -> Self {
151 Self(IndexMap::new())
152 }
153
154 pub fn get(&self, key: &str) -> Option<&GgufValue> {
155 self.0.get(key)
156 }
157
158 pub fn get_u64(&self, key: &str) -> Option<u64> {
159 self.get(key)?.as_u64()
160 }
161
162 pub fn get_u32(&self, key: &str) -> Option<u32> {
163 self.get(key)?.as_u32()
164 }
165
166 pub fn get_i64(&self, key: &str) -> Option<i64> {
167 self.get(key)?.as_i64()
168 }
169
170 pub fn get_i32(&self, key: &str) -> Option<i32> {
171 self.get(key)?.as_i32()
172 }
173
174 pub fn get_f32(&self, key: &str) -> Option<f32> {
175 self.get(key)?.as_f32()
176 }
177
178 pub fn get_f64(&self, key: &str) -> Option<f64> {
179 self.get(key)?.as_f64()
180 }
181
182 pub fn get_bool(&self, key: &str) -> Option<bool> {
183 self.get(key)?.as_bool()
184 }
185
186 pub fn get_str(&self, key: &str) -> Option<&str> {
187 self.get(key)?.as_str()
188 }
189
190 pub fn get_array(&self, key: &str) -> Option<&[GgufValue]> {
191 self.get(key)?.as_array()
192 }
193
194 /// Number of keys in the bag.
195 pub fn len(&self) -> usize {
196 self.0.len()
197 }
198
199 pub fn is_empty(&self) -> bool {
200 self.0.is_empty()
201 }
202
203 /// Iterates keys and values in file order.
204 pub fn iter(&self) -> impl Iterator<Item = (&str, &GgufValue)> {
205 self.0.iter().map(|(k, v)| (k.as_str(), v))
206 }
207}
208
209/// Model-level hyperparameters, promoted out of the raw metadata bag into
210/// named fields where the two supported formats (or at least GGUF, which is
211/// the only one of the two with a standardized hyperparameter vocabulary)
212/// agree on a concept.
213///
214/// # Why `Option` everywhere
215///
216/// Every field is optional because this struct has to represent both a
217/// fully-specified GGUF LLM export *and* a bare SafeTensors weight dump that
218/// carries no architecture metadata at all — SafeTensors' `__metadata__` is
219/// a free-form `string -> string` map with no standardized keys. Refusing
220/// to load a SafeTensors file just because it lacks `n_heads` would be
221/// wrong: the tensors are still perfectly loadable, and it is the
222/// architecture-specific consumer (in `kopitiam-runtime`, not here) that
223/// knows whether a missing field is fatal for the model it is trying to
224/// build.
225///
226/// # Why these fields specifically
227///
228/// This is the intersection of GGUF's `[llm].*` standardized keys (see
229/// `crates/kopitiam-ai/vendor/ggml/docs/gguf.md`) that essentially every
230/// dense transformer architecture needs to reconstruct its shape: layer
231/// count, head counts (including the separate KV head count for grouped-
232/// query attention), embedding/feed-forward widths, context length,
233/// vocabulary size, RoPE parameters, and normalization epsilon. Anything
234/// architecture-specific (MoE expert counts, SSM state sizes, ALiBi bias)
235/// is deliberately *not* promoted here — it stays reachable through `raw`
236/// so this struct does not have to grow a field for every architecture
237/// GGUF has ever described.
238#[derive(Debug, Clone, Default, PartialEq)]
239pub struct ModelMetadata {
240 /// `general.architecture` — e.g. `"llama"`, `"qwen2"`, `"gptneox"`.
241 /// `None` for formats or files that do not record one.
242 pub architecture: Option<String>,
243 /// `general.name`, if present.
244 pub name: Option<String>,
245 /// `[arch].block_count` — number of transformer blocks.
246 pub n_layers: Option<u64>,
247 /// `[arch].attention.head_count`.
248 pub n_heads: Option<u64>,
249 /// `[arch].attention.head_count_kv`. Distinct from `n_heads` only under
250 /// grouped-query or multi-query attention; `None` means "not recorded",
251 /// which the GGUF spec says should be read as "equal to `n_heads`" —
252 /// that fallback is a modeling decision for the consumer, not this
253 /// loader, so it is left as `None` rather than silently copied here.
254 pub n_kv_heads: Option<u64>,
255 /// `[arch].embedding_length` — the model's hidden/embedding width.
256 pub embedding_length: Option<u64>,
257 /// `[arch].feed_forward_length`.
258 pub feed_forward_length: Option<u64>,
259 /// `[arch].context_length` — the context window the model was trained
260 /// for.
261 pub context_length: Option<u64>,
262 /// Vocabulary size. GGUF has no dedicated key for this; it is derived
263 /// from the length of `tokenizer.ggml.tokens` when present.
264 pub vocab_size: Option<u64>,
265 /// `[arch].rope.freq_base` — the RoPE base frequency (`theta`).
266 pub rope_theta: Option<f32>,
267 /// `[arch].rope.dimension_count`.
268 pub rope_dimension_count: Option<u64>,
269 /// Normalization epsilon, from whichever of
270 /// `[arch].attention.layer_norm_rms_epsilon` or
271 /// `[arch].attention.layer_norm_epsilon` is present (RMS-norm is tried
272 /// first, since it is what every current GGUF LLM architecture uses).
273 pub norm_epsilon: Option<f32>,
274 /// `general.quantization_version`, present when any tensor is
275 /// quantized.
276 pub quantization_version: Option<u32>,
277 /// `general.file_type` — the GGUF enum describing the majority tensor
278 /// encoding. Left as the raw `u32` rather than decoded, since it is
279 /// advisory ("can be inferred from the tensor types") and this crate
280 /// already reports each tensor's real [`kopitiam_core::DType`].
281 pub file_type: Option<u32>,
282 /// Every metadata key as found in the file, verbatim. This is the
283 /// escape hatch for everything not promoted to a named field above:
284 /// tokenizer vocabularies, chat templates, MoE/SSM parameters,
285 /// community-namespaced keys, and SafeTensors' free-form
286 /// `__metadata__` strings.
287 pub raw: GgufMetadata,
288}