inferencelayer 0.2.3

//! GLiNER — zero-shot span NER on the engine's DeBERTa backbone.
//!
//! The backbone (disentangled attention) is the engine's; everything here is the head that turns
//! its per-token states into typed entity spans. The path, taken from the real checkpoint
//! (`urchade/gliner_small-v2.1`) and gated against the `gliner` package layer by layer:
//!
//! ```text
//! input:  <<ENT>> type₁ <<ENT>> type₂ … <<SEP>> word₁ word₂ …
//!
//! DeBERTa hidden states                                   [T, 768]
//!   → projection                       Linear(768 → 512)  [T, 512]
//!   → gather   prompts at the <<ENT>> positions           [C, 512]
//!              words at each word's FIRST subtoken        [W, 512]   (subtoken_pooling = "first")
//!   → BiLSTM   1 layer, 256/dir, concat                   [W, 512]
//!   → SpanMarkerV0:  start = project_start(words)
//!                    end   = project_end(words)
//!                    span[L][k] = out_project(relu(start[L] ‖ end[L+k]))
//!                                                         [W, max_width, 512]
//!   → prompt_rep(prompts)                                 [C, 512]
//!   → scores[L][k][c] = span[L][k] · prompt[c]            [W, max_width, C]
//!   → sigmoid, threshold, flat greedy overlap resolution  → entities
//! ```
//!
//! Every MLP is GLiNER's `create_projection_layer`: `Linear(d → 4·out) → ReLU → Dropout → Linear`.
//! Dropout is inference-time identity, which is why only the `.0.` and `.3.` weights exist.
//!
//! Span `(L, k)` means the word range `[L, L+k]` **inclusive**; it is scored only when `L+k < W`.

use anyhow::{Context, Result};
use std::path::Path;

use rayon::prelude::*;

use crate::EmbedEngine;
use crate::encoder_weights::EncBatch;
use crate::gliner_gpu::GlinerGpuHead;
use crate::weights::LazySt;

/// GPU scratch is sized once at load, so these bound what a single call may ask for. Generous:
/// GLiNER's backbone tops out at 512 positions anyway, and 25 types is its trained `max_types`.
const MAX_WORDS: usize = 512;
const MAX_TYPES: usize = 32;

/// Which device GLiNER runs on — backbone AND head together (they share one adapter, so the hidden
/// states never round-trip between devices).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum GlinerDevice {
    /// wgpu adapter when one exists (Metal/Vulkan/DX12/GL), CPU otherwise.
    Auto,
    /// Force the native CPU encoder.
    Cpu,
}

/// One `nn.Linear`: `y = x·Wᵀ + b`, HF layout (`W` is `[n, k]` row-major).
#[derive(Default)]
struct Linear {
    w: Vec<f32>,
    b: Vec<f32>,
    n: usize,
    k: usize,
    /// Packed on first use — see [`crate::cpu_gemm`]. The head's weights are immutable, and the
    /// `gemm` crate repacks them on every call.
    packed: std::sync::OnceLock<crate::cpu_gemm::PackedWeight>,
}

impl Linear {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        let w = st.tensor_f32(&format!("{prefix}.weight"))?;
        let b = st.tensor_f32(&format!("{prefix}.bias"))?;
        let n = b.len();
        let k = w.len() / n;
        Ok(Self {
            w,
            b,
            n,
            k,
            packed: std::sync::OnceLock::new(),
        })
    }

    /// `x` is `[rows, k]` → `[rows, n]`, through the same prepacked NEON GEMM the CPU encoder uses.
    ///
    /// A hand-rolled scalar dot-product loop here made the whole model **3.6× slower than torch**
    /// (which dispatches to a real BLAS): the head's matmuls — `out_project` over every span above
    /// all — are the bulk of GLiNER's compute, so this must be a real GEMM, and the caller must
    /// batch its rows into ONE call to give it an `m` worth parallelizing.
    fn forward(&self, x: &[f32]) -> Vec<f32> {
        let (n, k) = (self.n, self.k);
        let m = x.len() / k;
        assert_eq!(x.len(), m * k, "lhs shape");
        assert_eq!(self.w.len(), n * k, "weight shape");
        let mut out = vec![0f32; m * n];
        let packed = self
            .packed
            .get_or_init(|| crate::cpu_gemm::PackedWeight::new(&self.w, n, k));
        crate::cpu_gemm::gemm_packed(&mut out, x, packed, m, Some(&self.b));
        out
    }
}

/// GLiNER's `create_projection_layer`: `Linear(d → 4·out) → ReLU → Linear(4·out → out)`.
/// (The dropout between them is identity at inference — hence the `.0.` / `.3.` naming.)
struct Mlp {
    up: Linear,
    down: Linear,
}

