use std::borrow::Cow;
use candle_core::DType;
use crate::models::MultiVectorHead;
use crate::models::RerankHead;
use crate::models::SparseHead;
use crate::multi_vector::MaskedRowFill;
use crate::multi_vector::ScoreReduction;
use crate::pooling::Pooling;
use crate::registry::BackendSupport;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::registry::ModelDescriptor;
use crate::registry::WeightsFormat;
use crate::registry::multi_vector;
use crate::registry::multi_vector::MultiVectorDescriptor;
use crate::registry::rerank;
use crate::registry::rerank::RerankDescriptor;
use crate::registry::rerank::RerankInput;
use crate::registry::sparse::SparseComputation;
use crate::registry::sparse::SparseDescriptor;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PromptTemplate;
#[test]
fn a_new_model_descriptor_takes_the_snapshot_defaults() {
let expected = ModelDescriptor {
id: "org/dense".to_string(),
family: Family::Bert,
dim: 384,
default_max_length: 512,
pooling: Pooling::Cls,
normalize: true,
query_template: None,
passage_template: None,
default_instruction: None,
dense_dirs: Vec::new(),
mrl: false,
layer_norm_after_pooling: false,
license: "unknown".to_string(),
size_gb: 0.0,
gated: false,
phase: u8::MAX,
backends: BackendSupport::untested(),
default_dtype: DType::F32,
dtype_backends: DtypeBackendSupport::untested(),
};
let descriptor = ModelDescriptor::new("org/dense", Family::Bert, 384, Pooling::Cls);
assert_eq!(descriptor, expected);
}
#[test]
fn a_new_sparse_descriptor_is_an_encoded_head_with_the_snapshot_defaults() {
let expected = SparseDescriptor {
id: "org/sparse".to_string(),
family: Some(Family::Qwen3),
head: SparseHead::SpladeTied,
vocab_size: 151_936,
default_max_length: 512,
backends: BackendSupport::untested(),
dtype_backends: DtypeBackendSupport::untested(),
license: "unknown".to_string(),
size_gb: 0.0,
phase: u8::MAX,
weights: WeightsFormat::Safetensors,
gated: false,
head_source: None,
head_revision: None,
computation: SparseComputation::Encoded,
};
let descriptor =
SparseDescriptor::new("org/sparse", Family::Qwen3, SparseHead::SpladeTied, 151_936);
assert_eq!(descriptor, expected);
}
#[test]
fn a_new_multi_vector_descriptor_keeps_each_row_fact_in_its_own_field() {
let expected = MultiVectorDescriptor {
id: "org/colbert".to_string(),
family: Family::Bert,
head: MultiVectorHead::StanfordColbert,
default_max_length: 512,
keeps_cls_row: true,
projection_in: 768,
projection_out: 128,
projection_bias: false,
normalize: true,
score_reduction: ScoreReduction::Sum,
query_marker: Some("[unused0] ".to_string()),
document_marker: Some("[unused1] ".to_string()),
query_length: Some(32),
attends_to_mask_tokens: false,
masked_document_row_fill: MaskedRowFill::DtypeMinimum,
position_scheme: None,
backends: BackendSupport::untested(),
dtype_backends: DtypeBackendSupport::untested(),
license: "unknown".to_string(),
size_gb: 0.0,
phase: u8::MAX,
weights: WeightsFormat::Safetensors,
gated: false,
};
let head = MultiVectorHead::StanfordColbert;
let descriptor = MultiVectorDescriptor::new("org/colbert", Family::Bert, head, 768, 128);
assert_eq!(descriptor, expected);
}
#[test]
fn a_new_rerank_descriptor_takes_the_snapshot_defaults() {
let input = RerankInput::Prompt {
template: PromptTemplate::Qwen3Reranker,
default_instruction: Cow::Borrowed("judge it"),
};
let expected = RerankDescriptor {
id: "org/reranker".to_string(),
family: Family::Qwen3,
head: RerankHead::BertPooler,
default_max_length: 512,
score_activation: ScoreActivation::Sigmoid,
input: input.clone(),
alibi_variant: None,
position_scheme: None,
default_dtype: DType::F32,
backends: BackendSupport::untested(),
dtype_backends: DtypeBackendSupport::untested(),
license: "unknown".to_string(),
size_gb: 0.0,
phase: u8::MAX,
weights: WeightsFormat::Safetensors,
gated: false,
};
let head = RerankHead::BertPooler;
let descriptor = RerankDescriptor::new(
"org/reranker",
Family::Qwen3,
head,
ScoreActivation::Sigmoid,
input,
);
assert_eq!(descriptor, expected);
}
#[test]
fn a_new_multi_vector_descriptor_is_every_registered_entry_with_its_head_s_conventions() {
for registered in multi_vector::all() {
let mut built = MultiVectorDescriptor::new(
registered.id.clone(),
registered.family,
registered.head,
registered.projection_in,
registered.projection_out,
);
built.default_max_length = registered.default_max_length;
built.backends = registered.backends;
built.dtype_backends = registered.dtype_backends;
built.license.clone_from(®istered.license);
built.size_gb = registered.size_gb;
built.phase = registered.phase;
built.gated = registered.gated;
assert_eq!(built, registered, "{}", registered.id);
}
}
#[test]
fn a_new_rerank_descriptor_is_every_registered_entry_with_its_family_s_conventions() {
for registered in rerank::all() {
let mut built = RerankDescriptor::new(
registered.id.clone(),
registered.family,
registered.head,
registered.score_activation,
registered.input.clone(),
);
built.default_max_length = registered.default_max_length;
built.default_dtype = registered.default_dtype;
built.backends = registered.backends;
built.dtype_backends = registered.dtype_backends;
built.license.clone_from(®istered.license);
built.size_gb = registered.size_gb;
built.phase = registered.phase;
built.gated = registered.gated;
assert_eq!(built, registered, "{}", registered.id);
}
}
#[test]
fn every_dense_and_multi_vector_entry_normalizes() {
for descriptor in crate::registry::dense::all() {
assert!(descriptor.normalize, "{}", descriptor.id);
}
for descriptor in multi_vector::all() {
assert!(descriptor.normalize, "{}", descriptor.id);
}
}