frink-models 0.38.0

//! BERT: the encoder graph, transcribed from llama.cpp
//! `src/models/bert.cpp` (`llama_model_bert::graph::graph`).
//!
//! Loading lives next door in [`crate::bert_gguf_loader`]; pooling in
//! [`crate::pooling`]; the reason this is not an
//! [`crate::engine::Engine`] in [`crate::encoder`].
//!
//! # The graph, and the five places it is not a decoder
//!
//! ```text
//! h[i] = tok_embd[t[i]] + type_embd[seg[i]] + pos_embd[i] (1) (2)
//! h    = LayerNorm(h, token_embd_norm)                        (3)
//! for each layer:
//!     q,k,v = Wq h + bq,  Wk h + bk,  Wv h + bv
//!     a     = softmax(q·kᵀ / √head_dim) v                  (4)
//!     x     = LayerNorm(Wo a + bo + h,  attn_output_norm)      (3)
//!     f     = W_down · GELU(W_up x + b_up) + b_down        (5)
//!     h     = LayerNorm(f + x, layer_output_norm)              (3)
//! result = h                                              (6)
//! ```
//!
//! 1. **Learned position embeddings, added.** Not RoPE. `pos_embd` is a
//!    real `[n_ctx_train, n_embd]` table and position `i` is a row
//!    lookup. A learned table cannot be extrapolated, which is why
//!    [`crate::encoder::EncodeError::TooLong`] is an error and not a
//!    warning.
//! 2. **A token-type embedding, per position.** A single-sequence
//!    embedding pass is all "Sentence A" and uses row 0, which is what
//!    upstream hardcodes — `ggml_view_1d(ctx0, model.type_embd, n_embd,
//!    0)`, with the comment that token types are hardcoded to zero
//!    because `llama_batch` carries no segment ids. A cross-encoder
//!    PAIR is not that case: HuggingFace's `tokenizer(query, document)`
//!    emits `0…0 1…1` and `BertModel` adds row 1 to every position
//!    after the first `[SEP]`. Frink adds the row the caller names,
//!    which is row 0 for every embedding request and 0/1 for a rerank
//!    pair. Matching upstream here instead was measured, on
//!    `cross-encoder/ms-marco-MiniLM-L6-v2` against a NumPy transcription
//!    of `BertForSequenceClassification`, to put the RELEVANT document
//!    LAST in three of four rankings — see
//!    `tests/rerank_cross_encoder_ordering.rs`.
//! 3. **LayerNorm, not RMSNorm, at three sites per layer plus one on
//!    the input.** Mean-subtracting, and every one of them carries a
//!    `bias` tensor as well as a `weight`. Substituting RMSNorm here
//!    loads fine and produces a plausible-looking vector that is wrong.
//! 4. **No causal mask.** Row 0 attends to the last token. This is the
//!    single property that makes the whole model an encoder, and
//!    `attention_is_bidirectional_not_causal` below is the test that
//!    would go red if a mask ever appeared.
//! 5. **A plain GELU MLP, not a gated one.** Two matrices, not three,
//!    and both carry biases. `LLM_FFN_GELU, LLM_FFN_SEQ` upstream.
//! 6. **No output head and no logits.** The hidden states *are* the
//!    result (`res->t_embd`); this checkpoint has no `output.weight` at
//!    all.
//!
//! # What this module does not do
//!
//! Only `arch == "bert"`, and only its dense, non-RoPE, separate-QKV
//! shape. `nomic-bert` (RoPE + gated FFN), `jina-bert-v2` (GEGLU + a
//! second attention norm), `nomic-bert-moe` (expert layers) and
//! `modern-bert` all share `bert.cpp` upstream and are all refused by
//! name in the loader instead of being run through this graph.

use frink_core::matmul::{gelu, layer_norm};
use frink_core::weight_matrix::WeightMatrix;

use crate::encoder::{EncodeError, PairSequence, TextEncoder};
use crate::pooling::PoolingType;

/// `bert.*` metadata, after the loader has checked it.
#[derive(Debug, Clone)]
pub struct BertHparams {
    pub arch: String,
    pub n_layer: usize,
    pub n_embd: usize,
    pub n_ff: usize,
    pub n_head: usize,
    pub n_head_kv: usize,
    /// Height of the learned position table.
    pub n_ctx_train: usize,
    pub n_token_types: usize,
    pub layer_norm_eps: f32,
    /// NEOX RoPE on Q and K instead of a learned position table:
    /// `Some(theta)` for the architectures `bert.cpp:126-133` rotates
    /// (`nomic-bert`, `nomic-bert-moe`, `jina-bert-v3`), `None` for
    /// `bert` itself, which adds `position_embd` at `:90` instead.
    ///
    /// One field for both facts on purpose: a file cannot have a
    /// position table AND a rotation here, because upstream decides
    /// both from the architecture and never reads the table for a
    /// rotating one -- measured, libllama's load log never names
    /// `position_embd` for `nomic-bert`.
    pub rope_theta: Option<f32>,
    /// How many of each head's channels rotate
    /// (`{arch}.rope.dimension_count`), `head_dim` when the key is
    /// absent. Only read when [`Self::rope_theta`] is `Some`.
    pub rope_dim: usize,
    /// The FFN this architecture runs (`bert.cpp:179-201`).
    pub ffn: BertFfn,
    /// Where this architecture's norms sit and what they are.
    pub topology: BertTopology,
    /// `true` for the NORM (interleaved-pair) rotation, `false` for
    /// NEOX (split-half). `llama_model_rope_type` answers NORM for
    /// `neo-bert` alone among the encoders here.
    pub rope_interleaved: bool,
    /// ALiBi slopes, one per head, for the architecture whose graph
    /// carries a positional bias instead of a table or a rotation.
    ///
    /// `jina-bert-v2.cpp:5` sets `f_max_alibi_bias = 8.0f` as a
    /// LITERAL and `bert.cpp:78-80` builds no `inp_pos` for it at all,
    /// so this is the only place position enters that model. The bias
    /// is SYMMETRIC here -- `llama-graph.cpp:442` fills the mask with
    /// `-|p0 - p1|` for a non-causal model, where the decoder's is
    /// `p_key - p_query` -- which is why the encoder computes its own
    /// rather than calling the decoder's row helper.
    pub alibi_slopes: Option<Vec<f32>>,
    pub pooling: PoolingType,
    /// `[CLS]` / `[SEP]`, from `tokenizer.ggml.bos_token_id` and
    /// `tokenizer.ggml.seperator_token_id` (upstream's spelling of the
    /// key, typo included). See [`BertEncoder::wrap_special`].
    pub cls_id: u32,
    pub sep_id: u32,
}

