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

use candle_core::DType;
use candle_core::Module;
use candle_core::Tensor;
use candle_nn::Embedding;
use candle_nn::VarBuilder;
use candle_nn::embedding;

use super::config::Qwen3Config;
use super::layer::Qwen3Layer;
use crate::error::CanddingError;
use crate::error::Result;
use crate::layers::activation::Activation;
use crate::layers::mask::AttentionMask;
use crate::layers::rms_norm::RmsNorm;
use crate::layers::rms_norm::RmsNormStyle;
use crate::layers::rope::RotaryEmbedding;
use crate::models::Encoder;
use crate::models::EncoderInput;

pub struct Qwen3Model {
  embed_tokens: Embedding,
  layers: Vec<Qwen3Layer>,
  norm: RmsNorm,
  rope: RotaryEmbedding,
  config: Qwen3Config,
}

/// Manual because `RmsNorm`, `RotaryEmbedding` and `Qwen3Layer` don't derive `Debug`; tests need it for `unwrap_err`.
impl std::fmt::Debug for Qwen3Model {
  fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
    f.debug_struct("Qwen3Model").finish_non_exhaustive()
  }
}

impl Qwen3Model {
  /// Accepts the bare tensor names of the Qwen3-Embedding checkpoints and the `model.` prefix of a full causal LM.
  pub fn load(config: Qwen3Config, vb: VarBuilder) -> Result<Self> {
    if config.hidden_act != Activation::Silu {
      return Err(CanddingError::UnsupportedArchitecture(format!(
        "qwen3 with hidden_act {:?}",
        config.hidden_act
      )));
    }
    if config.use_sliding_window {
      return Err(CanddingError::UnsupportedArchitecture(
        "qwen3 with sliding window attention".to_string(),
      ));
    }
    if config.attention_bias {
      return Err(CanddingError::UnsupportedArchitecture(
        "qwen3 with attention bias".to_string(),
      ));
    }
    let vb = if vb.contains_tensor("embed_tokens.weight") {
      vb
    } else {
      vb.pp("model")
    };
    let embed_tokens = embedding(config.vocab_size, config.hidden_size, vb.pp("embed_tokens"))?;
    let layers = (0..config.num_hidden_layers)
      .map(|i| Qwen3Layer::load(&config, vb.pp(format!("layers.{i}"))))
      .collect::<Result<Vec<_>>>()?;
    let norm = RmsNorm::load(
      vb.pp("norm"),
      config.hidden_size,
      config.rms_norm_eps,
      RmsNormStyle::Scale,
    )?;
    let rope = RotaryEmbedding::new(
      config.head_dim,
      config.max_position_embeddings,
      config.rope_theta,
      vb.dtype(),
      vb.device(),
    )?;
    Ok(Self {
      embed_tokens,
      layers,
      norm,
      rope,
      config,
    })
  }

  /// The `(vocab, hidden)` token embeddings, reused by a tied sparse head instead of loading a second copy.
  pub(crate) fn embed_tokens(&self) -> &Tensor {
    self.embed_tokens.embeddings()
  }
}

impl Encoder for Qwen3Model {
  /// Causal attention needs no padding mask: real tokens never see padding and last-token pooling drops it.
  /// The bidirectional branch (`is_causal: false`) gets no such guarantee, so it builds a padding mask.
  fn forward(&self, input: &EncoderInput) -> Result<Tensor> {
    let mut hidden = self.embed_tokens.forward(input.input_ids)?;
    let (_, seq) = input.input_ids.dims2()?;
    let mask = if self.config.is_causal {
      AttentionMask::causal(seq, self.config.head_dim, hidden.dtype(), hidden.device())?
    } else {
      AttentionMask::padding(
        input.attention_mask,
        input.lengths,
        self.config.num_attention_heads,
        self.config.head_dim,
        hidden.dtype(),
      )?
    };
    for layer in &self.layers {
      hidden = layer.forward(&hidden, &self.rope, &mask)?;
    }
    self.norm.forward(&hidden)
  }

  fn hidden_size(&self) -> usize {
    self.config.hidden_size
  }

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

  fn weights_dtype(&self) -> DType {
    self.embed_tokens.embeddings().dtype()
  }
}

#[cfg(test)]
mod tests;