impl Mlp {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        Ok(Self {
            up: Linear::load(st, &format!("{prefix}.0"))?,
            down: Linear::load(st, &format!("{prefix}.3"))?,
        })
    }

    fn forward(&self, x: &[f32]) -> Vec<f32> {
        let mut h = self.up.forward(x);
        for v in h.iter_mut() {
            *v = v.max(0.0); // ReLU
        }
        self.down.forward(&h)
    }
}

/// A single-layer **bidirectional** LSTM (`nn.LSTM(input, hidden, bidirectional, batch_first)`).
/// PyTorch packs the four gates row-wise in the order **i, f, g, o**, and applies BOTH biases
/// (`b_ih` and `b_hh`) — getting either wrong yields plausible-but-wrong states, so both are pinned
/// by the parity gate.
struct BiLstm {
    /// `[w_ih (4H×I), w_hh (4H×H), b_ih (4H), b_hh (4H)]` for the forward direction…
    fwd: LstmDir,
    /// …and the reverse, which runs the sequence right-to-left.
    rev: LstmDir,
    hidden: usize,
}

struct LstmDir {
    /// `W_ih` + `b_ih` as a [`Linear`], so the state-independent half of the gates is one GEMM.
    in_proj: Linear,
    w_hh: Vec<f32>,
    b_hh: Vec<f32>,
    input: usize,
    hidden: usize,
}

impl LstmDir {
    fn load(st: &LazySt, prefix: &str, suffix: &str) -> Result<Self> {
        let w_ih = st.tensor_f32(&format!("{prefix}.weight_ih_l0{suffix}"))?;
        let w_hh = st.tensor_f32(&format!("{prefix}.weight_hh_l0{suffix}"))?;
        let b_ih = st.tensor_f32(&format!("{prefix}.bias_ih_l0{suffix}"))?;
        let b_hh = st.tensor_f32(&format!("{prefix}.bias_hh_l0{suffix}"))?;
        let hidden = b_ih.len() / 4;
        let input = w_ih.len() / (4 * hidden);
        Ok(Self {
            in_proj: Linear {
                n: 4 * hidden,
                k: input,
                w: w_ih,
                b: b_ih,
                ..Default::default()
            },
            w_hh,
            b_hh,
            input,
            hidden,
        })
    }

    /// Run one direction over `x` (`[t, input]`), visiting steps in `order`. Returns `[t, hidden]`
    /// indexed by the ORIGINAL step (so the reverse pass lands back in forward order).
    ///
    /// The **input** half of the gates (`W_ih·x + b_ih`) does not depend on the recurrent state, so
    /// it is one GEMM over all timesteps up front; only `W_hh·h` has to stay in the sequential loop.
    fn run(&self, x: &[f32], order: impl Iterator<Item = usize>) -> Vec<f32> {
        let (h_n, i_n) = (self.hidden, self.input);
        let t = x.len() / i_n;
        let xg = self.in_proj.forward(x); // [t, 4H] — W_ih·x + b_ih, batched
        let mut out = vec![0f32; t * h_n];
        let mut h = vec![0f32; h_n];
        let mut c = vec![0f32; h_n];
        let mut gates = vec![0f32; 4 * h_n];
        for step in order {
            let xr = &xg[step * 4 * h_n..(step + 1) * 4 * h_n];
            // gates = (W_ih·x + b_ih) + b_hh, then ACCUMULATE W_hh·h into it.
            //
            // This is the sequential half of the recurrence and it is a GEMV — one step's `h` against
            // a [4H, H] weight. It used to call `gemm::gemm` with m=1, which was a trap: on aarch64
            // `gemm` REPACKS its weight matrix on every call, and a 176-word input makes 352 of them,
            // so the layer moved ~352 MB just repacking 1 MB of `w_hh` over and over to perform
            // 92 MFLOP of arithmetic. A GEMV has no reuse to pack *for*. Hand-rolled below: 13 -> 2 ms.
            //
            // Still serial on purpose — rayon over the 4·H independent gates means ~400 spawns for
            // one input and swamps the work (it measured 4 ms one run, 49 ms the next). The two
            // DIRECTIONS run concurrently; that is the parallelism this layer admits.
            for (g, gate) in gates.iter_mut().enumerate() {
                *gate = xr[g] + self.b_hh[g];
            }
            gemv_acc(&mut gates, &self.w_hh, &h, h_n);
            for j in 0..h_n {
                let i_g = sigmoid(gates[j]);
                let f_g = sigmoid(gates[h_n + j]);
                let g_g = gates[2 * h_n + j].tanh();
                let o_g = sigmoid(gates[3 * h_n + j]);
                c[j] = f_g * c[j] + i_g * g_g;
                h[j] = o_g * c[j].tanh();
            }
            out[step * h_n..(step + 1) * h_n].copy_from_slice(&h);
        }
        out
    }
}

