ferrox_models/
pooling.rs

1//! `{arch}.pooling_type`: how a sequence of hidden states becomes one
2//! embedding vector.
3//!
4//! This is llama.cpp's `enum llama_pooling_type` and the switch in
5//! `llm_graph_context::build_pooling` (`src/llama-graph.cpp`),
6//! transcribed. The key is written by the HF converter
7//! (`gguf_writer.add_pooling_type`) as a **uint32** holding that enum's
8//! value, so the wire numbers below are load-bearing and are not
9//! ferrox's own invention.
10//!
11//! Its own module rather than a section of [`crate::bert_encoder`]
12//! because pooling is not a BERT fact: `llama-embed` and
13//! `gemma-embedding` are decoders that carry the same key, and
14//! `/v1/embeddings` needs to honour it for whatever produced the hidden
15//! states.
16//!
17//! # What is implemented, and what refuses
18//!
19//! `NONE`, `MEAN`, `CLS` and `LAST` are here. `RANK` is **not**: it is
20//! not a pooling rule at all but a classification head — upstream runs
21//! `cls`/`cls_out` matrices, a `tanh`, and an optional head norm over
22//! the pooled row, and reports the result through `/v1/rerank` against
23//! `classifier.output_labels`. None of that exists in ferrox yet, so
24//! [`PoolingType::Rank`] parses (so the refusal can name it) and
25//! [`pool`] returns [`PoolingError::Unimplemented`] rather than
26//! quietly handing back a CLS row that means something else.
27
28use thiserror::Error;
29
30/// `enum llama_pooling_type`, by its wire values.
31#[derive(Debug, Clone, Copy, PartialEq, Eq)]
32pub enum PoolingType {
33    /// No pooling: every token's hidden state is returned.
34    None,
35    /// Arithmetic mean over the sequence.
36    Mean,
37    /// The first token's row. What BERT/BGE checkpoints use, because
38    /// their `[CLS]` position is the one the model was trained to put
39    /// the sentence representation in.
40    Cls,
41    /// The last token's row. Decoder-style embedding models.
42    Last,
43    /// Not pooling — a reranker classification head. See module docs.
44    Rank,
45}
46
47#[derive(Debug, Error)]
48pub enum PoolingError {
49    #[error(
50        "{key} = {value} is not one of llama.cpp's llama_pooling_type values \
51         (-1 unspecified, 0 NONE, 1 MEAN, 2 CLS, 3 LAST, 4 RANK)"
52    )]
53    UnknownWireValue { key: String, value: i64 },
54    #[error("{key} is present but is not an integer: {value}")]
55    NotAnInteger { key: String, value: String },
56    #[error(
57        "pooling type RANK is a reranker classification head (cls / cls_out / \
58         classifier.output_labels), which ferrox does not implement — there is no /v1/rerank \
59         route and the head tensors are unread. Refusing rather than returning a CLS row \
60         that is not what RANK means"
61    )]
62    Unimplemented,
63    #[error("cannot pool an empty sequence")]
64    EmptySequence,
65    #[error("hidden states are {len} floats, which is not a whole number of {n_embd}-wide rows")]
66    RaggedHiddenStates { len: usize, n_embd: usize },
67}
68
69impl PoolingType {
70    /// The name upstream prints, so a refusal names the same thing the
71    /// user's `llama-embedding --pooling` flag does.
72    pub fn name(self) -> &'static str {
73        match self {
74            PoolingType::None => "NONE",
75            PoolingType::Mean => "MEAN",
76            PoolingType::Cls => "CLS",
77            PoolingType::Last => "LAST",
78            PoolingType::Rank => "RANK",
79        }
80    }
81
82    /// `None` for `-1` (`LLAMA_POOLING_TYPE_UNSPECIFIED`), which means
83    /// "the caller decides" and is not an error.
84    fn from_wire(key: &str, value: i64) -> Result<Option<Self>, PoolingError> {
85        Ok(match value {
86            -1 => None,
87            0 => Some(PoolingType::None),
88            1 => Some(PoolingType::Mean),
89            2 => Some(PoolingType::Cls),
90            3 => Some(PoolingType::Last),
91            4 => Some(PoolingType::Rank),
92            other => {
93                return Err(PoolingError::UnknownWireValue {
94                    key: key.to_string(),
95                    value: other,
96                })
97            }
98        })
99    }
100
101    /// Reads `{arch}.pooling_type`. `Ok(None)` means the key is absent
102    /// or explicitly unspecified — the caller picks a default and says
103    /// so, rather than this function inventing one.
104    pub fn from_gguf(
105        file: &impl ferrox_gguf::TensorSource,
106        arch: &str,
107    ) -> Result<Option<Self>, PoolingError> {
108        let key = format!("{arch}.pooling_type");
109        let Some(value) = file.metadata(&key) else {
110            return Ok(None);
111        };
112        let n = match value {
113            ferrox_gguf::GgufValue::U8(v) => i64::from(*v),
114            ferrox_gguf::GgufValue::I8(v) => i64::from(*v),
115            ferrox_gguf::GgufValue::U16(v) => i64::from(*v),
116            ferrox_gguf::GgufValue::I16(v) => i64::from(*v),
117            ferrox_gguf::GgufValue::U32(v) => i64::from(*v),
118            ferrox_gguf::GgufValue::I32(v) => i64::from(*v),
119            ferrox_gguf::GgufValue::U64(v) => *v as i64,
120            ferrox_gguf::GgufValue::I64(v) => *v,
121            other => {
122                return Err(PoolingError::NotAnInteger {
123                    key,
124                    value: format!("{other:?}"),
125                })
126            }
127        };
128        Self::from_wire(&key, n)
129    }
130}
131
132/// Pools `hidden` (`n_tokens` rows of `n_embd` floats, in row order)
133/// down to one vector — except for [`PoolingType::None`], which returns
134/// every row unchanged.
135///
136/// No L2 normalization happens here, because none happens in
137/// `build_pooling` either: upstream normalizes in the *caller*
138/// (`common_embd_normalize`, chosen by `llama-embedding --embd-normalize`
139/// and by the server's `/v1/embeddings`), and folding it in here would
140/// make MEAN and CLS silently return something the graph did not.
141pub fn pool(hidden: &[f32], n_embd: usize, ty: PoolingType) -> Result<Vec<f32>, PoolingError> {
142    if n_embd == 0 || hidden.is_empty() {
143        return Err(PoolingError::EmptySequence);
144    }
145    if !hidden.len().is_multiple_of(n_embd) {
146        return Err(PoolingError::RaggedHiddenStates {
147            len: hidden.len(),
148            n_embd,
149        });
150    }
151    let n_tokens = hidden.len() / n_embd;
152    Ok(match ty {
153        PoolingType::None => hidden.to_vec(),
154        PoolingType::Cls => hidden[..n_embd].to_vec(),
155        PoolingType::Last => hidden[(n_tokens - 1) * n_embd..].to_vec(),
156        PoolingType::Mean => {
157            let mut out = vec![0.0f32; n_embd];
158            for row in hidden.chunks_exact(n_embd) {
159                for (o, v) in out.iter_mut().zip(row) {
160                    *o += *v;
161                }
162            }
163            let inv = 1.0 / n_tokens as f32;
164            for o in out.iter_mut() {
165                *o *= inv;
166            }
167            out
168        }
169        PoolingType::Rank => return Err(PoolingError::Unimplemented),
170    })
171}
172
173/// L2-normalizes in place, which is what every BGE/E5/GTE consumer
174/// expects of an embedding and what `common_embd_normalize`'s default
175/// (`p == 2`) does. A zero vector is left alone, exactly as upstream
176/// leaves it (it divides by `norm > 0 ? 1/norm : 0`, i.e. it zeroes —
177/// and a zero vector is already zero).
178pub fn l2_normalize(v: &mut [f32]) {
179    let sum: f32 = v.iter().map(|x| x * x).sum();
180    if sum <= 0.0 {
181        return;
182    }
183    let inv = 1.0 / sum.sqrt();
184    for x in v.iter_mut() {
185        *x *= inv;
186    }
187}
188
189#[cfg(test)]
190mod tests {
191    use super::*;
192
193    #[test]
194    fn cls_takes_the_first_row_and_last_takes_the_last() {
195        let hidden = vec![1.0, 2.0, 10.0, 20.0, 100.0, 200.0];
196        assert_eq!(pool(&hidden, 2, PoolingType::Cls).unwrap(), vec![1.0, 2.0]);
197        assert_eq!(
198            pool(&hidden, 2, PoolingType::Last).unwrap(),
199            vec![100.0, 200.0]
200        );
201        assert_eq!(
202            pool(&hidden, 2, PoolingType::Mean).unwrap(),
203            vec![37.0, 74.0]
204        );
205        assert_eq!(pool(&hidden, 2, PoolingType::None).unwrap(), hidden);
206    }
207
208    /// RANK must refuse. If this ever starts returning a vector, the
209    /// caller is getting a CLS row labelled as a rerank score.
210    #[test]
211    fn rank_refuses_by_name() {
212        let err = pool(&[1.0, 2.0], 2, PoolingType::Rank).unwrap_err();
213        assert!(matches!(err, PoolingError::Unimplemented));
214        assert!(err.to_string().contains("RANK"));
215    }
216
217    #[test]
218    fn empty_and_ragged_inputs_refuse() {
219        assert!(matches!(
220            pool(&[], 4, PoolingType::Cls),
221            Err(PoolingError::EmptySequence)
222        ));
223        assert!(matches!(
224            pool(&[1.0, 2.0, 3.0], 2, PoolingType::Cls),
225            Err(PoolingError::RaggedHiddenStates { len: 3, n_embd: 2 })
226        ));
227    }
228
229    #[test]
230    fn wire_values_match_llama_pooling_type() {
231        let k = "bert.pooling_type";
232        assert_eq!(PoolingType::from_wire(k, -1).unwrap(), None);
233        assert_eq!(
234            PoolingType::from_wire(k, 0).unwrap(),
235            Some(PoolingType::None)
236        );
237        assert_eq!(
238            PoolingType::from_wire(k, 1).unwrap(),
239            Some(PoolingType::Mean)
240        );
241        assert_eq!(
242            PoolingType::from_wire(k, 2).unwrap(),
243            Some(PoolingType::Cls)
244        );
245        assert_eq!(
246            PoolingType::from_wire(k, 3).unwrap(),
247            Some(PoolingType::Last)
248        );
249        assert_eq!(
250            PoolingType::from_wire(k, 4).unwrap(),
251            Some(PoolingType::Rank)
252        );
253        assert!(PoolingType::from_wire(k, 5).is_err());
254    }
255
256    #[test]
257    fn l2_normalize_makes_a_unit_vector_and_leaves_zero_alone() {
258        let mut v = vec![3.0f32, 4.0];
259        l2_normalize(&mut v);
260        assert!((v[0] - 0.6).abs() < 1e-6 && (v[1] - 0.8).abs() < 1e-6);
261        let mut z = vec![0.0f32; 3];
262        l2_normalize(&mut z);
263        assert_eq!(z, vec![0.0, 0.0, 0.0]);
264    }
265}
ferrox_models/pooling.rs

ferrox_models/
pooling.rs
