nornir 0.3.1

Companion to cargo: dependency tracking, release gating, deploy, benchmarks, and documentation assembly. Project-agnostic.
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

//! Shared embedder support — model constants, `ModelProfile`, and the
//! tokenize-independent post-processing (mean-pool + L2-normalize) used by
//! **both** the tract (`embed`) and ort (`embed_ort`) backends running the
//! same `jina-embeddings-v2-base-code` ONNX model.
//!
//! Keeping this in one place guarantees the two runtimes produce the *same*
//! `model_profile` and the same vectors for the same text, so embeddings made
//! by one backend are reusable by the other (the warehouse dedup key is
//! backend-agnostic).
//!
//! Compiled when either embed feature is on.

use super::store::ModelProfile;

pub const MODEL_NAME: &str = "jinaai/jina-embeddings-v2-base-code";
pub const DIM: usize = 768;
/// jina v2 supports 8192 tokens; we truncate beyond this.
pub const MAX_TOKENS: usize = 8192;

/// Where `build.rs` placed the weights, and their content hashes.
pub const MODEL_DIR: &str = env!("NORNIR_MODEL_DIR");
pub const WEIGHTS_SHA: &str = env!("NORNIR_MODEL_WEIGHTS_SHA");
pub const TOKENIZER_SHA: &str = env!("NORNIR_MODEL_TOKENIZER_SHA");

/// The model identity shared by both backends. Note `weights_sha` is the SHA
/// of whichever model artifact build.rs fetched (the ONNX export) — the same
/// file feeds tract and ort, so the profile (and thus the dedup key) matches.
pub fn profile() -> ModelProfile {
    ModelProfile {
        model_name: MODEL_NAME.to_string(),
        weights_sha: WEIGHTS_SHA.to_string(),
        tokenizer_sha: TOKENIZER_SHA.to_string(),
        pooling: "mean".to_string(),
        normalize: true,
        dim: DIM,
        dtype: "f32".to_string(),
    }
}

/// Mean-pool a `[1, n_tokens, DIM]` row-major hidden state (flattened, so
/// element `(t, d)` is `hidden[t * DIM + d]`) over tokens, then L2-normalize.
/// Returns the `DIM`-length embedding.
pub fn pool_and_normalize(hidden: &[f32], n_tokens: usize) -> Vec<f32> {
    debug_assert_eq!(hidden.len(), n_tokens * DIM);
    let mut v = vec![0f32; DIM];
    for t in 0..n_tokens {
        let row = &hidden[t * DIM..(t + 1) * DIM];
        for (acc, &x) in v.iter_mut().zip(row) {
            *acc += x;
        }
    }
    let inv = 1.0 / n_tokens.max(1) as f32;
    for x in &mut v {
        *x *= inv;
    }
    l2_normalize(&mut v);
    v
}

/// L2-normalize in place. A zero vector is left unchanged.
pub fn l2_normalize(v: &mut [f32]) {
    let norm = v.iter().map(|x| x * x).sum::<f32>().sqrt();
    if norm > 0.0 {
        for x in v {
            *x /= norm;
        }
    }
}

/// Truncate raw token ids to `MAX_TOKENS` (and to at least one token), as
/// `i64` for the ONNX `input_ids` tensor, with an all-ones attention mask.
pub fn prepare_tokens(raw: &[u32]) -> (Vec<i64>, Vec<i64>) {
    let n = raw.len().clamp(1, MAX_TOKENS);
    let ids: Vec<i64> = raw[..n].iter().map(|&x| x as i64).collect();
    let mask = vec![1i64; n];
    (ids, mask)
}