candding 0.1.0

//! Model families implementing the `Encoder` trait.

pub mod bert;
pub mod gemma3;
pub mod gte_new;
pub mod jina_bert;
pub mod modernbert;
pub mod mpnet;
mod multi_vector;
pub mod nomic_bert;
pub mod qwen3;
mod rerank;
pub mod sparse;
pub mod xlm_roberta;
pub mod xlm_roberta_flash;

use std::path::Path;
use std::path::PathBuf;

use candle_core::DType;
use candle_core::Tensor;
use candle_nn::VarBuilder;
use serde::Deserialize;
use serde::Serialize;

use crate::error::CanddingError;
use crate::error::Result;
use crate::loader::SPARSE_LINEAR_FILENAME;
use crate::sparse::SparseEmbedding;

/// One padded batch; every tensor is `(batch, seq)` in `u32`. `lengths` holds one real-token count per row of
/// `input_ids`, which right padding makes the index of the last real token plus one; the mask builders assert it.
pub struct EncoderInput<'a> {
  pub input_ids: &'a Tensor,
  pub attention_mask: &'a Tensor,
  pub token_type_ids: &'a Tensor,
  pub lengths: &'a [usize],
}

/// A family produces per-token hidden states; pooling happens outside.
pub trait Encoder: Send + Sync {
  fn forward(&self, input: &EncoderInput) -> Result<Tensor>;
  fn hidden_size(&self) -> usize;
  fn max_position_embeddings(&self) -> usize;
  /// The dtype the weights arrived in, read from the token embedding table; later matmul weights must match it.
  fn weights_dtype(&self) -> DType;
}

pub fn load_encoder(
  model_type: &str,
  config_json: &str,
  vb: VarBuilder,
) -> Result<Box<dyn Encoder>> {
  match model_type {
    "bert" => Ok(Box::new(bert::BertModel::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    "qwen3" => Ok(Box::new(qwen3::Qwen3Model::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    "gemma3_text" => Ok(Box::new(gemma3::Gemma3Model::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    "jina_bert" => Ok(Box::new(jina_bert::JinaBertModel::load(
      serde_json::from_str(config_json)?,
      vb,
      jina_bert::AlibiSlopeVariant::Canonical,
    )?)),
    "modernbert" => Ok(Box::new(modernbert::ModernBertModel::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    "mpnet" => Ok(Box::new(mpnet::MPNetModel::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    "new" => Ok(Box::new(gte_new::GteNewModel::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    "nomic_bert" => Ok(Box::new(nomic_bert::NomicBertModel::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    "xlm-roberta" => Ok(Box::new(xlm_roberta::XlmRobertaModel::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    other => Err(CanddingError::UnsupportedArchitecture(other.to_string())),
  }
}

/// A sparse head: one `SparseEmbedding` per row, its pooling and postprocessing included.
pub trait SparseModel: Send + Sync {
  fn forward(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>>;
  /// Defaults to `forward`: only `MiniCoil` and `Bm42` weigh a query differently from a passage.
  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
    self.forward(input)
  }
  fn vocab_size(&self) -> usize;
  fn max_position_embeddings(&self) -> usize;
  /// `Encoder::weights_dtype` of the encoder the head runs on.
  fn weights_dtype(&self) -> DType;
}

/// Which sparse head a checkpoint runs; `Bm25` has no encoder and loads through `Bm25TextEmbedding` alone.
#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
#[serde(rename_all = "snake_case")]
#[non_exhaustive]
pub enum SparseHead {
  /// A masked-language-model head, max-pooled over the sequence: `prithivida/Splade_PP_en_v1`.
  SpladeMlm,
  /// A `Linear(hidden, 1)` per position, scattered by token id: `BAAI/bge-m3`.
  M3Linear,
  /// `SpladeMlm`'s pooling over a Qwen3 backbone's tied embedding table: `naver/splade-code-06B`.
  SpladeTied,
  /// A per-word matrix over a private vocabulary, with a hash for unknown words: `Qdrant/minicoil-v1`.
  MiniCoil,
  /// Last-layer attention summed per word and hashed: `sentence-transformers/all-MiniLM-L6-v2`.
  Bm42,
  /// No encoder: the term-frequency half of BM25 per hashed stem, `Qdrant/bm25`.
  Bm25,
}

/// The files `SparseHead::MiniCoil` needs beyond its encoder: the per-word matrix, vocabulary and stopwords.
/// All three come from `Qdrant/minicoil-v1`, not the encoder's repository.
pub struct MiniCoilFiles<'a> {
  pub matrix: &'a Path,
  pub vocab: &'a Path,
  pub stopwords: &'a Path,
}

/// Builds the sparse head `head` names; `sparse_linear` serves only `M3Linear` and `minicoil_files` only `MiniCoil`,
/// each refused by name when missing. `Bm25` has no weights and is always refused.
pub fn load_sparse(
  head: SparseHead,
  config_json: &str,
  vb: VarBuilder,
  tokenizer_path: &Path,
  sparse_linear: Option<&Path>,
  minicoil_files: Option<MiniCoilFiles<'_>>,
) -> Result<Box<dyn SparseModel>> {
  match head {
    SparseHead::SpladeMlm => Ok(Box::new(sparse::SpladeSparse::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    SparseHead::M3Linear => {
      let sparse_linear = sparse_linear
        .ok_or_else(|| CanddingError::MissingFile(PathBuf::from(SPARSE_LINEAR_FILENAME)))?;
      Ok(Box::new(sparse::M3Sparse::load(
        serde_json::from_str(config_json)?,
        vb,
        sparse_linear,
      )?))
    }
    SparseHead::SpladeTied => Ok(Box::new(sparse::SpladeTiedSparse::load(
      serde_json::from_str(config_json)?,
      vb,
    )?)),
    SparseHead::MiniCoil => {
      let files = minicoil_files.ok_or(CanddingError::MissingHeadSource)?;
      let encoder = load_encoder("jina_bert", config_json, vb)?;
      let matrix = crate::weights::npy::read(files.matrix)?;
      let vocab = sparse::minicoil::vocab::Vocab::parse(&std::fs::read_to_string(files.vocab)?)?;
      let stopwords =
        sparse::stopwords::Stopwords::parse(&std::fs::read_to_string(files.stopwords)?);
      let token_text = crate::tokenizer::token_text_from_file(tokenizer_path)?;
      Ok(Box::new(sparse::MiniCoilSparse::new(
        encoder, matrix, vocab, stopwords, token_text,
      )?))
    }
    SparseHead::Bm42 => {
      let config: bert::BertConfig = serde_json::from_str(config_json)?;
      let vocab_size = config.vocab_size;
      let bert = bert::BertModel::load(config, vb)?;
      let token_text = crate::tokenizer::token_text_from_file(tokenizer_path)?;
      Ok(Box::new(sparse::Bm42Sparse::new(
        bert, vocab_size, token_text,
      )))
    }
    SparseHead::Bm25 => Err(CanddingError::ComputedSparseModel {
      hint: "Bm25TextEmbedding::builder()".to_string(),
    }),
  }
}

pub use multi_vector::MultiVectorConventions;
pub use multi_vector::MultiVectorHead;
pub use multi_vector::MultiVectorModel;
pub use multi_vector::load_multi_vector;
pub use rerank::RerankConventions;
pub use rerank::RerankHead;
pub use rerank::RerankModel;
pub use rerank::Unembedding;
#[cfg(test)]
pub(crate) use rerank::flash_fixture;
pub use rerank::load_rerank;

#[cfg(test)]
mod tests;