candding 0.1.0

//! A zero-layer BERT whose only nonzero weight is the classifier bias, so every pair scores that bias.

use std::collections::HashMap;
use std::fs;
use std::path::PathBuf;

use candle_core::Tensor;
use tempfile::TempDir;
use tempfile::tempdir;

use super::*;
use crate::error::LengthKind;
use crate::models::RerankHead;
use crate::registry::BackendSupport;
use crate::registry::DtypeBackendSupport;
use crate::registry::Family;
use crate::registry::WeightsFormat;
use crate::registry::rerank::RerankInput;
use crate::rerank::ScoreActivation;
use crate::tokenizer::PairTruncationPolicy;

mod flash;
mod logit_score;

/// The classifier bias and its sigmoid, far enough apart that a head handed the other activation cannot pass.
const BIAS: f32 = 2.0;
const SIGMOID_OF_BIAS: f32 = 0.880_797_1;

/// `tiny_tokenizer.json` has 16 ids and adds two special tokens to a text and three to a pair.
const CONFIG: &str = r#"{"model_type":"bert","vocab_size":16,"hidden_size":4,"num_hidden_layers":0,
  "num_attention_heads":1,"intermediate_size":1,"max_position_embeddings":32,"type_vocab_size":2,
  "layer_norm_eps":1e-12,"pad_token_id":0}"#;

fn checkpoint() -> TempDir {
  let dir = tempdir().unwrap();
  fs::write(dir.path().join("config.json"), CONFIG).unwrap();
  fs::copy(
    PathBuf::from(env!("CARGO_MANIFEST_DIR")).join("tests/fixtures/tiny_tokenizer.json"),
    dir.path().join("tokenizer.json"),
  )
  .unwrap();
  let mut tensors: HashMap<String, Tensor> = [
    ("bert.embeddings.word_embeddings.weight", vec![16, 4]),
    ("bert.embeddings.position_embeddings.weight", vec![32, 4]),
    ("bert.embeddings.token_type_embeddings.weight", vec![2, 4]),
    ("bert.embeddings.LayerNorm.weight", vec![4]),
    ("bert.embeddings.LayerNorm.bias", vec![4]),
    ("bert.pooler.dense.weight", vec![4, 4]),
    ("bert.pooler.dense.bias", vec![4]),
    ("classifier.weight", vec![1, 4]),
  ]
  .into_iter()
  .map(|(name, shape)| {
    let zeros = Tensor::zeros(shape, DType::F32, &Device::Cpu).unwrap();
    (name.to_string(), zeros)
  })
  .collect();
  tensors.insert(
    "classifier.bias".to_string(),
    Tensor::new(&[BIAS], &Device::Cpu).unwrap(),
  );
  candle_core::safetensors::save(&tensors, dir.path().join("model.safetensors")).unwrap();
  dir
}

fn descriptor(score_activation: ScoreActivation) -> RerankDescriptor {
  RerankDescriptor {
    id: "synthetic/zero-bert-reranker".to_string(),
    family: Family::Bert,
    head: RerankHead::BertPooler,
    default_max_length: 512,
    score_activation,
    input: RerankInput::Pair {
      truncation: PairTruncationPolicy::LongestFirst,
    },
    alibi_variant: None,
    position_scheme: None,
    default_dtype: DType::F32,
    backends: BackendSupport::untested(),
    dtype_backends: DtypeBackendSupport::untested(),
    license: "unknown".to_string(),
    size_gb: 0.0,
    phase: u8::MAX,
    weights: WeightsFormat::Safetensors,
    gated: false,
  }
}

fn builder(dir: &TempDir, score_activation: ScoreActivation) -> TextCrossEncoderBuilder {
  TextCrossEncoder::builder(dir.path().to_str().unwrap())
    .descriptor(descriptor(score_activation))
    .device(Device::Cpu)
}

#[test]
fn default_batch_size_is_eight() {
  let builder = TextCrossEncoderBuilder::new("org/unregistered");
  assert_eq!(builder.batch_size, 8);
}

#[test]
fn build_hands_the_head_the_descriptors_activation() {
  let dir = checkpoint();
  for (activation, expected) in [
    (ScoreActivation::Sigmoid, SIGMOID_OF_BIAS),
    (ScoreActivation::Identity, BIAS),
  ] {
    let model = builder(&dir, activation).build().unwrap();
    let documents = ["a", "the cats sat on the mat", ""];
    let scores = model.rerank("a b", &documents, None).unwrap();
    assert_eq!(scores.len(), documents.len(), "{activation:?}");
    for score in scores {
      let value = score.value();
      assert!((value - expected).abs() < 1e-6, "{activation:?}: {value}");
    }
  }
}

/// A text tokenizer would refuse lengths below 3 with its own smaller minimum; a pair tokenizer refuses below 4.
#[test]
fn build_refuses_a_max_length_below_the_pair_minimum_with_the_pair_error() {
  let dir = checkpoint();
  for max_length in 0..4 {
    let built = builder(&dir, ScoreActivation::Sigmoid)
      .max_length(max_length)
      .build();
    match built {
      Err(CanddingError::LengthTooSmall {
        requested,
        minimum,
        kind,
      }) => assert_eq!(
        (requested, minimum, kind),
        (max_length, 4, LengthKind::Pair { added: 3 })
      ),
      Err(other) => panic!("max_length {max_length}: unexpected error: {other}"),
      Ok(_) => panic!("max_length {max_length} leaves a pair no room"),
    }
  }
  let model = builder(&dir, ScoreActivation::Sigmoid)
    .max_length(4)
    .build()
    .unwrap();
  assert_eq!(model.max_length(), 4);
}

/// On the CPU a BF16 default resolves to F32, the fallback `resolve_dtype` applies to every catalog's default.
#[test]
fn build_loads_the_descriptors_default_dtype_when_none_is_requested() {
  let dir = checkpoint();
  for (default, requested, expected) in [
    (DType::F16, None, DType::F16),
    (DType::F32, None, DType::F32),
    (DType::F16, Some(DType::F32), DType::F32),
    (DType::BF16, None, DType::F32),
  ] {
    let mut descriptor = descriptor(ScoreActivation::Identity);
    descriptor.default_dtype = default;
    let mut builder = TextCrossEncoder::builder(dir.path().to_str().unwrap())
      .descriptor(descriptor)
      .device(Device::Cpu);
    if let Some(dtype) = requested {
      builder = builder.dtype(dtype);
    }
    let model = builder.build().unwrap();
    let label = format!("default {default:?}, requested {requested:?}");
    assert_eq!(model.dtype(), expected, "{label}");
    assert_eq!(model.weights_dtype(), expected, "{label}");
    let scores = model.rerank("a b", &["a", ""], None).unwrap();
    assert!(scores.iter().all(|s| s.value() == BIAS), "{label}");
  }
}