1use std::collections::BTreeMap;
73use std::io::BufWriter;
74use std::path::{Path, PathBuf};
75
76use frink_gguf::{
77 GgmlType, GgufFile, GgufValue, GgufWriter, ShardedGguf, TensorPlan, TensorSource,
78};
79use frink_safetensors::SafetensorsFile;
80
81use crate::loader::{load_f32_vec_optional, load_weight_matrix, LoadError};
82use crate::rank_head::load_rank_head;
83use crate::safetensors_f32::widen_to_f32;
84
85pub const POOLER_SOURCE_KEY: &str = "frink.rerank.pooler_source";
88
89const CLS_W: &str = "cls.weight";
92const CLS_B: &str = "cls.bias";
93const CLS_OUT_W: &str = "cls.output.weight";
94const CLS_OUT_B: &str = "cls.output.bias";
95
96const HF_POOLER_W: [&str; 2] = ["bert.pooler.dense.weight", "pooler.dense.weight"];
100const HF_POOLER_B: [&str; 2] = ["bert.pooler.dense.bias", "pooler.dense.bias"];
101const HF_CLASSIFIER_W: &str = "classifier.weight";
102const HF_CLASSIFIER_B: &str = "classifier.bias";
103
104pub const IDENTITY_TOLERANCE: f32 = 1.0 / 128.0;
122
123pub const SPLICEABLE_HEAD_DTYPES: [GgmlType; 4] =
126 [GgmlType::F32, GgmlType::F16, GgmlType::BF16, GgmlType::Q8_0];
127
128#[derive(Debug, thiserror::Error)]
129pub enum SpliceError {
130 #[error(transparent)]
131 Gguf(#[from] frink_gguf::GgufError),
132 #[error(transparent)]
133 Load(#[from] LoadError),
134 #[error(transparent)]
135 Safetensors(#[from] frink_safetensors::SafetensorsError),
136 #[error("writing {path}: {source}")]
137 Write {
138 path: PathBuf,
139 #[source]
140 source: frink_gguf::GgufWriteError,
141 },
142 #[error("reopening the written file {path}: {source}")]
143 Reopen {
144 path: PathBuf,
145 #[source]
146 source: frink_gguf::ShardError,
147 },
148 #[error(
149 "{path} is a '{arch}' checkpoint; only a `bert` classification head is known to run \
150 classifier(tanh(pooler(cls))), so only a `bert` GGUF can take a pooler"
151 )]
152 NotBert { path: PathBuf, arch: String },
153 #[error(
154 "{path} is a split checkpoint ({shards} shards); merge it first (`frink gguf-split \
155 --merge`) so the pooler goes into one file"
156 )]
157 Split { path: PathBuf, shards: u64 },
158 #[error("{path} is missing `{key}`, which sizes the pooler")]
159 MissingHparam { path: PathBuf, key: String },
160 #[error(
161 "{path} already carries {CLS_W}{spliced_from}; splicing a second pooler over it would \
162 replace the head the file was converted with"
163 )]
164 AlreadyPooled { path: PathBuf, spliced_from: String },
165 #[error(
166 "{path} carries no {CLS_OUT_W}: there is no classifier for a pooler to feed, and \
167 no classifier to tie the pooler to. A plain embedding model has no rerank head"
168 )]
169 NoClassifier { path: PathBuf },
170 #[error(
171 "{path} stores {CLS_OUT_W} as {dtype:?}; the classifier identity check is derived \
172 for {allowed:?} and a coarser storage could pass it by accident"
173 )]
174 HeadDtype {
175 path: PathBuf,
176 dtype: GgmlType,
177 allowed: [GgmlType; 4],
178 },
179 #[error("{path} carries none of {tried:?}; it is not a BertForSequenceClassification export")]
180 MissingSafetensor {
181 path: PathBuf,
182 tried: Vec<&'static str>,
183 },
184 #[error("{path}: `{name}` is {dtype:?}, which is not a float type this splice reads")]
185 SafetensorDtype {
186 path: PathBuf,
187 name: String,
188 dtype: frink_safetensors::SafetensorsDtype,
189 },
190 #[error("{path}: `{name}` is {shape:?}, but the GGUF's encoder is {n_embd} wide so it must be {want:?}")]
191 Shape {
192 path: PathBuf,
193 name: String,
194 shape: Vec<usize>,
195 n_embd: usize,
196 want: Vec<usize>,
197 },
198 #[error(
199 "the pooler in {safetensors} does not belong to {gguf}: {mismatch}. A pooler from \
200 another checkpoint produces scores that look calibrated and are not, so nothing was \
201 written. Check that the safetensors is the exact HuggingFace repo this GGUF was \
202 converted from -- the GGUF's own `general.name` is not evidence, the published \
203 ms-marco-MiniLM-L6-v2 file names the L12 model"
204 )]
205 Mismatch {
206 gguf: PathBuf,
207 safetensors: PathBuf,
208 mismatch: IdentityMismatch,
209 },
210 #[error(
211 "the written file {path} loads without a pooler, which means the splice wrote the \
212 tensors under names the loader does not read; the file was removed"
213 )]
214 NotPooledAfterWrite { path: PathBuf },
215}
216
217#[derive(Debug, Clone, PartialEq)]
221pub struct IdentityMismatch {
222 pub tensor: &'static str,
224 pub index: usize,
226 pub gguf: f32,
227 pub reference: f32,
228 pub allowed: f32,
230}
231
232impl std::fmt::Display for IdentityMismatch {
233 fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
234 write!(
235 f,
236 "{} differs at element {} (GGUF {}, safetensors {}, allowed |diff| <= {:.3e})",
237 self.tensor, self.index, self.gguf, self.reference, self.allowed
238 )
239 }
240}
241
242#[derive(Debug, Clone, PartialEq)]
244pub struct SplicedPooler {
245 pub output: PathBuf,
246 pub n_embd: usize,
247 pub n_out: usize,
248 pub head_dtype: GgmlType,
251 pub classifier_max_abs_diff: f32,
255 pub classifier_allowed: f32,
257}
258
259pub fn classifier_matches(
268 tensor: &'static str,
269 gguf: &[f32],
270 reference: &[f32],
271) -> Result<(f32, f32), IdentityMismatch> {
272 let absmax = reference.iter().fold(0.0f32, |m, v| m.max(v.abs()));
273 let allowed = absmax * IDENTITY_TOLERANCE;
274 if gguf.len() != reference.len() {
275 return Err(IdentityMismatch {
276 tensor,
277 index: gguf.len().min(reference.len()),
278 gguf: f32::NAN,
279 reference: f32::NAN,
280 allowed,
281 });
282 }
283 let mut worst = (0usize, 0.0f32);
284 for (i, (g, r)) in gguf.iter().zip(reference).enumerate() {
285 let diff = (g - r).abs();
286 if diff > worst.1 || diff.is_nan() {
287 worst = (i, diff);
288 }
289 }
290 if worst.1 > allowed || worst.1.is_nan() {
291 return Err(IdentityMismatch {
292 tensor,
293 index: worst.0,
294 gguf: gguf[worst.0],
295 reference: reference[worst.0],
296 allowed,
297 });
298 }
299 Ok((worst.1, allowed))
300}
301
302fn read_hf(
305 file: &SafetensorsFile,
306 path: &Path,
307 names: &[&'static str],
308) -> Result<(Vec<usize>, Vec<f32>), SpliceError> {
309 let Some(name) = names.iter().find(|n| file.tensor_info(n).is_some()) else {
310 return Err(SpliceError::MissingSafetensor {
311 path: path.to_path_buf(),
312 tried: names.to_vec(),
313 });
314 };
315 let info = file.tensor_info(name).expect("found above");
316 let data = widen_to_f32(info.dtype, file.tensor_bytes(name)?).ok_or_else(|| {
317 SpliceError::SafetensorDtype {
318 path: path.to_path_buf(),
319 name: name.to_string(),
320 dtype: info.dtype,
321 }
322 })?;
323 Ok((info.shape.clone(), data))
324}
325
326fn want_shape(
327 path: &Path,
328 name: &str,
329 shape: &[usize],
330 want: &[usize],
331 n_embd: usize,
332) -> Result<(), SpliceError> {
333 if shape == want {
334 return Ok(());
335 }
336 Err(SpliceError::Shape {
337 path: path.to_path_buf(),
338 name: name.to_string(),
339 shape: shape.to_vec(),
340 n_embd,
341 want: want.to_vec(),
342 })
343}
344
345fn f32_bytes(v: &[f32]) -> Vec<u8> {
346 v.iter().flat_map(|x| x.to_le_bytes()).collect()
347}
348
349pub fn splice_pooler(
353 gguf: &Path,
354 safetensors: &Path,
355 out: &Path,
356) -> Result<SplicedPooler, SpliceError> {
357 let file = GgufFile::open(gguf)?;
358 let arch = file
359 .metadata_str("general.architecture")
360 .unwrap_or("")
361 .to_string();
362 if arch != crate::bert_gguf_loader::BERT_ARCH {
363 return Err(SpliceError::NotBert {
364 path: gguf.to_path_buf(),
365 arch,
366 });
367 }
368 if let Some(shards @ 2..) = file.metadata_u64("split.count") {
369 return Err(SpliceError::Split {
370 path: gguf.to_path_buf(),
371 shards,
372 });
373 }
374 let n_embd_key = format!("{arch}.embedding_length");
375 let n_embd = file
376 .metadata_u64(&n_embd_key)
377 .ok_or_else(|| SpliceError::MissingHparam {
378 path: gguf.to_path_buf(),
379 key: n_embd_key,
380 })? as usize;
381 if file.find_tensor(CLS_W).is_some() {
382 let source = file
383 .metadata_str(POOLER_SOURCE_KEY)
384 .map(|s| format!(" (spliced from {s})"))
385 .unwrap_or_default();
386 return Err(SpliceError::AlreadyPooled {
387 path: gguf.to_path_buf(),
388 spliced_from: source,
389 });
390 }
391 let Some(head_info) = file.find_tensor(CLS_OUT_W) else {
392 return Err(SpliceError::NoClassifier {
393 path: gguf.to_path_buf(),
394 });
395 };
396 let head_dtype = head_info.dtype;
397 if !SPLICEABLE_HEAD_DTYPES.contains(&head_dtype) {
398 return Err(SpliceError::HeadDtype {
399 path: gguf.to_path_buf(),
400 dtype: head_dtype,
401 allowed: SPLICEABLE_HEAD_DTYPES,
402 });
403 }
404
405 let head = load_weight_matrix(&file, CLS_OUT_W)?;
409 let n_out = head.rows();
410 let gguf_w: Vec<f32> = (0..n_out).flat_map(|r| head.dequant_row(r)).collect();
411 let gguf_b = load_f32_vec_optional(&file, CLS_OUT_B)?;
412
413 let hf = SafetensorsFile::open(safetensors)?;
414 let (cw_shape, hf_w) = read_hf(&hf, safetensors, &[HF_CLASSIFIER_W])?;
415 want_shape(
416 safetensors,
417 HF_CLASSIFIER_W,
418 &cw_shape,
419 &[n_out, n_embd],
420 n_embd,
421 )?;
422 let (pw_shape, pooler_w) = read_hf(&hf, safetensors, &HF_POOLER_W)?;
423 want_shape(
424 safetensors,
425 HF_POOLER_W[0],
426 &pw_shape,
427 &[n_embd, n_embd],
428 n_embd,
429 )?;
430 let (pb_shape, pooler_b) = read_hf(&hf, safetensors, &HF_POOLER_B)?;
431 want_shape(safetensors, HF_POOLER_B[0], &pb_shape, &[n_embd], n_embd)?;
432
433 let mismatch = |mismatch| SpliceError::Mismatch {
434 gguf: gguf.to_path_buf(),
435 safetensors: safetensors.to_path_buf(),
436 mismatch,
437 };
438 let (mut worst, allowed) = classifier_matches(CLS_OUT_W, &gguf_w, &hf_w).map_err(mismatch)?;
439 match (gguf_b, hf.tensor_info(HF_CLASSIFIER_B).is_some()) {
443 (Some(gguf_b), true) => {
444 let (_, hf_b) = read_hf(&hf, safetensors, &[HF_CLASSIFIER_B])?;
445 let (worst_b, _) = classifier_matches(CLS_OUT_B, &gguf_b, &hf_b).map_err(mismatch)?;
446 worst = worst.max(worst_b);
447 }
448 (None, false) => {}
449 (gguf_b, _) => {
450 return Err(mismatch(IdentityMismatch {
451 tensor: CLS_OUT_B,
452 index: 0,
453 gguf: gguf_b.map(|b| b[0]).unwrap_or(f32::NAN),
454 reference: f32::NAN,
455 allowed,
456 }));
457 }
458 }
459
460 let mut metadata: BTreeMap<String, GgufValue> = file
462 .metadata
463 .iter()
464 .map(|(k, v)| (k.clone(), v.clone()))
465 .collect();
466 metadata.insert(
467 POOLER_SOURCE_KEY.to_string(),
468 GgufValue::String(
469 safetensors
470 .file_name()
471 .map(|n| n.to_string_lossy().into_owned())
472 .unwrap_or_else(|| safetensors.display().to_string()),
473 ),
474 );
475 let mut plan: Vec<TensorPlan> = Vec::with_capacity(file.tensors.len() + 2);
476 for t in &file.tensors {
477 plan.push(TensorPlan {
478 name: t.name.clone(),
479 shape: t.shape.clone(),
480 dtype: t.dtype,
481 byte_len: file.tensor_bytes(&t.name)?.len(),
482 });
483 }
484 let pooler_w_bytes = f32_bytes(&pooler_w);
485 let pooler_b_bytes = f32_bytes(&pooler_b);
486 plan.push(TensorPlan {
491 name: CLS_W.to_string(),
492 shape: vec![n_embd as u64, n_embd as u64],
493 dtype: GgmlType::F32,
494 byte_len: pooler_w_bytes.len(),
495 });
496 plan.push(TensorPlan {
497 name: CLS_B.to_string(),
498 shape: vec![n_embd as u64],
499 dtype: GgmlType::F32,
500 byte_len: pooler_b_bytes.len(),
501 });
502
503 let write_err = |source| SpliceError::Write {
504 path: out.to_path_buf(),
505 source,
506 };
507 let sink = std::fs::File::create(out).map_err(|e| write_err(e.into()))?;
508 let mut w = GgufWriter::create(BufWriter::new(sink), &metadata, plan).map_err(write_err)?;
509 for t in &file.tensors {
510 w.write_tensor(&t.name, file.tensor_bytes(&t.name)?)
511 .map_err(write_err)?;
512 }
513 w.write_tensor(CLS_W, &pooler_w_bytes).map_err(write_err)?;
514 w.write_tensor(CLS_B, &pooler_b_bytes).map_err(write_err)?;
515 w.finish().map_err(write_err)?;
516
517 let eps = file
520 .metadata_f32(&format!("{arch}.attention.layer_norm_epsilon"))
521 .unwrap_or(1e-12);
522 let reopened = ShardedGguf::open(out).map_err(|source| SpliceError::Reopen {
523 path: out.to_path_buf(),
524 source,
525 })?;
526 let pooled = load_rank_head(&reopened, &arch, n_embd, eps)
527 .map(|h| h.is_some_and(|h| h.has_pooler()))
528 .unwrap_or(false);
529 if !pooled {
530 std::fs::remove_file(out).ok();
531 return Err(SpliceError::NotPooledAfterWrite {
532 path: out.to_path_buf(),
533 });
534 }
535
536 Ok(SplicedPooler {
537 output: out.to_path_buf(),
538 n_embd,
539 n_out,
540 head_dtype,
541 classifier_max_abs_diff: worst,
542 classifier_allowed: allowed,
543 })
544}
545
546#[cfg(test)]
547mod tests {
548 use super::*;
549
550 fn reference(n: usize) -> Vec<f32> {
554 (0..n)
555 .map(|i| {
556 let x = i as f32;
557 ((x * 0.37).sin() * 0.8 + (x * 0.011).cos() * 0.05)
558 * if i % 7 == 0 { 3.0 } else { 1.0 }
559 })
560 .collect()
561 }
562
563 #[test]
570 fn every_spliceable_storage_precision_passes_the_identity_bound() {
571 let r = reference(384);
572
573 let q8 = frink_quant::dequant_q8_0(&frink_quant::quantize_q8_0(&r)).unwrap();
574 let (worst, allowed) = classifier_matches("q8_0", &q8, &r).expect("Q8_0 round trip");
575 assert!(
576 worst > 0.0,
577 "the Q8_0 round trip must actually perturb something"
578 );
579 assert!(worst <= allowed);
580
581 let bf16: Vec<f32> = r
582 .iter()
583 .map(|x| {
584 let bits = x.to_bits();
585 let rounded = (bits.wrapping_add(0x7FFF + ((bits >> 16) & 1))) >> 16;
586 f32::from_bits(rounded << 16)
587 })
588 .collect();
589 let (worst, allowed) = classifier_matches("bf16", &bf16, &r).expect("BF16 round trip");
590 assert!(worst > 0.0);
591 assert!(worst <= allowed);
592
593 let f16: Vec<f32> = r.iter().map(|x| half::f16::from_f32(*x).to_f32()).collect();
594 classifier_matches("f16", &f16, &r).expect("F16 round trip");
595 classifier_matches("f32", &r, &r).expect("F32 is exact");
596 }
597
598 #[test]
604 fn a_classifier_off_by_more_than_the_files_own_rounding_is_refused_by_element() {
605 let r = reference(384);
606 let absmax = r.iter().fold(0.0f32, |m, v| m.max(v.abs()));
607 let mut other = r.clone();
608 other[200] += 2.0 * absmax * IDENTITY_TOLERANCE;
609 let err = classifier_matches(CLS_OUT_W, &other, &r).unwrap_err();
610 assert_eq!(err.tensor, CLS_OUT_W);
611 assert_eq!(err.index, 200);
612 assert_eq!(err.gguf, other[200]);
613 assert_eq!(err.reference, r[200]);
614 assert!(err.to_string().contains("element 200"), "{err}");
615 }
616
617 #[test]
619 fn a_classifier_of_another_width_is_refused_before_any_element_is_compared() {
620 let r = reference(384);
621 assert!(classifier_matches(CLS_OUT_W, &r[..383], &r).is_err());
622 assert!(classifier_matches(CLS_OUT_W, &r, &r[..383]).is_err());
623 }
624
625 #[test]
627 fn a_nan_never_matches() {
628 let r = reference(8);
629 let mut g = r.clone();
630 g[3] = f32::NAN;
631 assert!(classifier_matches(CLS_OUT_W, &g, &r).is_err());
632 }
633}