candding 0.1.0

use std::fs;

use tempfile::tempdir;

use super::*;

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

/// `register_marker` trims the trailing space before registering the marker.
#[test]
fn register_marker_turns_the_bare_token_into_one_piece() {
  let mut tokenizer =
    TextTokenizer::from_file(&markers_tokenizer_path(), 16, PadToken::Id(0)).unwrap();
  build::register_marker(&mut tokenizer, Some("[unused1] ")).unwrap();
  let ids = tokenizer
    .encode(&["[unused1] hello"], &Device::Cpu)
    .unwrap()
    .input_ids
    .to_vec2::<u32>()
    .unwrap();
  assert_eq!(
    ids[0],
    vec![2, 17, 4, 3],
    "[unused1] is id 17 in this fixture"
  );
}

#[test]
fn register_marker_is_a_no_op_for_none() {
  let mut tokenizer =
    TextTokenizer::from_file(&markers_tokenizer_path(), 16, PadToken::Id(0)).unwrap();
  build::register_marker(&mut tokenizer, None).unwrap();
  let ids = tokenizer
    .encode(&["[unused0] hello"], &Device::Cpu)
    .unwrap()
    .input_ids
    .to_vec2::<u32>()
    .unwrap();
  assert!(
    !ids[0].contains(&16),
    "nothing was registered: {:?}",
    ids[0]
  );
}

/// Pads with `[MASK]` (id 18 here); `formatted` prepends the marker, so this test prepends it itself.
#[test]
fn build_query_tokenizer_pads_to_a_fixed_width_when_the_descriptor_names_one() {
  let mut descriptor = stub_descriptor();
  descriptor.query_marker = Some("[unused0] ".to_string());
  descriptor.query_length = Some(8);
  let tokenizer = build::build_query_tokenizer(
    &markers_tokenizer_path(),
    16,
    PadToken::Id(0),
    &descriptor,
    512,
    None,
  )
  .unwrap();
  let batch = tokenizer
    .encode(&["[unused0] hello"], &Device::Cpu)
    .unwrap();
  assert_eq!(batch.input_ids.dims(), &[1, 8]);
  assert_eq!(
    batch.input_ids.to_vec2::<u32>().unwrap()[0],
    vec![2, 16, 4, 3, 18, 18, 18, 18],
    "CLS, [unused0], hello, SEP, then [MASK] filling the rest"
  );
  assert_eq!(
    batch.attention_mask.to_vec2::<u32>().unwrap()[0],
    vec![1, 1, 1, 1, 0, 0, 0, 0]
  );
}

/// The fallback is `[MASK]` (id 18), so naming `[unused1]` (id 17) proves the name is read from the file.
#[test]
fn build_query_tokenizer_resolves_the_mask_token_from_a_real_tokenizer_config() {
  let dir = tempdir().unwrap();
  let tokenizer_config = dir.path().join("tokenizer_config.json");
  fs::write(&tokenizer_config, r#"{"mask_token": "[unused1]"}"#).unwrap();
  let mut descriptor = stub_descriptor();
  descriptor.query_length = Some(6);
  let tokenizer = build::build_query_tokenizer(
    &markers_tokenizer_path(),
    16,
    PadToken::Id(0),
    &descriptor,
    512,
    Some(tokenizer_config.as_path()),
  )
  .unwrap();
  let batch = tokenizer.encode(&["hello"], &Device::Cpu).unwrap();
  assert_eq!(
    batch.input_ids.to_vec2::<u32>().unwrap()[0],
    vec![2, 4, 3, 17, 17, 17],
    "CLS, hello, SEP, then [unused1] (id 17) filling the rest — not [MASK] (id 18)"
  );
}

#[test]
fn build_query_tokenizer_matches_the_document_one_when_the_descriptor_names_no_query_length() {
  let descriptor = stub_descriptor();
  let path = PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/tiny_tokenizer.json");
  let tokenizer =
    build::build_query_tokenizer(&path, 16, PadToken::Id(0), &descriptor, 512, None).unwrap();
  let batch = tokenizer.encode(&["hello"], &Device::Cpu).unwrap();
  assert_eq!(
    batch.input_ids.dims(),
    &[1, 3],
    "batch-longest, not padded to a fixed width no descriptor asked for"
  );
}

/// Reports `[padded sequence width, real token count]`, which shows both the marker and the query padding.
struct SequenceShapeProbeModel;

impl MultiVectorModel for SequenceShapeProbeModel {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>> {
    let width = input.input_ids.dims()[1] as f32;
    Ok(
      input
        .lengths
        .iter()
        .map(|&len| MultiVectorEmbedding::from_rows(2, [vec![width, len as f32]]))
        .collect(),
    )
  }

  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<MultiVectorEmbedding>> {
    self.forward(input)
  }

  fn dim(&self) -> usize {
    2
  }

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

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

fn marker_descriptor() -> MultiVectorDescriptor {
  let mut descriptor = stub_descriptor();
  descriptor.query_marker = Some("[unused0] ".to_string());
  descriptor.document_marker = Some("[unused1] ".to_string());
  descriptor.query_length = Some(8);
  descriptor
}

fn marker_model(batch_size: usize) -> MultiVectorTextEmbedding {
  let path = markers_tokenizer_path();
  let mut tokenizer = TextTokenizer::from_file(&path, 16, PadToken::Id(0)).unwrap();
  build::register_marker(&mut tokenizer, Some("[unused1] ")).unwrap();
  let query_tokenizer =
    build::build_query_tokenizer(&path, 16, PadToken::Id(0), &marker_descriptor(), 512, None)
      .unwrap();
  MultiVectorTextEmbedding::new(
    Box::new(SequenceShapeProbeModel),
    tokenizer,
    query_tokenizer,
    marker_descriptor(),
    Device::Cpu,
    DType::F32,
    batch_size,
  )
}

#[test]
fn passage_embed_applies_the_document_marker_before_tokenizing() {
  let model = marker_model(8);
  let passage = model.passage_embed(&["hello"], None).unwrap();
  assert_eq!(model.embed(&["hello"], None).unwrap(), passage);
  let row: Vec<f32> = passage[0].rows().next().unwrap().to_vec();
  assert_eq!(
    row[1], 4.0,
    "CLS, [unused1], hello, SEP — the marker must reach the tokenizer through passage_embed()"
  );
}

#[test]
fn query_embed_applies_the_query_marker_and_pads_to_the_fixed_width() {
  let model = marker_model(8);
  let row: Vec<f32> = model.query_embed(&["hello"], None).unwrap()[0]
    .rows()
    .next()
    .unwrap()
    .to_vec();
  assert_eq!(
    row[0], 8.0,
    "the query tokenizer pads to query_length regardless of how short the real content is"
  );
  assert_eq!(
    row[1], 4.0,
    "CLS, [unused0], hello, SEP — the query marker must reach the tokenizer through query_embed()"
  );
}