candding 0.1.0

//! Tokenization into right-padded batches.

use std::path::Path;

use candle_core::Device;
use candle_core::Tensor;
use tokenizers::PaddingDirection;
use tokenizers::PaddingParams;
use tokenizers::PaddingStrategy;
use tokenizers::PostProcessor;
use tokenizers::Tokenizer;
use tokenizers::TruncationDirection;
use tokenizers::TruncationParams;
use tokenizers::TruncationStrategy;
use tokenizers::pad_encodings;

mod pair;
mod prompt;
mod special_tokens;
mod windowed;

pub use pair::PairTruncationPolicy;
pub use prompt::PromptTemplate;
pub use prompt::PromptTokenizer;
pub use special_tokens::PadToken;
pub use special_tokens::read_mask_token;
pub use special_tokens::read_pad_token;
pub use windowed::WindowedTokenizer;

use crate::error::CanddingError;
use crate::error::LengthKind;
use crate::error::Result;

/// `inner` never pads: `encode` pads its own batch and `count` must read unpadded lengths, yet a published
/// `tokenizer.json` can declare its own `padding`, which `Tokenizer::from_file` installs.
pub struct TextTokenizer {
  inner: Tokenizer,
  padding: PaddingParams,
  max_length: usize,
}

/// One padded batch on a device; every tensor is `(batch, seq)` in `u32`.
pub struct EncodedBatch {
  pub input_ids: Tensor,
  pub attention_mask: Tensor,
  pub token_type_ids: Tensor,
  pub lengths: Vec<usize>,
}

impl TextTokenizer {
  /// `max_length` includes special tokens; padding is on the right, to the batch's longest row. A `max_length`
  /// with no room for a real token is refused, since the tokenizers crate would underflow on it.
  pub fn from_file(path: &Path, max_length: usize, pad: PadToken) -> Result<Self> {
    Self::from_file_with_strategy(path, max_length, pad, PaddingStrategy::BatchLongest, false)
  }

  /// `from_file` for a tokenizer that encodes (query, document) pairs: a `max_length` too small for a pair is
  /// refused with the pair's own minimum, never with the smaller one a single text needs.
  pub fn from_file_for_pairs(path: &Path, max_length: usize, pad: PadToken) -> Result<Self> {
    Self::from_file_with_strategy(path, max_length, pad, PaddingStrategy::BatchLongest, true)
  }

  /// Pads every row to exactly `max_length`, not to the batch's longest: a ColBERT query is padded to its fixed
  /// width with `pad`, which for those checkpoints is the mask token the encoder embeds.
  pub fn from_file_fixed_length(path: &Path, max_length: usize, pad: PadToken) -> Result<Self> {
    Self::from_file_with_strategy(
      path,
      max_length,
      pad,
      PaddingStrategy::Fixed(max_length),
      false,
    )
  }

  /// `pair` picks the special-token count `max_length` must exceed: a pair's, which is never smaller than a
  /// text's, or a text's.
  fn from_file_with_strategy(
    path: &Path,
    max_length: usize,
    pad: PadToken,
    strategy: PaddingStrategy,
    pair: bool,
  ) -> Result<Self> {
    let mut inner = Tokenizer::from_file(path).map_err(tokenizer_error)?;
    inner.with_padding(None);
    let added = added_tokens(&inner, pair);
    if max_length <= added {
      return Err(too_small(max_length, added, pair));
    }
    let pad_id = pad_id(&inner, pad);
    let pad_token = inner
      .id_to_token(pad_id)
      .unwrap_or_else(|| "[PAD]".to_string());
    inner
      .with_truncation(Some(TruncationParams {
        max_length,
        strategy: TruncationStrategy::LongestFirst,
        stride: 0,
        direction: TruncationDirection::Right,
      }))
      .map_err(tokenizer_error)?;
    Ok(Self {
      inner,
      padding: PaddingParams {
        strategy,
        direction: PaddingDirection::Right,
        pad_to_multiple_of: None,
        pad_id,
        pad_type_id: 0,
        pad_token,
      },
      max_length,
    })
  }

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

  /// Registers `tokens` as special, non-normalized tokens, so a marker such as `[unused0]` stays one id instead of
  /// being split at its brackets before WordPiece runs.
  pub fn register_special_tokens(&mut self, tokens: &[&str]) -> Result<()> {
    let added: Vec<tokenizers::AddedToken> = tokens
      .iter()
      .map(|&t| tokenizers::AddedToken::from(t, true).normalized(false))
      .collect();
    self
      .inner
      .add_special_tokens(added)
      .map_err(tokenizer_error)?;
    Ok(())
  }

