candding 0.1.0

Pure candle embeddings for Rust with no ONNX: dense, sparse, multi-vector and rerank models today, image models on the roadmap
Documentation 
1
2
3
4
5
6
7
8
9
10
11
12
13
14
15
16
17
18
19
20
21
22
23
24
25
26
27
28
29
30
31
32
33
34
35
36
37
38
39
40
41
42
43
44
45
46
47
48
49
50
51
52
53
54
55
56
57
58
59
60
61
62
63
64
65
66
67
68
69
70
71
72
73
74
75
76
77
78
79
80
81
82
83
84
85
86
87
88
89
90
91
92
93
94
95
96
97
98
99
100
101
102
103
104
105
106
107
108
109
110
111
112
113
114
115
116
117
118
119
120
121
122
123
124
125
126
127
128

//! `RerankHead::LogitScore`: the `true_token_id` logit minus the `false_token_id` one at each prompt's last real
//! position. Only those two unembedding rows are kept, taken before the encoder loads.

use candle_core::DType;
use candle_core::Device;
use candle_core::Tensor;
use candle_nn::VarBuilder;

use super::RerankModel;
use super::Unembedding;
use crate::error::CanddingError;
use crate::error::Result;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::qwen3::Qwen3Config;
use crate::models::qwen3::Qwen3Model;
use crate::pooling;
use crate::pooling::Pooling;
use crate::rerank::RerankScore;
use crate::rerank::ScoreActivation;

pub struct LogitScoreRerank {
  encoder: Qwen3Model,
  /// `(hidden, 2)`: the true then the false token's unembedding row, transposed, in the model dtype.
  rows: Tensor,
  score_activation: ScoreActivation,
}

impl LogitScoreRerank {
  /// `token_ids` is `[true, false]`; `unembedding` is the entry's matrix, refused when `tie_word_embeddings`
  /// disagrees or when `lm_head.weight`'s presence contradicts it.
  pub fn load(
    model_type: &str,
    config_json: &str,
    vb: VarBuilder,
    token_ids: [u32; 2],
    unembedding: Unembedding,
    score_activation: ScoreActivation,
  ) -> Result<Self> {
    if model_type != "qwen3" {
      return Err(CanddingError::UnsupportedArchitecture(
        model_type.to_string(),
      ));
    }
    let config: Qwen3Config = serde_json::from_str(config_json)?;
    let (vocab_size, hidden) = (config.vocab_size, config.hidden_size);
    if let Some(&token_id) = token_ids.iter().find(|&&id| id as usize >= vocab_size) {
      return Err(CanddingError::LogitTokenOutOfRange {
        token_id,
        vocab_size,
      });
    }
    let tied = unembedding == Unembedding::Tied;
    if tied != config.tie_word_embeddings {
      return Err(CanddingError::UnembeddingDisagreesWithConfig {
        declared: unembedding.name().to_string(),
        tie_word_embeddings: config.tie_word_embeddings,
      });
    }
    let lm_head = vb.contains_tensor("lm_head.weight");
    let untied_rows = match (tied, lm_head) {
      (true, false) => None,
      (true, true) => return Err(CanddingError::TiedEmbeddingsWithLmHead),
      (false, true) => {
        let matrix = vb
          .clone()
          .set_device(Device::Cpu)
          .get((vocab_size, hidden), "lm_head.weight")?;
        Some(select_rows(&matrix, token_ids)?.to_device(vb.device())?)
      }
      (false, false) => return Err(CanddingError::MissingLmHead),
    };
    let encoder = Qwen3Model::load(config, vb.clone())?;
    let rows = match untied_rows {
      Some(rows) => rows,
      None => select_rows(encoder.embed_tokens(), token_ids)?,
    };
    Ok(Self {
      encoder,
      rows: rows.t()?.contiguous()?,
      score_activation,
    })
  }
}

/// Rows `token_ids` of a `(vocab, hidden)` matrix, as a `(2, hidden)` copy of its own.
fn select_rows(matrix: &Tensor, token_ids: [u32; 2]) -> Result<Tensor> {
  let index = Tensor::new(&token_ids, matrix.device())?;
  Ok(matrix.index_select(&index, 0)?)
}

impl RerankModel for LogitScoreRerank {
  /// The hidden state at each row's last real position; both logits and their difference in the model dtype,
  /// then the upcast, in `LogitScore`'s order; then the entry's activation.
  fn forward(&self, input: &EncoderInput) -> Result<Vec<RerankScore>> {
    let hidden = self.encoder.forward(input)?;
    let last = pooling::pool(
      Pooling::LastToken,
      &hidden,
      input.attention_mask,
      input.lengths,
    )?
    .contiguous()?;
    let logits = last.matmul(&self.rows)?;
    let difference = (logits.narrow(1, 0, 1)? - logits.narrow(1, 1, 1)?)?;
    let raw = difference
      .squeeze(1)?
      .to_dtype(DType::F32)?
      .to_vec1::<f32>()?;
    Ok(
      raw
        .into_iter()
        .map(|raw| RerankScore::new(self.score_activation.apply(raw)))
        .collect(),
    )
  }

  fn max_position_embeddings(&self) -> usize {
    self.encoder.max_position_embeddings()
  }

  fn weights_dtype(&self) -> DType {
    self.encoder.weights_dtype()
  }
}

#[cfg(test)]
mod tests;