impl BiLstm {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        let fwd = LstmDir::load(st, prefix, "")?;
        let rev = LstmDir::load(st, prefix, "_reverse")?;
        anyhow::ensure!(
            fwd.hidden == rev.hidden && fwd.input == rev.input,
            "BiLSTM directions disagree on geometry"
        );
        let hidden = fwd.hidden;
        Ok(Self { fwd, rev, hidden })
    }

    /// `[t, input]` → `[t, 2·hidden]`, forward states concatenated with reverse states per step.
    ///
    /// The two directions are INDEPENDENT recurrences, so they run concurrently — two coarse tasks,
    /// which is the parallelism this layer actually admits. (Not to be confused with parallelizing
    /// across the 4·H gates *within* a step: that was tried and is much worse — ~400 rayon spawns
    /// for one input, measured 4 ms one run and 49 ms the next. The recurrence itself is sequential
    /// and stays that way.)
    fn forward(&self, x: &[f32]) -> Vec<f32> {
        let t = x.len() / self.fwd.input;
        let (f, r) = rayon::join(|| self.fwd.run(x, 0..t), || self.rev.run(x, (0..t).rev()));
        let h = self.hidden;
        let mut out = vec![0f32; t * 2 * h];
        for step in 0..t {
            out[step * 2 * h..step * 2 * h + h].copy_from_slice(&f[step * h..(step + 1) * h]);
            out[step * 2 * h + h..(step + 1) * 2 * h].copy_from_slice(&r[step * h..(step + 1) * h]);
        }
        out
    }
}

/// `dst += W · x` for a row-major `W` of `dst.len()` rows and `k` columns.
///
/// The LSTM's recurrent weight is only ~1 MB and stays hot in L2, so this streams it once per step
/// with no packing at all — which is the whole point: the `gemm` crate it replaces repacked that
/// matrix on every single step. [`dot`] carries the vectorization (lane-wise accumulators; a plain
/// `.sum()` will not vectorize, because LLVM may not reassociate float addition).
#[inline]
fn gemv_acc(dst: &mut [f32], w: &[f32], x: &[f32], k: usize) {
    debug_assert_eq!(x.len(), k);
    crate::simd::gemv_acc(dst, w, x);
}

fn sigmoid(v: f32) -> f32 {
    if v >= 0.0 {
        1.0 / (1.0 + (-v).exp())
    } else {
        let e = v.exp();
        e / (1.0 + e)
    }
}

/// One predicted entity: a WORD range `[start, end]` (inclusive) with its type and probability.
#[derive(Debug, Clone, PartialEq)]
pub struct Entity {
    /// First word of the span (0-based index into the word sequence).
    pub start: usize,
    /// Last word of the span, inclusive.
    pub end: usize,
    /// Index into the entity-type list that was prompted.
    pub label: usize,
    /// Sigmoid of the span·type score.
    pub score: f32,
}

/// One predicted entity over the ORIGINAL text: byte offsets, the surface string, and its type.
#[derive(Debug, Clone, PartialEq)]
pub struct TextEntity {
    /// Byte offset of the first character (into the text passed to [`Gliner::predict_text`]).
    pub start: usize,
    /// Byte offset one past the last character.
    pub end: usize,
    /// `text[start..end]` — the matched surface form.
    pub text: String,
    /// The entity type, as spelled in the `labels` that were prompted.
    pub label: String,
    /// Sigmoid of the span·type score.
    pub score: f32,
}

/// GLiNER's word splitter: `\w+(?:[-_]\w+)*|\S` — a run of word characters (which may be joined
/// by `-`/`_`), or any single non-space character. Returns byte ranges into `text`.
///
/// Hand-rolled rather than pulling in `regex`: the pattern is small, and the WORD BOUNDARIES are
/// what the model's span indices are expressed in, so this has to be exact — a splitter that
/// disagreed with GLiNER's would silently shift every predicted offset.
fn split_words(text: &str) -> Vec<(usize, usize)> {
    fn is_word(c: char) -> bool {
        c.is_alphanumeric() || c == '_'
    }
    let chars: Vec<(usize, char)> = text.char_indices().collect();
    let end_of =
        |i: usize| -> usize { chars.get(i).map(|(o, _)| *o).unwrap_or_else(|| text.len()) };
    let mut out = Vec::new();
    let mut i = 0;
    while i < chars.len() {
        let (off, c) = chars[i];
        if c.is_whitespace() {
            i += 1;
        } else if is_word(c) {
            let mut j = i;
            while j < chars.len() && is_word(chars[j].1) {
                j += 1;
            }
            // (?:[-_]\w+)* — `_` is already a word char, so in practice this extends across `-`.
            while j + 1 < chars.len()
                && (chars[j].1 == '-' || chars[j].1 == '_')
                && is_word(chars[j + 1].1)
            {
                j += 1;
                while j < chars.len() && is_word(chars[j].1) {
                    j += 1;
                }
            }
            out.push((off, end_of(j)));
            i = j;
        } else {
            out.push((off, off + c.len_utf8()));
            i += 1;
        }
    }
    out
}

