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

//! `RerankHead::BertPooler`: `tanh(Linear(hidden, hidden))` over the first position, then `Linear(hidden, 1)`,
//! over a BERT or JinaBERT encoder under a `bert.` prefix, which the JinaBERT arm strips.

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::bert;
use crate::models::jina_bert;
use crate::pooling;
use crate::pooling::Pooling;
use crate::rerank::RerankScore;
use crate::rerank::ScoreActivation;

pub struct BertPoolerRerank {
  encoder: Box<dyn Encoder>,
  /// `bert.pooler.dense`, which no embedding model reads.
  pooler: Linear,
  /// `classifier`, `Linear(hidden, 1)`, shape-checked against the real tensor by `VarBuilder::get`.
  classifier: Linear,
  /// The entry's activation: a raw logit and a probability can look alike, so the registry decides.
  score_activation: ScoreActivation,
}

impl BertPoolerRerank {
  /// `model_type` is the resolved architecture key, `"bert"` or `"jina_bert"`; any other is refused by name.
  /// `jina_alibi_variant` is the entry's ALiBi slope formula, `None` for the canonical one; plain BERT ignores it.
  pub fn load(
    model_type: &str,
    config_json: &str,
    vb: VarBuilder,
    score_activation: ScoreActivation,
    jina_alibi_variant: Option<jina_bert::AlibiSlopeVariant>,
  ) -> Result<Self> {
    let (encoder, bert_vb): (Box<dyn Encoder>, VarBuilder) = match model_type {
      "bert" => {
        let config: bert::BertConfig = serde_json::from_str(config_json)?;
        let encoder = bert::BertModel::load(config, vb.clone())?;
        (Box::new(encoder), vb.pp("bert"))
      }
      "jina_bert" => {
        let config: jina_bert::JinaBertConfig = serde_json::from_str(config_json)?;
        let bert_vb = vb.pp("bert");
        let variant = jina_alibi_variant.unwrap_or(jina_bert::AlibiSlopeVariant::Canonical);
        let encoder = jina_bert::JinaBertModel::load(config, bert_vb.clone(), variant)?;
        (Box::new(encoder), bert_vb)
      }
      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 pooler_vb = bert_vb.pp("pooler").pp("dense");
    let pooler = Linear::new(
      pooler_vb.get((hidden, hidden), "weight")?,
      Some(pooler_vb.get(hidden, "bias")?),
    );
    let classifier_vb = vb.pp("classifier");
    let classifier = Linear::new(
      classifier_vb.get((1, hidden), "weight")?,
      Some(classifier_vb.get(1, "bias")?),
    );
    Ok(Self {
      encoder,
      pooler,
      classifier,
      score_activation,
    })
  }
}

impl RerankModel for BertPoolerRerank {
  /// `tanh(pooler(cls))` then `classifier` in the encoder's dtype, upcast to F32 only 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 pooled = self.pooler.forward(&cls)?.tanh()?;
    let logits: Vec<f32> = raw_logits(&self.classifier.forward(&pooled)?)?;
    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;