use super::*;
fn query_and_document() -> (MultiVectorEmbedding, MultiVectorEmbedding) {
let query = MultiVectorEmbedding::from_rows(2, [vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]]);
let document = MultiVectorEmbedding::from_rows(2, [vec![1.0, 0.0], vec![0.5, 0.5]]);
(query, document)
}
#[test]
fn sum_returns_the_raw_maxsim_sum_unchanged() {
let (query, document) = query_and_document();
let raw = late_interaction_score(&query, None, &document, None, MaskedRowFill::DtypeMinimum);
let reduced = late_interaction_score_for(
ScoreReduction::Sum,
&query,
&document,
MaskedRowFill::DtypeMinimum,
);
assert_eq!(reduced, raw);
}
#[test]
fn mean_by_query_length_divides_by_the_real_row_count() {
let (query, document) = query_and_document();
let raw = late_interaction_score(&query, None, &document, None, MaskedRowFill::DtypeMinimum);
let reduced = late_interaction_score_for(
ScoreReduction::MeanByQueryLength,
&query,
&document,
MaskedRowFill::DtypeMinimum,
);
assert!(
(reduced - raw / 3.0).abs() < 1e-6,
"{reduced} vs {}",
raw / 3.0
);
}
#[test]
fn sum_and_mean_by_query_length_disagree_on_the_same_pair() {
let (query, document) = query_and_document();
let sum = late_interaction_score_for(
ScoreReduction::Sum,
&query,
&document,
MaskedRowFill::DtypeMinimum,
);
let mean = late_interaction_score_for(
ScoreReduction::MeanByQueryLength,
&query,
&document,
MaskedRowFill::DtypeMinimum,
);
assert_ne!(
sum, mean,
"a 3-row query must not divide out to its own sum"
);
assert!((sum - mean * 3.0).abs() < 1e-5, "{sum} vs {}", mean * 3.0);
}
#[test]
fn mean_by_query_length_on_an_empty_query_stays_zero_not_nan() {
let empty_query = MultiVectorEmbedding::from_rows(2, Vec::<Vec<f32>>::new());
let document = MultiVectorEmbedding::from_rows(2, [vec![1.0, 0.0]]);
let reduced = late_interaction_score_for(
ScoreReduction::MeanByQueryLength,
&empty_query,
&document,
MaskedRowFill::DtypeMinimum,
);
assert_eq!(reduced, 0.0);
assert!(!reduced.is_nan());
}
#[test]
fn neither_reduction_touches_the_empty_document_sentinel() {
let query = MultiVectorEmbedding::from_rows(2, [vec![1.0, 0.0], vec![0.0, 1.0], vec![1.0, 1.0]]);
let empty_document = MultiVectorEmbedding::from_rows(2, Vec::<Vec<f32>>::new());
for reduction in [ScoreReduction::Sum, ScoreReduction::MeanByQueryLength] {
let reduced = late_interaction_score_for(
reduction,
&query,
&empty_document,
MaskedRowFill::DtypeMinimum,
);
assert_eq!(reduced, -1e9, "{reduction:?}");
}
}
#[test]
fn score_reduction_serializes_snake_case() {
assert_eq!(
serde_json::to_string(&ScoreReduction::MeanByQueryLength).unwrap(),
"\"mean_by_query_length\""
);
assert_eq!(
serde_json::from_str::<ScoreReduction>("\"sum\"").unwrap(),
ScoreReduction::Sum
);
}