use super::*;
use crate::models::xlm_roberta_flash::PositionScheme;
use crate::multi_vector::MaskedRowFill;
use crate::multi_vector::ScoreReduction;
mod status;
#[test]
fn bge_m3_is_registered_with_its_own_measured_shape() {
let all = all();
assert_eq!(all.len(), 4, "{all:?}");
let descriptor = lookup("BAAI/bge-m3").expect("BAAI/bge-m3 must be registered");
assert_eq!(descriptor.id, "BAAI/bge-m3");
assert_eq!(descriptor.head, crate::models::MultiVectorHead::M3Colbert);
assert!(
!descriptor.keeps_cls_row,
"bge-m3's own head drops the classification row"
);
assert_eq!(descriptor.projection_in, 1024);
assert_eq!(descriptor.projection_out, 1024);
assert!(
descriptor.projection_bias,
"colbert_linear.pt carries a bias"
);
assert!(descriptor.normalize);
assert_eq!(
descriptor.score_reduction,
ScoreReduction::MeanByQueryLength
);
assert_eq!(descriptor.query_marker, None);
assert_eq!(descriptor.document_marker, None);
assert_eq!(descriptor.query_length, None);
assert!(
!descriptor.attends_to_mask_tokens,
"bge-m3's own head carries no fixed-length query padding to attend to at all"
);
assert_eq!(
descriptor.masked_document_row_fill,
MaskedRowFill::Zero,
"bge-m3's own FlagEmbedding reference zeroes a masked document row rather than dtype-min-filling it"
);
assert_eq!(descriptor.phase, 4);
}
#[test]
fn jina_colbert_v2_is_registered_with_its_own_family_and_fourth_convention() {
let descriptor =
lookup("jinaai/jina-colbert-v2").expect("jinaai/jina-colbert-v2 must be registered");
assert_eq!(descriptor.id, "jinaai/jina-colbert-v2");
assert_eq!(descriptor.family, crate::registry::Family::XlmRobertaFlash);
assert_eq!(descriptor.head, crate::models::MultiVectorHead::JinaColbert);
assert!(
descriptor.keeps_cls_row,
"the same HF_ColBERT reference pipeline both StanfordColbert checkpoints load through keeps it too"
);
assert_eq!(descriptor.projection_in, 1024);
assert_eq!(descriptor.projection_out, 128);
assert!(
!descriptor.projection_bias,
"its own root linear.weight publishes no accompanying bias"
);
assert!(descriptor.normalize);
assert_eq!(descriptor.score_reduction, ScoreReduction::Sum);
assert_eq!(descriptor.query_marker.as_deref(), Some("[QueryMarker] "));
assert_eq!(
descriptor.document_marker.as_deref(),
Some("[DocumentMarker] ")
);
assert_eq!(descriptor.query_length, Some(32));
assert!(
descriptor.attends_to_mask_tokens,
"its self-attention sees its [MASK]-expansion query padding, where no StanfordColbert checkpoint's does"
);
assert_eq!(
descriptor.masked_document_row_fill,
MaskedRowFill::DtypeMinimum,
"shares sentence-transformers's own HF_ColBERT pipeline with both StanfordColbert entries, not bge-m3's"
);
assert_eq!(descriptor.default_max_length, 300);
assert_eq!(descriptor.license, "CC-BY-NC-4.0");
assert_eq!(descriptor.phase, 4);
}
#[test]
fn only_the_flash_entry_names_a_position_scheme_and_it_is_rotary() {
for descriptor in all() {
let json = serde_json::to_value(&descriptor).unwrap();
if descriptor.family == crate::registry::Family::XlmRobertaFlash {
assert_eq!(
descriptor.position_scheme,
Some(PositionScheme::Rotary),
"{}",
descriptor.id
);
assert_eq!(json["position_scheme"], "rotary", "{}", descriptor.id);
} else {
assert_eq!(descriptor.position_scheme, None, "{}", descriptor.id);
assert_eq!(
json["position_scheme"],
serde_json::Value::Null,
"{}",
descriptor.id
);
}
}
assert!(all().iter().any(|d| d.id == "jinaai/jina-colbert-v2"));
}
#[test]
fn only_jina_colbert_v2_attends_to_its_own_mask_tokens() {
for id in [
"BAAI/bge-m3",
"colbert-ir/colbertv2.0",
"answerdotai/answerai-colbert-small-v1",
] {
let descriptor = lookup(id).unwrap_or_else(|| panic!("{id} must be registered"));
assert!(!descriptor.attends_to_mask_tokens, "{id}");
}
assert!(
lookup("jinaai/jina-colbert-v2")
.unwrap()
.attends_to_mask_tokens
);
}
#[test]
fn the_two_stanford_colbert_entries_contradict_bge_m3s_three_conventions() {
for id in [
"colbert-ir/colbertv2.0",
"answerdotai/answerai-colbert-small-v1",
] {
let descriptor = lookup(id).unwrap_or_else(|| panic!("{id} must be registered"));
assert_eq!(
descriptor.head,
crate::models::MultiVectorHead::StanfordColbert
);
assert!(
descriptor.keeps_cls_row,
"{id}: both Stanford-ColBERT-style checkpoints keep the classification row"
);
assert_eq!(
descriptor.score_reduction,
ScoreReduction::Sum,
"{id}: both take the plain MaxSim sum, not bge-m3's mean-by-query-length"
);
assert_eq!(
descriptor.masked_document_row_fill,
MaskedRowFill::DtypeMinimum,
"{id}: both dtype-min-fill a masked document row, not bge-m3's zero-fill"
);
assert!(
!descriptor.projection_bias,
"{id}: neither checkpoint's own root linear.weight publishes a bias"
);
assert_eq!(descriptor.query_marker.as_deref(), Some("[unused0] "));
assert_eq!(descriptor.document_marker.as_deref(), Some("[unused1] "));
assert_eq!(descriptor.query_length, Some(32));
assert_eq!(descriptor.phase, 4);
}
}
#[test]
fn colbertv2_and_answerai_small_differ_only_in_projection_shape_and_document_length() {
let big = lookup("colbert-ir/colbertv2.0").unwrap();
let small = lookup("answerdotai/answerai-colbert-small-v1").unwrap();
assert_eq!(big.projection_in, 768);
assert_eq!(big.projection_out, 128);
assert_eq!(big.default_max_length, 180);
assert_eq!(small.projection_in, 384);
assert_eq!(small.projection_out, 96);
assert_eq!(small.default_max_length, 300);
assert_eq!(big.query_length, small.query_length);
assert_eq!(big.keeps_cls_row, small.keeps_cls_row);
assert_eq!(big.score_reduction, small.score_reduction);
assert_eq!(big.masked_document_row_fill, small.masked_document_row_fill);
assert_eq!(big.query_marker, small.query_marker);
assert_eq!(big.document_marker, small.document_marker);
}
#[test]
fn lookup_is_case_insensitive() {
assert!(lookup("baai/BGE-M3").is_some());
assert!(lookup("nonexistent/model").is_none());
}
#[test]
fn by_phase_finds_it_under_four_and_nowhere_else() {
assert_eq!(by_phase(4).len(), 4);
assert!(by_phase(3).is_empty());
}
#[test]
fn dtype_needs_accelerator_is_false_because_nothing_is_accelerator_only_here() {
assert!(!dtype_needs_accelerator(
"BAAI/bge-m3",
candle_core::DType::F16
));
}
#[test]
fn the_descriptor_round_trips_through_owned_json_deserialization() {
let descriptor = lookup("colbert-ir/colbertv2.0").unwrap();
let json = serde_json::to_string(&descriptor).unwrap();
let restored: MultiVectorDescriptor = serde_json::from_str(&json).unwrap();
assert_eq!(restored, descriptor);
}