/// `out_project`'s first layer, **split down the concat seam** — the head's one real algorithmic
/// win, and the reason the engine beats torch on long text instead of merely matching it.
///
/// SpanMarkerV0 feeds `up` a concatenation, `cat = relu(start[l]) ‖ relu(end[l+k])`, so
///
/// ```text
///   W_up · cat  =  W_left · relu(start[l])  +  W_right · relu(end[l+k])
/// ```
///
/// The naive form (what torch does) materializes an `[n_spans, 1024]` matrix and burns a
/// `[n_spans, 1024] × [1024, 2048]` GEMM — and `n_spans ≈ 12·W`, so it grows with the TEXT, not the
/// word count. Projecting each half **per word** instead costs two `[W, 512] × [512, 2048]` GEMMs
/// and turns the per-span work into a vector ADD. On a 176-word text that is ~2,050 spans: 4.3
/// GMAC of GEMM becomes 0.37 GMAC + a few million adds.
///
/// Only the FIRST layer collapses this way. The ReLU after it is non-linear, so `down` still has to
/// see every span.
struct SpanUp {
    /// `W_up[:, :512]` — applied to `relu(start)`. Bias lives on the right half (added once).
    left: Linear,
    /// `W_up[:, 512:]` + `b_up` — applied to `relu(end)`.
    right: Linear,
}

impl SpanUp {
    /// Load `out_project.0` (`[4·d, 2·d]`) and slice each row at the concat seam.
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        let up = Linear::load(st, prefix)?;
        let (n, k) = (up.n, up.k); // n = 4·d (2048), k = 2·d (1024)
        anyhow::ensure!(
            k % 2 == 0,
            "out_project input {k} is not a concat of two halves"
        );
        let half = k / 2;
        let mut left = Vec::with_capacity(n * half);
        let mut right = Vec::with_capacity(n * half);
        for row in up.w.chunks_exact(k) {
            left.extend_from_slice(&row[..half]);
            right.extend_from_slice(&row[half..]);
        }
        Ok(Self {
            left: Linear {
                w: left,
                b: vec![0.0; n], // the bias is applied once, on the right half
                n,
                k: half,
                ..Default::default()
            },
            right: Linear {
                w: right,
                b: up.b,
                n,
                k: half,
                ..Default::default()
            },
        })
    }
}

/// GLiNER: the engine's DeBERTa backbone plus the span head.
pub struct Gliner {
    backbone: EmbedEngine,
    projection: Linear,
    rnn: BiLstm,
    project_start: Mlp,
    project_end: Mlp,
    /// `out_project`, with its first layer pre-split across the concat seam (see [`SpanUp`]).
    span_up: SpanUp,
    /// `out_project`'s second layer (`2048 → 512`); must see every span (the ReLU is non-linear).
    span_down: Linear,
    prompt_rep: Mlp,
    /// The head on wgpu — present whenever the backbone is (they share one device and one adapter).
    /// `None` on CPU, where the head runs through `gemm`.
    gpu_head: Option<GlinerGpuHead>,
    max_width: usize,
    hidden: usize,
    ent_token_id: u32,
    /// The checkpoint's tokenizer (with `<<ENT>>`/`<<SEP>>` as added tokens) — `None` when the dir
    /// ships no `tokenizer.json`, in which case only the id-level API ([`Gliner::scores`],
    /// [`Gliner::predict`]) is available. Behind `cli`, which is this crate's gate for the
    /// `tokenizers` dependency (it pulls a C library with no wasm build).
    #[cfg(feature = "cli")]
    tokenizer: Option<tokenizers::Tokenizer>,
    #[cfg(feature = "cli")]
    ent_token: String,
    #[cfg(feature = "cli")]
    sep_token: String,
}

impl Gliner {
    /// Load a GLiNER checkpoint (one `model.safetensors` holding backbone AND head — see
    /// `tests/fixtures/export_gliner.py`), running the backbone on the best available device.
    /// `ent_token_id` is the `<<ENT>>` token that marks each entity-type slot in the prompt.
    pub fn load(dir: &Path, ent_token_id: u32) -> Result<Self> {
        Self::load_on(dir, ent_token_id, GlinerDevice::Auto)
    }

