candding 0.1.0

Pure candle embeddings for Rust with no ONNX: dense, sparse, multi-vector and rerank models today, image models on the roadmap
Documentation
use super::*;

/// Chosen so the raw MaxSim sum is not a round number.
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
  );
}

/// Code that ignored `reduction` passes at most one assertion.
#[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
  );
}