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

//! `RerankHead::RobertaClassificationHead`: `Linear`, `tanh`, `Linear(hidden, 1)` on the first position from a root
//! `classifier` beside a `roberta.`-prefixed encoder; serves the bge rerankers and jina-reranker-v2.

use candle_core::DType;
use candle_core::Tensor;
use candle_nn::Linear;
use candle_nn::Module;
use candle_nn::VarBuilder;

use super::RerankModel;
use crate::error::CanddingError;
use crate::error::Result;
use crate::models::Encoder;
use crate::models::EncoderInput;
use crate::models::xlm_roberta;
use crate::models::xlm_roberta_flash;
use crate::models::xlm_roberta_flash::PositionScheme;
use crate::pooling;
use crate::pooling::Pooling;
use crate::rerank::RerankScore;
use crate::rerank::ScoreActivation;

pub struct RobertaClassificationRerank {
  encoder: Box<dyn Encoder>,
  /// `classifier.dense`, `Linear(hidden, hidden)`, the head's first projection over the raw first position.
  dense: Linear,
  /// `classifier.out_proj`, `Linear(hidden, 1)`, shape-checked against the real tensor by `VarBuilder::get`.
  out_proj: Linear,
  /// The entry's activation, as in `BertPoolerRerank`.
  score_activation: ScoreActivation,
}

impl RobertaClassificationRerank {
  /// Takes the resolved key, `xlm-roberta` or `xlm_roberta_flash::ARCHITECTURE`, refusing any other by name;
  /// `position_scheme` is the entry's, checked by the flash encoder against its `config.json`.
  pub fn load(
    model_type: &str,
    config_json: &str,
    vb: VarBuilder,
    score_activation: ScoreActivation,
    position_scheme: Option<PositionScheme>,
  ) -> Result<Self> {
    let encoder: Box<dyn Encoder> = match model_type {
      "xlm-roberta" => {
        let config: xlm_roberta::XlmRobertaConfig = serde_json::from_str(config_json)?;
        Box::new(xlm_roberta::XlmRobertaModel::load(
          config,
          vb.pp("roberta"),
        )?)
      }
      xlm_roberta_flash::ARCHITECTURE => {
        let config: xlm_roberta_flash::XlmRobertaFlashConfig = serde_json::from_str(config_json)?;
        Box::new(xlm_roberta_flash::XlmRobertaFlashModel::load(
          config,
          vb.clone(),
          position_scheme,
        )?)
      }
      other => {
        return Err(CanddingError::UnsupportedArchitecture(other.to_string()));
      }
    };
    // The encoder's real width; `VarBuilder::get` below fails if the head's tensors disagree.
    let hidden = encoder.hidden_size();
    let head_vb = vb.pp("classifier");
    let dense_vb = head_vb.pp("dense");
    let dense = Linear::new(
      dense_vb.get((hidden, hidden), "weight")?,
      Some(dense_vb.get(hidden, "bias")?),
    );
    let out_proj_vb = head_vb.pp("out_proj");
    let out_proj = Linear::new(
      out_proj_vb.get((1, hidden), "weight")?,
      Some(out_proj_vb.get(1, "bias")?),
    );
    Ok(Self {
      encoder,
      dense,
      out_proj,
      score_activation,
    })
  }
}

impl RerankModel for RobertaClassificationRerank {
  /// `tanh(dense(cls))` then `out_proj` in the encoder's dtype, no dropout at inference, upcast to F32 for the scores.
  fn forward(&self, input: &EncoderInput) -> Result<Vec<RerankScore>> {
    let hidden = self.encoder.forward(input)?;
    // `pool(Cls)` is a view; Metal's matmul refuses it ("Invalid matmul arguments"), so materialize it.
    let cls =
      pooling::pool(Pooling::Cls, &hidden, input.attention_mask, input.lengths)?.contiguous()?;
    let projected = self.dense.forward(&cls)?.tanh()?;
    let logits: Vec<f32> = raw_logits(&self.out_proj.forward(&projected)?)?;
    Ok(
      logits
        .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()
  }
}

/// `logits` is `(batch, 1)`; one raw value per row, in F32 regardless of the model's own runtime dtype.
fn raw_logits(logits: &Tensor) -> Result<Vec<f32>> {
  Ok(logits.squeeze(1)?.to_dtype(DType::F32)?.to_vec1::<f32>()?)
}

#[cfg(test)]
mod tests;