impl BertHparams {
    pub fn head_dim(&self) -> usize {
        self.n_embd / self.n_head
    }
}

/// One transformer block's weights. Biases that llama.cpp marks
/// `TENSOR_NOT_REQUIRED` are `Option`, so a checkpoint without them is
/// run without them rather than with a silently fabricated zero vector.
/// Where an encoder layer's norms sit, and which function they are.
///
/// `bert.cpp`'s graph is POST-norm with LayerNorm: attention, residual,
/// norm, FFN, residual, norm. `neo-bert.cpp:59-118` and
/// `eurobert.cpp:55-114` are the other shape -- RMSNorm BEFORE each
/// block and a bare residual after it, with one final norm at the end
/// -- and they are one topology with four table columns between them,
/// not two graphs.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BertTopology {
    /// `bert.cpp`: LayerNorm with biases, after each residual.
    PostNormLayerNorm,
    /// `neo-bert` / `eurobert`: RMSNorm with a weight and no bias,
    /// before each block, and one final norm.
    PreNormRms,
}

/// The two FFN shapes `bert.cpp` builds on this graph for the
/// architectures frink serves.
///
/// `bert` and `nomic-bert-moe`'s dense layers take the ungated GELU
/// with both biases (`:179-187`); `nomic-bert` takes the gated SiLU
/// with none (`:195-201`, the final `else`). The variant is the
/// architecture's, read once at load, so a layer body cannot ask
/// "is there a gate tensor" and answer differently on two files.
#[derive(Debug, Clone, Copy, PartialEq, Eq)]
pub enum BertFfn {
    /// `down(gelu(up(x)))`, biases where the file has them.
    GeluSeq,
    /// `down(silu(gate(x)) * up(x))`, no biases.
    SwigluPar,
    /// `down(gelu(gate(x)) * up(x))`: `jina-bert-v2` with a separate
    /// `ffn_gate` (`bert.cpp:188-194` picks `LLM_FFN_GELU` under
    /// `LLM_FFN_PAR`, which `build_ffn` turns into `ggml_geglu_split`).
    GegluPar,
    /// The same function with the gate FUSED into `ffn_up`: the matrix
    /// is `2 * n_ff` rows and `ggml_geglu` splits it, the FIRST half
    /// being the gate (`bert.cpp:189`, `up_contains_gate`).
    GegluFusedUp,
    /// SwiGLU with the gate fused into `ffn_up` the same way:
    /// `neo-bert.cpp:35,110-115` creates a `2 * n_ff`-wide matrix and
    /// passes `LLM_FFN_SWIGLU` under `LLM_FFN_SEQ`, which is
    /// `ggml_swiglu` over the doubled row.
    SwigluFusedUp,
}

/// One layer's Q/K LayerNorm pair, weights and biases.
#[derive(Debug, Clone)]
pub struct QkLayerNorm {
    pub q_w: Vec<f32>,
    pub q_b: Vec<f32>,
    pub k_w: Vec<f32>,
    pub k_b: Vec<f32>,
}

pub struct BertLayer {
    pub wq: WeightMatrix,
    pub bq: Option<Vec<f32>>,
    pub wk: WeightMatrix,
    pub bk: Option<Vec<f32>>,
    pub wv: WeightMatrix,
    pub bv: Option<Vec<f32>>,
    pub wo: WeightMatrix,
    pub bo: Option<Vec<f32>>,
    /// A LayerNorm over the WHOLE Q / K projection, with biases, when
    /// the file carries one (`bert.cpp:109-123`; `jina-bert-v2.cpp:
    /// 30-35` creates the pair optional). Not per head: the reshape at
    /// `:110` is `n_embd_head * n_head` wide, so one norm covers every
    /// head's channels together.
    pub qk_norm: Option<QkLayerNorm>,
    /// `attn_norm` / `ffn_norm`: the PRE-norm weights, `Some` exactly
    /// for [`BertTopology::PreNormRms`].
    pub pre_attn_norm: Option<Vec<f32>>,
    pub pre_ffn_norm: Option<Vec<f32>>,
    /// `attn_output_norm`, applied after the attention residual.
    /// `Some` exactly for [`BertTopology::PostNormLayerNorm`].
    pub attn_out_norm_w: Option<Vec<f32>>,
    pub attn_out_norm_b: Option<Vec<f32>>,
    /// `attn_norm_2`, `jina-bert-v2`'s second attention norm
    /// (`bert.cpp:156-159`): the LAYER INPUT is re-added and normed a
    /// second time before the FFN reads it.
    pub attn_norm_2: Option<(Vec<f32>, Vec<f32>)>,
    pub ffn_up: WeightMatrix,
    pub ffn_up_b: Option<Vec<f32>>,
    /// `ffn_gate`, present exactly for [`BertFfn::SwigluPar`].
    pub ffn_gate: Option<WeightMatrix>,
    pub ffn_down: WeightMatrix,
    pub ffn_down_b: Option<Vec<f32>>,
    /// `layer_output_norm`, applied after the FFN residual. `Some`
    /// exactly for [`BertTopology::PostNormLayerNorm`].
    pub layer_out_norm_w: Option<Vec<f32>>,
    pub layer_out_norm_b: Option<Vec<f32>>,
}

