use std::collections::HashMap;
use crate::registry;
use crate::registry::WeightsFormat;
use crate::registry::multi_vector;
use crate::registry::rerank;
use crate::registry::sparse;
fn every_entry() -> Vec<(&'static str, String, WeightsFormat)> {
let mut entries: Vec<_> = registry::dense::all()
.into_iter()
.map(|d| {
let weights = registry::entry_for(&d.id).expect("registered").weights;
("dense", d.id, weights)
})
.collect();
entries.extend(
sparse::all()
.into_iter()
.map(|d| ("sparse", d.id, d.weights)),
);
entries.extend(
multi_vector::all()
.into_iter()
.map(|d| ("multi_vector", d.id, d.weights)),
);
entries.extend(
rerank::all()
.into_iter()
.map(|d| ("rerank", d.id, d.weights)),
);
entries
}
#[test]
fn a_case_variant_of_a_rerank_only_id_canonicalizes_to_the_registry_spelling() {
assert!(registry::dense::lookup("cross-encoder/ms-marco-MiniLM-L6-v2").is_none());
assert!(sparse::lookup("cross-encoder/ms-marco-MiniLM-L6-v2").is_none());
assert!(multi_vector::lookup("cross-encoder/ms-marco-MiniLM-L6-v2").is_none());
assert_eq!(
registry::canonical_id("CROSS-ENCODER/ms-marco-minilm-l6-v2"),
Some("cross-encoder/ms-marco-MiniLM-L6-v2")
);
for d in rerank::all() {
assert_eq!(
registry::canonical_id(&d.id.to_ascii_uppercase()),
Some(d.id.as_str()),
"{}",
d.id
);
}
}
#[test]
fn a_case_variant_of_a_sparse_only_or_multi_vector_only_id_canonicalizes_too() {
assert_eq!(
registry::canonical_id("PRITHIVIDA/splade_pp_en_v1"),
Some("prithivida/Splade_PP_en_v1")
);
assert_eq!(
registry::canonical_id("COLBERT-IR/ColBERTv2.0"),
Some("colbert-ir/colbertv2.0")
);
assert_eq!(registry::canonical_id("nope/nothing"), None);
}
#[test]
fn weights_format_returns_each_rerank_and_multi_vector_entrys_own_format() {
for d in rerank::all() {
assert_eq!(
registry::registered_weights_format(&d.id),
Some(d.weights),
"{}",
d.id
);
assert_eq!(registry::weights_format(&d.id), d.weights, "{}", d.id);
}
for d in multi_vector::all() {
assert_eq!(
registry::registered_weights_format(&d.id),
Some(d.weights),
"{}",
d.id
);
}
assert_eq!(registry::registered_weights_format("nope/nothing"), None);
}
#[test]
fn an_id_in_several_registries_has_one_spelling_and_one_weights_format() {
let mut seen: HashMap<String, Vec<(&'static str, String, WeightsFormat)>> = HashMap::new();
for (catalog, id, weights) in every_entry() {
seen
.entry(id.to_ascii_lowercase())
.or_default()
.push((catalog, id, weights));
}
let shared: Vec<_> = seen.values().filter(|hits| hits.len() > 1).collect();
assert!(
!shared.is_empty(),
"BAAI/bge-m3 is carried by three catalogs"
);
for hits in shared {
let (_, first_id, first_weights) = &hits[0];
for (catalog, id, weights) in hits {
assert_eq!(id, first_id, "{catalog} spells {first_id} differently");
assert_eq!(weights, first_weights, "{catalog} fetches {id} differently");
}
}
}