  pub fn encode(&self, texts: &[&str], device: &Device) -> Result<EncodedBatch> {
    let mut encodings = self
      .inner
      .encode_batch(texts.to_vec(), true)
      .map_err(tokenizer_error)?;
    pad_encodings(&mut encodings, &self.padding).map_err(tokenizer_error)?;
    encodings_to_batch(&encodings, device)
  }

  /// Token counts per text after truncation, special tokens included; never padded, since it sees the whole list.
  pub fn count(&self, texts: &[&str]) -> Result<Vec<usize>> {
    let encodings = self
      .inner
      .encode_batch(texts.to_vec(), true)
      .map_err(tokenizer_error)?;
    Ok(encodings.iter().map(|e| e.len()).collect())
  }
}

/// The id `pad` names in this tokenizer's vocabulary, 0 when it names none.
fn pad_id(tokenizer: &Tokenizer, pad: PadToken) -> u32 {
  match pad {
    PadToken::Id(id) => id,
    PadToken::Named(name) => tokenizer.token_to_id(&name).unwrap_or(0),
    PadToken::Unknown => 0,
  }
}

/// How many special tokens the published post-processor adds to one text, or to one pair.
fn added_tokens(tokenizer: &Tokenizer, pair: bool) -> usize {
  tokenizer
    .get_post_processor()
    .map(|processor| processor.added_tokens(pair))
    .unwrap_or(0)
}

/// The refusal for a `max_length` at or below `added` special tokens, for a pair or for a text.
fn too_small(requested: usize, added: usize, pair: bool) -> CanddingError {
  let kind = if pair {
    LengthKind::Pair { added }
  } else {
    LengthKind::Text { added }
  };
  CanddingError::LengthTooSmall {
    requested,
    minimum: added + 1,
    kind,
  }
}

fn real_tokens(mask: &[u32]) -> usize {
  mask.iter().filter(|m| **m == 1).count()
}

/// Stacks an already-padded batch of `Encoding`s into the tensors `EncodedBatch` carries.
fn encodings_to_batch(encodings: &[tokenizers::Encoding], device: &Device) -> Result<EncodedBatch> {
  let batch = encodings.len();
  let seq = encodings.first().map(|e| e.len()).unwrap_or(0);
  let mut ids = Vec::with_capacity(batch * seq);
  let mut mask = Vec::with_capacity(batch * seq);
  let mut types = Vec::with_capacity(batch * seq);
  let mut lengths = Vec::with_capacity(batch);
  for encoding in encodings {
    ids.extend_from_slice(encoding.get_ids());
    mask.extend_from_slice(encoding.get_attention_mask());
    types.extend_from_slice(encoding.get_type_ids());
    lengths.push(real_tokens(encoding.get_attention_mask()));
  }
  Ok(EncodedBatch {
    input_ids: Tensor::from_vec(ids, (batch, seq), device)?,
    attention_mask: Tensor::from_vec(mask, (batch, seq), device)?,
    token_type_ids: Tensor::from_vec(types, (batch, seq), device)?,
    lengths,
  })
}

/// An id-indexed token-text table from `tokenizer.json` (`##` marks a WordPiece continuation), for the sparse
/// heads' word reconstruction; read again because the pipeline's tokenizer is built after the model.
pub fn token_text_from_file(path: &Path) -> Result<Vec<String>> {
  let inner = Tokenizer::from_file(path).map_err(tokenizer_error)?;
  let size = inner.get_vocab_size(true);
  Ok(
    (0..size as u32)
      .map(|id| inner.id_to_token(id).unwrap_or_default())
      .collect(),
  )
}

/// Python's `string.punctuation`, the ColBERT skiplist, each character resolved through the vocabulary on its own as
/// the reference does; `$+<=>^|~` are Unicode symbols, which a punctuation-category check would miss.
const ASCII_PUNCTUATION: &str = "!\"#$%&'()*+,-./:;<=>?@[\\]^_`{|}~";

/// Vocabulary ids for `ASCII_PUNCTUATION`, the ColBERT document skiplist, read before the pipeline's tokenizer
/// exists; a character with no single-token id is left out.
pub fn ascii_punctuation_ids(path: &Path) -> Result<std::collections::HashSet<u32>> {
  let inner = Tokenizer::from_file(path).map_err(tokenizer_error)?;
  Ok(
    ASCII_PUNCTUATION
      .chars()
      .filter_map(|c| inner.token_to_id(&c.to_string()))
      .collect(),
  )
}

fn tokenizer_error(err: impl std::fmt::Display) -> CanddingError {
  CanddingError::Tokenizer(err.to_string())
}

#[cfg(test)]
mod tests;