candding 0.1.0

//! `tests/fixtures/pair_truncation.json` comes from `scripts/pair_fixture.py`, running the pair encoders the rerank goldens use.

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::PairTruncationPolicy;
use crate::tokenizer::TextTokenizer;

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

#[derive(Deserialize)]
struct Case {
  tokenizer: String,
  policy: PairTruncationPolicy,
  max_length: usize,
  query: String,
  document: String,
  ids: Vec<u32>,
  type_ids: Vec<u32>,
  attention_mask: Vec<u32>,
}

/// `(fixture, pair special tokens, smallest accepted budget, largest budget the query's share and the specials fill)`.
const TOKENIZERS: [(&str, usize, usize, usize); 2] = [
  ("tiny_tokenizer.json", 3, 4, 12),
  ("tiny_tokenizer_double_separator.json", 4, 5, 16),
];

const POLICIES: [PairTruncationPolicy; 2] = [
  PairTruncationPolicy::LongestFirst,
  PairTruncationPolicy::PreserveQueryToThreeQuarters,
];

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

fn cases() -> Vec<Case> {
  let text = fs::read_to_string(fixtures().join("pair_truncation.json")).unwrap();
  serde_json::from_str::<Fixture>(&text).unwrap().cases
}

/// Includes the budgets where the document is cut to nothing and where the encoder's pair runs past `max_length`.
#[test]
fn encode_pairs_matches_the_generators_pair_encoders_at_small_budgets() {
  let cases = cases();
  assert_eq!(cases.len(), 144);
  for case in &cases {
    let path = fixtures().join(&case.tokenizer);
    let tokenizer =
      TextTokenizer::from_file_for_pairs(&path, case.max_length, PadToken::Id(0)).unwrap();
    let label = format!(
      "{} {:?} max_length {} ({:?}, {:?})",
      case.tokenizer, case.policy, case.max_length, case.query, case.document
    );
    let batch = tokenizer
      .encode_pairs(
        &[(case.query.as_str(), case.document.as_str())],
        case.policy,
        &Device::Cpu,
      )
      .unwrap_or_else(|err| panic!("{label}: {err}"));
    assert_eq!(
      batch.input_ids.to_vec2::<u32>().unwrap()[0],
      case.ids,
      "{label}: ids"
    );
    assert_eq!(
      batch.token_type_ids.to_vec2::<u32>().unwrap()[0],
      case.type_ids,
      "{label}: type ids"
    );
    assert_eq!(
      batch.attention_mask.to_vec2::<u32>().unwrap()[0],
      case.attention_mask,
      "{label}: attention mask"
    );
  }
}

#[test]
fn the_fixture_covers_each_tokenizers_minimum_and_edge() {
  let cases = cases();
  for (tokenizer, _, minimum, edge) in TOKENIZERS {
    for policy in POLICIES {
      for budget in [minimum, edge] {
        assert!(
          cases
            .iter()
            .any(|c| c.tokenizer == tokenizer && c.policy == policy && c.max_length == budget),
          "{tokenizer} {policy:?} {budget}"
        );
      }
    }
    assert!(
      cases.iter().any(|c| {
        c.tokenizer == tokenizer
          && c.policy == PairTruncationPolicy::PreserveQueryToThreeQuarters
          && c.max_length == edge
          && !c.document.is_empty()
          && c.ids.len() == edge
      }),
      "{tokenizer}: the edge with a document cut to zero"
    );
  }
}

/// The tokenizers crate subtracts the special tokens unchecked: at four, a budget of 3 panics in a debug build.
#[test]
fn a_budget_with_no_room_for_a_pair_is_refused_with_the_pair_minimum() {
  for (tokenizer, added, minimum, _) in TOKENIZERS {
    let path = fixtures().join(tokenizer);
    for max_length in 0..minimum {
      let refused = TextTokenizer::from_file_for_pairs(&path, max_length, PadToken::Id(0));
      assert!(
        matches!(refused, Err(CanddingError::LengthTooSmall { requested, minimum: m, kind: LengthKind::Pair { added: a } })
          if requested == max_length && m == minimum && a == added),
        "{tokenizer} {max_length}: {:?}",
        refused.err()
      );
      let Ok(text_tokenizer) = TextTokenizer::from_file(&path, max_length, PadToken::Id(0)) else {
        continue;
      };
      for policy in POLICIES {
        let refused = text_tokenizer.encode_pairs(&[("a b", "c d")], policy, &Device::Cpu);
        assert!(
          matches!(refused, Err(CanddingError::LengthTooSmall { minimum: m, kind: LengthKind::Pair { .. }, .. }) if m == minimum),
          "{tokenizer} {max_length} {policy:?}: {:?}",
          refused.err()
        );
      }
    }
    assert!(TextTokenizer::from_file_for_pairs(&path, minimum, PadToken::Id(0)).is_ok());
  }
}