candding 0.1.0

//! A pair too long for one sequence, cut into windows each scored as its own pair, reproducing jina-reranker-v2's
//! `reranker_tokenize_preproc` id for id: a capped query plus document windows of `max_length - query - 2`.

use std::path::Path;

use candle_core::Device;
use tokenizers::Encoding;
use tokenizers::PaddingDirection;
use tokenizers::PaddingParams;
use tokenizers::PaddingStrategy;
use tokenizers::Token;
use tokenizers::Tokenizer;
use tokenizers::TruncationDirection;
use tokenizers::pad_encodings;

use super::EncodedBatch;
use super::PadToken;
use super::added_tokens;
use super::encodings_to_batch;
use super::pad_id;
use super::pair::with_special_tokens;
use super::tokenizer_error;
use crate::error::CanddingError;
use crate::error::LengthKind;
use crate::error::Result;

/// Encodes (query, document) pairs as windows of the document, each a pair with the query.
pub struct WindowedTokenizer {
  /// Pads and truncates nothing: the query cap and the windows are this tokenizer's cuts, never the tokenizer's own.
  inner: Tokenizer,
  padding: PaddingParams,
  max_length: usize,
  query_cap: usize,
  overlap: usize,
  /// The special tokens the published post-processor adds to a text alone, which the query cap counts.
  single: usize,
}

impl WindowedTokenizer {
  /// `query_cap` counts the query's ids with its special tokens, `overlap` what consecutive windows share; refuses
  /// caps and lengths with no room for a token, and a pair template other than the text's with two separators.
  pub fn from_file(
    path: &Path,
    max_length: usize,
    pad: PadToken,
    query_cap: usize,
    overlap: usize,
  ) -> Result<Self> {
    let mut inner = Tokenizer::from_file(path).map_err(tokenizer_error)?;
    inner.with_padding(None);
    inner.with_truncation(None).map_err(tokenizer_error)?;
    let (single, pair) = (added_tokens(&inner, false), added_tokens(&inner, true));
    if pair != single + 2 {
      return Err(CanddingError::Tokenizer(format!(
        "a windowed input needs a pair template that adds two separators to a text's: this tokenizer adds {single} special tokens to a text and {pair} to a pair"
      )));
    }
    if query_cap <= single {
      return Err(CanddingError::LengthTooSmall {
        requested: query_cap,
        minimum: single + 1,
        kind: LengthKind::QueryCap { added: single },
      });
    }
    let minimum = (2 * query_cap).max(query_cap + 3);
    if max_length < minimum {
      return Err(CanddingError::LengthTooSmall {
        requested: max_length,
        minimum,
        kind: LengthKind::Windowed { query_cap },
      });
    }
    let pad_id = pad_id(&inner, pad);
    let pad_token = inner
      .id_to_token(pad_id)
      .unwrap_or_else(|| "[PAD]".to_string());
    Ok(Self {
      inner,
      padding: PaddingParams {
        strategy: PaddingStrategy::BatchLongest,
        direction: PaddingDirection::Right,
        pad_to_multiple_of: None,
        pad_id,
        pad_type_id: 0,
        pad_token,
      },
      max_length,
      query_cap,
      overlap,
      single,
    })
  }

  pub fn max_length(&self) -> usize {
    self.max_length
  }

  /// Every window of `document` as a pair with `query`, in document order.
  pub fn windows(&self, query: &str, document: &str) -> Result<Vec<Encoding>> {
    let mut query = self.inner.encode(query, false).map_err(tokenizer_error)?;
    query.truncate(self.query_cap - self.single, 0, TruncationDirection::Right);
    query.take_overflowing();
    let budget = self.max_length - (query.len() + self.single) - 2;
    let document = self
      .inner
      .encode(document, false)
      .map_err(tokenizer_error)?;
    window_spans(document.len(), budget, self.overlap.min(budget / 4))
      .into_iter()
      .map(|(start, end)| {
        with_special_tokens(&self.inner, query.clone(), slice(&document, start, end))
      })
      .collect()
  }

  /// `windows`' ids, one list per window.
  pub fn window_ids(&self, query: &str, document: &str) -> Result<Vec<Vec<u32>>> {
    Ok(
      self
        .windows(query, document)?
        .iter()
        .map(|window| window.get_ids().to_vec())
        .collect(),
    )
  }

  /// One right-padded batch of windows, from any documents.
  pub fn batch(&self, windows: &[Encoding], device: &Device) -> Result<EncodedBatch> {
    let mut windows = windows.to_vec();
    pad_encodings(&mut windows, &self.padding).map_err(tokenizer_error)?;
    encodings_to_batch(&windows, device)
  }
}

/// `[start, end)` of each window: `budget` tokens every `budget - overlap`, then the document's last `budget`
/// tokens, so the last window overlaps the previous one by `overlap` or more.
fn window_spans(tokens: usize, budget: usize, overlap: usize) -> Vec<(usize, usize)> {
  if tokens <= budget {
    return vec![(0, tokens)];
  }
  let mut spans = Vec::new();
  let mut start = 0;
  loop {
    let end = start + budget;
    if end >= tokens {
      spans.push((tokens - budget, tokens));
      return spans;
    }
    spans.push((start, end));
    start = end - overlap;
  }
}

/// Tokens `start..end` of a document's encoding as an encoding of their own.
fn slice(document: &Encoding, start: usize, end: usize) -> Encoding {
  let tokens = (start..end)
    .map(|index| {
      Token::new(
        document.get_ids()[index],
        document.get_tokens()[index].clone(),
        document.get_offsets()[index],
      )
    })
    .collect();
  Encoding::from_tokens(tokens, 0)
}