candding 0.1.0

//! Query cap 6, overlap 2, maximum length 16: the query `a b` is 4 ids, so a window holds 10 document tokens and one starts every 8.
//! The stub models score each row from its ids, so a document's score shows which window it came from.

use std::sync::Arc;
use std::sync::Mutex;

use super::*;

const QUERY: &str = "a b";
const WINDOWED: RerankInput = RerankInput::Windowed {
  query_cap: 6,
  overlap: 2,
};

/// Scores each row by the sum of its real ids, and records how many rows each forward pass read.
#[derive(Default)]
struct SumModel {
  batches: Arc<Mutex<Vec<usize>>>,
}

impl RerankModel for SumModel {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<RerankScore>> {
    let ids = input.input_ids.to_vec2::<u32>()?;
    self.batches.lock().unwrap().push(ids.len());
    Ok(
      ids
        .iter()
        .zip(input.lengths)
        .map(|(row, &len)| RerankScore::new(row[..len].iter().sum::<u32>() as f32))
        .collect(),
    )
  }

  fn max_position_embeddings(&self) -> usize {
    512
  }

  fn weights_dtype(&self) -> DType {
    DType::F32
  }
}

fn double_separator_path() -> PathBuf {
  PathBuf::from(env!("CARGO_MANIFEST_DIR"))
    .join("tests/fixtures/tiny_tokenizer_double_separator.json")
}

fn windowed_model(input: RerankInput, batch_size: usize) -> TextCrossEncoder {
  windowed_with(Box::new(SumModel::default()), input, batch_size)
}

fn windowed_with(
  model: Box<dyn RerankModel>,
  input: RerankInput,
  batch_size: usize,
) -> TextCrossEncoder {
  let mut descriptor = stub_descriptor();
  descriptor.input = input;
  cross_encoder(model, descriptor, &double_separator_path(), 16, batch_size)
}

fn words(word: &str, n: usize) -> String {
  vec![word; n].join(" ")
}

/// 30 tokens cut into windows at 0..10, 8..18, 16..26 and 20..30; only the second holds every `d`.
fn peaked_document() -> String {
  [words("hello", 8), words("d", 10), words("hello", 12)].join(" ")
}

fn window_sums(document: &str) -> Vec<f32> {
  let tokenizer = crate::tokenizer::WindowedTokenizer::from_file(
    &double_separator_path(),
    16,
    PadToken::Id(0),
    6,
    2,
  )
  .unwrap();
  tokenizer
    .window_ids(QUERY, document)
    .unwrap()
    .iter()
    .map(|ids| ids.iter().sum::<u32>() as f32)
    .collect()
}

#[test]
fn a_windowed_document_scores_its_best_window() {
  let document = peaked_document();
  let sums = window_sums(&document);
  assert_eq!(sums.len(), 4, "{sums:?}");
  let best = sums.iter().copied().fold(f32::MIN, f32::max);
  let mean = sums.iter().sum::<f32>() / sums.len() as f32;
  assert!(best > mean && best > sums[0] && best > sums[3], "{sums:?}");
  assert_eq!(best, sums[1], "{sums:?}");
  let model = windowed_model(WINDOWED, 8);
  let scores = model.rerank(QUERY, &[document.as_str()], None).unwrap();
  assert_eq!(scores[0].value(), best, "window sums {sums:?}");
}

#[test]
fn windows_are_batched_across_documents_by_the_batch_size() {
  let documents = ["c".to_string(), peaked_document(), words("world", 12)];
  let sums: Vec<Vec<f32>> = documents.iter().map(|d| window_sums(d)).collect();
  assert_eq!(sums.iter().map(Vec::len).collect::<Vec<_>>(), [1, 4, 2]);
  let best: Vec<f32> = sums
    .iter()
    .map(|s| s.iter().copied().fold(f32::MIN, f32::max))
    .collect();
  for (batch_size, expected) in [(4, vec![4, 3]), (1, vec![1; 7]), (8, vec![7])] {
    let model = SumModel::default();
    let batches = Arc::clone(&model.batches);
    let encoder = windowed_with(Box::new(model), WINDOWED, batch_size);
    let scores: Vec<f32> = encoder
      .rerank(QUERY, &documents, None)
      .unwrap()
      .iter()
      .map(RerankScore::value)
      .collect();
    assert_eq!(scores, best, "batch size {batch_size}");
    assert_eq!(
      *batches.lock().unwrap(),
      expected,
      "batch size {batch_size}"
    );
  }
}

#[test]
fn a_windowed_documents_counts_total_its_windows() {
  let document = peaked_document();
  let model = windowed_model(WINDOWED, 8);
  let documents = [document.as_str(), "", "c"];
  assert_eq!(
    model.pair_window_count(QUERY, &documents).unwrap(),
    [4, 1, 1]
  );
  assert_eq!(
    model.pair_token_count(QUERY, &documents).unwrap(),
    [4 * 16, 4 + 2, 4 + 3]
  );
  let pair = stub_model(Box::new(StubRerankModel), 512, 8);
  assert_eq!(
    pair.pair_window_count("a", &["a", "a b c"]).unwrap(),
    [1, 1]
  );
}

#[test]
fn the_descriptors_query_cap_and_overlap_reach_the_windows() {
  let document = peaked_document();
  let long_query = "a b c d a b";
  let capped = windowed_model(
    RerankInput::Windowed {
      query_cap: 4,
      overlap: 2,
    },
    8,
  );
  let uncapped = windowed_model(WINDOWED, 8);
  let short = capped
    .pair_token_count(QUERY, &[document.as_str()])
    .unwrap();
  let long = capped
    .pair_token_count(long_query, &[document.as_str()])
    .unwrap();
  assert_eq!(short, long, "a cap of 4 keeps `a b` of either query");
  assert_ne!(
    uncapped
      .pair_token_count(long_query, &[document.as_str()])
      .unwrap(),
    long
  );
  let no_overlap = windowed_model(
    RerankInput::Windowed {
      query_cap: 6,
      overlap: 0,
    },
    8,
  );
  assert_eq!(
    no_overlap
      .pair_window_count(QUERY, &[document.as_str()])
      .unwrap(),
    [3]
  );
}

#[test]
fn a_windowed_entry_refuses_an_instruction() {
  let model = windowed_model(WINDOWED, 8);
  let refused = model.rerank_with_instruction(QUERY, &["c"], "judge it", None);
  assert!(matches!(
    refused,
    Err(CanddingError::InstructionOnPairInput { .. })
  ));
  let refused = model.pair_token_count_with_instruction(QUERY, &["c"], "judge it");
  assert!(matches!(
    refused,
    Err(CanddingError::InstructionOnPairInput { .. })
  ));
}