    /// [`Self::load`] with an explicit device for the backbone (the benchmark drives both).
    pub fn load_on(dir: &Path, ent_token_id: u32, device: GlinerDevice) -> Result<Self> {
        let backbone = match device {
            // The DeBERTa GPU path is the disentangled-attention kernel (`ENC_ATTN_DISENT`).
            GlinerDevice::Auto => EmbedEngine::auto(dir, 8192)?,
            GlinerDevice::Cpu => EmbedEngine::cpu(dir)?,
        };
        let st = LazySt::open(dir)?;
        let p = "span_rep_layer.span_rep_layer";
        let projection = Linear::load(&st, "token_rep_layer.projection")?;
        let rnn = BiLstm::load(&st, "rnn.lstm")?;
        // The head's own params (`gliner_head.json`); the shipped checkpoints all use these values,
        // so a missing file is not an error.
        let head_cfg: serde_json::Value = std::fs::read(dir.join("gliner_head.json"))
            .ok()
            .and_then(|b| serde_json::from_slice(&b).ok())
            .unwrap_or(serde_json::Value::Null);
        #[cfg(feature = "cli")]
        let s = |k: &str, d: &str| -> String {
            head_cfg
                .get(k)
                .and_then(|x| x.as_str())
                .unwrap_or(d)
                .to_string()
        };
        // No truncation here: GLiNER's window is the backbone's, and silently dropping the tail of
        // a prompt would drop entity TYPES.
        #[cfg(feature = "cli")]
        let tokenizer = tokenizers::Tokenizer::from_file(dir.join("tokenizer.json")).ok();
        let head = Self {
            projection,
            hidden: rnn.hidden * 2,
            rnn,
            project_start: Mlp::load(&st, &format!("{p}.project_start"))?,
            project_end: Mlp::load(&st, &format!("{p}.project_end"))?,
            span_up: SpanUp::load(&st, &format!("{p}.out_project.0"))?,
            span_down: Linear::load(&st, &format!("{p}.out_project.3"))?,
            prompt_rep: Mlp::load(&st, "prompt_rep_layer")?,
            max_width: head_cfg
                .get("max_width")
                .and_then(|x| x.as_u64())
                .unwrap_or(12) as usize,
            backbone,
            ent_token_id,
            #[cfg(feature = "cli")]
            tokenizer,
            #[cfg(feature = "cli")]
            ent_token: s("ent_token", "<<ENT>>"),
            #[cfg(feature = "cli")]
            sep_token: s("sep_token", "<<SEP>>"),
            gpu_head: None,
        };
        anyhow::ensure!(
            head.projection.n == head.hidden,
            "projection outputs {} but the BiLSTM is {}-wide",
            head.projection.n,
            head.hidden
        );

        // The head goes wherever the backbone went — same device, same adapter, so the hidden
        // states never round-trip. The head is O(n_spans) and n_spans ≈ 12·W, so on long text it
        // outgrows the CPU: this is what keeps the engine ahead of torch at every size.
        let mut head = head;
        if let Some(ctx) = head.backbone.gpu_ctx() {
            let (d, up_w) = (head.hidden, head.span_up.left.n);
            head.gpu_head = Some(GlinerGpuHead::new(
                ctx,
                (
                    &head.project_start.up.w,
                    &head.project_start.up.b,
                    &head.project_start.down.w,
                    &head.project_start.down.b,
                ),
                (
                    &head.project_end.up.w,
                    &head.project_end.up.b,
                    &head.project_end.down.w,
                    &head.project_end.down.b,
                ),
                (&head.span_up.left.w, &head.span_up.left.b),
                (&head.span_up.right.w, &head.span_up.right.b),
                (&head.span_down.w, &head.span_down.b),
                d,
                up_w,
                head.max_width,
                MAX_WORDS,
                MAX_TYPES,
            )?);
        }
        Ok(head)
    }

    /// Override `max_width` (GLiNER's `gliner_head.json`; the shipped checkpoints use 12).
    pub fn with_max_width(mut self, max_width: usize) -> Self {
        self.max_width = max_width;
        self
    }

