candding 0.1.0

//! A reranker input rendered as one prompt: the Qwen3 rerankers tokenize prefix, content and suffix as one string
//! through the published `tokenizer.json`, with no pair post-processor and no token types.

use std::path::Path;

use candle_core::Device;
use candle_core::Tensor;
use serde::Deserialize;
use serde::Serialize;
use tokenizers::Tokenizer;

use super::EncodedBatch;
use super::PadToken;
use super::pad_id;
use super::tokenizer_error;
use crate::error::CanddingError;
use crate::error::LengthKind;
use crate::error::Result;

/// A named prompt template: the fixed text a prompt reranker's references render around its three slots, plain
/// strings because the checkpoints' `chat_template.jinja` is byte-identical across the three sizes.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum PromptTemplate {
  /// The Qwen3 rerankers' template: a system turn, a user turn with the three slots, an empty think block; it ends
  /// in two newlines after `</think>`, as Jinja drops the template file's third.
  Qwen3Reranker,
}

const QWEN3_RERANKER_PREFIX: &str = "<|im_start|>system\nJudge whether the Document meets the requirements based on the Query and the Instruct provided. Note that the answer can only be \"yes\" or \"no\".<|im_end|>\n<|im_start|>user\n";
const QWEN3_RERANKER_SUFFIX: &str = "<|im_end|>\n<|im_start|>assistant\n<think>\n\n</think>\n\n";

impl PromptTemplate {
  /// The serde name, `qwen3_reranker`, for error messages.
  pub fn name(self) -> &'static str {
    match self {
      PromptTemplate::Qwen3Reranker => "qwen3_reranker",
    }
  }

  /// The fixed text before the content.
  pub fn prefix(self) -> &'static str {
    match self {
      PromptTemplate::Qwen3Reranker => QWEN3_RERANKER_PREFIX,
    }
  }

  /// The fixed text after the content, which the cut keeps whatever else it drops.
  pub fn suffix(self) -> &'static str {
    match self {
      PromptTemplate::Qwen3Reranker => QWEN3_RERANKER_SUFFIX,
    }
  }

  /// The three slots, filled verbatim. A control-token string typed inside a slot becomes that control token when
  /// the prompt is tokenized, as on both references.
  pub fn content(self, instruction: &str, query: &str, document: &str) -> String {
    match self {
      PromptTemplate::Qwen3Reranker => {
        format!("<Instruct>: {instruction}\n<Query>: {query}\n<Document>: {document}")
      }
    }
  }

  /// The whole prompt: the prefix, the content and the suffix.
  pub fn render(self, instruction: &str, query: &str, document: &str) -> String {
    [
      self.prefix(),
      &self.content(instruction, query, document),
      self.suffix(),
    ]
    .concat()
  }
}

/// Encodes (query, document) pairs as prompts of one template; prefix and suffix lengths are measured with the
/// loaded tokenizer, and a prompt over `max_length` keeps its first `max_length - suffix` ids, then the suffix.
pub struct PromptTokenizer {
  /// Pads and truncates nothing: `encode` pads the batch itself, and the cut is `cut`'s, not the tokenizer's.
  inner: Tokenizer,
  template: PromptTemplate,
  max_length: usize,
  pad_id: u32,
  prefix_len: usize,
  suffix_ids: Vec<u32>,
}

