candding 0.1.0

//! `TextCrossEncoder`: one relevance score per (query, document) pair, through the checkpoint's own rerank head.

mod build;
mod builder;
mod input;
mod windows;

pub use builder::TextCrossEncoderBuilder;
use candle_core::DType;
use candle_core::Device;
pub(crate) use input::InputTokenizer;
use input::RowTokenizer;

use super::METAL_FORWARD_LOCK;
use crate::error::Result;
use crate::models::EncoderInput;
use crate::models::RerankModel;
use crate::registry::rerank::RerankDescriptor;
use crate::rerank::RerankScore;
use crate::tokenizer::EncodedBatch;

/// A loaded reranker: one `RerankScore` per (query, document) pair, read as a segment pair, a rendered prompt or
/// one pair per window of a long document, as the descriptor's `RerankInput` says.
pub struct TextCrossEncoder {
  model: Box<dyn RerankModel>,
  tokenizer: InputTokenizer,
  descriptor: RerankDescriptor,
  device: Device,
  dtype: DType,
  batch_size: usize,
}

impl TextCrossEncoder {
  pub fn builder(spec: impl Into<String>) -> TextCrossEncoderBuilder {
    TextCrossEncoderBuilder::new(spec)
  }

  /// `tokenizer` is `InputTokenizer::from_file` over `descriptor.input`, which has already refused a `max_length`
  /// too small for that input kind, so a builder handed one fails when it builds rather than on the first call.
  pub(crate) fn new(
    model: Box<dyn RerankModel>,
    tokenizer: InputTokenizer,
    descriptor: RerankDescriptor,
    device: Device,
    dtype: DType,
    batch_size: usize,
  ) -> Self {
    Self {
      model,
      tokenizer,
      descriptor,
      device,
      dtype,
      batch_size,
    }
  }

  pub fn descriptor(&self) -> &RerankDescriptor {
    &self.descriptor
  }

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

  pub fn device(&self) -> &Device {
    &self.device
  }

  pub fn dtype(&self) -> DType {
    self.dtype
  }

  /// The dtype the loaded weights are in, read back from the encoder's token embedding table
  /// (`RerankModel::weights_dtype`); `dtype` is the one the builder resolved and asked the loader for.
  pub fn weights_dtype(&self) -> DType {
    self.model.weights_dtype()
  }

  /// One `RerankScore` per document, in input order and never sorted; a prompt entry renders its default instruction.
  pub fn rerank<S: AsRef<str>>(
    &self,
    query: &str,
    documents: &[S],
    batch_size: Option<usize>,
  ) -> Result<Vec<RerankScore>> {
    self.score(query, documents, None, batch_size)
  }

  /// `rerank` with `instruction` in a prompt entry's instruction slot, where an empty one renders empty; a pair or
  /// windowed entry refuses any instruction with `CanddingError::InstructionOnPairInput`.
  pub fn rerank_with_instruction<S: AsRef<str>>(
    &self,
    query: &str,
    documents: &[S],
    instruction: &str,
    batch_size: Option<usize>,
  ) -> Result<Vec<RerankScore>> {
    self.score(query, documents, Some(instruction), batch_size)
  }

  fn score<S: AsRef<str>>(
    &self,
    query: &str,
    documents: &[S],
    instruction: Option<&str>,
    batch_size: Option<usize>,
  ) -> Result<Vec<RerankScore>> {
    self
      .descriptor
      .input
      .check_instruction(&self.descriptor.id, instruction)?;
    if documents.is_empty() {
      return Ok(Vec::new());
    }
    let batch_size = batch_size.unwrap_or(self.batch_size).max(1);
    let rows = match &self.tokenizer {
      InputTokenizer::Rows(rows) => rows,
      InputTokenizer::Windowed(tokenizer) => {
        return self.score_windows(tokenizer, query, documents, batch_size);
      }
    };
    let mut scores = Vec::with_capacity(documents.len());
    for chunk in documents.chunks(batch_size) {
      let pairs: Vec<(&str, &str)> = chunk.iter().map(|d| (query, d.as_ref())).collect();
      scores.extend(self.score_chunk(rows, instruction, &pairs)?);
    }
    Ok(scores)
  }

  /// Each pair's token count after truncation, the query and every special or template token included; a windowed
  /// entry counts the total over the document's windows.
  pub fn pair_token_count<S: AsRef<str>>(
    &self,
    query: &str,
    documents: &[S],
  ) -> Result<Vec<usize>> {
    self.count(query, documents, None)
  }

  /// `pair_token_count` for `rerank_with_instruction`'s input: the same counts that call's rows have, and the same
  /// refusal of an instruction for a pair entry.
  pub fn pair_token_count_with_instruction<S: AsRef<str>>(
    &self,
    query: &str,
    documents: &[S],
    instruction: &str,
  ) -> Result<Vec<usize>> {
    self.count(query, documents, Some(instruction))
  }

  fn count<S: AsRef<str>>(
    &self,
    query: &str,
    documents: &[S],
    instruction: Option<&str>,
  ) -> Result<Vec<usize>> {
    self
      .descriptor
      .input
      .check_instruction(&self.descriptor.id, instruction)?;
    let pairs: Vec<(&str, &str)> = documents.iter().map(|d| (query, d.as_ref())).collect();
    self.tokenizer.count(instruction, &pairs)
  }

  /// How many rows the model reads for each document: one for a pair or a prompt entry, and for a windowed entry the
  /// document's window count, the windows `rerank` scores and `pair_token_count` totals.
  pub fn pair_window_count<S: AsRef<str>>(
    &self,
    query: &str,
    documents: &[S],
  ) -> Result<Vec<usize>> {
    let pairs: Vec<(&str, &str)> = documents.iter().map(|d| (query, d.as_ref())).collect();
    self.tokenizer.window_count(&pairs)
  }

  fn score_chunk(
    &self,
    rows: &RowTokenizer,
    instruction: Option<&str>,
    pairs: &[(&str, &str)],
  ) -> Result<Vec<RerankScore>> {
    self.forward(|device| rows.encode(instruction, pairs, device))
  }

  /// One batch through the model; on Metal the whole body runs under `METAL_FORWARD_LOCK`, as in every pipeline.
  fn forward(
    &self,
    batch: impl FnOnce(&Device) -> Result<EncodedBatch>,
  ) -> Result<Vec<RerankScore>> {
    let _guard = self.device.is_metal().then(|| {
      METAL_FORWARD_LOCK
        .lock()
        .unwrap_or_else(|poisoned| poisoned.into_inner())
    });
    let batch = batch(&self.device)?;
    let input = EncoderInput {
      input_ids: &batch.input_ids,
      attention_mask: &batch.attention_mask,
      token_type_ids: &batch.token_type_ids,
      lengths: &batch.lengths,
    };
    self.model.forward(&input)
  }
}

#[cfg(test)]
mod tests;