candding 0.1.0

//! Compact static form of a sparse catalog entry, and its public runtime form.

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

use crate::models::SparseHead;
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;

/// Whether a head runs an encoder (`Encoded`) or computes its weights from the text alone (`Computed`, `Qdrant/bm25`).
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
pub enum SparseComputation {
  Encoded,
  Computed,
}

/// Compact static form of a sparse catalog entry; no pooling, templates or `dense_dirs`, the head does it all.
pub(crate) struct SparseEntry {
  pub id: &'static str,
  /// The encoder's family, or `None` for a `Computed` entry, which has no encoder at all.
  pub family: Option<Family>,
  pub head: SparseHead,
  /// `0` for a `Computed` entry (an unbounded hash space); `describe` and `list-models` print `n/a`.
  pub vocab_size: usize,
  /// `usize::MAX` for a `Computed` entry, which has no length cap; `describe` and `list-models` print `n/a`,
  /// while the JSON form keeps the raw number.
  pub default_max_length: usize,
  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,
  /// Dtypes correct on an accelerator but not on the CPU (bge-m3's and MiniCOIL's F16, where candle's CPU f16
  /// matmul misses the tolerance); `DtypeBackendSupport::derive` turns them into the row the builder refuses from.
  pub accel_only_dtypes: &'static [DType],
  /// A second repository holding the head's extra files, `Qdrant/minicoil-v1` for MiniCOIL; `None` elsewhere.
  /// The builder resolves it as a second download, and `fetch phase:<n>` includes it.
  pub head_source: Option<&'static str>,
  /// `head_source`'s pinned commit, as the reference generator pins it: the shape guard cannot catch a same-shape
  /// retrain or a reordered vocabulary. `None` without a `head_source`.
  pub head_revision: Option<&'static str>,
  /// `Encoded` for every entry but `Qdrant/bm25`.
  pub computation: SparseComputation,
  /// Cells a run decided that the derivation cannot state, written over its rows (`CellException`).
  pub cell_exceptions: &'static [CellException],
}

impl SparseEntry {
  pub fn descriptor(&self) -> SparseDescriptor {
    let backends = BackendSupport {
      cpu: self.cpu,
      metal: self.metal,
      cuda: self.cuda,
    };
    SparseDescriptor {
      id: self.id.to_string(),
      family: self.family,
      head: self.head,
      vocab_size: self.vocab_size,
      default_max_length: self.default_max_length,
      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,
      head_source: self.head_source.map(str::to_string),
      head_revision: self.head_revision.map(str::to_string),
      computation: self.computation,
    }
  }

  /// 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 sparse model beyond its weights; a row's width is `vocab_size`, and the
/// head decides pooling and normalization.
#[derive(Clone, Debug, PartialEq, Serialize, Deserialize)]
#[non_exhaustive]
pub struct SparseDescriptor {
  pub id: String,
  /// The encoder's family; `None` for a `Computed` entry.
  pub family: Option<Family>,
  pub head: SparseHead,
  /// The index width; `0` for a `Computed` entry.
  pub vocab_size: usize,
  /// The default truncation; `usize::MAX` for a `Computed` entry.
  pub default_max_length: usize,
  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,
  /// The second repository a head reads extra files from: `Qdrant/minicoil-v1` for the miniCOIL entry.
  #[serde(default, skip_serializing_if = "Option::is_none")]
  pub head_source: Option<String>,
  /// The commit `head_source` is pinned to.
  #[serde(default, skip_serializing_if = "Option::is_none")]
  pub head_revision: Option<String>,
  /// Whether the head is encoded or computed.
  pub computation: SparseComputation,
}

impl SparseDescriptor {
  /// An encoded head's four facts with no safe default; the rest defaults as in `ModelDescriptor::new`, with
  /// safetensors weights and no second repository.
  pub fn new(id: impl Into<String>, family: Family, head: SparseHead, vocab_size: usize) -> Self {
    Self {
      id: id.into(),
      family: Some(family),
      head,
      vocab_size,
      default_max_length: 512,
      backends: BackendSupport::untested(),
      dtype_backends: DtypeBackendSupport::untested(),
      license: "unknown".to_string(),
      size_gb: 0.0,
      phase: u8::MAX,
      weights: WeightsFormat::Safetensors,
      gated: false,
      head_source: None,
      head_revision: None,
      computation: SparseComputation::Encoded,
    }
  }
}