    /// **The end-to-end API**: raw text + entity types → entities with byte offsets into `text`.
    ///
    /// Builds GLiNER's prompt (`<<ENT>> type₁ <<ENT>> type₂ … <<SEP>> word₁ word₂ …`), tokenizes it
    /// PRE-SPLIT so each subtoken carries its word index, marks each word's first subtoken
    /// (`subtoken_pooling: "first"`), runs the model, and maps the predicted WORD ranges back to
    /// byte offsets in the original text.
    #[cfg(feature = "cli")]
    pub fn predict_text(
        &mut self,
        text: &str,
        labels: &[impl AsRef<str>],
        threshold: f32,
    ) -> Result<Vec<TextEntity>> {
        let tokenizer = self
            .tokenizer
            .as_ref()
            .context("checkpoint has no tokenizer.json — use `predict` with your own token ids")?;
        anyhow::ensure!(!labels.is_empty(), "no entity types given");

        let spans = split_words(text);
        // The tokenizer sees the prompt words and the text words as ONE pre-split sequence, so a
        // subtoken's word id indexes into `seq` — the text words start at `prompt_len`.
        let mut seq: Vec<String> = Vec::with_capacity(2 * labels.len() + 1 + spans.len());
        for l in labels {
            seq.push(self.ent_token.clone());
            seq.push(l.as_ref().to_string());
        }
        seq.push(self.sep_token.clone());
        let prompt_len = seq.len();
        seq.extend(spans.iter().map(|&(a, b)| text[a..b].to_string()));

        let enc = tokenizer
            .encode(tokenizers::InputSequence::from(seq), true)
            .map_err(|e| anyhow::anyhow!("gliner tokenize: {e}"))?;
        let ids = enc.get_ids().to_vec();

        // words_mask: the 1-based TEXT-word index on each word's FIRST subtoken, 0 elsewhere
        // (special tokens, the prompt region, and continuation subtokens).
        let mut words_mask = vec![0u32; ids.len()];
        let mut prev: Option<u32> = None;
        for (i, wid) in enc.get_word_ids().iter().enumerate() {
            if let Some(w) = wid {
                if prev != Some(*w) && (*w as usize) >= prompt_len {
                    words_mask[i] = *w + 1 - prompt_len as u32;
                }
                prev = Some(*w);
            }
        }

        let hits = self.predict(&ids, &words_mask, threshold)?;
        Ok(hits
            .into_iter()
            .map(|e| {
                let (start, _) = spans[e.start];
                let (_, end) = spans[e.end];
                TextEntity {
                    start,
                    end,
                    text: text[start..end].to_string(),
                    label: labels[e.label].as_ref().to_string(),
                    score: e.score,
                }
            })
            .collect())
    }

    /// The device the backbone runs on (`"Metal/Apple M4 Max (f16)"`, `"cpu"`, …).
    pub fn device(&self) -> String {
        self.backbone.device()
    }

    /// The span head's width (512 on the shipped checkpoints — NOT the backbone's 768).
    pub fn hidden(&self) -> usize {
        self.hidden
    }

