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

ferrox_models/
pooling.rs
