candding/models/
mod.rs

1//! Model families implementing the `Encoder` trait.
2
3pub mod bert;
4pub mod gemma3;
5pub mod gte_new;
6pub mod jina_bert;
7pub mod modernbert;
8pub mod mpnet;
9mod multi_vector;
10pub mod nomic_bert;
11pub mod qwen3;
12mod rerank;
13pub mod sparse;
14pub mod xlm_roberta;
15pub mod xlm_roberta_flash;
16
17use std::path::Path;
18use std::path::PathBuf;
19
20use candle_core::DType;
21use candle_core::Tensor;
22use candle_nn::VarBuilder;
23use serde::Deserialize;
24use serde::Serialize;
25
26use crate::error::CanddingError;
27use crate::error::Result;
28use crate::loader::SPARSE_LINEAR_FILENAME;
29use crate::sparse::SparseEmbedding;
30
31/// One padded batch; every tensor is `(batch, seq)` in `u32`. `lengths` holds one real-token count per row of
32/// `input_ids`, which right padding makes the index of the last real token plus one; the mask builders assert it.
33pub struct EncoderInput<'a> {
34  pub input_ids: &'a Tensor,
35  pub attention_mask: &'a Tensor,
36  pub token_type_ids: &'a Tensor,
37  pub lengths: &'a [usize],
38}
39
40/// A family produces per-token hidden states; pooling happens outside.
41pub trait Encoder: Send + Sync {
42  fn forward(&self, input: &EncoderInput) -> Result<Tensor>;
43  fn hidden_size(&self) -> usize;
44  fn max_position_embeddings(&self) -> usize;
45  /// The dtype the weights arrived in, read from the token embedding table; later matmul weights must match it.
46  fn weights_dtype(&self) -> DType;
47}
48
49pub fn load_encoder(
50  model_type: &str,
51  config_json: &str,
52  vb: VarBuilder,
53) -> Result<Box<dyn Encoder>> {
54  match model_type {
55    "bert" => Ok(Box::new(bert::BertModel::load(
56      serde_json::from_str(config_json)?,
57      vb,
58    )?)),
59    "qwen3" => Ok(Box::new(qwen3::Qwen3Model::load(
60      serde_json::from_str(config_json)?,
61      vb,
62    )?)),
63    "gemma3_text" => Ok(Box::new(gemma3::Gemma3Model::load(
64      serde_json::from_str(config_json)?,
65      vb,
66    )?)),
67    "jina_bert" => Ok(Box::new(jina_bert::JinaBertModel::load(
68      serde_json::from_str(config_json)?,
69      vb,
70      jina_bert::AlibiSlopeVariant::Canonical,
71    )?)),
72    "modernbert" => Ok(Box::new(modernbert::ModernBertModel::load(
73      serde_json::from_str(config_json)?,
74      vb,
75    )?)),
76    "mpnet" => Ok(Box::new(mpnet::MPNetModel::load(
77      serde_json::from_str(config_json)?,
78      vb,
79    )?)),
80    "new" => Ok(Box::new(gte_new::GteNewModel::load(
81      serde_json::from_str(config_json)?,
82      vb,
83    )?)),
84    "nomic_bert" => Ok(Box::new(nomic_bert::NomicBertModel::load(
85      serde_json::from_str(config_json)?,
86      vb,
87    )?)),
88    "xlm-roberta" => Ok(Box::new(xlm_roberta::XlmRobertaModel::load(
89      serde_json::from_str(config_json)?,
90      vb,
91    )?)),
92    other => Err(CanddingError::UnsupportedArchitecture(other.to_string())),
93  }
94}
95
96/// A sparse head: one `SparseEmbedding` per row, its pooling and postprocessing included.
97pub trait SparseModel: Send + Sync {
98  fn forward(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>>;
99  /// Defaults to `forward`: only `MiniCoil` and `Bm42` weigh a query differently from a passage.
100  fn forward_query(&self, input: &EncoderInput) -> Result<Vec<SparseEmbedding>> {
101    self.forward(input)
102  }
103  fn vocab_size(&self) -> usize;
104  fn max_position_embeddings(&self) -> usize;
105  /// `Encoder::weights_dtype` of the encoder the head runs on.
106  fn weights_dtype(&self) -> DType;
107}
108
109/// Which sparse head a checkpoint runs; `Bm25` has no encoder and loads through `Bm25TextEmbedding` alone.
110#[derive(Clone, Copy, Debug, PartialEq, Eq, Serialize, Deserialize)]
111#[serde(rename_all = "snake_case")]
112#[non_exhaustive]
113pub enum SparseHead {
114  /// A masked-language-model head, max-pooled over the sequence: `prithivida/Splade_PP_en_v1`.
115  SpladeMlm,
116  /// A `Linear(hidden, 1)` per position, scattered by token id: `BAAI/bge-m3`.
117  M3Linear,
118  /// `SpladeMlm`'s pooling over a Qwen3 backbone's tied embedding table: `naver/splade-code-06B`.
119  SpladeTied,
120  /// A per-word matrix over a private vocabulary, with a hash for unknown words: `Qdrant/minicoil-v1`.
121  MiniCoil,
122  /// Last-layer attention summed per word and hashed: `sentence-transformers/all-MiniLM-L6-v2`.
123  Bm42,
124  /// No encoder: the term-frequency half of BM25 per hashed stem, `Qdrant/bm25`.
125  Bm25,
126}
127
128/// The files `SparseHead::MiniCoil` needs beyond its encoder: the per-word matrix, vocabulary and stopwords.
129/// All three come from `Qdrant/minicoil-v1`, not the encoder's repository.
130pub struct MiniCoilFiles<'a> {
131  pub matrix: &'a Path,
132  pub vocab: &'a Path,
133  pub stopwords: &'a Path,
134}
135
136/// Builds the sparse head `head` names; `sparse_linear` serves only `M3Linear` and `minicoil_files` only `MiniCoil`,
137/// each refused by name when missing. `Bm25` has no weights and is always refused.
138pub fn load_sparse(
139  head: SparseHead,
140  config_json: &str,
141  vb: VarBuilder,
142  tokenizer_path: &Path,
143  sparse_linear: Option<&Path>,
144  minicoil_files: Option<MiniCoilFiles<'_>>,
145) -> Result<Box<dyn SparseModel>> {
146  match head {
147    SparseHead::SpladeMlm => Ok(Box::new(sparse::SpladeSparse::load(
148      serde_json::from_str(config_json)?,
149      vb,
150    )?)),
151    SparseHead::M3Linear => {
152      let sparse_linear = sparse_linear
153        .ok_or_else(|| CanddingError::MissingFile(PathBuf::from(SPARSE_LINEAR_FILENAME)))?;
154      Ok(Box::new(sparse::M3Sparse::load(
155        serde_json::from_str(config_json)?,
156        vb,
157        sparse_linear,
158      )?))
159    }
160    SparseHead::SpladeTied => Ok(Box::new(sparse::SpladeTiedSparse::load(
161      serde_json::from_str(config_json)?,
162      vb,
163    )?)),
164    SparseHead::MiniCoil => {
165      let files = minicoil_files.ok_or(CanddingError::MissingHeadSource)?;
166      let encoder = load_encoder("jina_bert", config_json, vb)?;
167      let matrix = crate::weights::npy::read(files.matrix)?;
168      let vocab = sparse::minicoil::vocab::Vocab::parse(&std::fs::read_to_string(files.vocab)?)?;
169      let stopwords =
170        sparse::stopwords::Stopwords::parse(&std::fs::read_to_string(files.stopwords)?);
171      let token_text = crate::tokenizer::token_text_from_file(tokenizer_path)?;
172      Ok(Box::new(sparse::MiniCoilSparse::new(
173        encoder, matrix, vocab, stopwords, token_text,
174      )?))
175    }
176    SparseHead::Bm42 => {
177      let config: bert::BertConfig = serde_json::from_str(config_json)?;
178      let vocab_size = config.vocab_size;
179      let bert = bert::BertModel::load(config, vb)?;
180      let token_text = crate::tokenizer::token_text_from_file(tokenizer_path)?;
181      Ok(Box::new(sparse::Bm42Sparse::new(
182        bert, vocab_size, token_text,
183      )))
184    }
185    SparseHead::Bm25 => Err(CanddingError::ComputedSparseModel {
186      hint: "Bm25TextEmbedding::builder()".to_string(),
187    }),
188  }
189}
190
191pub use multi_vector::MultiVectorConventions;
192pub use multi_vector::MultiVectorHead;
193pub use multi_vector::MultiVectorModel;
194pub use multi_vector::load_multi_vector;
195pub use rerank::RerankConventions;
196pub use rerank::RerankHead;
197pub use rerank::RerankModel;
198pub use rerank::Unembedding;
199#[cfg(test)]
200pub(crate) use rerank::flash_fixture;
201pub use rerank::load_rerank;
202
203#[cfg(test)]
204mod tests;
candding/models/mod.rs

candding/models/
mod.rs
