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

//! What `MultiVectorTextEmbeddingBuilder::build` decides from the descriptor: the conventions, the weight files
//! and the query-side tokenizer.

use std::path::Path;
use std::path::PathBuf;

use candle_core::DType;

use crate::embed::fit;
use crate::error::CanddingError;
use crate::error::Result;
use crate::loader::ModelFiles;
use crate::models::MultiVectorConventions;
use crate::models::MultiVectorHead;
use crate::registry::multi_vector::MultiVectorDescriptor;
use crate::tokenizer::PadToken;
use crate::tokenizer::TextTokenizer;

/// Kept out of `build`, which reaches it only after a download, so a test can catch a hardcoded or swapped field.
pub(super) fn conventions_from_descriptor(
  descriptor: &MultiVectorDescriptor,
) -> MultiVectorConventions {
  MultiVectorConventions {
    keeps_cls_row: descriptor.keeps_cls_row,
    projection_in: descriptor.projection_in,
    projection_out: descriptor.projection_out,
    projection_bias: descriptor.projection_bias,
    normalize: descriptor.normalize,
    attends_to_mask_tokens: descriptor.attends_to_mask_tokens,
    position_scheme: descriptor.position_scheme,
  }
}

/// The files the build reads, all in the run's dtype: the encoder's, and bge-m3's `colbert_linear.pt` for its head.
pub(super) fn weight_files(
  descriptor: &MultiVectorDescriptor,
  files: &ModelFiles,
  dtype: DType,
) -> Vec<(PathBuf, DType)> {
  let mut weights = fit::at(&files.weights, dtype);
  if descriptor.head == MultiVectorHead::M3Colbert
    && let Some(head) = &files.colbert_linear
  {
    weights.push((head.clone(), dtype));
  }
  weights
}

/// Refuses a `StanfordColbert` descriptor without `query_length`: its query keeps every position, so `[PAD]` rows
/// would silently count toward MaxSim.
pub(super) fn require_fixed_query_length(
  model: &str,
  descriptor: &MultiVectorDescriptor,
) -> Result<()> {
  if descriptor.head == MultiVectorHead::StanfordColbert && descriptor.query_length.is_none() {
    return Err(CanddingError::MissingQueryLength {
      model: model.to_string(),
    });
  }
  Ok(())
}

/// Registers the bare marker as a special token; its trailing space belongs to the prompt string, not the token.
pub(super) fn register_marker(tokenizer: &mut TextTokenizer, marker: Option<&str>) -> Result<()> {
  if let Some(marker) = marker {
    tokenizer.register_special_tokens(&[marker.trim()])?;
  }
  Ok(())
}

/// The query-side tokenizer: the document one's twin without a `query_length`, else fixed-length and filled with
/// the mask token `tokenizer_config.json` names (`<mask>` for `jinaai/jina-colbert-v2`, `[MASK]` for the others).
pub(super) fn build_query_tokenizer(
  tokenizer_path: &Path,
  max_length: usize,
  pad: PadToken,
  descriptor: &MultiVectorDescriptor,
  max_position_embeddings: usize,
  tokenizer_config: Option<&Path>,
) -> Result<TextTokenizer> {
  let mut tokenizer = match descriptor.query_length {
    Some(query_length) => {
      let mask_token = crate::tokenizer::read_mask_token(tokenizer_config)?;
      TextTokenizer::from_file_fixed_length(
        tokenizer_path,
        query_length.min(max_position_embeddings),
        PadToken::Named(mask_token),
      )?
    }
    None => TextTokenizer::from_file(tokenizer_path, max_length, pad)?,
  };
  register_marker(&mut tokenizer, descriptor.query_marker.as_deref())?;
  Ok(tokenizer)
}