candding 0.1.0

//! The rerank score: one number for a (query, document) pair, and how a checkpoint's head activates it.

use std::cmp::Ordering;

use serde::Deserialize;
use serde::Serialize;

/// How a checkpoint's head turns its final projection into the reported score; a logit and a probability can look
/// alike, so this is per-checkpoint registry data.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum ScoreActivation {
  /// The raw projection output, an unbounded logit.
  Identity,
  /// A sigmoid applied to the raw projection, bounding the reported score to `(0, 1)`.
  Sigmoid,
}

impl ScoreActivation {
  /// Applies the activation to a head's raw output.
  pub fn apply(self, raw: f32) -> f32 {
    match self {
      ScoreActivation::Identity => raw,
      ScoreActivation::Sigmoid => 1.0 / (1.0 + (-raw).exp()),
    }
  }
}

/// One pair's score, a logit or a probability per the checkpoint's `ScoreActivation`, aligned with the input order.
/// `Ord` puts every `NaN` below every number and orders numbers by `f32::total_cmp`, so a sort never panics.
#[derive(Clone, Copy, Debug)]
pub struct RerankScore {
  value: f32,
}

impl RerankScore {
  /// Wraps a raw score.
  pub fn new(value: f32) -> Self {
    Self { value }
  }

  /// The score as an `f32`; the checkpoint's `ScoreActivation` says whether it is a logit or a probability.
  pub fn value(&self) -> f32 {
    self.value
  }
}

impl PartialEq for RerankScore {
  fn eq(&self, other: &Self) -> bool {
    self.cmp(other) == Ordering::Equal
  }
}

impl Eq for RerankScore {}

impl PartialOrd for RerankScore {
  fn partial_cmp(&self, other: &Self) -> Option<Ordering> {
    Some(self.cmp(other))
  }
}

impl Ord for RerankScore {
  /// Every `NaN` equals every other `NaN`, whatever the sign bit or payload, and sorts below every number, `-inf`
  /// included; two numbers compare by `f32::total_cmp`. The type's own doc says why a best-first sort needs this.
  fn cmp(&self, other: &Self) -> Ordering {
    match (self.value.is_nan(), other.value.is_nan()) {
      (true, true) => Ordering::Equal,
      (true, false) => Ordering::Less,
      (false, true) => Ordering::Greater,
      (false, false) => self.value.total_cmp(&other.value),
    }
  }
}

#[cfg(test)]
mod tests;