candding 0.1.0

//! Compact static form of a rerank catalog entry and its public runtime form; a rerank head outputs one number.

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

use crate::models::RerankHead;
use crate::models::jina_bert::AlibiSlopeVariant;
use crate::models::xlm_roberta_flash::PositionScheme;
use crate::registry::BackendSupport;
use crate::registry::CellException;
use crate::registry::DtypeBackendSupport;
use crate::registry::DtypeSupport;
use crate::registry::Family;
use crate::registry::Support;
use crate::registry::WeightsFormat;
use crate::registry::rerank::RerankInput;
use crate::registry::support::dtype_serde;
use crate::rerank::ScoreActivation;

/// Compact static form of a rerank catalog entry. It stores no projection width: every rerank head reads the
/// encoder's hidden size and writes one number, from tensors inside the checkpoint's own weights.
pub(crate) struct RerankEntry {
  pub id: &'static str,
  pub family: Family,
  pub head: RerankHead,
  pub default_max_length: usize,
  /// Raw logit or sigmoid, as the checkpoint's reference reports it; the golden suite checks it per fixture.
  pub score_activation: ScoreActivation,
  /// What the model reads for a pair (segment pair, prompt or windows); the golden suite checks it per fixture.
  pub input: RerankInput,
  /// The ALiBi slope formula of a `Family::JinaBert` encoder: `None` for the canonical one and for every encoder
  /// without ALiBi, `HalvedInterpolatedTail` for `jinaai/jina-reranker-v1-tiny-en`.
  pub alibi_variant: Option<AlibiSlopeVariant>,
  /// How a flash XLM-RoBERTa checkpoint numbers its tokens, `None` for every other family; stated per entry because
  /// the two implementation repositories disagree, and the encoder's load checks it against `config.json`.
  pub position_scheme: Option<PositionScheme>,
  /// The dtype `TextCrossEncoderBuilder::build` loads when the caller sets none, through the same
  /// `resolve_dtype` the dense builder calls, so a BF16 default on the CPU loads F32 instead.
  pub default_dtype: DType,
  pub cpu: Support,
  pub metal: Support,
  pub cuda: Support,
  pub dtypes: DtypeSupport,
  pub license: &'static str,
  pub size_gb: f32,
  pub phase: u8,
  pub weights: WeightsFormat,
  pub gated: bool,
  pub accel_only_dtypes: &'static [DType],
  /// Cells a run decided that the derivation cannot state, written over its rows (`CellException`).
  pub cell_exceptions: &'static [CellException],
}

impl RerankEntry {
  pub fn descriptor(&self) -> RerankDescriptor {
    let backends = BackendSupport {
      cpu: self.cpu,
      metal: self.metal,
      cuda: self.cuda,
    };
    RerankDescriptor {
      id: self.id.to_string(),
      family: self.family,
      head: self.head,
      default_max_length: self.default_max_length,
      score_activation: self.score_activation,
      input: self.input.clone(),
      alibi_variant: self.alibi_variant,
      position_scheme: self.position_scheme,
      default_dtype: self.default_dtype,
      backends,
      dtype_backends: DtypeBackendSupport::derive(backends, self.dtypes, self.accel_only_dtypes)
        .with_exceptions(self.cell_exceptions),
      license: self.license.to_string(),
      size_gb: self.size_gb,
      phase: self.phase,
      weights: self.weights,
      gated: self.gated,
    }
  }

  /// Whether `dtype` is correct only on an accelerator, from the same `dtype_backends` row the builder refuses from.
  pub fn needs_accelerator(&self, dtype: DType) -> bool {
    self
      .descriptor()
      .dtype_backends
      .row(dtype)
      .needs_accelerator()
  }

  /// This entry as `registry::support::rules` reads it.
  #[cfg(test)]
  pub(crate) fn cells(&self) -> crate::registry::support::rules::Cells<'static> {
    let descriptor = self.descriptor();
    crate::registry::support::rules::Cells {
      id: self.id,
      backends: descriptor.backends,
      dtypes: self.dtypes,
      accel_only: self.accel_only_dtypes,
      exceptions: self.cell_exceptions,
      rows: descriptor.dtype_backends,
    }
  }
}

/// Everything the library knows about a rerank model beyond its weights, as `describe --json` and
/// `list-models --json` print it.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct RerankDescriptor {
  pub id: String,
  pub family: Family,
  pub head: RerankHead,
  pub default_max_length: usize,
  pub score_activation: ScoreActivation,
  /// What the model reads for a query and a document, serialized as a `kind`-tagged object.
  pub input: RerankInput,
  pub alibi_variant: Option<AlibiSlopeVariant>,
  /// How a flash XLM-RoBERTa encoder numbers tokens, `rotary` or `absolute`; `null` for every other family.
  pub position_scheme: Option<PositionScheme>,
  /// The dtype the builder loads when none is set, serialized as `f32`, `f16` or `bf16`.
  #[serde(with = "dtype_serde")]
  pub default_dtype: DType,
  pub backends: BackendSupport,
  /// Per-dtype support broken down by device, the same shape `ModelDescriptor::dtype_backends` carries.
  pub dtype_backends: DtypeBackendSupport,
  pub license: String,
  pub size_gb: f32,
  pub phase: u8,
  pub weights: WeightsFormat,
  pub gated: bool,
}

impl RerankDescriptor {
  /// The facts that decide a reranker's score. A JinaBERT reranker takes the halved ALiBi tail both published ones
  /// bundle, a flash one absolute positions, which its load checks; the rest as in `ModelDescriptor::new`.
  pub fn new(
    id: impl Into<String>,
    family: Family,
    head: RerankHead,
    score_activation: ScoreActivation,
    input: RerankInput,
  ) -> Self {
    Self {
      id: id.into(),
      family,
      head,
      default_max_length: 512,
      score_activation,
      input,
      alibi_variant: (family == Family::JinaBert)
        .then_some(AlibiSlopeVariant::HalvedInterpolatedTail),
      position_scheme: (family == Family::XlmRobertaFlash).then_some(PositionScheme::Absolute),
      default_dtype: DType::F32,
      backends: BackendSupport::untested(),
      dtype_backends: DtypeBackendSupport::untested(),
      license: "unknown".to_string(),
      size_gb: 0.0,
      phase: u8::MAX,
      weights: WeightsFormat::Safetensors,
      gated: false,
    }
  }
}