pub struct BertEncoder {
    pub hp: BertHparams,
    pub tok_embd: WeightMatrix,
    /// `token_types.weight`, **every** row: `[n_token_types, n_embd]`.
    /// Row 0 is "Sentence A" and row 1 "Sentence B". `None` when the
    /// checkpoint carries no table at all, which upstream allows
    /// (`TENSOR_NOT_REQUIRED`) and which means no segment embedding is
    /// added anywhere. Loading only row 0 — what this held before — is
    /// what made a rerank pair score both halves as Sentence A.
    pub type_embd: Option<Vec<Vec<f32>>>,
    /// The learned position table, `None` for a rotating architecture
    /// (see [`BertHparams::rope_theta`]).
    pub pos_embd: Option<WeightMatrix>,
    /// The embedding LayerNorm, `Some` exactly for
    /// [`BertTopology::PostNormLayerNorm`]; a pre-norm encoder feeds
    /// the raw embeddings into layer 0.
    pub tok_norm_w: Option<Vec<f32>>,
    pub tok_norm_b: Option<Vec<f32>>,
    /// The final RMSNorm weight, `Some` exactly for
    /// [`BertTopology::PreNormRms`] (`enc.output_norm` for
    /// `neo-bert`, `output_norm` for `eurobert`).
    pub final_norm: Option<Vec<f32>>,
    pub layers: Vec<BertLayer>,
}

/// Adds `bias` to every `width`-wide row of `rows`, when there is one.
fn add_bias_rows(rows: &mut [f32], width: usize, bias: Option<&Vec<f32>>) {
    let Some(b) = bias else { return };
    debug_assert_eq!(b.len(), width);
    for row in rows.chunks_exact_mut(width) {
        for (x, bv) in row.iter_mut().zip(b.iter()) {
            *x += bv;
        }
    }
}

/// RMSNorm applied independently to each `width`-wide row, returning a
/// new buffer: the pre-norm shape needs the normed value AND the
/// unnormed residual, so this one does not work in place.
fn rms_norm_rows(rows: &[f32], width: usize, weight: &[f32], eps: f32) -> Vec<f32> {
    debug_assert_eq!(weight.len(), width);
    rows.chunks_exact(width)
        .flat_map(|row| frink_core::matmul::rms_norm(row, weight, eps))
        .collect()
}

/// LayerNorm applied independently to each `width`-wide row, in place.
fn layer_norm_rows(rows: &mut [f32], width: usize, weight: &[f32], bias: &[f32], eps: f32) {
    for row in rows.chunks_exact_mut(width) {
        let normed = layer_norm(row, weight, bias, eps);
        row.copy_from_slice(&normed);
    }
}

/// In-place softmax over one score row, max-shifted.
fn softmax_row(scores: &mut [f32]) {
    let max = scores.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let mut sum = 0.0f32;
    for s in scores.iter_mut() {
        *s = (*s - max).exp();
        sum += *s;
    }
    let inv = 1.0 / sum;
    for s in scores.iter_mut() {
        *s *= inv;
    }
}

/// Full bidirectional multi-head attention over `n` positions.
///
/// `q` is `[n][n_head * head_dim]`; `k` and `v` are
/// `[n][n_head_kv * head_dim]`. **Every query row attends to every key
/// row** — there is no mask argument here on purpose, so a causal mask
/// cannot be added by accident.
#[allow(clippy::too_many_arguments)]
fn bidirectional_attention(
    q: &[f32],
    k: &[f32],
    v: &[f32],
    n: usize,
    n_head: usize,
    n_head_kv: usize,
    head_dim: usize,
    // One slope per head, or `None`. The bias is `-slope * |i - j|`:
    // SYMMETRIC, because `llama-graph.cpp:442` fills a non-causal
    // model's mask with `-|p0 - p1|` and `ggml_soft_max_ext`
    // multiplies it by the head's slope.
    alibi_slopes: Option<&[f32]>,
) -> Vec<f32> {
    let q_width = n_head * head_dim;
    let kv_width = n_head_kv * head_dim;
    let heads_per_kv = n_head / n_head_kv;
    let scale = 1.0 / (head_dim as f32).sqrt();
    let mut out = vec![0.0f32; n * q_width];
    let mut scores = vec![0.0f32; n];
    for h in 0..n_head {
        let kv_h = h / heads_per_kv;
        let q_off = h * head_dim;
        let kv_off = kv_h * head_dim;
        let slope = alibi_slopes.map(|s| s[h]);
        for i in 0..n {
            let qi = &q[i * q_width + q_off..i * q_width + q_off + head_dim];
            for (j, s) in scores.iter_mut().enumerate() {
                let kj = &k[j * kv_width + kv_off..j * kv_width + kv_off + head_dim];
                *s = qi.iter().zip(kj).map(|(a, b)| a * b).sum::<f32>() * scale;
                if let Some(slope) = slope {
                    *s += slope * -((i as f32 - j as f32).abs());
                }
            }
            softmax_row(&mut scores);
            let dst = &mut out[i * q_width + q_off..i * q_width + q_off + head_dim];
            for (j, &p) in scores.iter().enumerate() {
                let vj = &v[j * kv_width + kv_off..j * kv_width + kv_off + head_dim];
                for (o, &vv) in dst.iter_mut().zip(vj) {
                    *o += p * vv;
                }
            }
        }
    }
    out
}

impl BertEncoder {
    pub fn vocab_size(&self) -> usize {
        self.tok_embd.rows()
    }
}

impl TextEncoder for BertEncoder {
    fn bert_hparams(&self) -> Option<&BertHparams> {
        Some(&self.hp)
    }

    fn n_embd(&self) -> usize {
        self.hp.n_embd
    }

    fn n_ctx_train(&self) -> usize {
        self.hp.n_ctx_train
    }

    fn pooling_type(&self) -> PoolingType {
        self.hp.pooling
    }

