candding 0.1.0

//! The tokenizer a `TextCrossEncoder` encodes with, one per `RerankInput` kind, built from the descriptor's input:
//! one row per pair for a pair or prompt input, one row per document window for a windowed one.

use std::path::Path;

use candle_core::Device;
use tokenizers::Encoding;

use crate::error::Result;
use crate::registry::rerank::RerankInput;
use crate::tokenizer::EncodedBatch;
use crate::tokenizer::PadToken;
use crate::tokenizer::PairTruncationPolicy;
use crate::tokenizer::PromptTokenizer;
use crate::tokenizer::TextTokenizer;
use crate::tokenizer::WindowedTokenizer;

pub(crate) enum InputTokenizer {
  /// A pair or a prompt input: one row per (query, document) pair.
  Rows(RowTokenizer),
  /// A windowed input: one row per window of each document, a document scored by its best window.
  Windowed(WindowedTokenizer),
}

pub(crate) enum RowTokenizer {
  Pair {
    tokenizer: TextTokenizer,
    truncation: PairTruncationPolicy,
  },
  Prompt {
    tokenizer: PromptTokenizer,
    default_instruction: String,
  },
}

impl InputTokenizer {
  /// A `max_length` too small for the input kind is refused here, with that kind's own minimum: a pair's special
  /// tokens plus one, a prompt template's prefix and suffix plus one, or twice a windowed input's query cap.
  pub(crate) fn from_file(
    input: &RerankInput,
    path: &Path,
    max_length: usize,
    pad: PadToken,
  ) -> Result<Self> {
    Ok(match input {
      RerankInput::Pair { truncation } => InputTokenizer::Rows(RowTokenizer::Pair {
        tokenizer: TextTokenizer::from_file_for_pairs(path, max_length, pad)?,
        truncation: *truncation,
      }),
      RerankInput::Prompt {
        template,
        default_instruction,
      } => InputTokenizer::Rows(RowTokenizer::Prompt {
        tokenizer: PromptTokenizer::from_file(path, max_length, pad, *template)?,
        default_instruction: default_instruction.to_string(),
      }),
      RerankInput::Windowed { query_cap, overlap } => InputTokenizer::Windowed(
        WindowedTokenizer::from_file(path, max_length, pad, *query_cap, *overlap)?,
      ),
    })
  }

  pub(crate) fn max_length(&self) -> usize {
    match self {
      InputTokenizer::Rows(RowTokenizer::Pair { tokenizer, .. }) => tokenizer.max_length(),
      InputTokenizer::Rows(RowTokenizer::Prompt { tokenizer, .. }) => tokenizer.max_length(),
      InputTokenizer::Windowed(tokenizer) => tokenizer.max_length(),
    }
  }

  /// Each pair's token count as the model reads it, with nothing padded: its row's length, or the total over a
  /// windowed document's windows.
  pub(crate) fn count(
    &self,
    instruction: Option<&str>,
    pairs: &[(&str, &str)],
  ) -> Result<Vec<usize>> {
    match self {
      InputTokenizer::Rows(rows) => rows.count(instruction, pairs),
      InputTokenizer::Windowed(tokenizer) => pairs
        .iter()
        .map(|&(query, document)| {
          Ok(
            tokenizer
              .windows(query, document)?
              .iter()
              .map(Encoding::len)
              .sum(),
          )
        })
        .collect(),
    }
  }

  /// Each pair's row count: one for a pair or a prompt, a document's window count for a windowed input.
  pub(crate) fn window_count(&self, pairs: &[(&str, &str)]) -> Result<Vec<usize>> {
    match self {
      InputTokenizer::Rows(_) => Ok(vec![1; pairs.len()]),
      InputTokenizer::Windowed(tokenizer) => pairs
        .iter()
        .map(|&(query, document)| Ok(tokenizer.windows(query, document)?.len()))
        .collect(),
    }
  }
}

impl RowTokenizer {
  /// One padded batch of (query, document) pairs; a prompt input renders `instruction` or its default, and an
  /// empty one renders empty. Pair inputs never get an instruction here.
  pub(crate) fn encode(
    &self,
    instruction: Option<&str>,
    pairs: &[(&str, &str)],
    device: &Device,
  ) -> Result<EncodedBatch> {
    match self {
      RowTokenizer::Pair {
        tokenizer,
        truncation,
      } => tokenizer.encode_pairs(pairs, *truncation, device),
      RowTokenizer::Prompt {
        tokenizer,
        default_instruction,
      } => tokenizer.encode(instruction.unwrap_or(default_instruction), pairs, device),
    }
  }

  /// Each pair's token count as `encode` builds it for the same instruction, with nothing padded.
  fn count(&self, instruction: Option<&str>, pairs: &[(&str, &str)]) -> Result<Vec<usize>> {
    match self {
      RowTokenizer::Pair {
        tokenizer,
        truncation,
      } => tokenizer.count_pairs(pairs, *truncation),
      RowTokenizer::Prompt {
        tokenizer,
        default_instruction,
      } => tokenizer.count(instruction.unwrap_or(default_instruction), pairs),
    }
  }
}