candding 0.1.0

//! The rerank trait and its dispatch, a fourth output kind: one score per (query, document) pair. `load_rerank`
//! builds the head a `RerankHead` names: `BertPooler`, `RobertaClassificationHead` or `LogitScore`.

mod bert_pooler;
#[cfg(test)]
pub(crate) mod flash_fixture;
mod logit_score;
mod roberta_classification;

use candle_core::DType;
use serde::Deserialize;
use serde::Serialize;

use super::EncoderInput;
use super::jina_bert;
use super::xlm_roberta_flash;
use crate::error::Result;
use crate::rerank::RerankScore;
use crate::rerank::ScoreActivation;

/// A rerank head: one `RerankScore` per row of a pair-tokenized batch, with no separate query entry point, since the
/// tokenizer, not the model, tells the query from the document.
pub trait RerankModel: Send + Sync {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<RerankScore>>;
  fn max_position_embeddings(&self) -> usize;
  /// `Encoder::weights_dtype` of the encoder the head runs on.
  fn weights_dtype(&self) -> DType;
}

/// Which rerank head a checkpoint runs, serialized with a `kind` tag such as `{"kind": "bert_pooler"}`.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(tag = "kind", rename_all = "snake_case")]
#[non_exhaustive]
pub enum RerankHead {
  /// `tanh(Linear(hidden, hidden))` over the first position, then `Linear(hidden, 1)`, from the checkpoint's
  /// `bert.pooler.dense` and `classifier`.
  BertPooler,
  /// The same arithmetic from `classifier.dense` and `classifier.out_proj`, with no pooler tensor.
  RobertaClassificationHead,
  /// The logit of `true_token_id` minus that of `false_token_id` at the prompt's last position: the Qwen3 rerankers.
  LogitScore {
    true_token_id: u32,
    false_token_id: u32,
    unembedding: Unembedding,
  },
}

/// Which matrix a `LogitScore` head reads its two rows from: the tied embedding table, or `lm_head.weight` for
/// `Qwen/Qwen3-Reranker-8B`; a `config.json` that disagrees is refused at load.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum Unembedding {
  Tied,
  Untied,
}

impl Unembedding {
  /// The serde name, `tied` or `untied`, which errors print.
  pub fn name(self) -> &'static str {
    match self {
      Unembedding::Tied => "tied",
      Unembedding::Untied => "untied",
    }
  }
}

/// The per-checkpoint facts `load_rerank` builds a head from: a `models`-side copy of the registry entry's
/// fields, since `models` cannot depend on `registry`.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub struct RerankConventions {
  pub head: RerankHead,
  /// `registry::rerank::RerankEntry::score_activation`: whether the head returns its raw logit or a sigmoid of it.
  pub score_activation: ScoreActivation,
  /// `registry::rerank::RerankEntry::alibi_variant`: the ALiBi slope formula of a `Family::JinaBert` encoder,
  /// `None` for the family's canonical one and for every encoder without ALiBi.
  pub alibi_variant: Option<jina_bert::AlibiSlopeVariant>,
  /// How a flash XLM-RoBERTa encoder numbers its tokens, `None` for every other family; the encoder's load
  /// refuses a `config.json` naming another scheme or none.
  pub position_scheme: Option<xlm_roberta_flash::PositionScheme>,
}

/// Builds a head from the resolved architecture key, a string because `models` cannot depend on the registry, handing
/// each field of `conventions` to the part that reads it; the match stays exhaustive over `RerankHead`.
pub fn load_rerank(
  model_type: &str,
  config_json: &str,
  vb: candle_nn::VarBuilder,
  conventions: RerankConventions,
) -> Result<Box<dyn RerankModel>> {
  match conventions.head {
    RerankHead::BertPooler => Ok(Box::new(bert_pooler::BertPoolerRerank::load(
      model_type,
      config_json,
      vb,
      conventions.score_activation,
      conventions.alibi_variant,
    )?)),
    RerankHead::RobertaClassificationHead => Ok(Box::new(
      roberta_classification::RobertaClassificationRerank::load(
        model_type,
        config_json,
        vb,
        conventions.score_activation,
        conventions.position_scheme,
      )?,
    )),
    RerankHead::LogitScore {
      true_token_id,
      false_token_id,
      unembedding,
    } => Ok(Box::new(logit_score::LogitScoreRerank::load(
      model_type,
      config_json,
      vb,
      [true_token_id, false_token_id],
      unembedding,
      conventions.score_activation,
    )?)),
  }
}

#[cfg(test)]
mod tests;