use std::collections::BTreeMap;
use std::path::{Path, PathBuf};
use ferrox_gguf::{GgmlType, GgufValue, GgufWriter, ShardedGguf, TensorPlan};
use ferrox_models::{load_rank_head, pool, EmbeddingModel, PoolingType};
fn checkpoint() -> PathBuf {
let path =
Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models/ms-marco-MiniLM-L6-v2-Q8_0.gguf");
assert!(
path.exists(),
"{} is missing. These tests are `#[ignore]`d precisely so that running them means \
you want them to run, so this fails instead of passing in 0.00s.\n \
ferrox download sinjab/ms-marco-MiniLM-L6-v2-Q8_0-GGUF --local-dir models\n\
(a git worktree has no models/ of its own — symlink the one in the main checkout)",
path.display()
);
path
}
const QUERY: &str = "How many people live in Berlin?";
const DOCUMENTS: [&str; 5] = [
"Berlin is well known for its museums.",
"Berlin had a population of 3,520,031 registered inhabitants in an area of 891.82 square kilometers.",
"The capital of France is Paris.",
"Elephants are the largest land animals.",
"Berlin is the capital and largest city of Germany by both area and population.",
];
const REFERENCE_SCORES: [f32; 5] = [-0.027210, 0.073057, -0.172285, -0.245995, 0.010248];
const REFERENCE_ORDER: [usize; 5] = [1, 4, 0, 2, 3];
const HF_REFERENCE_SCORES: [f32; 5] = [-4.320078, 8.60714, -11.101217, -11.18758, 0.636921];
const TOLERANCE: f32 = 4e-3;
fn ranking(scores: &[f32]) -> Vec<usize> {
let mut order: Vec<usize> = (0..scores.len()).collect();
order.sort_by(|&a, &b| scores[b].total_cmp(&scores[a]));
order
}
#[test]
#[ignore = "needs models/ms-marco-MiniLM-L6-v2-Q8_0.gguf"]
fn a_real_reranker_checkpoint_loads_with_its_classification_head() {
let model = EmbeddingModel::from_gguf_path(checkpoint()).expect("load the reranker");
assert_eq!(model.architecture(), "bert");
assert_eq!(model.n_embd(), 384);
let head = model
.rank_head()
.expect("a reranker checkpoint with no head is not a reranker");
assert_eq!(head.n_cls_out(), 1);
assert_eq!(head.labels(), ["LABEL_0"]);
}
#[test]
#[ignore = "needs models/ms-marco-MiniLM-L6-v2-Q8_0.gguf"]
fn the_pair_is_cls_query_sep_document_sep_with_the_document_on_segment_one() {
let model = EmbeddingModel::from_gguf_path(checkpoint()).expect("load the reranker");
let pair = model.rerank_input(QUERY, DOCUMENTS[0]).expect("pair input");
assert_eq!(
pair.tokens,
vec![
101, 2129, 2116, 2111, 2444, 1999, 4068, 1029, 102, 4068, 2003, 2092, 2124, 2005, 2049,
9941, 1012, 102
]
);
assert_eq!(
pair.segments,
vec![0, 0, 0, 0, 0, 0, 0, 0, 0, 1, 1, 1, 1, 1, 1, 1, 1, 1]
);
assert_eq!(pair.tokens.len(), pair.segments.len());
assert_eq!(pair.tokens[8], 102);
assert_eq!(pair.segments[8], 0);
assert_eq!(pair.segments[9], 1);
}
#[test]
#[ignore = "needs models/ms-marco-MiniLM-L6-v2-Q8_0.gguf"]
fn the_relevant_document_ranks_first_and_the_scores_match_the_numpy_reference() {
let model = EmbeddingModel::from_gguf_path(checkpoint()).expect("load the reranker");
let scores: Vec<f32> = DOCUMENTS
.iter()
.map(|d| {
let pair = model.rerank_input(QUERY, d).expect("pair input");
model.rerank_score(&pair).expect("score")
})
.collect();
let order = ranking(&scores);
assert_eq!(
order[0],
1,
"the document that answers the query ranked {} of {}: {scores:?}",
order.iter().position(|&i| i == 1).unwrap() + 1,
DOCUMENTS.len()
);
assert_eq!(
order,
REFERENCE_ORDER.to_vec(),
"ferrox ordered {order:?}, HuggingFace {REFERENCE_ORDER:?}, from {scores:?}"
);
for (i, (got, want)) in scores.iter().zip(REFERENCE_SCORES).enumerate() {
assert!(
(got - want).abs() < TOLERANCE,
"document {i}: ferrox {got}, NumPy reference {want}"
);
}
assert!(scores.iter().any(|s| *s < 0.0) && scores.iter().any(|s| *s > 0.0));
}
#[test]
#[ignore = "needs models/ms-marco-MiniLM-L6-v2-Q8_0.gguf"]
fn the_reported_score_is_the_head_run_on_the_cls_row() {
let model = EmbeddingModel::from_gguf_path(checkpoint()).expect("load the reranker");
let head = model.rank_head().expect("head is present");
let pair = model.rerank_input(QUERY, DOCUMENTS[1]).expect("pair input");
let hidden = model.pair_hidden_states(&pair).expect("hidden states");
assert_eq!(hidden.len(), pair.tokens.len() * model.n_embd());
let cls = pool(&hidden, model.n_embd(), PoolingType::Cls).expect("cls row");
let by_hand = head.score(&cls);
let by_route = model.rerank_score(&pair).expect("score");
assert_eq!(
by_hand, by_route,
"the score the model reports is not the head applied to the CLS row"
);
assert_eq!(cls.len(), 384);
assert_ne!(by_route, cls[0]);
}
#[test]
#[ignore = "needs models/ms-marco-MiniLM-L6-v2-Q8_0.gguf"]
fn scoring_both_halves_as_segment_zero_ranks_the_relevant_document_last() {
let model = EmbeddingModel::from_gguf_path(checkpoint()).expect("load the reranker");
let head = model.rank_head().expect("head is present");
let segment_blind: Vec<f32> = DOCUMENTS
.iter()
.map(|d| {
let mut pair = model.rerank_input(QUERY, d).expect("pair input");
pair.segments.fill(0);
let hidden = model.pair_hidden_states(&pair).expect("hidden states");
let cls = pool(&hidden, model.n_embd(), PoolingType::Cls).expect("cls row");
head.score(&cls)
})
.collect();
let order = ranking(&segment_blind);
assert_eq!(
*order.last().unwrap(),
1,
"the segment-blind graph no longer ranks the relevant document last ({order:?} from \
{segment_blind:?}); if the reference changed, re-run \
scripts/rerank_reference_ms_marco.py before relaxing this"
);
assert_ne!(order, REFERENCE_ORDER.to_vec());
}
fn pooler_weights() -> (usize, usize, Vec<u8>, Vec<u8>) {
let path =
Path::new(env!("CARGO_MANIFEST_DIR")).join("../../models/ms-marco-MiniLM-L6-v2-pooler.bin");
assert!(
path.exists(),
"{} is missing. It is the checkpoint's own pooler, which the GGUF does not carry \
(issue #82), dumped from the safetensors by the same NumPy reference this file's \
golden values come from:\n \
python3 scripts/rerank_reference_ms_marco.py --dump-pooler \
models/ms-marco-MiniLM-L6-v2-pooler.bin",
path.display()
);
let raw = std::fs::read(&path).expect("read the pooler dump");
assert!(
raw.len() >= 16 && &raw[..8] == b"FXPOOL01",
"{} is not a --dump-pooler file (magic is {:?})",
path.display(),
&raw[..raw.len().min(8)]
);
let u32_at = |o: usize| u32::from_le_bytes(raw[o..o + 4].try_into().unwrap()) as usize;
let (out, inp) = (u32_at(8), u32_at(12));
let w_end = 16 + out * inp * 4;
let b_end = w_end + out * 4;
assert_eq!(
raw.len(),
b_end,
"{} is {} bytes but [{out}, {inp}] + [{out}] needs {b_end}",
path.display(),
raw.len()
);
(
out,
inp,
raw[16..w_end].to_vec(),
raw[w_end..b_end].to_vec(),
)
}
fn write_pooled_head_gguf(
file: &ShardedGguf,
out: usize,
inp: usize,
w: &[u8],
b: &[u8],
) -> PathBuf {
let mut metadata = BTreeMap::new();
metadata.insert(
"general.architecture".to_string(),
GgufValue::String("bert".to_string()),
);
metadata.insert(
"bert.classifier.output_labels".to_string(),
GgufValue::Array(vec![GgufValue::String("LABEL_0".to_string())]),
);
let mut plan = vec![
TensorPlan {
name: "cls.weight".to_string(),
shape: vec![inp as u64, out as u64],
dtype: GgmlType::F32,
byte_len: w.len(),
},
TensorPlan {
name: "cls.bias".to_string(),
shape: vec![out as u64],
dtype: GgmlType::F32,
byte_len: b.len(),
},
];
let copied: Vec<(&str, Vec<u8>)> = ["cls.output.weight", "cls.output.bias"]
.iter()
.map(|name| {
let info = file
.find_tensor(name)
.unwrap_or_else(|| panic!("the checkpoint carries {name}"));
let bytes = file.tensor_bytes(name).expect("read the tensor").to_vec();
plan.push(TensorPlan {
name: (*name).to_string(),
shape: info.shape.clone(),
dtype: info.dtype,
byte_len: bytes.len(),
});
(*name, bytes)
})
.collect();
let dir = PathBuf::from(env!("CARGO_TARGET_TMPDIR"));
std::fs::create_dir_all(&dir).expect("create the fixture directory");
let path = dir.join(format!("ms-marco-head-pooled-{}.gguf", std::process::id()));
let sink = std::fs::File::create(&path).expect("create the fixture");
let mut writer = GgufWriter::create(sink, &metadata, plan).expect("write the header");
writer.write_tensor("cls.weight", w).expect("write cls");
writer.write_tensor("cls.bias", b).expect("write cls.bias");
for (name, bytes) in &copied {
writer.write_tensor(name, bytes).expect("copy cls.output");
}
writer.finish().expect("finish the fixture");
path
}
#[test]
#[ignore = "needs models/ms-marco-MiniLM-L6-v2-Q8_0.gguf"]
fn the_shipped_checkpoint_is_the_unpooled_regime_and_says_so() {
let model = EmbeddingModel::from_gguf_path(checkpoint()).expect("load the reranker");
let head = model.rank_head().expect("head is present");
assert!(
!head.has_pooler(),
"llama.cpp's converter deletes pooler.dense, so this file cannot carry a cls.weight; \
if it now does, the GGUF changed and HF_REFERENCE_SCORES is the golden set"
);
assert_eq!(head.graph(), "classifier(cls)");
let scores: Vec<f32> = DOCUMENTS
.iter()
.map(|d| {
let pair = model.rerank_input(QUERY, d).expect("pair input");
model.rerank_score(&pair).expect("score")
})
.collect();
for (i, (got, want)) in scores.iter().zip(REFERENCE_SCORES).enumerate() {
assert!(
(got - want).abs() < TOLERANCE,
"document {i}: ferrox {got}, NumPy reference {want} -- the unpooled path moved"
);
}
}
#[test]
#[ignore = "needs models/ms-marco-MiniLM-L6-v2-Q8_0.gguf and its --dump-pooler sidecar"]
fn the_head_reproduces_huggingface_when_the_gguf_carries_the_pooler() {
let model = EmbeddingModel::from_gguf_path(checkpoint()).expect("load the reranker");
let (out, inp, w, b) = pooler_weights();
assert_eq!((out, inp), (model.n_embd(), model.n_embd()));
let source = ShardedGguf::open(checkpoint()).expect("reopen the checkpoint");
let path = write_pooled_head_gguf(&source, out, inp, &w, &b);
let spliced = ShardedGguf::open(&path).expect("open the head fixture");
let head = load_rank_head(&spliced, "bert", model.n_embd(), 1e-12)
.expect("the head loads")
.expect("the fixture carries a head");
assert!(head.has_pooler());
assert_eq!(head.graph(), "classifier(tanh(pooler(cls)))");
let pooled: Vec<f32> = DOCUMENTS
.iter()
.map(|d| {
let pair = model.rerank_input(QUERY, d).expect("pair input");
let hidden = model.pair_hidden_states(&pair).expect("hidden states");
let cls = pool(&hidden, model.n_embd(), PoolingType::Cls).expect("cls row");
head.score(&cls)
})
.collect();
std::fs::remove_file(&path).ok();
const POOLED_TOLERANCE: f32 = 0.02;
for (i, (got, want)) in pooled.iter().zip(HF_REFERENCE_SCORES).enumerate() {
assert!(
(got - want).abs() < POOLED_TOLERANCE,
"document {i}: ferrox with the pooler {got}, HuggingFace {want}"
);
}
assert_eq!(ranking(&pooled), REFERENCE_ORDER.to_vec());
let widest = |v: &[f32]| v.iter().fold(0.0f32, |m, s| m.max(s.abs()));
let unpooled = widest(&REFERENCE_SCORES);
assert!(
widest(&pooled) > 20.0 * unpooled,
"the pooled head's range ({}) is not the calibration difference #82 describes \
against the unpooled {unpooled}",
widest(&pooled)
);
}