    /// Score every span against every prompted type.
    ///
    /// `input_ids` is the tokenized prompt+text; `words_mask[i]` is the **1-based** word index when
    /// subtoken `i` is a word's FIRST subtoken and `0` otherwise (GLiNER's `subtoken_pooling:
    /// "first"`). Returns `[W][max_width][C]` raw scores — `(L, k)` is the word range `[L, L+k]`,
    /// and entries with `L + k >= W` are meaningless (the caller/decoder skips them).
    pub fn scores(&mut self, input_ids: &[u32], words_mask: &[u32]) -> Result<Vec<Vec<Vec<f32>>>> {
        anyhow::ensure!(
            input_ids.len() == words_mask.len(),
            "words_mask has {} entries for {} tokens",
            words_mask.len(),
            input_ids.len()
        );
        let d = self.hidden;

        // `GLINER_PROFILE=1` prints the per-phase budget. Guessing which phase dominates is how you
        // spend an afternoon optimizing the wrong one.
        let prof = std::env::var("GLINER_PROFILE").is_ok();
        let mut t = std::time::Instant::now();
        let mut lap = |name: &str| {
            if prof {
                eprintln!(
                    "  [profile] {name:<12} {:6.1} ms",
                    t.elapsed().as_secs_f64() * 1e3
                );
            }
            t = std::time::Instant::now();
        };

        // 1. backbone → 2. projection to the head's width.
        let states = self
            .backbone
            .forward_hidden(&EncBatch::from_seqs([input_ids.to_vec()]))
            .context("gliner backbone")?;
        lap("backbone");
        let tokens = self.projection.forward(&states);
        lap("projection");

        // 3a. prompts: the <<ENT>> markers, in order — one per entity type.
        let mut prompts: Vec<f32> = Vec::new();
        for (i, &id) in input_ids.iter().enumerate() {
            if id == self.ent_token_id {
                prompts.extend_from_slice(&tokens[i * d..(i + 1) * d]);
            }
        }
        let n_types = prompts.len() / d;
        anyhow::ensure!(n_types > 0, "no <<ENT>> markers in input_ids");

        // 3b. words: each word's FIRST subtoken (words_mask is 1-based; 0 = prompt/pad region).
        let n_words = *words_mask.iter().max().unwrap_or(&0) as usize;
        anyhow::ensure!(n_words > 0, "words_mask marks no words");
        let mut words = vec![0f32; n_words * d];
        for (i, &w) in words_mask.iter().enumerate() {
            if w > 0 {
                let dst = (w as usize - 1) * d;
                words[dst..dst + d].copy_from_slice(&tokens[i * d..(i + 1) * d]);
            }
        }

        // 4. BiLSTM over the word sequence. Inherently sequential, so it stays on the CPU.
        let words = self.rnn.forward(&words);
        lap("bilstm");

        // 5/6. SpanMarkerV0 + scoring. On the GPU this is one submit and one small readback.
        let prompt = self.prompt_rep.forward(&prompts);
        if let (Some(gh), Some(ctx)) = (&self.gpu_head, self.backbone.gpu_ctx())
            && n_words <= MAX_WORDS
            && n_types <= MAX_TYPES
        {
            let flat = gh.scores(ctx, &words, &prompt, n_words, n_types)?;
            lap("head(gpu)");
            // [n_spans, C] with span = l·max_width + k → the [W][max_width][C] the caller expects.
            let mut out = vec![vec![vec![0f32; n_types]; self.max_width]; n_words];
            for (l, per_width) in out.iter_mut().enumerate() {
                for (k, per_type) in per_width.iter_mut().enumerate() {
                    if l + k >= n_words {
                        continue; // ran past the text — the shader computed it against a clamped
                        // row; GLiNER's span_mask discards it too
                    }
                    let base = (l * self.max_width + k) * n_types;
                    per_type.copy_from_slice(&flat[base..base + n_types]);
                }
            }
            return Ok(out);
        }

        let start = self.project_start.forward(&words);
        let end = self.project_end.forward(&words);

        // Spans running past the text are skipped, exactly as GLiNER's span_mask does, so nothing
        // is computed that would be thrown away.
        let valid: Vec<(usize, usize)> = (0..n_words)
            .flat_map(|l| (0..self.max_width).map(move |k| (l, k)))
            .filter(|&(l, k)| l + k < n_words)
            .collect();

        // `out_project`'s first layer, SPLIT AT THE CONCAT SEAM (see `SpanUp`): project each half
        // once PER WORD, then the per-span work is a vector add rather than a GEMM row. The naive
        // form costs an [n_spans, 1024] × [1024, 2048] GEMM, and n_spans ≈ 12·W — so it grows with
        // the text. This is what makes long inputs cheap.
        let relu = |v: &[f32]| -> Vec<f32> { v.iter().map(|x| x.max(0.0)).collect() };
        let a = self.span_up.left.forward(&relu(&start)); // [W, 4d]
        let b = self.span_up.right.forward(&relu(&end)); // [W, 4d] (carries the bias)
        let up_w = self.span_up.left.n; // 4·d

        // h = relu(A[l] + B[l+k]) for every span → [n_spans, 4d]; `down` must still see them all
        // (the ReLU is non-linear, so it cannot be folded back into the per-word projections).
        let mut h = vec![0f32; valid.len() * up_w];
        h.par_chunks_exact_mut(up_w)
            .zip(valid.par_iter())
            .for_each(|(dst, &(l, k))| {
                let (ar, br) = (&a[l * up_w..], &b[(l + k) * up_w..]);
                for j in 0..up_w {
                    dst[j] = (ar[j] + br[j]).max(0.0);
                }
            });
        let span_reps = self.span_down.forward(&h);
        lap("head(cpu)");

        // 6. score every span against every type. This is an [n_spans, d] × [d, n_types] matmul —
        // ~2,300 × 512 × 20 at 176 words — and it was a scalar dot loop. `prompt` is [C, d], i.e.
        // exactly a weight matrix in HF layout, so it goes through the same GEMM as everything
        // else (no bias: the score is a bare inner product).
        let scorer = Linear {
            w: prompt,
            b: vec![0.0; n_types],
            n: n_types,
            k: d,
            ..Default::default()
        };
        let flat = scorer.forward(&span_reps);

        let mut out = vec![vec![vec![0f32; n_types]; self.max_width]; n_words];
        for (row, &(l, k)) in valid.iter().enumerate() {
            out[l][k].copy_from_slice(&flat[row * n_types..(row + 1) * n_types]);
        }
        Ok(out)
    }

