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

use std::collections::HashMap;

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

use super::*;

fn zeros(shape: impl Into<candle_core::Shape>) -> Tensor {
  Tensor::zeros(shape, DType::F32, &Device::Cpu).unwrap()
}

fn tensor_map_without_a_position_table(config: &XlmRobertaFlashConfig) -> HashMap<String, Tensor> {
  let h = config.hidden_size;
  let inter = config.intermediate_size;
  let mut map = HashMap::new();
  map.insert(
    "embeddings.word_embeddings.weight".to_string(),
    zeros((config.vocab_size, h)),
  );
  map.insert(
    "embeddings.token_type_embeddings.weight".to_string(),
    zeros((config.type_vocab_size, h)),
  );
  map.insert("emb_ln.weight".to_string(), zeros(h));
  map.insert("emb_ln.bias".to_string(), zeros(h));
  for l in 0..config.num_hidden_layers {
    let p = format!("encoder.layers.{l}");
    map.insert(format!("{p}.mixer.Wqkv.weight"), zeros((3 * h, h)));
    map.insert(format!("{p}.mixer.Wqkv.bias"), zeros(3 * h));
    map.insert(format!("{p}.mixer.out_proj.weight"), zeros((h, h)));
    map.insert(format!("{p}.mixer.out_proj.bias"), zeros(h));
    map.insert(format!("{p}.norm1.weight"), zeros(h));
    map.insert(format!("{p}.norm1.bias"), zeros(h));
    map.insert(format!("{p}.mlp.fc1.weight"), zeros((inter, h)));
    map.insert(format!("{p}.mlp.fc1.bias"), zeros(inter));
    map.insert(format!("{p}.mlp.fc2.weight"), zeros((h, inter)));
    map.insert(format!("{p}.mlp.fc2.bias"), zeros(h));
    map.insert(format!("{p}.norm2.weight"), zeros(h));
    map.insert(format!("{p}.norm2.bias"), zeros(h));
  }
  map
}

/// `VarBuilder::from_tensors` fails a `.get` for a missing name, so the load must never ask for a position tensor.
#[test]
fn loading_succeeds_from_tensors_that_include_no_position_table() {
  let config = tiny_config();
  let map = tensor_map_without_a_position_table(&config);
  assert!(
    !map.keys().any(|k| k.contains("position")),
    "fixture must not smuggle in the thing this test proves absent"
  );
  let vb = VarBuilder::from_tensors(map, DType::F32, &Device::Cpu);
  assert!(XlmRobertaFlashModel::load(config, vb, Some(PositionScheme::Rotary)).is_ok());
}

/// Token 5 at positions 0 and 1 differs only through rotary, so a table built and never applied fails here.
#[test]
fn rotary_positions_make_an_identical_token_diverge_by_position() {
  let (_varmap, model) = tiny_model();
  let out = run(&model, &[vec![5, 5, 9]]);
  let shift = max_shift(&[vec![out[0][0].clone()]], &[vec![out[0][1].clone()]]);
  assert!(
    shift > 1e-3,
    "token 5 at positions 0 and 1 produced the same output, so nothing is reading position: {shift}"
  );
}