impl PromptTokenizer {
  /// Refuses a `max_length` that leaves no room for one content token between the template's prefix and suffix,
  /// 49 for the Qwen3 template with its published tokenizer.
  pub fn from_file(
    path: &Path,
    max_length: usize,
    pad: PadToken,
    template: PromptTemplate,
  ) -> 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 prefix_len = encode(&inner, template.prefix())?.len();
    let suffix_ids = encode(&inner, template.suffix())?;
    let minimum = prefix_len + suffix_ids.len() + 1;
    if max_length < minimum {
      return Err(CanddingError::LengthTooSmall {
        requested: max_length,
        minimum,
        kind: LengthKind::Prompt {
          template,
          prefix: prefix_len,
          suffix: suffix_ids.len(),
        },
      });
    }
    let pad_id = pad_id(&inner, pad);
    Ok(Self {
      inner,
      template,
      max_length,
      pad_id,
      prefix_len,
      suffix_ids,
    })
  }

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

  pub fn template(&self) -> PromptTemplate {
    self.template
  }

  /// How many ids the template's prefix encodes to with this tokenizer.
  pub fn prefix_len(&self) -> usize {
    self.prefix_len
  }

  /// The ids the template's suffix encodes to with this tokenizer, which every prompt ends in.
  pub fn suffix_ids(&self) -> &[u32] {
    &self.suffix_ids
  }

  /// One prompt's ids, cut to `max_length`; refused when they do not end in the suffix ids, which would mean the
  /// tokenizer merged the suffix into the text or the template and rendering disagree.
  pub fn prompt_ids(&self, instruction: &str, query: &str, document: &str) -> Result<Vec<u32>> {
    self.rendered_ids(&self.template.render(instruction, query, document))
  }

  /// `prompt_ids` for an already rendered prompt, the step the suffix check guards.
  pub(super) fn rendered_ids(&self, rendered: &str) -> Result<Vec<u32>> {
    let ids = encode(&self.inner, rendered)?;
    if !ids.ends_with(&self.suffix_ids) {
      let tail = ids.len().saturating_sub(self.suffix_ids.len());
      return Err(CanddingError::PromptSuffixMismatch {
        template: self.template.name().to_string(),
        expected: self.suffix_ids.clone(),
        found: ids[tail..].to_vec(),
      });
    }
    Ok(cut(ids, self.max_length, &self.suffix_ids))
  }

  /// One right-padded batch: `lengths` is each prompt's real length and `token_type_ids` is all zeros.
  pub fn encode(
    &self,
    instruction: &str,
    pairs: &[(&str, &str)],
    device: &Device,
  ) -> Result<EncodedBatch> {
    let rows = pairs
      .iter()
      .map(|&(query, document)| self.prompt_ids(instruction, query, document))
      .collect::<Result<Vec<_>>>()?;
    rows_to_batch(&rows, self.pad_id, device)
  }

  /// Each prompt's length after the cut, from the same `prompt_ids` `encode` stacks, with nothing padded.
  pub fn count(&self, instruction: &str, pairs: &[(&str, &str)]) -> Result<Vec<usize>> {
    pairs
      .iter()
      .map(|&(query, document)| Ok(self.prompt_ids(instruction, query, document)?.len()))
      .collect()
  }
}

/// No special tokens: the Qwen3 rerankers' post-processor is a bare `ByteLevel` that adds none, and a
/// post-processor that did add some would put tokens outside the prompt both references tokenize.
fn encode(tokenizer: &Tokenizer, text: &str) -> Result<Vec<u32>> {
  let encoding = tokenizer.encode(text, false).map_err(tokenizer_error)?;
  Ok(encoding.get_ids().to_vec())
}

/// Both references' cut: the first `max_length - suffix` ids, then the suffix; `PromptTokenizer::from_file`
/// guarantees `max_length` is larger than the suffix.
fn cut(mut ids: Vec<u32>, max_length: usize, suffix: &[u32]) -> Vec<u32> {
  if ids.len() > max_length {
    ids.truncate(max_length - suffix.len());
    ids.extend_from_slice(suffix);
  }
  ids
}

/// Stacks rows of unequal length, right-padded with `pad_id` to the longest.
fn rows_to_batch(rows: &[Vec<u32>], pad_id: u32, device: &Device) -> Result<EncodedBatch> {
  let batch = rows.len();
  let seq = rows.iter().map(Vec::len).max().unwrap_or(0);
  let mut ids = Vec::with_capacity(batch * seq);
  let mut mask = Vec::with_capacity(batch * seq);
  for row in rows {
    let padding = seq - row.len();
    ids.extend_from_slice(row);
    ids.extend(std::iter::repeat_n(pad_id, padding));
    mask.extend(std::iter::repeat_n(1u32, row.len()));
    mask.extend(std::iter::repeat_n(0u32, padding));
  }
  let input_ids = Tensor::from_vec(ids, (batch, seq), device)?;
  Ok(EncodedBatch {
    token_type_ids: input_ids.zeros_like()?,
    attention_mask: Tensor::from_vec(mask, (batch, seq), device)?,
    input_ids,
    lengths: rows.iter().map(Vec::len).collect(),
  })
}