candding 0.1.0

use candle_core::Device;

use super::tiny;
use crate::tokenizer::PairTruncationPolicy;

fn words(n: usize) -> String {
  ["a", "b", "c", "d"]
    .iter()
    .cycle()
    .take(n)
    .copied()
    .collect::<Vec<_>>()
    .join(" ")
}

/// `[CLS] query [SEP]` is the `0` run and `document [SEP]` the `1` run in `BertProcessing`'s layout.
fn query_and_document_real_lengths(type_ids: &[u32]) -> (usize, usize) {
  let zeros = type_ids.iter().filter(|&&t| t == 0).count();
  let ones = type_ids.iter().filter(|&&t| t == 1).count();
  (zeros - 2, ones - 1)
}

/// No concatenation through the single-sequence path produces `token_type_ids` `[0,..,0,1,..,1]`.
#[test]
fn a_pair_encodes_as_a_pair_not_a_concatenation() {
  let tokenizer = tiny(64);
  let pair = tokenizer
    .encode_pairs(
      &[("a b", "c d")],
      PairTruncationPolicy::LongestFirst,
      &Device::Cpu,
    )
    .unwrap();
  let pair_types = pair.token_type_ids.to_vec2::<u32>().unwrap()[0].clone();
  assert!(
    pair_types.contains(&1),
    "a real pair's document segment must carry type id 1: {pair_types:?}"
  );
  assert_eq!(
    pair.input_ids.dims(),
    &[1, 7],
    "[CLS] a b [SEP] c d [SEP] — two SEPs, one closing each segment"
  );

  let concatenated = tokenizer.encode(&["a b c d"], &Device::Cpu).unwrap();
  let concat_types = concatenated.token_type_ids.to_vec2::<u32>().unwrap()[0].clone();
  assert!(
    concat_types.iter().all(|&t| t == 0),
    "plain concatenation as one sequence can only ever produce an all-zero token_type_ids row: {concat_types:?}"
  );
  assert_eq!(
    concatenated.input_ids.dims(),
    &[1, 6],
    "[CLS] a b c d [SEP] — one SEP; concatenation has no second segment to close"
  );
}

#[test]
fn longest_first_trims_whichever_side_is_currently_longer() {
  let tokenizer = tiny(10);
  let query = words(10);
  let document = words(2);
  let batch = tokenizer
    .encode_pairs(
      &[(query.as_str(), document.as_str())],
      PairTruncationPolicy::LongestFirst,
      &Device::Cpu,
    )
    .unwrap();
  let type_ids = &batch.token_type_ids.to_vec2::<u32>().unwrap()[0];
  assert_eq!(
    query_and_document_real_lengths(type_ids),
    (5, 2),
    "the 10-token query is the longer side and is the one trimmed, down to 5; \
     the already-short 2-token document is untouched"
  );
}

#[test]
fn preserve_query_to_three_quarters_only_trims_the_document() {
  let tokenizer = tiny(20);
  let query = words(16);
  let document = words(10);
  let batch = tokenizer
    .encode_pairs(
      &[(query.as_str(), document.as_str())],
      PairTruncationPolicy::PreserveQueryToThreeQuarters,
      &Device::Cpu,
    )
    .unwrap();
  let type_ids = &batch.token_type_ids.to_vec2::<u32>().unwrap()[0];
  assert_eq!(
    query_and_document_real_lengths(type_ids),
    (15, 2),
    "the 16-token query is capped at three quarters of 20 (15) and never reduced further; \
     the document gives up the rest of the budget, trimmed from 10 down to 2"
  );
}

#[test]
fn swapping_the_policy_swaps_which_side_survives_truncation() {
  let tokenizer = tiny(20);
  let query = words(16);
  let document = words(10);

  let longest_first = tokenizer
    .encode_pairs(
      &[(query.as_str(), document.as_str())],
      PairTruncationPolicy::LongestFirst,
      &Device::Cpu,
    )
    .unwrap();
  let preserve_query = tokenizer
    .encode_pairs(
      &[(query.as_str(), document.as_str())],
      PairTruncationPolicy::PreserveQueryToThreeQuarters,
      &Device::Cpu,
    )
    .unwrap();

  let longest_first_lengths =
    query_and_document_real_lengths(&longest_first.token_type_ids.to_vec2::<u32>().unwrap()[0]);
  let preserve_query_lengths =
    query_and_document_real_lengths(&preserve_query.token_type_ids.to_vec2::<u32>().unwrap()[0]);

  assert_ne!(
    longest_first_lengths, preserve_query_lengths,
    "the same over-long pair must truncate differently under the two policies: \
     LongestFirst {longest_first_lengths:?}, PreserveQueryToThreeQuarters {preserve_query_lengths:?}"
  );
  assert_eq!(longest_first_lengths, (9, 8));
  assert_eq!(preserve_query_lengths, (15, 2));
}

/// The long query against `a` counts 20 under `LongestFirst` and 19 under `PreserveQueryToThreeQuarters`.
#[test]
fn count_pairs_equals_the_lengths_of_the_padded_batch_under_both_policies() {
  let tokenizer = tiny(20);
  let long = words(24);
  let pairs = [
    ("a", "a"),
    ("", ""),
    (long.as_str(), "a"),
    ("a b", long.as_str()),
    (long.as_str(), long.as_str()),
  ];
  for policy in [
    PairTruncationPolicy::LongestFirst,
    PairTruncationPolicy::PreserveQueryToThreeQuarters,
  ] {
    let counts = tokenizer.count_pairs(&pairs, policy).unwrap();
    let batch = tokenizer
      .encode_pairs(&pairs, policy, &Device::Cpu)
      .unwrap();
    assert_eq!(counts, batch.lengths, "{policy:?}");
    assert_eq!(
      batch.input_ids.dims()[1],
      *counts.iter().max().unwrap(),
      "{policy:?}: the batch is padded to its longest pair"
    );
    assert!(
      counts.iter().min() < counts.iter().max(),
      "{policy:?}: {counts:?}"
    );
  }
  assert_eq!(
    tokenizer
      .count_pairs(&pairs[2..3], PairTruncationPolicy::LongestFirst)
      .unwrap(),
    vec![20]
  );
  assert_eq!(
    tokenizer
      .count_pairs(
        &pairs[2..3],
        PairTruncationPolicy::PreserveQueryToThreeQuarters
      )
      .unwrap(),
    vec![19]
  );
}