candding 0.1.0

//! `tests/fixtures/windowed_pairs.json` comes from the checkpoint's own `reranker_tokenize_preproc`, via `scripts/window_fixture.py`.

use std::fs;
use std::path::PathBuf;

use candle_core::Device;
use serde::Deserialize;

use crate::error::CanddingError;
use crate::error::LengthKind;
use crate::tokenizer::PadToken;
use crate::tokenizer::WindowedTokenizer;

#[derive(Deserialize)]
struct Fixture {
  tokenizer: String,
  cases: Vec<Case>,
}

#[derive(Deserialize)]
struct Case {
  max_length: usize,
  query_cap: usize,
  overlap: usize,
  query: String,
  document: String,
  query_ids: Vec<u32>,
  windows: Vec<Vec<u32>>,
}

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

fn fixture() -> Fixture {
  let text = fs::read_to_string(fixtures().join("windowed_pairs.json")).unwrap();
  serde_json::from_str(&text).unwrap()
}

fn tokenizer(fixture: &Fixture, case: &Case) -> WindowedTokenizer {
  WindowedTokenizer::from_file(
    &fixtures().join(&fixture.tokenizer),
    case.max_length,
    PadToken::Id(0),
    case.query_cap,
    case.overlap,
  )
  .unwrap()
}

fn label(case: &Case) -> String {
  format!(
    "max_length {} query cap {} overlap {}, query of {} words, document of {} words",
    case.max_length,
    case.query_cap,
    case.overlap,
    case.query.split_whitespace().count(),
    case.document.split_whitespace().count()
  )
}

#[test]
fn every_window_is_the_references_id_for_id() {
  let fixture = fixture();
  assert_eq!(fixture.cases.len(), 56);
  for case in &fixture.cases {
    let tokenizer = tokenizer(&fixture, case);
    let windows = tokenizer.windows(&case.query, &case.document).unwrap();
    let ids: Vec<Vec<u32>> = windows.iter().map(|w| w.get_ids().to_vec()).collect();
    assert_eq!(ids, case.windows, "{}", label(case));
    assert!(
      ids.iter().all(|w| w.starts_with(&case.query_ids)),
      "{}",
      label(case)
    );
    for window in &windows {
      assert!(
        window.get_type_ids().iter().all(|&t| t == 0),
        "{}",
        label(case)
      );
      assert!(
        window.get_attention_mask().iter().all(|&m| m == 1),
        "{}",
        label(case)
      );
      assert!(window.len() <= case.max_length, "{}", label(case));
    }
  }
}

#[test]
fn the_fixture_covers_every_edge_of_the_windowing() {
  let fixture = fixture();
  let has = |check: &dyn Fn(&Case) -> bool| fixture.cases.iter().any(check);
  assert!(has(&|c| c.document.is_empty() && c.windows.len() == 1));
  assert!(has(&|c| !c.document.is_empty()
    && c.document.trim().is_empty()
    && c.windows.len() == 1));
  let budget = |c: &Case| c.max_length - c.query_ids.len() - 2;
  let content = |c: &Case| {
    c.windows
      .iter()
      .map(|w| w.len() - c.query_ids.len() - 2)
      .collect::<Vec<_>>()
  };
  assert!(has(&|c| content(c) == [budget(c)]));
  assert!(has(
    &|c| c.windows.len() == 2 && content(c) == [budget(c), budget(c)]
  ));
  assert!(has(
    &|c| c.query_ids.len() == c.query_cap && c.query.split_whitespace().count() > c.query_cap
  ));
  assert!(has(&|c| c.overlap > budget(c) / 4 && c.windows.len() > 2));
  assert!(has(&|c| c.overlap < budget(c) / 4 && c.windows.len() > 2));
  let aligned_to_the_end = |c: &Case| {
    let [.., previous, last] = c.windows.as_slice() else {
      return false;
    };
    let content = |w: &[u32]| w[c.query_ids.len() + 1..w.len() - 1].to_vec();
    let (previous, last) = (content(previous), content(last));
    c.windows.len() >= 3 && previous[1..] == last[..budget(c) - 1]
  };
  assert!(has(&aligned_to_the_end));
  assert!(has(&|c| c.max_length == 1024
    && c.query_cap == 512
    && c.overlap == 80
    && c.windows.len() >= 3));
}

/// The reference's own `rerank` asserts a `max_length` of at least twice the query cap.
#[test]
fn a_windowed_tokenizer_refuses_lengths_and_templates_it_cannot_cut() {
  let path = fixtures().join("tiny_tokenizer_double_separator.json");
  for (max_length, query_cap, expected) in [
    (16, 2, (2, 3, LengthKind::QueryCap { added: 2 })),
    (11, 6, (11, 12, LengthKind::Windowed { query_cap: 6 })),
    (4, 3, (4, 6, LengthKind::Windowed { query_cap: 3 })),
  ] {
    let refused = WindowedTokenizer::from_file(&path, max_length, PadToken::Id(0), query_cap, 2);
    assert!(
      matches!(
        &refused,
        Err(CanddingError::LengthTooSmall { requested, minimum, kind })
          if (*requested, *minimum, *kind) == expected
      ),
      "{max_length} {query_cap}: {:?}",
      refused.err()
    );
  }
  assert!(WindowedTokenizer::from_file(&path, 12, PadToken::Id(0), 6, 2).is_ok());
  let bert = fixtures().join("tiny_tokenizer.json");
  let refused = WindowedTokenizer::from_file(&bert, 16, PadToken::Id(0), 6, 2);
  assert!(
    matches!(&refused, Err(CanddingError::Tokenizer(m)) if m.contains("two separators")),
    "{:?}",
    refused.err()
  );
}

#[test]
fn a_batch_of_windows_pads_to_the_longest() {
  let fixture = fixture();
  let case = &fixture.cases[2];
  let tokenizer = tokenizer(&fixture, case);
  let mut windows = tokenizer.windows(&case.query, "").unwrap();
  windows.extend(tokenizer.windows(&case.query, &case.document).unwrap());
  let batch = tokenizer.batch(&windows, &Device::Cpu).unwrap();
  let ids = batch.input_ids.to_vec2::<u32>().unwrap();
  let mask = batch.attention_mask.to_vec2::<u32>().unwrap();
  let longest = windows.iter().map(|w| w.len()).max().unwrap();
  assert_eq!(
    batch.lengths,
    windows.iter().map(|w| w.len()).collect::<Vec<_>>()
  );
  for (row, window) in windows.iter().enumerate() {
    assert_eq!(ids[row].len(), longest);
    assert_eq!(&ids[row][..window.len()], window.get_ids());
    assert!(ids[row][window.len()..].iter().all(|&id| id == 0));
    assert_eq!(mask[row].iter().filter(|&&m| m == 1).count(), window.len());
  }
}