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
//! Pooling & embedding-output semantics shared by the GPU encoder and the CPU oracle/runtime.
//!
//! The encoder forward produces per-token hidden states; an embedding MODEL is defined by what
//! happens after — the pooling reduction and the L2 normalization. Those semantics live here once,
//! CPU-side (the WGSL twins in `encoder` implement the same math and are parity-tested against
//! these), so every consumer — GPU path, CPU fallback, tests — agrees on what "the embedding" is.
/// How per-token hidden states reduce to the final embedding(s) of a sequence.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Pooling {
/// Mask-aware mean over the sequence's tokens (BERT-family sentence-transformers default).
Mean,
/// The first token's hidden state (CLS).
Cls,
/// The last token's hidden state (decoder-based embedders, e.g. Qwen3-Embedding).
LastToken,
/// SigLIP multi-head attention-pooling head: a learned probe query cross-attends the tokens.
MapHead,
/// No reduction: one vector per token, each projected to `dim` and L2-normalized
/// (late-interaction / ColBERT).
PerToken {
/// Per-token output dimension after the projection.
dim: usize,
},
/// `*ForSequenceClassification` scoring head (sentence-transformers CrossEncoder): the CLS
/// hidden state → `pooler.dense` + tanh → `classifier` → `n_labels` logits. NOT normalized —
/// the output is a raw score (relevance/verification), one `[n_labels]` vector per input pair.
CrossEncoder {
/// Classifier output width (1 for a relevance scorer).
n_labels: usize,
},
}
/// The embedding output of one encode call.
#[derive(Clone, Debug, PartialEq)]
pub enum EmbedOut {
/// One vector per input sequence (`Pooling::Mean`/`Cls`/`LastToken`/`MapHead`).
Pooled(Vec<Vec<f32>>),
/// One `[tokens, dim]` matrix per input sequence (`Pooling::PerToken`), flattened as a
/// per-sequence list of per-token vectors.
PerToken(Vec<Vec<Vec<f32>>>),
}
/// Mask-aware mean pooling over a ragged batch: `hidden` is `[total_tokens, h]` row-major,
/// `seq_starts` (length `B+1`) delimits sequences. Exactness comes from raggedness — there are no
/// padding rows to mask out.
pub fn mean_pool(hidden: &[f32], h: usize, seq_starts: &[u32]) -> Vec<Vec<f32>> {
let mut out = Vec::with_capacity(seq_starts.len().saturating_sub(1));
for w in seq_starts.windows(2) {
let (a, b) = (w[0] as usize, w[1] as usize);
let mut acc = vec![0f32; h];
for row in hidden[a * h..b * h].chunks_exact(h) {
for (s, x) in acc.iter_mut().zip(row) {
*s += x;
}
}
let inv = 1.0 / (b - a).max(1) as f32;
for s in &mut acc {
*s *= inv;
}
out.push(acc);
}
out
}
/// L2-normalize in place (cosine == dot for the consumers). A zero vector is left unchanged —
/// normalizing it has no defined direction.
pub fn l2_normalize(v: &mut [f32]) {
let n: f32 = v.iter().map(|x| x * x).sum::<f32>().sqrt();
if n > 0.0 {
for x in v {
*x /= n;
}
}
}
/// Matryoshka truncation: keep the first `dim` components, then re-L2-normalize (the
/// nomic-embed / OpenAI `dimensions` semantics — truncate-then-renormalize, never just slice).
pub fn matryoshka_truncate(v: &[f32], dim: usize) -> Vec<f32> {
let mut out = v[..dim.min(v.len())].to_vec();
l2_normalize(&mut out);
out
}
/// MaxSim late-interaction relevance (ColBERT): for each query token vector, its best dot
/// product over the document's token vectors, summed. Token vectors are expected L2-normalized
/// (the `PerToken` output guarantees it), so dot == cosine. Empty query or document scores 0.0
/// — the empty sum, plain math.
pub fn maxsim(query_tokens: &[Vec<f32>], doc_tokens: &[Vec<f32>]) -> f32 {
if doc_tokens.is_empty() {
return 0.0;
}
query_tokens
.iter()
.map(|q| {
doc_tokens
.iter()
.map(|d| q.iter().zip(d).map(|(a, b)| a * b).sum::<f32>())
.fold(f32::NEG_INFINITY, f32::max)
})
.sum()
}