    /// `[CLS] … [SEP]`, which is what llama.cpp's WPM branch adds when
    /// `add_special` is set: it pushes `special_bos_id` before the
    /// pieces and `special_sep_id` after, unconditionally — the
    /// `add_bos`/`add_eos` flags are not consulted on that path
    /// (`llama-vocab.cpp`, `case LLAMA_VOCAB_TYPE_WPM`).
    fn wrap_special(&self, pieces: &[u32]) -> Vec<u32> {
        let mut out = Vec::with_capacity(pieces.len() + 2);
        out.push(self.hp.cls_id);
        out.extend_from_slice(pieces);
        out.push(self.hp.sep_id);
        out
    }

    /// The height of `token_types.weight`, or 1 when the checkpoint
    /// carries no table (nothing is added at any position, which is
    /// what a one-row table would do anyway).
    fn n_segments(&self) -> usize {
        self.type_embd.as_ref().map(Vec::len).unwrap_or(1)
    }

    /// `[CLS] a [SEP] b [SEP]` with segments `0…0 1…1` — what
    /// HuggingFace's `tokenizer(query, document)` builds for a BERT
    /// cross-encoder, which is the input these checkpoints were
    /// trained on.
    ///
    /// The boundary is defined once, here, and both vectors are cut on
    /// it: the first `[SEP]` closes segment 0 (HF counts it as part of
    /// the first half) and everything after it is segment 1. Returning
    /// the ids alone and letting the graph assume a segment is how the
    /// document half came to be scored as "Sentence A".
    fn wrap_special_pair(&self, a: &[u32], b: &[u32]) -> Option<PairSequence> {
        let mut tokens = Vec::with_capacity(a.len() + b.len() + 3);
        tokens.push(self.hp.cls_id);
        tokens.extend_from_slice(a);
        tokens.push(self.hp.sep_id);
        let first_half = tokens.len();
        tokens.extend_from_slice(b);
        tokens.push(self.hp.sep_id);
        let mut segments = vec![0u32; tokens.len()];
        for s in segments[first_half..].iter_mut() {
            *s = 1;
        }
        Some(PairSequence { tokens, segments })
    }