    /// Predict entities: sigmoid the scores, keep those over `threshold`, then resolve overlaps
    /// greedily by score (GLiNER's `flat_ner`) so the result is a non-overlapping set.
    pub fn predict(
        &mut self,
        input_ids: &[u32],
        words_mask: &[u32],
        threshold: f32,
    ) -> Result<Vec<Entity>> {
        let scores = self.scores(input_ids, words_mask)?;
        let n_words = scores.len();
        let mut hits: Vec<Entity> = Vec::new();
        for (l, per_width) in scores.iter().enumerate() {
            for (k, per_type) in per_width.iter().enumerate() {
                if l + k >= n_words {
                    continue;
                }
                for (c, &raw) in per_type.iter().enumerate() {
                    let p = sigmoid(raw);
                    if p > threshold {
                        hits.push(Entity {
                            start: l,
                            end: l + k,
                            label: c,
                            score: p,
                        });
                    }
                }
            }
        }
        // Flat NER: highest score first, drop anything overlapping an already-taken span.
        hits.sort_by(|a, b| {
            b.score
                .partial_cmp(&a.score)
                .unwrap_or(std::cmp::Ordering::Equal)
                .then(a.start.cmp(&b.start))
        });
        let mut kept: Vec<Entity> = Vec::new();
        for h in hits {
            if kept.iter().any(|k| h.start <= k.end && k.start <= h.end) {
                continue;
            }
            kept.push(h);
        }
        kept.sort_by_key(|e| (e.start, e.end));
        Ok(kept)
    }
}

#[cfg(test)]
mod lstm_probe {
    use super::*;

    /// Where does the recurrent step actually go — the GEMV, or the gate epilogue?
    ///
    /// The GEMV moves 1 MB of `w_hh` per step for 0.5 MFLOP of work (0.5 FLOP/byte), so on paper it
    /// is bandwidth-bound and should cost ~1 ms over a 176-word input. The epilogue does ~1280
    /// transcendentals per step. One of those two is the 11 ms; guessing which is how you spend an
    /// afternoon optimizing the wrong one.
    #[test]
    #[ignore = "perf probe: cargo test --release lstm_probe -- --ignored --nocapture"]
    fn split_the_recurrent_step() {
        const H: usize = 256; // GLiNER: 256 per direction
        const T: usize = 176; // the losing row's word count
        let w: Vec<f32> = (0..4 * H * H)
            .map(|i| ((i % 19) as f32 - 9.0) * 0.01)
            .collect();
        let x: Vec<f32> = (0..H).map(|i| ((i % 23) as f32 - 11.0) * 0.01).collect();
        let mut gates = vec![0.1f32; 4 * H];
        let (mut c, mut h) = (vec![0.05f32; H], vec![0f32; H]);

        let (mut gemv, mut epi) = (f64::MAX, f64::MAX);
        for _ in 0..20 {
            let t0 = std::time::Instant::now();
            for _ in 0..T {
                gemv_acc(&mut gates, &w, &x, H);
            }
            gemv = gemv.min(t0.elapsed().as_secs_f64());

            let t0 = std::time::Instant::now();
            for _ in 0..T {
                for j in 0..H {
                    let i_g = sigmoid(gates[j]);
                    let f_g = sigmoid(gates[H + j]);
                    let g_g = gates[2 * H + j].tanh();
                    let o_g = sigmoid(gates[3 * H + j]);
                    c[j] = f_g * c[j] + i_g * g_g;
                    h[j] = o_g * c[j].tanh();
                }
            }
            epi = epi.min(t0.elapsed().as_secs_f64());
        }
        // Pure-read ceiling: touch the same 1 MB the same number of times, doing almost no work.
        // If the GEMV already sits at this number it is at the machine's streaming limit and only
        // moving FEWER BYTES (f16 weights, or the GPU's bandwidth) can help; if it is far below,
        // the kernel is the problem.
        let mut read = f64::MAX;
        let mut sink = 0f32;
        for _ in 0..20 {
            let t0 = std::time::Instant::now();
            for _ in 0..T {
                let mut acc = [0f32; 4];
                for ch in w.chunks_exact(4) {
                    for l in 0..4 {
                        acc[l] += ch[l];
                    }
                }
                sink += acc[0] + acc[1] + acc[2] + acc[3];
            }
            read = read.min(t0.elapsed().as_secs_f64());
        }
        let read_gbs = (T * 4 * H * H * 4) as f64 / read / 1e9;
        eprintln!(
            "  pure-read ceiling: {:6.2} ms ({read_gbs:5.1} GB/s)  [sink {sink:.1}]",
            read * 1e3
        );

        let gf = 2.0 * (T * 4 * H * H) as f64 / gemv / 1e9;
        let gbs = (T * 4 * H * H * 4) as f64 / gemv / 1e9;
        eprintln!(
            "\n  one direction, T={T}:  gemv {:6.2} ms ({gf:5.1} GF/s, {gbs:5.1} GB/s) | epilogue {:6.2} ms\n",
            gemv * 1e3,
            epi * 1e3,
        );
        assert!(h.iter().any(|v| *v != 0.0));
    }
}