candding 0.1.0

//! Pair encoding: `Tokenizer::encode((query, document), true)` already builds the two-segment sequence through the
//! published post-processor; this module adds the per-checkpoint truncation policy.

use candle_core::Device;
use tokenizers::Encoding;
use tokenizers::PostProcessor;
use tokenizers::Tokenizer;
use tokenizers::TruncationDirection;
use tokenizers::pad_encodings;

use super::EncodedBatch;
use super::TextTokenizer;
use super::added_tokens;
use super::encodings_to_batch;
use super::real_tokens;
use super::tokenizer_error;
use super::too_small;
use crate::error::CanddingError;
use crate::error::Result;

/// How a pair longer than `max_length` is cut to fit; each rerank entry names its own.
#[derive(Clone, Copy, Debug, PartialEq, Eq, serde::Serialize, serde::Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum PairTruncationPolicy {
  /// The tokenizer's default: remove one token at a time from whichever side is longer until the pair fits.
  LongestFirst,
  /// FlagEmbedding's policy: the query alone is cut to three quarters of `max_length`, then only the document, to
  /// zero tokens if need be; at a `max_length` of 12 or less the pair can still exceed it, as in the reference.
  PreserveQueryToThreeQuarters,
}

impl TextTokenizer {
  /// Refuses a `max_length` with no room for one real token beside a pair's special tokens (three for BERT, four
  /// for RoBERTa and XLM-RoBERTa), more than the two `from_file` counts for a text.
  pub fn require_pair_room(&self) -> Result<()> {
    let added = added_tokens(&self.inner, true);
    if self.max_length <= added {
      return Err(too_small(self.max_length, added, true));
    }
    Ok(())
  }

  /// Encodes (query, document) pairs into one padded batch; `policy` is per call because one tokenizer serves
  /// every entry's `default_max_length` and each entry names its own policy.
  pub fn encode_pairs(
    &self,
    pairs: &[(&str, &str)],
    policy: PairTruncationPolicy,
    device: &Device,
  ) -> Result<EncodedBatch> {
    self.require_pair_room()?;
    let mut encodings = pairs
      .iter()
      .map(|&(query, document)| encode_pair(&self.inner, self.max_length, query, document, policy))
      .collect::<Result<Vec<_>>>()?;
    pad_encodings(&mut encodings, &self.padding).map_err(tokenizer_error)?;
    encodings_to_batch(&encodings, device)
  }

  /// Token counts per pair after `policy`, special tokens included: `encode_pairs`' lengths, with nothing padded.
  pub fn count_pairs(
    &self,
    pairs: &[(&str, &str)],
    policy: PairTruncationPolicy,
  ) -> Result<Vec<usize>> {
    self.require_pair_room()?;
    pairs
      .iter()
      .map(|&(query, document)| {
        encode_pair(&self.inner, self.max_length, query, document, policy)
          .map(|encoding| real_tokens(encoding.get_attention_mask()))
      })
      .collect()
  }
}

/// Encodes one (query, document) pair per `policy`; `encode_pairs` maps it over a batch.
pub(super) fn encode_pair(
  tokenizer: &Tokenizer,
  max_length: usize,
  query: &str,
  document: &str,
  policy: PairTruncationPolicy,
) -> Result<Encoding> {
  match policy {
    PairTruncationPolicy::LongestFirst => tokenizer
      .encode((query, document), true)
      .map_err(tokenizer_error),
    PairTruncationPolicy::PreserveQueryToThreeQuarters => {
      encode_pair_preserving_query(tokenizer, max_length, query, document)
    }
  }
}

/// FlagEmbedding's three stages: the query cut alone, the document alone, then the document cut to what the budget
/// leaves, by hand because `truncate_encodings` with `OnlySecond` refuses to cut a document to zero.
fn encode_pair_preserving_query(
  tokenizer: &Tokenizer,
  max_length: usize,
  query: &str,
  document: &str,
) -> Result<Encoding> {
  let mut query_encoding = tokenizer.encode(query, false).map_err(tokenizer_error)?;
  query_encoding.truncate((max_length * 3) / 4, 0, TruncationDirection::Right);
  let mut document_encoding = tokenizer.encode(document, false).map_err(tokenizer_error)?;
  let budget = max_length.saturating_sub(added_tokens(tokenizer, true));
  document_encoding.truncate(
    budget.saturating_sub(query_encoding.len()),
    0,
    TruncationDirection::Right,
  );
  with_special_tokens(tokenizer, query_encoding, document_encoding)
}

/// The post-processor's special tokens around an already cut pair, without the installed truncation, which at
/// small `max_length` would cut the query below its three-quarter share where the reference does not.
pub(super) fn with_special_tokens(
  tokenizer: &Tokenizer,
  query: Encoding,
  document: Encoding,
) -> Result<Encoding> {
  match tokenizer.get_post_processor() {
    Some(processor) => processor
      .process(query, Some(document), true)
      .map_err(tokenizer_error),
    None => <dyn PostProcessor>::default_process(vec![query, document], true)
      .map_err(tokenizer_error)?
      .pop()
      .ok_or_else(|| {
        CanddingError::Tokenizer("the default post-processor merged no encoding".to_string())
      }),
  }
}