use std::collections::BTreeMap;
use std::io::BufWriter;
use std::path::{Path, PathBuf};
use frink_gguf::{
GgmlType, GgufFile, GgufValue, GgufWriter, ShardedGguf, TensorPlan, TensorSource,
};
use frink_safetensors::SafetensorsFile;
use crate::loader::{load_f32_vec_optional, load_weight_matrix, LoadError};
use crate::rank_head::load_rank_head;
use crate::safetensors_f32::widen_to_f32;
pub const POOLER_SOURCE_KEY: &str = "frink.rerank.pooler_source";
const CLS_W: &str = "cls.weight";
const CLS_B: &str = "cls.bias";
const CLS_OUT_W: &str = "cls.output.weight";
const CLS_OUT_B: &str = "cls.output.bias";
const HF_POOLER_W: [&str; 2] = ["bert.pooler.dense.weight", "pooler.dense.weight"];
const HF_POOLER_B: [&str; 2] = ["bert.pooler.dense.bias", "pooler.dense.bias"];
const HF_CLASSIFIER_W: &str = "classifier.weight";
const HF_CLASSIFIER_B: &str = "classifier.bias";
pub const IDENTITY_TOLERANCE: f32 = 1.0 / 128.0;
pub const SPLICEABLE_HEAD_DTYPES: [GgmlType; 4] =
[GgmlType::F32, GgmlType::F16, GgmlType::BF16, GgmlType::Q8_0];
#[derive(Debug, thiserror::Error)]
pub enum SpliceError {
#[error(transparent)]
Gguf(#[from] frink_gguf::GgufError),
#[error(transparent)]
Load(#[from] LoadError),
#[error(transparent)]
Safetensors(#[from] frink_safetensors::SafetensorsError),
#[error("writing {path}: {source}")]
Write {
path: PathBuf,
#[source]
source: frink_gguf::GgufWriteError,
},
#[error("reopening the written file {path}: {source}")]
Reopen {
path: PathBuf,
#[source]
source: frink_gguf::ShardError,
},
#[error(
"{path} is a '{arch}' checkpoint; only a `bert` classification head is known to run \
classifier(tanh(pooler(cls))), so only a `bert` GGUF can take a pooler"
)]
NotBert { path: PathBuf, arch: String },
#[error(
"{path} is a split checkpoint ({shards} shards); merge it first (`frink gguf-split \
--merge`) so the pooler goes into one file"
)]
Split { path: PathBuf, shards: u64 },
#[error("{path} is missing `{key}`, which sizes the pooler")]
MissingHparam { path: PathBuf, key: String },
#[error(
"{path} already carries {CLS_W}{spliced_from}; splicing a second pooler over it would \
replace the head the file was converted with"
)]
AlreadyPooled { path: PathBuf, spliced_from: String },
#[error(
"{path} carries no {CLS_OUT_W}: there is no classifier for a pooler to feed, and \
no classifier to tie the pooler to. A plain embedding model has no rerank head"
)]
NoClassifier { path: PathBuf },
#[error(
"{path} stores {CLS_OUT_W} as {dtype:?}; the classifier identity check is derived \
for {allowed:?} and a coarser storage could pass it by accident"
)]
HeadDtype {
path: PathBuf,
dtype: GgmlType,
allowed: [GgmlType; 4],
},
#[error("{path} carries none of {tried:?}; it is not a BertForSequenceClassification export")]
MissingSafetensor {
path: PathBuf,
tried: Vec<&'static str>,
},
#[error("{path}: `{name}` is {dtype:?}, which is not a float type this splice reads")]
SafetensorDtype {
path: PathBuf,
name: String,
dtype: frink_safetensors::SafetensorsDtype,
},
#[error("{path}: `{name}` is {shape:?}, but the GGUF's encoder is {n_embd} wide so it must be {want:?}")]
Shape {
path: PathBuf,
name: String,
shape: Vec<usize>,
n_embd: usize,
want: Vec<usize>,
},
#[error(
"the pooler in {safetensors} does not belong to {gguf}: {mismatch}. A pooler from \
another checkpoint produces scores that look calibrated and are not, so nothing was \
written. Check that the safetensors is the exact HuggingFace repo this GGUF was \
converted from -- the GGUF's own `general.name` is not evidence, the published \
ms-marco-MiniLM-L6-v2 file names the L12 model"
)]
Mismatch {
gguf: PathBuf,
safetensors: PathBuf,
mismatch: IdentityMismatch,
},
#[error(
"the written file {path} loads without a pooler, which means the splice wrote the \
tensors under names the loader does not read; the file was removed"
)]
NotPooledAfterWrite { path: PathBuf },
}
#[derive(Debug, Clone, PartialEq)]
pub struct IdentityMismatch {
pub tensor: &'static str,
pub index: usize,
pub gguf: f32,
pub reference: f32,
pub allowed: f32,
}
impl std::fmt::Display for IdentityMismatch {
fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
write!(
f,
"{} differs at element {} (GGUF {}, safetensors {}, allowed |diff| <= {:.3e})",
self.tensor, self.index, self.gguf, self.reference, self.allowed
)
}
}
#[derive(Debug, Clone, PartialEq)]
pub struct SplicedPooler {
pub output: PathBuf,
pub n_embd: usize,
pub n_out: usize,
pub head_dtype: GgmlType,
pub classifier_max_abs_diff: f32,
pub classifier_allowed: f32,
}
pub fn classifier_matches(
tensor: &'static str,
gguf: &[f32],
reference: &[f32],
) -> Result<(f32, f32), IdentityMismatch> {
let absmax = reference.iter().fold(0.0f32, |m, v| m.max(v.abs()));
let allowed = absmax * IDENTITY_TOLERANCE;
if gguf.len() != reference.len() {
return Err(IdentityMismatch {
tensor,
index: gguf.len().min(reference.len()),
gguf: f32::NAN,
reference: f32::NAN,
allowed,
});
}
let mut worst = (0usize, 0.0f32);
for (i, (g, r)) in gguf.iter().zip(reference).enumerate() {
let diff = (g - r).abs();
if diff > worst.1 || diff.is_nan() {
worst = (i, diff);
}
}
if worst.1 > allowed || worst.1.is_nan() {
return Err(IdentityMismatch {
tensor,
index: worst.0,
gguf: gguf[worst.0],
reference: reference[worst.0],
allowed,
});
}
Ok((worst.1, allowed))
}
fn read_hf(
file: &SafetensorsFile,
path: &Path,
names: &[&'static str],
) -> Result<(Vec<usize>, Vec<f32>), SpliceError> {
let Some(name) = names.iter().find(|n| file.tensor_info(n).is_some()) else {
return Err(SpliceError::MissingSafetensor {
path: path.to_path_buf(),
tried: names.to_vec(),
});
};
let info = file.tensor_info(name).expect("found above");
let data = widen_to_f32(info.dtype, file.tensor_bytes(name)?).ok_or_else(|| {
SpliceError::SafetensorDtype {
path: path.to_path_buf(),
name: name.to_string(),
dtype: info.dtype,
}
})?;
Ok((info.shape.clone(), data))
}
fn want_shape(
path: &Path,
name: &str,
shape: &[usize],
want: &[usize],
n_embd: usize,
) -> Result<(), SpliceError> {
if shape == want {
return Ok(());
}
Err(SpliceError::Shape {
path: path.to_path_buf(),
name: name.to_string(),
shape: shape.to_vec(),
n_embd,
want: want.to_vec(),
})
}
fn f32_bytes(v: &[f32]) -> Vec<u8> {
v.iter().flat_map(|x| x.to_le_bytes()).collect()
}
pub fn splice_pooler(
gguf: &Path,
safetensors: &Path,
out: &Path,
) -> Result<SplicedPooler, SpliceError> {
let file = GgufFile::open(gguf)?;
let arch = file
.metadata_str("general.architecture")
.unwrap_or("")
.to_string();
if arch != crate::bert_gguf_loader::BERT_ARCH {
return Err(SpliceError::NotBert {
path: gguf.to_path_buf(),
arch,
});
}
if let Some(shards @ 2..) = file.metadata_u64("split.count") {
return Err(SpliceError::Split {
path: gguf.to_path_buf(),
shards,
});
}
let n_embd_key = format!("{arch}.embedding_length");
let n_embd = file
.metadata_u64(&n_embd_key)
.ok_or_else(|| SpliceError::MissingHparam {
path: gguf.to_path_buf(),
key: n_embd_key,
})? as usize;
if file.find_tensor(CLS_W).is_some() {
let source = file
.metadata_str(POOLER_SOURCE_KEY)
.map(|s| format!(" (spliced from {s})"))
.unwrap_or_default();
return Err(SpliceError::AlreadyPooled {
path: gguf.to_path_buf(),
spliced_from: source,
});
}
let Some(head_info) = file.find_tensor(CLS_OUT_W) else {
return Err(SpliceError::NoClassifier {
path: gguf.to_path_buf(),
});
};
let head_dtype = head_info.dtype;
if !SPLICEABLE_HEAD_DTYPES.contains(&head_dtype) {
return Err(SpliceError::HeadDtype {
path: gguf.to_path_buf(),
dtype: head_dtype,
allowed: SPLICEABLE_HEAD_DTYPES,
});
}
let head = load_weight_matrix(&file, CLS_OUT_W)?;
let n_out = head.rows();
let gguf_w: Vec<f32> = (0..n_out).flat_map(|r| head.dequant_row(r)).collect();
let gguf_b = load_f32_vec_optional(&file, CLS_OUT_B)?;
let hf = SafetensorsFile::open(safetensors)?;
let (cw_shape, hf_w) = read_hf(&hf, safetensors, &[HF_CLASSIFIER_W])?;
want_shape(
safetensors,
HF_CLASSIFIER_W,
&cw_shape,
&[n_out, n_embd],
n_embd,
)?;
let (pw_shape, pooler_w) = read_hf(&hf, safetensors, &HF_POOLER_W)?;
want_shape(
safetensors,
HF_POOLER_W[0],
&pw_shape,
&[n_embd, n_embd],
n_embd,
)?;
let (pb_shape, pooler_b) = read_hf(&hf, safetensors, &HF_POOLER_B)?;
want_shape(safetensors, HF_POOLER_B[0], &pb_shape, &[n_embd], n_embd)?;
let mismatch = |mismatch| SpliceError::Mismatch {
gguf: gguf.to_path_buf(),
safetensors: safetensors.to_path_buf(),
mismatch,
};
let (mut worst, allowed) = classifier_matches(CLS_OUT_W, &gguf_w, &hf_w).map_err(mismatch)?;
match (gguf_b, hf.tensor_info(HF_CLASSIFIER_B).is_some()) {
(Some(gguf_b), true) => {
let (_, hf_b) = read_hf(&hf, safetensors, &[HF_CLASSIFIER_B])?;
let (worst_b, _) = classifier_matches(CLS_OUT_B, &gguf_b, &hf_b).map_err(mismatch)?;
worst = worst.max(worst_b);
}
(None, false) => {}
(gguf_b, _) => {
return Err(mismatch(IdentityMismatch {
tensor: CLS_OUT_B,
index: 0,
gguf: gguf_b.map(|b| b[0]).unwrap_or(f32::NAN),
reference: f32::NAN,
allowed,
}));
}
}
let mut metadata: BTreeMap<String, GgufValue> = file
.metadata
.iter()
.map(|(k, v)| (k.clone(), v.clone()))
.collect();
metadata.insert(
POOLER_SOURCE_KEY.to_string(),
GgufValue::String(
safetensors
.file_name()
.map(|n| n.to_string_lossy().into_owned())
.unwrap_or_else(|| safetensors.display().to_string()),
),
);
let mut plan: Vec<TensorPlan> = Vec::with_capacity(file.tensors.len() + 2);
for t in &file.tensors {
plan.push(TensorPlan {
name: t.name.clone(),
shape: t.shape.clone(),
dtype: t.dtype,
byte_len: file.tensor_bytes(&t.name)?.len(),
});
}
let pooler_w_bytes = f32_bytes(&pooler_w);
let pooler_b_bytes = f32_bytes(&pooler_b);
plan.push(TensorPlan {
name: CLS_W.to_string(),
shape: vec![n_embd as u64, n_embd as u64],
dtype: GgmlType::F32,
byte_len: pooler_w_bytes.len(),
});
plan.push(TensorPlan {
name: CLS_B.to_string(),
shape: vec![n_embd as u64],
dtype: GgmlType::F32,
byte_len: pooler_b_bytes.len(),
});
let write_err = |source| SpliceError::Write {
path: out.to_path_buf(),
source,
};
let sink = std::fs::File::create(out).map_err(|e| write_err(e.into()))?;
let mut w = GgufWriter::create(BufWriter::new(sink), &metadata, plan).map_err(write_err)?;
for t in &file.tensors {
w.write_tensor(&t.name, file.tensor_bytes(&t.name)?)
.map_err(write_err)?;
}
w.write_tensor(CLS_W, &pooler_w_bytes).map_err(write_err)?;
w.write_tensor(CLS_B, &pooler_b_bytes).map_err(write_err)?;
w.finish().map_err(write_err)?;
let eps = file
.metadata_f32(&format!("{arch}.attention.layer_norm_epsilon"))
.unwrap_or(1e-12);
let reopened = ShardedGguf::open(out).map_err(|source| SpliceError::Reopen {
path: out.to_path_buf(),
source,
})?;
let pooled = load_rank_head(&reopened, &arch, n_embd, eps)
.map(|h| h.is_some_and(|h| h.has_pooler()))
.unwrap_or(false);
if !pooled {
std::fs::remove_file(out).ok();
return Err(SpliceError::NotPooledAfterWrite {
path: out.to_path_buf(),
});
}
Ok(SplicedPooler {
output: out.to_path_buf(),
n_embd,
n_out,
head_dtype,
classifier_max_abs_diff: worst,
classifier_allowed: allowed,
})
}
#[cfg(test)]
mod tests {
use super::*;
fn reference(n: usize) -> Vec<f32> {
(0..n)
.map(|i| {
let x = i as f32;
((x * 0.37).sin() * 0.8 + (x * 0.011).cos() * 0.05)
* if i % 7 == 0 { 3.0 } else { 1.0 }
})
.collect()
}
#[test]
fn every_spliceable_storage_precision_passes_the_identity_bound() {
let r = reference(384);
let q8 = frink_quant::dequant_q8_0(&frink_quant::quantize_q8_0(&r)).unwrap();
let (worst, allowed) = classifier_matches("q8_0", &q8, &r).expect("Q8_0 round trip");
assert!(
worst > 0.0,
"the Q8_0 round trip must actually perturb something"
);
assert!(worst <= allowed);
let bf16: Vec<f32> = r
.iter()
.map(|x| {
let bits = x.to_bits();
let rounded = (bits.wrapping_add(0x7FFF + ((bits >> 16) & 1))) >> 16;
f32::from_bits(rounded << 16)
})
.collect();
let (worst, allowed) = classifier_matches("bf16", &bf16, &r).expect("BF16 round trip");
assert!(worst > 0.0);
assert!(worst <= allowed);
let f16: Vec<f32> = r.iter().map(|x| half::f16::from_f32(*x).to_f32()).collect();
classifier_matches("f16", &f16, &r).expect("F16 round trip");
classifier_matches("f32", &r, &r).expect("F32 is exact");
}
#[test]
fn a_classifier_off_by_more_than_the_files_own_rounding_is_refused_by_element() {
let r = reference(384);
let absmax = r.iter().fold(0.0f32, |m, v| m.max(v.abs()));
let mut other = r.clone();
other[200] += 2.0 * absmax * IDENTITY_TOLERANCE;
let err = classifier_matches(CLS_OUT_W, &other, &r).unwrap_err();
assert_eq!(err.tensor, CLS_OUT_W);
assert_eq!(err.index, 200);
assert_eq!(err.gguf, other[200]);
assert_eq!(err.reference, r[200]);
assert!(err.to_string().contains("element 200"), "{err}");
}
#[test]
fn a_classifier_of_another_width_is_refused_before_any_element_is_compared() {
let r = reference(384);
assert!(classifier_matches(CLS_OUT_W, &r[..383], &r).is_err());
assert!(classifier_matches(CLS_OUT_W, &r, &r[..383]).is_err());
}
#[test]
fn a_nan_never_matches() {
let r = reference(8);
let mut g = r.clone();
g[3] = f32::NAN;
assert!(classifier_matches(CLS_OUT_W, &g, &r).is_err());
}
}