use candle_core::DType;
use super::*;
use crate::models::RerankHead;
use crate::registry::BackendSupport;
use crate::registry::Family;
use crate::registry::Support;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PairTruncationPolicy;
mod cells;
mod jina_v2;
mod qwen3;
#[test]
fn ms_marco_is_registered_with_its_own_measured_shape() {
let descriptor = lookup("cross-encoder/ms-marco-MiniLM-L6-v2")
.expect("cross-encoder/ms-marco-MiniLM-L6-v2 must be registered");
assert_eq!(descriptor.id, "cross-encoder/ms-marco-MiniLM-L6-v2");
assert_eq!(descriptor.family, Family::Bert);
assert_eq!(descriptor.head, RerankHead::BertPooler);
assert_eq!(
descriptor.score_activation,
ScoreActivation::Identity,
"sbert_ce_default_activation_function: Identity"
);
assert_eq!(
descriptor.input,
RerankInput::Pair {
truncation: PairTruncationPolicy::LongestFirst
},
"a long pair truncates to a near-even split"
);
assert_eq!(descriptor.default_max_length, 512);
assert_eq!(descriptor.phase, 5);
assert!(!descriptor.gated);
assert_eq!(descriptor.license, "Apache-2.0");
}
#[test]
fn jina_reranker_v1_tiny_en_is_registered_with_its_own_measured_shape() {
let descriptor = lookup("jinaai/jina-reranker-v1-tiny-en")
.expect("jinaai/jina-reranker-v1-tiny-en must be registered");
assert_eq!(descriptor.id, "jinaai/jina-reranker-v1-tiny-en");
assert_eq!(descriptor.family, Family::JinaBert);
assert_eq!(descriptor.head, RerankHead::BertPooler);
assert_eq!(
descriptor.score_activation,
ScoreActivation::Sigmoid,
"compute_score applies sigmoid, though the bare forward() does not"
);
assert_eq!(
descriptor.input,
RerankInput::Pair {
truncation: PairTruncationPolicy::LongestFirst
},
"the checkpoint truncates longest first"
);
assert_eq!(descriptor.default_max_length, 512);
assert_eq!(descriptor.phase, 5);
assert!(!descriptor.gated);
assert_eq!(descriptor.license, "Apache-2.0");
}
#[test]
fn bge_reranker_base_is_registered_with_its_own_measured_shape() {
let descriptor =
lookup("BAAI/bge-reranker-base").expect("BAAI/bge-reranker-base must be registered");
assert_eq!(descriptor.id, "BAAI/bge-reranker-base");
assert_eq!(descriptor.family, Family::XlmRoberta);
assert_eq!(descriptor.head, RerankHead::RobertaClassificationHead);
assert_eq!(
descriptor.score_activation,
ScoreActivation::Identity,
"FlagEmbedding's compute_score defaults normalize=False, a raw logit"
);
assert_eq!(
descriptor.input,
RerankInput::Pair {
truncation: PairTruncationPolicy::PreserveQueryToThreeQuarters
},
"query capped at three quarters of max_length, only the document truncated further"
);
assert_eq!(descriptor.default_max_length, 512);
assert_eq!(descriptor.phase, 5);
assert!(!descriptor.gated);
assert_eq!(descriptor.license, "MIT");
}
#[test]
fn bge_reranker_v2_m3_is_registered_with_its_own_measured_shape() {
let descriptor =
lookup("BAAI/bge-reranker-v2-m3").expect("BAAI/bge-reranker-v2-m3 must be registered");
assert_eq!(descriptor.id, "BAAI/bge-reranker-v2-m3");
assert_eq!(descriptor.family, Family::XlmRoberta);
assert_eq!(descriptor.head, RerankHead::RobertaClassificationHead);
assert_eq!(
descriptor.score_activation,
ScoreActivation::Identity,
"the same FlagReranker default as bge-reranker-base, a raw logit"
);
assert_eq!(
descriptor.input,
RerankInput::Pair {
truncation: PairTruncationPolicy::PreserveQueryToThreeQuarters
},
"the same truncation as bge-reranker-base"
);
assert_eq!(
descriptor.default_max_length, 512,
"the class-level FlagEmbedding default, not this checkpoint's own wider 8192 tokenizer limit"
);
assert_eq!(descriptor.phase, 5);
assert!(!descriptor.gated);
assert_eq!(descriptor.license, "Apache-2.0");
}
#[test]
fn every_entry_names_a_position_scheme_exactly_when_its_family_is_the_flash_one() {
for descriptor in all() {
let flash = descriptor.family == Family::XlmRobertaFlash;
assert_eq!(
descriptor.position_scheme.is_some(),
flash,
"{}",
descriptor.id
);
let json = serde_json::to_value(&descriptor).unwrap();
assert_eq!(
json["position_scheme"].is_null(),
!flash,
"{}",
descriptor.id
);
}
}
const DEFAULT_BF16: &[&str] = &["Qwen/Qwen3-Reranker-4B", "Qwen/Qwen3-Reranker-8B"];
#[test]
fn every_entry_defaults_to_f32_but_the_bf16_ones() {
for descriptor in all() {
let (dtype, name) = if DEFAULT_BF16.contains(&descriptor.id.as_str()) {
(DType::BF16, "bf16")
} else {
(DType::F32, "f32")
};
assert_eq!(descriptor.default_dtype, dtype, "{}", descriptor.id);
let json = serde_json::to_value(&descriptor).unwrap();
assert_eq!(json["default_dtype"], name, "{}", descriptor.id);
}
}
#[test]
fn all_and_by_phase_carry_exactly_these_eight_entries() {
let all = all();
assert_eq!(all.len(), 8, "{all:?}");
let phase5 = by_phase(5);
assert_eq!(phase5.len(), 8, "{phase5:?}");
assert!(!dtype_needs_accelerator(
"cross-encoder/ms-marco-MiniLM-L6-v2",
DType::F32
));
}
const AWAITING_CPU_F32: &[&str] = &["Qwen/Qwen3-Reranker-4B"];
#[test]
fn every_registered_rerank_entry_has_already_flipped_past_pending() {
for id in AWAITING_CPU_F32 {
assert!(
lookup(id).is_some(),
"{id} awaits a run but is not registered"
);
}
for d in all() {
let awaiting = AWAITING_CPU_F32.contains(&d.id.as_str());
assert_eq!(d.backends.cpu == Support::Pending, awaiting, "{} cpu", d.id);
assert_ne!(d.backends.metal, Support::Pending, "{}", d.id);
assert_ne!(d.backends.cuda, Support::Pending, "{}", d.id);
for (dtype, row) in [
(DType::F32, d.dtype_backends.f32),
(DType::F16, d.dtype_backends.f16),
(DType::BF16, d.dtype_backends.bf16),
] {
for (device, status) in [("cpu", row.cpu), ("metal", row.metal), ("cuda", row.cuda)] {
let pending_here = awaiting && dtype == DType::F32 && device == "cpu";
assert_eq!(
status == Support::Pending,
pending_here,
"{} {dtype:?} on {device}: {status:?}",
d.id
);
}
}
}
}
fn assert_accelerator_only_f16_and_bf16(id: &str, f16_on_cuda: Support) {
let descriptor = lookup(id).expect("registered");
let accelerator_only = |cuda| BackendSupport {
cpu: Support::Unsupported,
metal: Support::Verified,
cuda,
};
assert_eq!(
descriptor.dtype_backends.f16,
accelerator_only(f16_on_cuda),
"{id} f16"
);
assert_eq!(
descriptor.dtype_backends.bf16,
accelerator_only(Support::Untested),
"{id} bf16"
);
assert!(dtype_needs_accelerator(id, DType::F16), "{id}");
}
#[test]
fn bge_reranker_base_f16_is_accelerator_only_and_bf16_is_verified_on_metal() {
assert_accelerator_only_f16_and_bf16("BAAI/bge-reranker-base", Support::Verified);
}
#[test]
fn the_other_three_rerank_entries_are_f16_accelerator_only_and_bf16_verified_on_metal() {
for (id, f16_on_cuda) in [
("cross-encoder/ms-marco-MiniLM-L6-v2", Support::Verified),
("jinaai/jina-reranker-v1-tiny-en", Support::Verified),
("BAAI/bge-reranker-v2-m3", Support::Verified),
] {
assert_accelerator_only_f16_and_bf16(id, f16_on_cuda);
}
}
#[test]
fn accelerator_only_f16_leaves_f32_the_same_on_cpu_and_metal() {
for descriptor in all() {
let f32 = descriptor.dtype_backends.f32;
assert!(
!dtype_needs_accelerator(&descriptor.id, DType::F32),
"{}",
descriptor.id
);
assert_ne!(f32.cpu, Support::Unsupported, "{}", descriptor.id);
if descriptor.backends.cpu == Support::Verified
&& descriptor.backends.metal == Support::Verified
{
assert_eq!(f32.cpu, f32.metal, "{}", descriptor.id);
assert_eq!(f32.cpu, Support::Verified, "{}", descriptor.id);
}
}
}