    fn encode_on_worker(
        &self,
        tokens: &[u32],
        segments: Option<&[u32]>,
    ) -> Result<Vec<f32>, EncodeError> {
        let n = tokens.len();
        if n == 0 {
            return Err(EncodeError::EmptySequence);
        }
        if let Some(seg) = segments {
            if seg.len() != n {
                return Err(EncodeError::RaggedSegments {
                    tokens: n,
                    segments: seg.len(),
                });
            }
        }
        if n > self.hp.n_ctx_train {
            return Err(EncodeError::TooLong {
                got: n,
                max: self.hp.n_ctx_train,
                arch: self.hp.arch.clone(),
            });
        }
        let d = self.hp.n_embd;
        let vocab_size = self.vocab_size();

        // (1)(2) token + type + position, then the input LayerNorm.
        let mut h = vec![0.0f32; n * d];
        for (i, &t) in tokens.iter().enumerate() {
            if t as usize >= vocab_size {
                return Err(EncodeError::TokenOutOfRange { id: t, vocab_size });
            }
            let tok = self.tok_embd.dequant_row(t as usize);
            let row = &mut h[i * d..(i + 1) * d];
            match &self.pos_embd {
                Some(table) => {
                    let pos = table.dequant_row(i);
                    for (j, slot) in row.iter_mut().enumerate() {
                        *slot = tok[j] + pos[j];
                    }
                }
                // A rotating architecture adds no table here;
                // `bert.cpp:90` is gated on `arch == LLM_ARCH_BERT`.
                None => row.copy_from_slice(&tok),
            }
            if let Some(table) = &self.type_embd {
                let seg = segments.map(|s| s[i]).unwrap_or(0);
                let ty = table
                    .get(seg as usize)
                    .ok_or(EncodeError::SegmentOutOfRange {
                        id: seg,
                        pos: i,
                        n_segments: table.len(),
                    })?;
                for (slot, tv) in row.iter_mut().zip(ty.iter()) {
                    *slot += tv;
                }
            }
        }
        if let (Some(w), Some(b)) = (&self.tok_norm_w, &self.tok_norm_b) {
            layer_norm_rows(&mut h, d, w, b, self.hp.layer_norm_eps);
        }

        let head_dim = self.hp.head_dim();
        for layer in &self.layers {
            // A pre-norm layer normalises what the block reads and
            // leaves the residual alone; a post-norm one reads the
            // residual directly and norms after each add.
            // `neo-bert.cpp:62-65` against `bert.cpp:103`.
            let block_in = match &layer.pre_attn_norm {
                None => h.clone(),
                Some(w) => rms_norm_rows(&h, d, w, self.hp.layer_norm_eps),
            };
            let mut q = layer.wq.apply_batch(&block_in, n);
            let mut k = layer.wk.apply_batch(&block_in, n);
            let mut v = layer.wv.apply_batch(&block_in, n);
            add_bias_rows(&mut q, self.hp.n_head * head_dim, layer.bq.as_ref());
            add_bias_rows(&mut k, self.hp.n_head_kv * head_dim, layer.bk.as_ref());
            add_bias_rows(&mut v, self.hp.n_head_kv * head_dim, layer.bv.as_ref());

            if let Some(qk) = &layer.qk_norm {
                // `bert.cpp:109-123`: a LayerNorm over the whole
                // projection, BEFORE the rotation below, which is the
                // order the graph builds them in.
                layer_norm_rows(
                    &mut q,
                    self.hp.n_head * head_dim,
                    &qk.q_w,
                    &qk.q_b,
                    self.hp.layer_norm_eps,
                );
                layer_norm_rows(
                    &mut k,
                    self.hp.n_head_kv * head_dim,
                    &qk.k_w,
                    &qk.k_b,
                    self.hp.layer_norm_eps,
                );
            }

            if let Some(theta) = self.hp.rope_theta {
                // `bert.cpp:126-133`, NEOX (`llama_model_rope_type`),
                // over the first `rope_dim` channels of every head and
                // at the row's own position -- the same rotation the
                // decoder path applies, on both Q and K.
                let rot = self.hp.rope_dim.min(head_dim);
                // NORM (interleaved pairs) for `neo-bert`, NEOX
                // (split-half) for the others: `llama_model_rope_type`
                // is the table and `BertHparams::rope_interleaved`
                // carries its answer.
                let rotate: fn(&mut [f32], usize, f32) = if self.hp.rope_interleaved {
                    frink_core::attention::apply_rope_interleaved
                } else {
                    frink_core::attention::apply_rope
                };
                for (pos, row) in q.chunks_exact_mut(self.hp.n_head * head_dim).enumerate() {
                    for head in row.chunks_exact_mut(head_dim) {
                        rotate(&mut head[..rot], pos, theta);
                    }
                }
                for (pos, row) in k.chunks_exact_mut(self.hp.n_head_kv * head_dim).enumerate() {
                    for head in row.chunks_exact_mut(head_dim) {
                        rotate(&mut head[..rot], pos, theta);
                    }
                }
            }

            let attn = bidirectional_attention(
                &q,
                &k,
                &v,
                n,
                self.hp.n_head,
                self.hp.n_head_kv,
                head_dim,
                self.hp.alibi_slopes.as_deref(),
            );

            let mut x = layer.wo.apply_batch(&attn, n);
            add_bias_rows(&mut x, d, layer.bo.as_ref());
            // Residual over the *layer input*; the post-norm shape
            // then norms it, the pre-norm shape does not.
            for (xv, hv) in x.iter_mut().zip(h.iter()) {
                *xv += hv;
            }
            if let (Some(w), Some(b)) = (&layer.attn_out_norm_w, &layer.attn_out_norm_b) {
                layer_norm_rows(&mut x, d, w, b, self.hp.layer_norm_eps);
            }

            // `bert.cpp:156-159`: the layer INPUT is re-added and
            // normed a second time. Only `jina-bert-v2` carries the
            // tensor, and only on the layers that have it.
            if let Some((w, b)) = &layer.attn_norm_2 {
                for (xv, hv) in x.iter_mut().zip(h.iter()) {
                    *xv += hv;
                }
                layer_norm_rows(&mut x, d, w, b, self.hp.layer_norm_eps);
            }

            // (5) the architecture's MLP; the FFN residual is over
            // `x`, i.e. over the post-norm value, not over the layer
            // input.
            let ffn_in = match &layer.pre_ffn_norm {
                None => x.clone(),
                Some(w) => rms_norm_rows(&x, d, w, self.hp.layer_norm_eps),
            };
            let mut up = layer.ffn_up.apply_batch(&ffn_in, n);
            add_bias_rows(&mut up, self.hp.n_ff, layer.ffn_up_b.as_ref());
            match (self.hp.ffn, &layer.ffn_gate) {
                (BertFfn::GeluSeq, _) => {
                    for a in up.iter_mut() {
                        *a = gelu(*a);
                    }
                }
                (BertFfn::SwigluPar, Some(gate)) => {
                    let g = gate.apply_batch(&ffn_in, n);
                    for (a, gv) in up.iter_mut().zip(g.iter()) {
                        *a *= frink_core::matmul::silu(*gv);
                    }
                }
                (BertFfn::GegluPar, Some(gate)) => {
                    let g = gate.apply_batch(&ffn_in, n);
                    for (a, gv) in up.iter_mut().zip(g.iter()) {
                        *a *= gelu(*gv);
                    }
                }
                (BertFfn::SwigluFusedUp, _) => {
                    // `ggml_swiglu` over a `2 * n_ff`-wide row: the
                    // first half is the gate.
                    let wide = self.hp.n_ff * 2;
                    debug_assert_eq!(up.len(), n * wide);
                    let mut folded = vec![0.0f32; n * self.hp.n_ff];
                    for (row, out) in up
                        .chunks_exact(wide)
                        .zip(folded.chunks_exact_mut(self.hp.n_ff))
                    {
                        let (gate, rest) = row.split_at(self.hp.n_ff);
                        for ((o, g), u) in out.iter_mut().zip(gate).zip(rest) {
                            *o = frink_core::matmul::silu(*g) * u;
                        }
                    }
                    up = folded;
                }
                (BertFfn::GegluFusedUp, _) => {
                    // `ggml_geglu` over a `2 * n_ff`-wide row: the
                    // first half is the gate and the second the up, per
                    // row, and the result is `n_ff` wide.
                    let wide = self.hp.n_ff * 2;
                    debug_assert_eq!(up.len(), n * wide);
                    let mut folded = vec![0.0f32; n * self.hp.n_ff];
                    for (row, out) in up
                        .chunks_exact(wide)
                        .zip(folded.chunks_exact_mut(self.hp.n_ff))
                    {
                        let (gate, rest) = row.split_at(self.hp.n_ff);
                        for ((o, g), u) in out.iter_mut().zip(gate).zip(rest) {
                            *o = gelu(*g) * u;
                        }
                    }
                    up = folded;
                }
                // Unreachable through the loader, which builds the
                // pair together; spelled so a third FFN has to answer
                // here rather than silently running GELU.
                (BertFfn::SwigluPar | BertFfn::GegluPar, None) => {
                    return Err(EncodeError::MissingGate { layer: 0 });
                }
            }
            let mut down = layer.ffn_down.apply_batch(&up, n);
            add_bias_rows(&mut down, d, layer.ffn_down_b.as_ref());
            for (dv, xv) in down.iter_mut().zip(x.iter()) {
                *dv += xv;
            }
            if let (Some(w), Some(b)) = (&layer.layer_out_norm_w, &layer.layer_out_norm_b) {
                layer_norm_rows(&mut down, d, w, b, self.hp.layer_norm_eps);
            }
            h = down;
        }
        // One final norm for the pre-norm shape, which has normed
        // nothing since the last block read its input
        // (`neo-bert.cpp:122-125`, `eurobert.cpp:116-119`).
        if let Some(w) = &self.final_norm {
            h = rms_norm_rows(&h, d, w, self.hp.layer_norm_eps);
        }
        Ok(h)
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use frink_core::tensor::Tensor;

    /// Deterministic pseudo-random weights: a small LCG, so the fixture
    /// is reproducible without pulling in a dependency.
    struct Lcg(u64);
    impl Lcg {
        fn next_f32(&mut self) -> f32 {
            self.0 = self.0.wrapping_mul(6364136223846793005).wrapping_add(1);
            ((self.0 >> 33) as f32 / (1u64 << 31) as f32) - 0.5
        }
        fn vec(&mut self, n: usize) -> Vec<f32> {
            (0..n).map(|_| self.next_f32()).collect()
        }
        fn matrix(&mut self, rows: usize, cols: usize) -> WeightMatrix {
            WeightMatrix::F32(Tensor::new(self.vec(rows * cols), vec![rows, cols]))
        }
    }

    const D: usize = 8;
    const FF: usize = 16;
    const HEADS: usize = 2;
    const VOCAB: usize = 20;
    const CTX: usize = 12;
    const EPS: f32 = 1e-12;

    fn fixture(n_layer: usize) -> BertEncoder {
        let mut r = Lcg(0x5EED);
        let tok_embd = r.matrix(VOCAB, D);
        let pos_embd = r.matrix(CTX, D);
        // Two rows, like every real BERT: "Sentence A" and "Sentence B".
        let type_embd = Some(vec![r.vec(D), r.vec(D)]);
        let tok_norm_w = r.vec(D);
        let tok_norm_b = r.vec(D);
        let layers = (0..n_layer)
            .map(|_| BertLayer {
                ffn_gate: None,
                qk_norm: None,
                attn_norm_2: None,
                pre_attn_norm: None,
                pre_ffn_norm: None,
                wq: r.matrix(D, D),
                bq: Some(r.vec(D)),
                wk: r.matrix(D, D),
                bk: Some(r.vec(D)),
                wv: r.matrix(D, D),
                bv: Some(r.vec(D)),
                wo: r.matrix(D, D),
                bo: Some(r.vec(D)),
                attn_out_norm_w: Some(r.vec(D)),
                attn_out_norm_b: Some(r.vec(D)),
                ffn_up: r.matrix(FF, D),
                ffn_up_b: Some(r.vec(FF)),
                ffn_down: r.matrix(D, FF),
                ffn_down_b: Some(r.vec(D)),
                layer_out_norm_w: Some(r.vec(D)),
                layer_out_norm_b: Some(r.vec(D)),
            })
            .collect();
        BertEncoder {
            hp: BertHparams {
                topology: BertTopology::PostNormLayerNorm,
                rope_interleaved: false,
                alibi_slopes: None,
                rope_theta: None,
                rope_dim: 0,
                ffn: BertFfn::GeluSeq,
                arch: "bert".into(),
                n_layer,
                n_embd: D,
                n_ff: FF,
                n_head: HEADS,
                n_head_kv: HEADS,
                n_ctx_train: CTX,
                n_token_types: 2,
                layer_norm_eps: EPS,
                pooling: PoolingType::Cls,
                cls_id: 1,
                sep_id: 2,
            },
            tok_embd,
            type_embd,
            pos_embd: Some(pos_embd),
            tok_norm_w: Some(tok_norm_w),
            tok_norm_b: Some(tok_norm_b),
            final_norm: None,
            layers,
        }
    }

    /// An f64 transcription of the graph in the module docs, written
    /// the slowest possible way: no `apply_batch`, no shared buffers,
    /// one scalar loop per matrix element. It exists to disagree with
    /// [`BertEncoder::encode`] if the fast path transposes a matrix,
    /// drops a bias, norms the wrong residual, reuses a buffer it
    /// should not, or reads the wrong row of the token-type table.
    fn reference_forward(m: &BertEncoder, tokens: &[u32], segments: &[u32]) -> Vec<f64> {
        let d = m.hp.n_embd;
        let n = tokens.len();
        let hd = m.hp.head_dim();

        let dense = |w: &WeightMatrix| -> Vec<Vec<f64>> {
            (0..w.rows())
                .map(|r| w.dequant_row(r).iter().map(|&v| v as f64).collect())
                .collect()
        };
        let matvec = |w: &Vec<Vec<f64>>, x: &[f64]| -> Vec<f64> {
            w.iter()
                .map(|row| row.iter().zip(x).map(|(a, b)| a * b).sum())
                .collect()
        };
        let ln = |x: &[f64], wt: &[f32], b: &[f32]| -> Vec<f64> {
            let mean = x.iter().sum::<f64>() / x.len() as f64;
            let var = x.iter().map(|v| (v - mean).powi(2)).sum::<f64>() / x.len() as f64;
            let inv = 1.0 / (var + m.hp.layer_norm_eps as f64).sqrt();
            x.iter()
                .zip(wt)
                .zip(b)
                .map(|((v, w), bb)| (v - mean) * inv * (*w as f64) + (*bb as f64))
                .collect()
        };

        let mut h: Vec<Vec<f64>> = tokens
            .iter()
            .enumerate()
            .map(|(i, &t)| {
                let tok = m.tok_embd.dequant_row(t as usize);
                let pos = m
                    .pos_embd
                    .as_ref()
                    .expect("the reference model has a table")
                    .dequant_row(i);
                let ty = m
                    .type_embd
                    .as_ref()
                    .map(|t| t[segments[i] as usize].clone())
                    .unwrap_or_else(|| vec![0.0; d]);
                let row: Vec<f64> = (0..d)
                    .map(|j| tok[j] as f64 + pos[j] as f64 + ty[j] as f64)
                    .collect();
                // The naive reference covers the POST-norm topology,
                // which is the one it was written for; a pre-norm
                // model has no embedding norm at all.
                let (tw, tb) = (
                    m.tok_norm_w.as_ref().expect("post-norm reference"),
                    m.tok_norm_b.as_ref().expect("post-norm reference"),
                );
                ln(&row, tw, tb)
            })
            .collect();

        for layer in &m.layers {
            let (wq, wk, wv, wo) = (
                dense(&layer.wq),
                dense(&layer.wk),
                dense(&layer.wv),
                dense(&layer.wo),
            );
            let (wu, wd) = (dense(&layer.ffn_up), dense(&layer.ffn_down));
            let bias = |v: &mut Vec<f64>, b: &Option<Vec<f32>>| {
                if let Some(b) = b {
                    for (x, bb) in v.iter_mut().zip(b) {
                        *x += *bb as f64;
                    }
                }
            };
            let mut q = Vec::new();
            let mut k = Vec::new();
            let mut v = Vec::new();
            for row in &h {
                let mut a = matvec(&wq, row);
                bias(&mut a, &layer.bq);
                q.push(a);
                let mut a = matvec(&wk, row);
                bias(&mut a, &layer.bk);
                k.push(a);
                let mut a = matvec(&wv, row);
                bias(&mut a, &layer.bv);
                v.push(a);
            }
            let mut attn = vec![vec![0.0f64; d]; n];
            for head in 0..m.hp.n_head {
                let off = head * hd;
                for i in 0..n {
                    let raw: Vec<f64> = (0..n)
                        .map(|j| {
                            (0..hd).map(|c| q[i][off + c] * k[j][off + c]).sum::<f64>()
                                / (hd as f64).sqrt()
                        })
                        .collect();
                    let mx = raw.iter().cloned().fold(f64::NEG_INFINITY, f64::max);
                    let ex: Vec<f64> = raw.iter().map(|s| (s - mx).exp()).collect();
                    let sum: f64 = ex.iter().sum();
                    for j in 0..n {
                        let p = ex[j] / sum;
                        for c in 0..hd {
                            attn[i][off + c] += p * v[j][off + c];
                        }
                    }
                }
            }
            let mut next = Vec::new();
            for i in 0..n {
                let mut o = matvec(&wo, &attn[i]);
                bias(&mut o, &layer.bo);
                for (x, hv) in o.iter_mut().zip(&h[i]) {
                    *x += hv;
                }
                let x = ln(
                    &o,
                    layer.attn_out_norm_w.as_ref().expect("post-norm reference"),
                    layer.attn_out_norm_b.as_ref().expect("post-norm reference"),
                );
                let mut up = matvec(&wu, &x);
                bias(&mut up, &layer.ffn_up_b);
                let act: Vec<f64> = up
                    .iter()
                    .map(|&u| {
                        const K: f64 = 0.797_884_560_802_865_4;
                        const C: f64 = 0.044_715;
                        0.5 * u * (1.0 + (K * (u + C * u * u * u)).tanh())
                    })
                    .collect();
                let mut down = matvec(&wd, &act);
                bias(&mut down, &layer.ffn_down_b);
                for (dv, xv) in down.iter_mut().zip(&x) {
                    *dv += xv;
                }
                next.push(ln(
                    &down,
                    layer
                        .layer_out_norm_w
                        .as_ref()
                        .expect("post-norm reference"),
                    layer
                        .layer_out_norm_b
                        .as_ref()
                        .expect("post-norm reference"),
                ));
            }
            h = next;
        }
        h.into_iter().flatten().collect()
    }

    #[test]
    fn matches_an_independent_f64_transcription_of_the_graph() {
        let m = fixture(3);
        let tokens = [1u32, 7, 13, 4, 9, 2];
        let got = m.encode_tokens(&tokens).unwrap();
        let want = reference_forward(&m, &tokens, &[0; 6]);
        assert_eq!(got.len(), want.len());
        for (i, (g, w)) in got.iter().zip(&want).enumerate() {
            assert!(
                (*g as f64 - w).abs() < 2e-4,
                "element {i}: {g} vs reference {w}"
            );
        }
    }

    /// The same transcription, driven with a real `0 0 0 1 1 1` split.
    /// The point is not that segments *do something* — it is that the
    /// fast path reads the SAME row the reference does at every
    /// position, so an off-by-one on the boundary or a table indexed
    /// with the token id would show up here.
    #[test]
    fn the_segment_id_selects_the_token_type_row_at_every_position() {
        let m = fixture(3);
        let tokens = [1u32, 7, 13, 4, 9, 2];
        let segments = [0u32, 0, 0, 1, 1, 1];
        let got = m.encode(&tokens, Some(&segments)).unwrap();
        let want = reference_forward(&m, &tokens, &segments);
        for (i, (g, w)) in got.iter().zip(&want).enumerate() {
            assert!(
                (*g as f64 - w).abs() < 2e-4,
                "element {i}: {g} vs reference {w}"
            );
        }
        // And it is genuinely a different graph from the all-zeros one,
        // which is the whole of issue #44: scoring the second half as
        // "Sentence A" is not a rounding difference.
        let all_zero = m.encode_tokens(&tokens).unwrap();
        let moved: f32 = all_zero
            .iter()
            .zip(&got)
            .map(|(x, y)| (x - y).abs())
            .sum::<f32>();
        assert!(moved > 1e-3, "segment 1 changed nothing ({moved})");
    }

    /// A segment id with no row, and a segment list that is not one per
    /// token, are refusals rather than a panic or a silently wrong row.
    #[test]
    fn a_segment_id_off_the_table_and_a_ragged_segment_list_are_refused() {
        let m = fixture(1);
        assert!(matches!(
            m.encode(&[1, 7, 2], Some(&[0, 2, 0])),
            Err(EncodeError::SegmentOutOfRange { id: 2, pos: 1, .. })
        ));
        assert!(matches!(
            m.encode(&[1, 7, 2], Some(&[0, 0])),
            Err(EncodeError::RaggedSegments {
                tokens: 3,
                segments: 2
            })
        ));
    }

    /// The property that makes this an encoder. Row 0's output must
    /// change when the *last* token changes; under a causal mask it
    /// could not, because position 0 would attend only to itself.
    #[test]
    fn attention_is_bidirectional_not_causal() {
        let m = fixture(2);
        let a = m.encode_tokens(&[5u32, 6, 7, 8]).unwrap();
        let b = m.encode_tokens(&[5u32, 6, 7, 19]).unwrap();
        let moved: f32 = a[..D].iter().zip(&b[..D]).map(|(x, y)| (x - y).abs()).sum();
        assert!(
            moved > 1e-3,
            "row 0 barely moved ({moved}) when the last token changed — \
             attention is behaving causally"
        );
    }

    /// Position is a learned table lookup, so the same token at a
    /// different index must land somewhere else.
    #[test]
    fn position_embeddings_make_the_same_token_differ_by_index() {
        let m = fixture(1);
        let out = m.encode_tokens(&[11u32, 11]).unwrap();
        let delta: f32 = out[..D]
            .iter()
            .zip(&out[D..2 * D])
            .map(|(x, y)| (x - y).abs())
            .sum();
        assert!(
            delta > 1e-3,
            "identical tokens gave identical rows: {delta}"
        );
    }

    /// The graph ends on a LayerNorm: with unit weight and zero bias
    /// each output row is mean-zero and unit-variance. An RMSNorm in
    /// that slot would leave the mean wherever it was.
    #[test]
    fn the_last_op_is_a_mean_subtracting_layer_norm() {
        let mut m = fixture(2);
        let last = m.layers.last_mut().unwrap();
        last.layer_out_norm_w = Some(vec![1.0; D]);
        last.layer_out_norm_b = Some(vec![0.0; D]);
        let out = m.encode_tokens(&[3u32, 4, 5]).unwrap();
        for row in out.as_chunks::<D>().0 {
            let mean: f32 = row.iter().sum::<f32>() / D as f32;
            let var: f32 = row.iter().map(|v| (v - mean).powi(2)).sum::<f32>() / D as f32;
            assert!(mean.abs() < 1e-4, "row mean {mean} is not zero");
            assert!((var - 1.0).abs() < 1e-3, "row variance {var} is not one");
        }
    }

    #[test]
    fn refuses_an_empty_sequence_and_one_past_the_position_table() {
        let m = fixture(1);
        assert!(matches!(
            m.encode_tokens(&[]),
            Err(EncodeError::EmptySequence)
        ));
        let long: Vec<u32> = (0..CTX as u32 + 1).map(|i| i % VOCAB as u32).collect();
        let err = m.encode_tokens(&long).unwrap_err();
        assert!(
            matches!(err, EncodeError::TooLong { got, max, .. } if got == CTX + 1 && max == CTX)
        );
        assert!(matches!(
            m.encode_tokens(&[VOCAB as u32]),
            Err(EncodeError::TokenOutOfRange { .. })
        ));
    }

    #[test]
    fn wrap_special_brackets_the_pieces_with_cls_and_sep() {
        let m = fixture(1);
        assert_eq!(m.wrap_special(&[7, 8]), vec![1, 7, 8, 2]);
        assert_eq!(m.wrap_special(&[]), vec![1, 2]);
    }

    /// The cross-encoder input is `[CLS] a [SEP] b [SEP]` with segments
    /// `0 0 0 0 1 1` — the boundary between the two halves is the whole
    /// reason a reranker scores differently from an embedding model.
    /// Concatenating without it, dropping the trailing `[SEP]`, or
    /// leaving every segment at 0, produces a perfectly plausible
    /// ranking that is not the model's, so both vectors are asserted
    /// exactly rather than by length.
    ///
    /// The first `[SEP]` belongs to segment 0, which is what
    /// HuggingFace's `tokenizer(query, document)` emits: an off-by-one
    /// there is a one-position difference that no shape check catches.
    #[test]
    fn the_pair_form_separates_the_two_halves_and_labels_each_one() {
        let m = fixture(1);
        let pair = m.wrap_special_pair(&[7, 8], &[9]).unwrap();
        assert_eq!(pair.tokens, vec![1, 7, 8, 2, 9, 2]);
        assert_eq!(pair.segments, vec![0, 0, 0, 0, 1, 1]);
        // An empty half is still a half: the boundary stays.
        let empty = m.wrap_special_pair(&[], &[]).unwrap();
        assert_eq!(empty.tokens, vec![1, 2, 2]);
        assert_eq!(empty.segments, vec![0, 0, 1]);
        // And it is NOT the single-sequence form of the two texts run
        // together, which is what a defaulted implementation would give.
        assert_ne!(pair.tokens, m.wrap_special(&[7, 8, 9]));
    }

    /// A checkpoint with no "Sentence B" row cannot express a pair, and
    /// [`crate::EmbeddingModel`] refuses one at load. This is the value
    /// that refusal reads.
    #[test]
    fn n_segments_is_the_height_of_the_token_type_table() {
        let mut m = fixture(1);
        assert_eq!(m.n_segments(), 2);
        m.type_embd = Some(vec![vec![0.0; D]]);
        assert_eq!(m.n_segments(), 1);
        m.type_embd = None;
        assert_eq!(m.n_segments(), 1);
    }

    /// `embed_tokens` must return the CLS row of the hidden states this
    /// checkpoint's `pooling_type` names, not the mean and not the last.
    #[test]
    fn embed_tokens_pools_the_way_the_hparams_say() {
        let m = fixture(2);
        let tokens = [1u32, 9, 4, 2];
        let hidden = m.encode_tokens(&tokens).unwrap();
        assert_eq!(m.embed_tokens(&tokens).unwrap(), hidden[..D].to_vec());
    }
}