inferencelayer 0.2.4

//! gliner2 (`fastino/gliner2-multi-v1`) — zero-shot RELATION extraction on the engine.
//!
//! Same DeBERTa-v2 backbone and SpanMarkerV0 span head as GLiNER (`gliner.rs`), but a *different*
//! head turns per-token states into relations. A relation `R` is a structure with two span fields
//! (`head`, `tail`); the model is prompted with a schema and, per predicted instance, emits a span
//! score over the text for each field. The path, gated against the `gliner2` package layer by layer
//! (`tests/fixtures/export_gliner2.py` → `gliner2_relation_oracle.json`):
//!
//! ```text
//! input (NO CLS/SEP):  per relation  ( [P] name ( [R] head [R] tail ) )  … [SEP_TEXT]  word₁ word₂ …
//!
//! DeBERTa hidden states                                        [T, 768]
//!   → gather   schema markers  [P],[R],[R]  per relation block
//!              words at each word's FIRST subtoken (token_pooling="first")   [L, 768]
//!   → span_rep (SpanMarkerV0): span[l][k] = out_project(relu(start[l] ‖ end[l+k]))  [L, max_width=8, 768]
//!   per relation block:
//!     → pred_count = argmax(count_pred([P]))                    (0..19; ≤0 ⇒ no relations)
//!     → struct_proj = GRU_unroll(count_embed, [head,tail], pred_count)          [count, 2, 768]
//!     → span_scores = sigmoid( span_rep · struct_projᵀ )        [count, 2, L, 8]
//!     → per instance, per field: earliest-then-narrowest span ≥ threshold; emit head/tail pairs
//! ```
//!
//! The whole thing runs on the engine's backbone (CPU or wgpu) with a CPU head. Text is LOWERCASED
//! for word-splitting and tokenization (gliner2's `WhitespaceTokenSplitter`), but returned spans are
//! sliced from the ORIGINAL text at CHARACTER offsets — matching gliner2's Python indices.

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

use rayon::prelude::*;

use crate::EmbedEngine;
use crate::encoder_weights::EncBatch;
use crate::gliner::GlinerDevice;
use crate::weights::LazySt;

// ── primitives (kept local: a self-contained module; `gliner.rs`'s are private and shaped for v1) ──

/// One `nn.Linear`: `y = x·Wᵀ + b`, HF layout (`W` is `[n, k]` row-major). Prepacked NEON GEMM.
#[derive(Default)]
struct Linear {
    w: Vec<f32>,
    b: Vec<f32>,
    n: usize,
    k: usize,
    packed: std::sync::OnceLock<crate::cpu_gemm::PackedWeight>,
    /// Device-side (weight, bias), uploaded on first GPU use and reused for the model's life.
    /// Uploading them per call is what made the naive GPU path only ~3x the CPU one instead of
    /// ~30x (tests/gemm_probe.rs); at these shapes the weights are ~9 MB per GEMM.
    gpu: std::sync::OnceLock<(wgpu::Buffer, wgpu::Buffer)>,
}

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(),
            gpu: std::sync::OnceLock::new(),
        })
    }

    /// The device-side (weight, bias) pair, uploaded once and kept for the model's life.
    fn gpu_weights(&self, ctx: &crate::GpuCtx) -> (&wgpu::Buffer, &wgpu::Buffer) {
        let (w, b) = self
            .gpu
            .get_or_init(|| (ctx.storage(&self.w), ctx.storage(&self.b)));
        (w, b)
    }

    /// `x` is `[rows, k]` → `[rows, n]`, on the DEVICE, reusing weights uploaded once.
    /// `None` when there is no GPU — callers fall back to `forward`.
    fn forward_gpu(
        &self,
        gpu: Option<(&crate::GpuCtx, &crate::encoder::EncKernels)>,
        x: &[f32],
    ) -> Option<Vec<f32>> {
        let (ctx, kern) = gpu?;
        let (wb, bb) = self
            .gpu
            .get_or_init(|| (ctx.storage(&self.w), ctx.storage(&self.b)));
        let m = x.len() / self.k;
        kern.gemm_resident(ctx, x, wb, bb, m, self.n, self.k, None)
            .ok()
    }

    /// `x` is `[rows, k]` → `[rows, n]`.
    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");
        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
    }
}

/// gliner2's `create_mlp` / gliner's `create_projection_layer`: `Linear → ReLU → Linear`.
/// The dropout between is identity at inference; gliner2's own heads name the two Linears `.0`/`.2`
/// (dropout=0., no module), while the span_rep MLPs keep `.0`/`.3` (gliner's layer). Hence `down_idx`.
struct Mlp {
    up: Linear,
    down: Linear,
}

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

    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);
        }
        self.down.forward(&h)
    }

    /// Same math, both GEMMs on the device when one is available (see `Linear::forward_gpu`).
    /// The ReLU between stays on the host: it is elementwise over a small buffer and moving it
    /// would need a kernel for no measurable gain.
    fn forward_on(
        &self,
        gpu: Option<(&crate::GpuCtx, &crate::encoder::EncKernels)>,
        x: &[f32],
    ) -> Vec<f32> {
        let mut h = self
            .up
            .forward_gpu(gpu, x)
            .unwrap_or_else(|| self.up.forward(x));
        for v in h.iter_mut() {
            *v = v.max(0.0);
        }
        self.down
            .forward_gpu(gpu, &h)
            .unwrap_or_else(|| self.down.forward(&h))
    }
}

/// `out_project.0` split at the concat seam: `W_up·relu(start ‖ end) = W_left·relu(start) +
/// W_right·relu(end)` (relu is elementwise, so it commutes with the concat). Projecting each half
/// per WORD turns the per-span first layer into a vector add. Same trick as `gliner.rs`.
struct SpanUp {
    left: Linear,
    right: Linear,
}

impl SpanUp {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        let up = Linear::load(st, prefix)?;
        let (n, k) = (up.n, up.k);
        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],
                n,
                k: half,
                ..Default::default()
            },
            right: Linear {
                w: right,
                b: up.b,
                n,
                k: half,
                ..Default::default()
            },
        })
    }
}

/// gliner2's `count_embed` — a single-layer unidirectional **GRU** (`nn.GRU`, despite the config's
/// "count_lstm" name), then a projector MLP. It unrolls `count` steps over a learned positional
/// table, seeding the hidden state with the field (head/tail) marker embeddings, to produce a
/// distinct query vector for each (instance, field). No GRU existed in the engine before this.
///
/// PyTorch GRU (gate order r, z, n): `r=σ(Wir·x+bir+Whr·h+bhr)`, `z=σ(Wiz·x+biz+Whz·h+bhz)`,
/// `n=tanh(Win·x+bin + r·(Whn·h+bhn))`, `h=(1-z)·n + z·h`. The new-gate hidden bias `bhn` sits
/// INSIDE the reset-gated term — get that wrong and it's plausible-but-off.
struct Gru {
    w_ih: Vec<f32>, // [3H, H]
    w_hh: Vec<f32>, // [3H, H]
    b_ih: Vec<f32>, // [3H]
    b_hh: Vec<f32>, // [3H]
    pos: Vec<f32>,  // [max_count, H] learned positional embedding (the per-step input)
    projector: Mlp, // create_mlp(2H → 4H → H), cat(gru_out, field)
    h: usize,
    max_count: usize,
}

impl Gru {
    fn load(st: &LazySt, prefix: &str) -> Result<Self> {
        let w_ih = st.tensor_f32(&format!("{prefix}.gru.weight_ih_l0"))?;
        let w_hh = st.tensor_f32(&format!("{prefix}.gru.weight_hh_l0"))?;
        let b_ih = st.tensor_f32(&format!("{prefix}.gru.bias_ih_l0"))?;
        let b_hh = st.tensor_f32(&format!("{prefix}.gru.bias_hh_l0"))?;
        let pos = st.tensor_f32(&format!("{prefix}.pos_embedding.weight"))?;
        let h = b_ih.len() / 3;
        let max_count = pos.len() / h;
        anyhow::ensure!(
            w_ih.len() == 3 * h * h && w_hh.len() == 3 * h * h,
            "GRU weight geometry"
        );
        Ok(Self {
            w_ih,
            w_hh,
            b_ih,
            b_hh,
            pos,
            projector: Mlp::load(st, &format!("{prefix}.projector"), 2)?,
            h,
            max_count,
        })
    }

    /// `fields` = `[M, H]` marker embeddings (head, tail). Returns `struct_proj` `[count·M, H]` laid
    /// out in `(instance·M + field)` row order — i.e. the "prompt reps" the scorer contracts against.
    fn forward(&self, fields: &[f32], count: usize) -> Vec<f32> {
        let h = self.h;
        let m = fields.len() / h;
        let count = count.min(self.max_count);
        // Input projection W_ih·pos[t] + b_ih for each step (independent of field). Precompute here.
        let mut ih = vec![0f32; count * 3 * h];
        for t in 0..count {
            let dst = &mut ih[t * 3 * h..(t + 1) * 3 * h];
            dst.copy_from_slice(&self.b_ih);
            crate::simd::gemv_acc(dst, &self.w_ih, &self.pos[t * h..(t + 1) * h]);
        }
        // Per field: unroll the recurrence, then project cat(h_t, field).
        let mut out = vec![0f32; count * m * h];
        let mut hh = vec![0f32; 3 * h];
        for f in 0..m {
            let field = &fields[f * h..(f + 1) * h];
            let mut hs = field.to_vec();
            for t in 0..count {
                let ihs = &ih[t * 3 * h..(t + 1) * 3 * h];
                hh.copy_from_slice(&self.b_hh);
                crate::simd::gemv_acc(&mut hh, &self.w_hh, &hs);
                for j in 0..h {
                    let r = sigmoid(ihs[j] + hh[j]);
                    let z = sigmoid(ihs[h + j] + hh[h + j]);
                    let n = (ihs[2 * h + j] + r * hh[2 * h + j]).tanh();
                    hs[j] = (1.0 - z) * n + z * hs[j];
                }
                // struct_proj row (t·M + f): projector(cat(h_t, field)).
                let mut cat = Vec::with_capacity(2 * h);
                cat.extend_from_slice(&hs);
                cat.extend_from_slice(field);
                let proj = self.projector.forward(&cat);
                out[(t * m + f) * h..(t * m + f + 1) * h].copy_from_slice(&proj);
            }
        }
        out
    }
}

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

/// Numerically stable softmax (matches `torch.softmax(logits, dim=-1)` on a classification block).
fn softmax(logits: &[f32]) -> Vec<f32> {
    let max = logits.iter().copied().fold(f32::NEG_INFINITY, f32::max);
    let exps: Vec<f32> = logits.iter().map(|&v| (v - max).exp()).collect();
    let sum: f32 = exps.iter().sum();
    exps.iter().map(|&e| e / sum).collect()
}

// ── output types ────────────────────────────────────────────────────────────────

/// One end of a relation: the span text, its CHARACTER offsets into the original text, and the
/// sigmoid score of the span. (The gliner2-ner service discards the score — its `Relation.confidence`
/// is always 1.0 — and keeps only text + offsets.)
#[derive(Debug, Clone, PartialEq)]
pub struct RelEnd {
    pub text: String,
    pub start: usize,
    pub end: usize,
    pub confidence: f32,
}

/// One extracted relation: `head` → `tail` under `label`, in the order gliner2 emits them.
#[derive(Debug, Clone, PartialEq)]
pub struct Relation {
    pub label: String,
    pub head: RelEnd,
    pub tail: RelEnd,
}

/// One entity from the gliner2 schema-based entity path (`/extract/combined`). Distinct from the
/// engine GLiNER `/ner` entities: this is the gliner2 model's own `extract_entities`.
#[derive(Debug, Clone, PartialEq)]
pub struct Ner2Entity {
    pub text: String,
    pub label: String,
    pub start: usize,
    pub end: usize,
    pub confidence: f32,
}

/// One label's probability from a classification task. `scores` in [`ClsResult`] carries every
/// label (input order); `chosen` carries gliner2's decode.
#[derive(Debug, Clone, PartialEq)]
pub struct ClsLabelScore {
    pub label: String,
    pub score: f32,
}

/// One classification task's outcome (gliner2 `classify_text`): all per-label probabilities plus
/// the chosen set — single-label → argmax; multi-label → every label ≥ `cls_threshold`, falling
/// back to the argmax when none clears it (engine.py `_extract_classification_result`).
#[derive(Debug, Clone, PartialEq)]
pub struct ClsResult {
    pub task: String,
    pub scores: Vec<ClsLabelScore>,
    pub chosen: Vec<ClsLabelScore>,
}

/// Activation over a classification block's logits. gliner2's `class_act` config: `Auto` (the
/// default) is sigmoid for multi-label tasks and softmax otherwise.
#[derive(Debug, Clone, Copy, PartialEq, Eq, Default)]
pub enum ClsActivation {
    #[default]
    Auto,
    Sigmoid,
    Softmax,
}

/// One classification task, mirroring gliner2's `Schema.classification(...)` config. The optional
/// `prompt`, ordered `label_descriptions` and `examples` extend the `[P]` prompt string exactly as
/// `SchemaTransformer._transform_schema` does at inference (`example_mode="both"`).
#[derive(Debug, Clone, PartialEq)]
pub struct ClsTask {
    pub name: String,
    pub labels: Vec<String>,
    pub multi_label: bool,
    pub cls_threshold: f32,
    pub class_act: ClsActivation,
    pub prompt: Option<String>,
    pub label_descriptions: Vec<(String, String)>,
    pub examples: Vec<(String, String)>,
}

impl ClsTask {
    /// A task with gliner2's defaults (`multi_label=false`, `cls_threshold=0.5`, auto activation).
    pub fn new(name: impl Into<String>, labels: Vec<String>) -> Self {
        Self {
            name: name.into(),
            labels,
            multi_label: false,
            cls_threshold: 0.5,
            class_act: ClsActivation::Auto,
            prompt: None,
            label_descriptions: Vec::new(),
            examples: Vec::new(),
        }
    }
}

/// Every stage of the classification forward pass, for the layer-by-layer parity gate.
pub struct ClsIntermediates {
    pub input_ids: Vec<u32>,
    pub last_hidden_state: Vec<f32>, // [T, H]
    pub schema_block0: Vec<f32>,     // [1 + n_labels, H] first task: [P], then its [L] markers
    pub logits0: Vec<f32>,           // [n_labels] first task
    pub probs0: Vec<f32>,            // [n_labels] first task, post-activation
}

/// Every stage of the forward pass, for the layer-by-layer parity gate.
pub struct RelIntermediates {
    pub input_ids: Vec<u32>,
    pub last_hidden_state: Vec<f32>, // [T, H]
    pub token_embs: Vec<f32>,        // [L, H]
    pub schema_block0: Vec<f32>,     // [3, H] first block: [P], head-[R], tail-[R]
    pub count_logits: Vec<f32>,      // [num_count] first block
    pub pred_count: usize,           // first block
    pub struct_proj: Vec<f32>,       // [count·M, H] first block, (inst·M+field) order
    pub span_rep: Vec<f32>,          // [L, K, H]
    pub span_scores: Vec<f32>,       // [count, M, L, K] first block
    pub relations: Vec<Relation>,
    pub n_words: usize,
    pub t: usize,
}

// ── the model ────────────────────────────────────────────────────────────────

/// A schema block for one relation: the positions (in `input_ids`) of its `[P]` and the two `[R]`
/// (head, tail) markers.
struct Block {
    label: String,
    p_pos: usize,
    field_pos: [usize; 2],
}

/// A text word: the position of its FIRST subtoken and its CHARACTER offsets in the lowercased text.
struct Word {
    first_tok: usize,
    char_start: usize,
    char_end: usize,
}

/// A text word tokenized ONCE (its subtoken ids + global CHARACTER offsets), reused across every
/// chunk window; `first_tok` is assigned per window when the window's `input_ids` are assembled.
#[cfg(feature = "cli")]
struct WordTok {
    ids: Vec<u32>,
    char_start: usize,
    char_end: usize,
}

/// A classification schema block: the positions (in `input_ids`) of its `[P]` and its `[L]`
/// label markers (one per label, in label order).
struct ClsBlock {
    p_pos: usize,
    l_pos: Vec<usize>,
}

/// The first relation block's intermediates, captured for the layer-by-layer parity gate.
struct Block0 {
    schema: Vec<f32>,
    count_logits: Vec<f32>,
    pred_count: usize,
    struct_proj: Vec<f32>,
    span_scores: Vec<f32>,
}

/// The span_rep "fill": `hbuf[row] = relu(a[li] + b[lj])`, one output element per thread.
///
/// The only kernel gliner2 owns. It exists so the ~22.8 MB `hbuf` never crosses the bus: with the
/// projections and the down-projection already on the device, the fill was the one step forcing a
/// round-trip, and that transfer — not the math — was the remaining cost.
///
/// The dispatch is 2D on purpose. One thread per element at 64/group would need `spans*4H/64`
/// ≈ 89k workgroups for a full window, past the 65535-per-dimension limit; `x` walks the 3072-wide
/// row and `y` walks the spans, so both stay far under it.
const FILL_WGSL: &str = r#"
@group(0) @binding(0) var<storage, read> a: array<f32>;
@group(0) @binding(1) var<storage, read> b: array<f32>;
@group(0) @binding(2) var<storage, read> idx: array<u32>;
@group(0) @binding(3) var<storage, read_write> out: array<f32>;
@group(0) @binding(4) var<uniform> dims: vec4<u32>;

@compute @workgroup_size(64)
fn main(@builtin(global_invocation_id) g: vec3<u32>) {
    let rows = dims.x;
    let w = dims.y;
    let j = g.x;
    let row = g.y;
    if (j >= w || row >= rows) { return; }
    let li = idx[row * 2u];
    let lj = idx[row * 2u + 1u];
    out[row * w + j] = max(a[li * w + j] + b[lj * w + j], 0.0);
}
"#;

pub struct Gliner2 {
    /// The fill kernel, compiled on first device use. Built with an auto bind-group layout, so it
    /// needs NO change to the shared `EncKernels` pipeline set (that struct is on the hot path of
    /// every embed/score/rerank call — keeping this local is deliberate).
    fill_pl: std::sync::OnceLock<wgpu::ComputePipeline>,
    backbone: EmbedEngine,
    project_start: Mlp, // span_rep_layer.project_start (.0/.3)
    project_end: Mlp,
    span_up: SpanUp,   // out_project.0, split
    span_down: Linear, // out_project.3
    count_pred: Mlp,   // count_pred (.0/.2) → num_count logits
    /// `None` when the checkpoint's counting layer isn't the v1 GRU (`count_lstm`) — e.g. the
    /// GLiGuard family ships `count_lstm_v2` (GRU + a transformer refiner, different tensors).
    /// Classification never touches it; the span paths error explicitly instead of failing load.
    count_embed: Option<Gru>,
    counting_layer: String,
    classifier: Mlp, // classifier (.0/.2) → one logit per [L] label marker
    max_width: usize,
    hidden: usize,
    #[cfg(feature = "cli")]
    tokenizer: tokenizers::Tokenizer,
    #[cfg(feature = "cli")]
    splitter: regex::Regex,
}

impl Gliner2 {
    /// Load a gliner2 checkpoint (`tests/fixtures/export_gliner2.py` layout) on the best device.
    pub fn load(dir: &Path) -> Result<Self> {
        Self::load_on(dir, GlinerDevice::Auto)
    }

    /// [`Self::load`] with an explicit backbone device (the parity gate drives CPU for determinism).
    pub fn load_on(dir: &Path, device: GlinerDevice) -> Result<Self> {
        let backbone = match device {
            GlinerDevice::Auto => EmbedEngine::auto(dir, 8192)?,
            GlinerDevice::Cpu => EmbedEngine::cpu(dir)?,
        };
        let st = LazySt::open(dir)?;
        let head_cfg: serde_json::Value = std::fs::read(dir.join("gliner2_head.json"))
            .ok()
            .and_then(|b| serde_json::from_slice(&b).ok())
            .unwrap_or(serde_json::Value::Null);
        let hidden = head_cfg
            .get("hidden_size")
            .and_then(|x| x.as_u64())
            .unwrap_or(768) as usize;
        let max_width = head_cfg
            .get("max_width")
            .and_then(|x| x.as_u64())
            .unwrap_or(8) as usize;
        let counting_layer = head_cfg
            .get("counting_layer")
            .and_then(|x| x.as_str())
            .unwrap_or("count_lstm")
            .to_string();
        // Only the v1 GRU (`count_lstm`) is ported; other counting layers load classify-only.
        let count_embed = if counting_layer == "count_lstm" {
            Some(Gru::load(&st, "count_embed")?)
        } else {
            None
        };

        let sp = "span_rep.span_rep_layer";

        Ok(Self {
            fill_pl: std::sync::OnceLock::new(),
            project_start: Mlp::load(&st, &format!("{sp}.project_start"), 3)?,
            project_end: Mlp::load(&st, &format!("{sp}.project_end"), 3)?,
            span_up: SpanUp::load(&st, &format!("{sp}.out_project.0"))?,
            span_down: Linear::load(&st, &format!("{sp}.out_project.3"))?,
            count_pred: Mlp::load(&st, "count_pred", 2)?,
            count_embed,
            counting_layer,
            classifier: Mlp::load(&st, "classifier", 2)?,
            max_width,
            hidden,
            backbone,
            #[cfg(feature = "cli")]
            tokenizer: tokenizers::Tokenizer::from_file(dir.join("tokenizer.json"))
                .map_err(|e| anyhow::anyhow!("gliner2 tokenizer: {e}"))?,
            // gliner2's WhitespaceTokenSplitter regex (applied to the LOWERCASED text).
            #[cfg(feature = "cli")]
            splitter: regex::Regex::new(
                r"(?ix)(?:https?://[^\s]+|www\.[^\s]+)|[a-z0-9._%+-]+@[a-z0-9.-]+\.[a-z]{2,}|@[a-z0-9_]+|\w+(?:[-_]\w+)*|\S",
            )?,
        })
    }

    /// The backbone's device string (e.g. `"Metal/Apple M4 Max (f16)"`, `"cpu"`).
    pub fn device(&self) -> String {
        self.backbone.device()
    }

    /// **End-to-end**: raw text + relation labels → relations with character offsets into `text`.
    #[cfg(feature = "cli")]
    pub fn predict_relations(
        &mut self,
        text: &str,
        labels: &[impl AsRef<str>],
        threshold: f32,
    ) -> Result<Vec<Relation>> {
        Ok(self.relations_debug(text, labels, threshold)?.0)
    }

    /// [`Self::predict_relations`] that also returns every intermediate, for the parity gate.
    #[cfg(feature = "cli")]
    pub fn relations_debug(
        &mut self,
        text: &str,
        labels: &[impl AsRef<str>],
        threshold: f32,
    ) -> Result<(Vec<Relation>, RelIntermediates)> {
        anyhow::ensure!(!labels.is_empty(), "no relation labels given");
        let Some(count_embed) = &self.count_embed else {
            anyhow::bail!(
                "counting layer {:?} is not ported — this checkpoint serves classification only \
                 (classify_text)",
                self.counting_layer
            );
        };
        let text = normalize_text(text);
        let text = text.as_str();
        let (input_ids, blocks, words) = self.tokenize(text, labels);
        let orig: Vec<char> = text.chars().collect();

        // 1. backbone → per-token hidden states [T, H].
        let states = self
            .backbone
            .forward_hidden(&EncBatch::from_seqs([input_ids.clone()]))
            .context("gliner2 backbone")?;
        let h = self.hidden;
        let t = states.len() / h;

        // 2. gather word embeddings (first subtoken) → token_embs [L, H].
        let l = words.len();
        let mut token_embs = vec![0f32; l * h];
        for (wi, w) in words.iter().enumerate() {
            token_embs[wi * h..(wi + 1) * h]
                .copy_from_slice(&states[w.first_tok * h..(w.first_tok + 1) * h]);
        }

        // 3. span representations (SpanMarkerV0), shared across all relation blocks → [L, K, H] dense.
        let span_rep = self.span_rep(&token_embs, l);

        // 4. per relation block: count → GRU → score → decode.
        let mut relations = Vec::new();
        let mut dbg_block0: Option<Block0> = None;
        for blk in &blocks {
            // schema embeddings for this block: [P], head-[R], tail-[R].
            let mut schema = vec![0f32; 3 * h];
            schema[0..h].copy_from_slice(&states[blk.p_pos * h..(blk.p_pos + 1) * h]);
            schema[h..2 * h]
                .copy_from_slice(&states[blk.field_pos[0] * h..(blk.field_pos[0] + 1) * h]);
            schema[2 * h..3 * h]
                .copy_from_slice(&states[blk.field_pos[1] * h..(blk.field_pos[1] + 1) * h]);

            // count = argmax(count_pred([P])).
            let count_logits = self.count_pred.forward(&schema[0..h]);
            let pred_count = argmax(&count_logits);

            let (struct_proj, span_scores) = if pred_count == 0 {
                (Vec::new(), Vec::new())
            } else {
                // struct_proj [count·M, H] from the GRU over the two field embeddings.
                let struct_proj = count_embed.forward(&schema[h..3 * h], pred_count);
                let bp = struct_proj.len() / h; // count·M
                // score[span, bp] = sigmoid(span_rep[span] · struct_proj[bp]). Score the DENSE span_rep
                // [L·K, H] (invalid spans carry the (0,0) value, as gliner2's safe-index does), so
                // span_scores matches the oracle everywhere; the decoder reads only valid spans.
                let scorer = Linear {
                    w: struct_proj.clone(),
                    b: vec![0.0; bp],
                    n: bp,
                    k: h,
                    ..Default::default()
                };
                let t_score = std::time::Instant::now();
                let mut flat = scorer.forward(&span_rep); // [L·K, bp]
                if std::env::var("OSFKB_GLINER2_TRACE").is_ok_and(|v| v != "0") {
                    eprintln!(
                        "  scorer: {:.3}s  (rows {}, k {}, n {})",
                        t_score.elapsed().as_secs_f64(),
                        span_rep.len() / h,
                        h,
                        bp,
                    );
                }
                for v in flat.iter_mut() {
                    *v = sigmoid(*v);
                }
                // reshape [L·K, bp] → [count, M, L, K] (M=2), einsum index order bplk.
                let count = pred_count.min(count_embed.max_count);
                let m = bp / count;
                let kmax = self.max_width;
                let mut ss = vec![0f32; count * m * l * kmax];
                for li in 0..l {
                    for ki in 0..kmax {
                        let srow = (li * kmax + ki) * bp;
                        for b in 0..count {
                            for p in 0..m {
                                ss[((b * m + p) * l + li) * kmax + ki] = flat[srow + b * m + p];
                            }
                        }
                    }
                }
                (struct_proj, ss)
            };

            // decode: per instance, head+tail earliest-narrowest span ≥ threshold.
            if !span_scores.is_empty() {
                let count = pred_count.min(count_embed.max_count);
                let m = 2usize;
                let kmax = self.max_width;
                for b in 0..count {
                    let head = find_span(&span_scores, b, 0, m, l, kmax, threshold, &words, &orig);
                    let tail = find_span(&span_scores, b, 1, m, l, kmax, threshold, &words, &orig);
                    if let (Some(head), Some(tail)) = (head, tail) {
                        relations.push(Relation {
                            label: blk.label.clone(),
                            head,
                            tail,
                        });
                    }
                }
            }

            if dbg_block0.is_none() {
                dbg_block0 = Some(Block0 {
                    schema,
                    count_logits,
                    pred_count,
                    struct_proj,
                    span_scores,
                });
            }
        }

        let b0 = dbg_block0.unwrap_or(Block0 {
            schema: vec![],
            count_logits: vec![],
            pred_count: 0,
            struct_proj: vec![],
            span_scores: vec![],
        });
        let inter = RelIntermediates {
            input_ids,
            last_hidden_state: states,
            token_embs,
            schema_block0: b0.schema,
            count_logits: b0.count_logits,
            pred_count: b0.pred_count,
            struct_proj: b0.struct_proj,
            span_rep,
            span_scores: b0.span_scores,
            relations: relations.clone(),
            n_words: l,
            t,
        };
        Ok((relations, inter))
    }

    /// SpanMarkerV0: `span[l][k] = out_project(relu(project_start(tok)[l] ‖ project_end(tok)[l+k]))`.
    /// Returns the dense `[L, K, H]` tensor (invalid spans `l+k>=L` carry the `(0,0)` value, matching
    /// gliner2's safe-index behaviour). The per-word `SpanUp` seam split keeps this O(L), not O(L·K).
    /// span_rep with the WHOLE chain resident on the device: proj → fill → down, one download.
    ///
    /// `Some` only when a GPU is present and the shapes are sane; callers fall back to the CPU/
    /// per-GEMM path otherwise. Returns `span_reps_flat` (`[n_valid, H]`), i.e. exactly what the
    /// CPU path produces before the dense expansion.
    ///
    /// The two ReLUs are folded into the preceding GEMMs (`Act::Relu`): the one inside each `Mlp`,
    /// and the `relu(&start)`/`relu(&end)` the CPU path applies before `span_up` — algebraically the
    /// same, and it keeps the chain from breaking for an elementwise op.
    #[allow(clippy::too_many_arguments)]
    fn span_rep_chained(
        &self,
        ctx: &crate::GpuCtx,
        kern: &crate::encoder::EncKernels,
        token_embs: &[f32],
        l: usize,
        valid: &[(usize, usize)],
    ) -> Option<Vec<f32>> {
        use crate::encoder_weights::Act;
        use crate::forward::{make_bg_pub, pipeline_pub, uni_pub};
        let h = self.hidden;
        let up_w = self.span_up.left.n;
        let rows = valid.len();
        if rows == 0 || l == 0 {
            return None;
        }

        // (li, li+ki) per valid span — the gather the fill kernel indexes with. Tiny (~15 KB).
        let idx: Vec<u32> = valid
            .iter()
            .flat_map(|&(li, ki)| [li as u32, (li + ki) as u32])
            .collect();

        let xb = ctx.storage(token_embs);
        let mid = ctx.storage(&vec![0f32; l * self.project_start.up.n]);
        let startb = ctx.storage(&vec![0f32; l * h]);
        let endb = ctx.storage(&vec![0f32; l * h]);
        let ab = ctx.storage(&vec![0f32; l * up_w]);
        let bbuf = ctx.storage(&vec![0f32; l * up_w]);
        let hbuf = ctx.storage(&vec![0f32; rows * up_w]);
        let outb = ctx.storage(&vec![0f32; rows * h]);
        let idxb = ctx.storage_bytes(bytemuck::cast_slice(&idx));

        // project_start: up(Relu) → down(Relu, folding the caller's relu(&start))
        let (w, b) = &self.project_start.up.gpu_weights(ctx);
        kern.gemm_resident_into(
            ctx,
            &xb,
            w,
            b,
            &mid,
            l,
            self.project_start.up.n,
            self.project_start.up.k,
            Some(Act::Relu),
        )
        .ok()?;
        let (w, b) = &self.project_start.down.gpu_weights(ctx);
        kern.gemm_resident_into(
            ctx,
            &mid,
            w,
            b,
            &startb,
            l,
            h,
            self.project_start.down.k,
            Some(Act::Relu),
        )
        .ok()?;
        // project_end, same shape
        let (w, b) = &self.project_end.up.gpu_weights(ctx);
        kern.gemm_resident_into(
            ctx,
            &xb,
            w,
            b,
            &mid,
            l,
            self.project_end.up.n,
            self.project_end.up.k,
            Some(Act::Relu),
        )
        .ok()?;
        let (w, b) = &self.project_end.down.gpu_weights(ctx);
        kern.gemm_resident_into(
            ctx,
            &mid,
            w,
            b,
            &endb,
            l,
            h,
            self.project_end.down.k,
            Some(Act::Relu),
        )
        .ok()?;
        // span_up halves
        let (w, b) = &self.span_up.left.gpu_weights(ctx);
        kern.gemm_resident_into(ctx, &startb, w, b, &ab, l, up_w, self.span_up.left.k, None)
            .ok()?;
        let (w, b) = &self.span_up.right.gpu_weights(ctx);
        kern.gemm_resident_into(ctx, &endb, w, b, &bbuf, l, up_w, self.span_up.right.k, None)
            .ok()?;

        // the fill — this is what keeps hbuf off the bus
        let pl = self
            .fill_pl
            .get_or_init(|| pipeline_pub(ctx, "gliner2_span_fill", FILL_WGSL));
        let meta = uni_pub(
            ctx,
            bytemuck::cast_slice(&[rows as u32, up_w as u32, 0u32, 0u32]),
        );
        let bg = make_bg_pub(ctx, pl, &[&ab, &bbuf, &idxb, &hbuf], &meta);
        crate::encoder::dispatch(ctx, pl, &bg, (up_w as u32).div_ceil(64), rows as u32);

        let (w, b) = &self.span_down.gpu_weights(ctx);
        kern.gemm_resident_into(ctx, &hbuf, w, b, &outb, rows, h, self.span_down.k, None)
            .ok()?;
        ctx.read(&outb, rows * h).ok()
    }

    /// `span_reps_flat` (valid rows) → the dense `[L, K, H]` gliner2 expects. Out-of-range spans
    /// (l+k ≥ L) take the SAFE (0,0) value, not zero — matching gliner2's own masking, which the
    /// oracle dumps pre-mask. Shared by the device and CPU paths so they cannot diverge here.
    fn densify(
        flat: &[f32],
        valid: &[(usize, usize)],
        l: usize,
        kmax: usize,
        h: usize,
    ) -> Vec<f32> {
        let zero = flat[0..h].to_vec(); // the (0,0) span, always valid row 0
        let mut dense = vec![0f32; l * kmax * h];
        for cell in dense.chunks_exact_mut(h) {
            cell.copy_from_slice(&zero);
        }
        for (row, &(li, ki)) in valid.iter().enumerate() {
            dense[(li * kmax + ki) * h..(li * kmax + ki + 1) * h]
                .copy_from_slice(&flat[row * h..(row + 1) * h]);
        }
        dense
    }

    fn span_rep(&self, token_embs: &[f32], l: usize) -> Vec<f32> {
        let tr = std::env::var("OSFKB_GLINER2_TRACE").is_ok_and(|v| v != "0");
        let t0 = std::time::Instant::now();
        let h = self.hidden;
        let kmax = self.max_width;

        // Fully-on-device chain first (proj → fill → down, ONE download). Falls through to the
        // per-GEMM path below when there is no GPU, when it is switched off, or if anything in the
        // chain declines — so the CPU behaviour is unchanged and always reachable.
        if !std::env::var("OSFKB_GLINER2_CHAIN").is_ok_and(|v| v == "0") {
            if let Some((ctx, enc)) = self.backbone.gpu_parts() {
                let valid: Vec<(usize, usize)> = (0..l)
                    .flat_map(|li| (0..kmax).map(move |ki| (li, ki)))
                    .filter(|&(li, ki)| li + ki < l)
                    .collect();
                if let Some(flat) = self.span_rep_chained(ctx, enc.kernels(), token_embs, l, &valid)
                {
                    if tr {
                        eprintln!(
                            "  span_rep {:.3}s (device chain, spans {})",
                            t0.elapsed().as_secs_f64(),
                            valid.len()
                        );
                    }
                    return Self::densify(&flat, &valid, l, kmax, h);
                }
            }
        }
        // Head GEMMs on the DEVICE when the backbone is there. MEASURED (fleet trace): the heads are
        // ~75 % of a GPU call and span_rep is ~85 % of the heads, all of it CPU `Linear::forward` —
        // the Amdahl ceiling on the whole NER GPU port. Weights are uploaded once (`Linear::gpu`),
        // because re-uploading them per call is what held the naive path to ~3x instead of ~30x.
        // OSFKB_GLINER2_HEAD_GPU=0 reverts to the CPU heads without a rebuild.
        let gpu = (!std::env::var("OSFKB_GLINER2_HEAD_GPU").is_ok_and(|v| v == "0"))
            .then(|| self.backbone.gpu_parts())
            .flatten()
            .map(|(ctx, enc)| (ctx, enc.kernels()));
        let lin = |l: &Linear, x: &[f32]| -> Vec<f32> {
            l.forward_gpu(gpu, x).unwrap_or_else(|| l.forward(x))
        };
        let start = self.project_start.forward_on(gpu, token_embs); // [L, H]
        let end = self.project_end.forward_on(gpu, token_embs); // [L, H]
        let relu = |v: &[f32]| -> Vec<f32> { v.iter().map(|x| x.max(0.0)).collect() };
        let a = lin(&self.span_up.left, &relu(&start)); // [L, 4H]
        let b = lin(&self.span_up.right, &relu(&end)); // [L, 4H] (bias here)
        let up_w = self.span_up.left.n;
        let t_proj = t0.elapsed();

        let valid: Vec<(usize, usize)> = (0..l)
            .flat_map(|li| (0..kmax).map(move |ki| (li, ki)))
            .filter(|&(li, ki)| li + ki < l)
            .collect();
        let t1 = std::time::Instant::now();
        // relu(a[li] ‖ b[li+ki]) for every valid span. MEASURED on the fleet: this loop is
        // 0.034-0.048 s of a 0.27 s call — ~20 % of span_rep — because it touches
        // spans x 4H ≈ 5.7 M floats scalar. Each row writes a DISJOINT slice and the op is
        // elementwise, so parallelising it is bit-exact: no reduction is reordered, unlike the
        // GEMMs around it. (The trace that found this: OSFKB_GLINER2_TRACE=1.)
        let mut hbuf = vec![0f32; valid.len() * up_w];
        hbuf.par_chunks_mut(up_w)
            .zip(valid.par_iter())
            .for_each(|(row, &(li, ki))| {
                let (ar, br) = (&a[li * up_w..], &b[(li + ki) * up_w..]);
                for j in 0..up_w {
                    row[j] = (ar[j] + br[j]).max(0.0);
                }
            });
        let t_fill = t1.elapsed();
        let t_down = std::time::Instant::now();
        let span_reps_flat = lin(&self.span_down, &hbuf); // [n_valid, H]
        let t_dn = t_down.elapsed();
        let t2 = std::time::Instant::now();
        // dense [L, K, H]. gliner2 replaces out-of-range spans (l+k>=L) with the SAFE index (0,0), so
        // every invalid entry equals the (0,0) span's value (= valid row 0), NOT zero. Reproduce that
        // so the oracle matches exactly (it dumps the raw, pre-mask span_rep).
        let dense = Self::densify(&span_reps_flat, &valid, l, kmax, h);
        if tr {
            eprintln!(
                "  span_rep {:.3}s = proj {:.3} + fill {:.3} + down {:.3} + dense {:.3}  \
                 (L {}, spans {}, up_w {})",
                t0.elapsed().as_secs_f64(),
                t_proj.as_secs_f64(),
                t_fill.as_secs_f64(),
                t_dn.as_secs_f64(),
                t2.elapsed().as_secs_f64(),
                l,
                valid.len(),
                up_w,
            );
        }
        dense
    }

    /// Build `input_ids` (schema blocks + `[SEP_TEXT]` + lowercased text words, each element tokenized
    /// independently, NO CLS/SEP) plus the schema-marker positions and word first-subtoken/offsets.
    #[cfg(feature = "cli")]
    fn tokenize(
        &self,
        text: &str,
        labels: &[impl AsRef<str>],
    ) -> (Vec<u32>, Vec<Block>, Vec<Word>) {
        let mut ids: Vec<u32> = Vec::new();
        let mut blocks: Vec<Block> = Vec::new();
        // encode one combined element (add_special_tokens=false) and append; return its token range.
        let push = |ids: &mut Vec<u32>, s: &str| -> (usize, usize) {
            let start = ids.len();
            let enc = self
                .tokenizer
                .encode(s, false)
                .expect("gliner2 tokenize element");
            ids.extend_from_slice(enc.get_ids());
            (start, ids.len())
        };

        for (i, lab) in labels.iter().enumerate() {
            if i > 0 {
                push(&mut ids, "[SEP_STRUCT]");
            }
            push(&mut ids, "(");
            let (ps, _) = push(&mut ids, "[P]");
            push(&mut ids, lab.as_ref());
            push(&mut ids, "(");
            let (r1, _) = push(&mut ids, "[R]");
            push(&mut ids, "head");
            let (r2, _) = push(&mut ids, "[R]");
            push(&mut ids, "tail");
            push(&mut ids, ")");
            push(&mut ids, ")");
            blocks.push(Block {
                label: lab.as_ref().to_string(),
                p_pos: ps,
                field_pos: [r1, r2],
            });
        }
        push(&mut ids, "[SEP_TEXT]");

        // text words: lowercased split; each word tokenized independently; first subtoken pooled.
        let low = text.to_lowercase();
        let mut words: Vec<Word> = Vec::new();
        for m in self.splitter.find_iter(&low) {
            let char_start = low[..m.start()].chars().count();
            let char_end = char_start + m.as_str().chars().count();
            let (first_tok, _) = push(&mut ids, m.as_str());
            words.push(Word {
                first_tok,
                char_start,
                char_end,
            });
        }
        (ids, blocks, words)
    }

    /// The ENTITIES schema prefix `( [P] entities ( [E] t1 [E] t2 … ) ) [SEP_TEXT]` — the part BEFORE
    /// the text, fixed across every chunk window. Returns the ids, the `[P]` position, and the `[E]`
    /// marker positions (one per type). Text words are tokenized separately by [`Self::word_toks`].
    #[cfg(feature = "cli")]
    fn entity_schema_ids(&self, labels: &[impl AsRef<str>]) -> (Vec<u32>, usize, Vec<usize>) {
        let mut ids: Vec<u32> = Vec::new();
        let push = |ids: &mut Vec<u32>, s: &str| -> usize {
            let start = ids.len();
            let enc = self
                .tokenizer
                .encode(s, false)
                .expect("gliner2 tokenize element");
            ids.extend_from_slice(enc.get_ids());
            start
        };
        push(&mut ids, "(");
        let p_pos = push(&mut ids, "[P]");
        push(&mut ids, "entities");
        push(&mut ids, "(");
        let mut e_pos = Vec::with_capacity(labels.len());
        for lab in labels {
            e_pos.push(push(&mut ids, "[E]"));
            push(&mut ids, lab.as_ref());
        }
        push(&mut ids, ")");
        push(&mut ids, ")");
        push(&mut ids, "[SEP_TEXT]");
        (ids, p_pos, e_pos)
    }

    /// Tokenize each text word ONCE — its subtoken ids and CHARACTER offsets into the (lowercased)
    /// text. Reused by every window; the offsets are global, so decoded spans always point back into
    /// the original document regardless of which window found them.
    #[cfg(feature = "cli")]
    fn word_toks(&self, text: &str) -> Vec<WordTok> {
        let low = text.to_lowercase();
        let mut out = Vec::new();
        for m in self.splitter.find_iter(&low) {
            let char_start = low[..m.start()].chars().count();
            let char_end = char_start + m.as_str().chars().count();
            let enc = self
                .tokenizer
                .encode(m.as_str(), false)
                .expect("gliner2 tokenize word");
            out.push(WordTok {
                ids: enc.get_ids().to_vec(),
                char_start,
                char_end,
            });
        }
        out
    }

    /// **Entity extraction** (`extract_entities`; for `/extract/combined`) — CHUNKED so a document of
    /// ANY length works. The mDeBERTa backbone caps at `max_pos` (512) tokens; rather than error on a
    /// long document, the label schema is held fixed and the text words are packed into windows that
    /// fit `max_pos − schema`, overlapping by `max_width` words so no span (≤ `max_width` words wide)
    /// is ever cut at a boundary. Each window decodes independently (words carry GLOBAL char offsets),
    /// then spans are merged with the SAME per-type greedy overlap removal. A document that fits in
    /// one window is bit-identical to the unchunked decode (the merge is idempotent on one window).
    #[cfg(feature = "cli")]
    pub fn predict_entities(
        &mut self,
        text: &str,
        labels: &[impl AsRef<str>],
        threshold: f32,
    ) -> Result<Vec<Ner2Entity>> {
        self.predict_entities_windowed(text, labels, threshold, true)
    }

    /// `predict_entities`, with the window batching selectable. **Test/benchmark hook only** —
    /// production always takes `batched = true` (the default path above).
    ///
    /// `batched = false` reproduces the pre-2026-07-24 behaviour: one backbone forward per window.
    /// Keeping it callable buys two things that are otherwise impossible:
    ///
    /// * **an equivalence proof** — both paths run against the same weights in the same process, so
    ///   a test can assert the entities are identical rather than merely plausible, and
    /// * **an honest A/B** — absolute latencies on a shared dev box are not comparable across runs
    ///   (a build or a browser moves them by 5x), so the only trustworthy comparison interleaves
    ///   the two paths inside ONE process.
    ///
    /// Both are in `tests/gliner2_batched_windows.rs`.
    #[cfg(feature = "cli")]
    pub fn predict_entities_windowed(
        &mut self,
        text: &str,
        labels: &[impl AsRef<str>],
        threshold: f32,
        batched: bool,
    ) -> Result<Vec<Ner2Entity>> {
        anyhow::ensure!(!labels.is_empty(), "no entity labels given");
        anyhow::ensure!(
            self.count_embed.is_some(),
            "counting layer {:?} is not ported — this checkpoint serves classification only \
             (classify_text)",
            self.counting_layer
        );
        let label_strs: Vec<String> = labels.iter().map(|l| l.as_ref().to_string()).collect();
        let (schema_ids, p_pos, e_pos) = self.entity_schema_ids(labels);
        let (orig, win_ids, win_words) = self.entity_windows(text, &schema_ids);

        // Run the backbone ONCE over every window (see `entity_windows` for the assembly).
        //
        // This used to be one `forward_hidden` per window, each an `EncBatch` of a single sequence.
        // `EncBatch` is ragged (concatenated tokens + `seq_starts`), so batching costs no padding —
        // it just turns N small GEMMs into one large one. The math per window is untouched: the
        // backbone attends within a sequence, so window i's states are bit-identical to what a solo
        // forward produced. Gated by `tests/gliner2_batched_windows.rs`.
        let t_bb = std::time::Instant::now();
        let states_all: Vec<f32> = if batched {
            self.backbone
                .forward_hidden(&EncBatch::from_seqs(win_ids.iter().cloned()))
                .context("gliner2 backbone (batched windows)")?
        } else {
            // The pre-2026-07-24 path: one forward per window, each an EncBatch of one sequence.
            let mut v = Vec::new();
            for ids in &win_ids {
                v.extend(
                    self.backbone
                        .forward_hidden(&EncBatch::from_seqs([ids.clone()]))
                        .context("gliner2 backbone (per-window)")?,
                );
            }
            v
        };
        let t_backbone = t_bb.elapsed();

        // Decode each window from its slice of the batched states; spans already carry global char
        // offsets. `seq_starts` is implicit here: sequences are concatenated in order, so window i
        // begins at the running token offset.
        //
        // OSFKB_GLINER2_TRACE=1 splits the call into BACKBONE (device) vs HEADS (always CPU:
        // span_rep / count_pred / the scorer GEMM). Without that split it is easy to "optimise" the
        // backbone and measure nothing, because on a GPU the heads can be the whole cost.
        let trace = std::env::var("OSFKB_GLINER2_TRACE").is_ok_and(|v| v != "0");
        let t_heads = std::time::Instant::now();
        let h = self.hidden;
        let mut all: Vec<Ner2Entity> = Vec::new();
        let mut tok_off = 0usize;
        for (ids, wds) in win_ids.iter().zip(&win_words) {
            let states = &states_all[tok_off * h..(tok_off + ids.len()) * h];
            tok_off += ids.len();
            let ents =
                self.entities_window(states, p_pos, &e_pos, wds, &orig, &label_strs, threshold)?;
            all.extend(ents);
        }

        if trace {
            eprintln!(
                "gliner2 trace: backbone {:.3}s | heads {:.3}s | windows {}",
                t_backbone.as_secs_f64(),
                t_heads.elapsed().as_secs_f64(),
                win_ids.len(),
            );
        }
        Ok(Self::merge_windows(all, &label_strs))
    }

    /// Entity extraction over MANY texts in ONE backbone forward.
    ///
    /// The single-text path already batches a document's own windows; this batches across
    /// *requests* as well, which is what actually fills a GPU. Measured motivation: on the fleet's
    /// V100 a single call leaves the device ~idle (0 % utilisation between calls, and the GPU is
    /// only ~4-5x a fast CPU per call) — the work per launch is simply too small. The obvious
    /// alternative, running N model instances concurrently, is ruled out where VRAM is tight: each
    /// instance duplicates the weights, whereas batching reuses ONE copy and pays only for
    /// activations. On the reasoninglayer fleet that is decisive — the V100 is time-sliced with
    /// glmocr/docextract and has ~2 GB free.
    ///
    /// Returns one entity list per input text, in order. Spans/labels are identical to calling
    /// `predict_entities` per text, confidences within float noise, and a batch of ONE is
    /// bit-exact — all gated by `tests/gliner2_batched_windows.rs`.
    ///
    /// **MEASURED NEUTRAL TODAY — do not reach for this expecting a speedup.** Interleaved A/B on
    /// Metal: 0.92x / 0.79x / 0.94x at n = 2 / 4 / 8. The reason is in the trace below: with the
    /// backbone on a GPU the *heads* (span_rep, count_pred, the scorer GEMM — all CPU) are ~80 % of
    /// the call, so batching the backbone optimises the small end. Kept because it is correct,
    /// gated and additive, and it becomes the right lever the moment the heads move to the device.
    #[cfg(feature = "cli")]
    pub fn predict_entities_batch(
        &mut self,
        texts: &[impl AsRef<str>],
        labels: &[impl AsRef<str>],
        threshold: f32,
    ) -> Result<Vec<Vec<Ner2Entity>>> {
        anyhow::ensure!(!labels.is_empty(), "no entity labels given");
        anyhow::ensure!(
            self.count_embed.is_some(),
            "counting layer {:?} is not ported — this checkpoint serves classification only \
             (classify_text)",
            self.counting_layer
        );
        if texts.is_empty() {
            return Ok(Vec::new());
        }
        let label_strs: Vec<String> = labels.iter().map(|l| l.as_ref().to_string()).collect();
        // One schema for the whole batch: the label set is shared, so the prompt prefix is too.
        let (schema_ids, p_pos, e_pos) = self.entity_schema_ids(labels);

        // Assemble every text's windows up front; remember each text's slice of the window list.
        let mut origs: Vec<Vec<char>> = Vec::with_capacity(texts.len());
        let mut spans: Vec<(usize, usize)> = Vec::with_capacity(texts.len()); // window range per text
        let mut all_ids: Vec<Vec<u32>> = Vec::new();
        let mut all_words: Vec<Vec<Word>> = Vec::new();
        for text in texts {
            let (orig, ids, wds) = self.entity_windows(text.as_ref(), &schema_ids);
            let start = all_ids.len();
            all_ids.extend(ids);
            all_words.extend(wds);
            spans.push((start, all_ids.len()));
            origs.push(orig);
        }

        let states_all = self
            .backbone
            .forward_hidden(&EncBatch::from_seqs(all_ids.iter().cloned()))
            .context("gliner2 backbone (batched texts)")?;

        // Decode per window, then merge per text. Sequences are concatenated in order, so each
        // window begins at the running token offset.
        let h = self.hidden;
        let mut offsets: Vec<usize> = Vec::with_capacity(all_ids.len() + 1);
        let mut run = 0usize;
        for ids in &all_ids {
            offsets.push(run);
            run += ids.len();
        }
        let mut out: Vec<Vec<Ner2Entity>> = Vec::with_capacity(texts.len());
        for (ti, (wa, wb)) in spans.into_iter().enumerate() {
            let mut found: Vec<Ner2Entity> = Vec::new();
            for wi in wa..wb {
                let off = offsets[wi];
                let states = &states_all[off * h..(off + all_ids[wi].len()) * h];
                let ents = self.entities_window(
                    states,
                    p_pos,
                    &e_pos,
                    &all_words[wi],
                    &origs[ti],
                    &label_strs,
                    threshold,
                )?;
                found.extend(ents);
            }
            out.push(Self::merge_windows(found, &label_strs));
        }
        Ok(out)
    }

    /// Assemble one text's windows: `(original chars, per-window token ids, per-window words)`.
    /// Each window is `schema ‖ window words`, greedily filled to the backbone's position budget and
    /// overlapping by `max_width` words so no entity is cut at a boundary. Shared by the single-text
    /// and batched paths so the two cannot drift apart.
    #[cfg(feature = "cli")]
    fn entity_windows(
        &mut self,
        text: &str,
        schema_ids: &[u32],
    ) -> (Vec<char>, Vec<Vec<u32>>, Vec<Vec<Word>>) {
        let text = normalize_text(text);
        let orig: Vec<char> = text.chars().collect();
        let words = self.word_toks(&text);

        // Per-window token budget for the TEXT words = max_pos − schema − a small margin (covers the
        // learned-position offset, if any). The backbone enforces `len + offset ≤ max_pos`.
        let budget = self
            .backbone
            .config()
            .max_pos
            .saturating_sub(schema_ids.len())
            .saturating_sub(4)
            .max(1);
        let overlap = self.max_width.max(1);

        // Greedy word windows ≤ budget subtokens, overlapping by `overlap` words. `j == a` guarantees
        // at least one word per window even if a single word tokenizes beyond the budget.
        let mut windows: Vec<(usize, usize)> = Vec::new();
        let mut a = 0usize;
        while a < words.len() {
            let mut j = a;
            let mut toks = 0usize;
            while j < words.len() && (j == a || toks + words[j].ids.len() <= budget) {
                toks += words[j].ids.len();
                j += 1;
            }
            windows.push((a, j));
            if j >= words.len() {
                break;
            }
            a = j.saturating_sub(overlap).max(a + 1);
        }
        if windows.is_empty() {
            windows.push((0, 0)); // empty text → one empty window (matches the unchunked empty case)
        }

        let mut win_ids: Vec<Vec<u32>> = Vec::with_capacity(windows.len());
        let mut win_words: Vec<Vec<Word>> = Vec::with_capacity(windows.len());
        for (wa, wb) in windows {
            let mut ids = schema_ids.to_vec();
            let mut wds: Vec<Word> = Vec::with_capacity(wb - wa);
            for w in &words[wa..wb] {
                let first_tok = ids.len();
                ids.extend_from_slice(&w.ids);
                wds.push(Word {
                    first_tok,
                    char_start: w.char_start,
                    char_end: w.char_end,
                });
            }
            win_ids.push(ids);
            win_words.push(wds);
        }
        (orig, win_ids, win_words)
    }

    /// Merge one text's per-window spans: per label, greedy overlap removal by char offset, highest
    /// confidence first (dedups the overlap regions AND resolves cross-window disagreement
    /// consistently).
    #[cfg(feature = "cli")]
    fn merge_windows(all: Vec<Ner2Entity>, label_strs: &[String]) -> Vec<Ner2Entity> {
        let mut out: Vec<Ner2Entity> = Vec::new();
        for lab in label_strs {
            let mut spans: Vec<Ner2Entity> =
                all.iter().filter(|e| &e.label == lab).cloned().collect();
            spans.sort_by(|a, b| {
                b.confidence
                    .partial_cmp(&a.confidence)
                    .unwrap_or(std::cmp::Ordering::Equal)
            });
            let mut selected: Vec<Ner2Entity> = Vec::new();
            for sp in spans {
                let ov = selected
                    .iter()
                    .any(|s| !(sp.end <= s.start || sp.start >= s.end));
                if !ov {
                    selected.push(sp);
                }
            }
            out.extend(selected);
        }
        out
    }

    /// Decode entities for ONE assembled window (`input_ids` = schema ‖ window words). This is
    /// gliner2's `extract_entities` math verbatim: backbone → span_rep → count → score → per-type
    /// greedy select; the words carry global char offsets so the returned spans need no adjustment.
    #[cfg(feature = "cli")]
    #[allow(clippy::too_many_arguments)]
    /// `states` is this window's backbone output (`[len × hidden]`), sliced out of ONE batched
    /// forward over every window — see `predict_entities`. Taking states rather than token ids is
    /// what lets the caller batch; the head math below is unchanged.
    fn entities_window(
        &mut self,
        states: &[f32],
        p_pos: usize,
        e_pos: &[usize],
        words: &[Word],
        orig: &[char],
        labels: &[String],
        threshold: f32,
    ) -> Result<Vec<Ner2Entity>> {
        let l = words.len();
        if l == 0 {
            return Ok(Vec::new());
        }
        let count_embed = self
            .count_embed
            .as_ref()
            .expect("count_embed presence checked by predict_entities");
        let h = self.hidden;
        let mut token_embs = vec![0f32; l * h];
        for (wi, w) in words.iter().enumerate() {
            token_embs[wi * h..(wi + 1) * h]
                .copy_from_slice(&states[w.first_tok * h..(w.first_tok + 1) * h]);
        }
        let span_rep = self.span_rep(&token_embs, l);

        // count = argmax(count_pred([P])); ≤0 ⇒ no entities.
        let count_logits = self.count_pred.forward(&states[p_pos * h..(p_pos + 1) * h]);
        let pred_count = argmax(&count_logits);
        if pred_count == 0 {
            return Ok(Vec::new());
        }
        // struct_proj [count·n, H] from the GRU over the n [E] field markers.
        let n = e_pos.len();
        let mut fields = vec![0f32; n * h];
        for (i, &ep) in e_pos.iter().enumerate() {
            fields[i * h..(i + 1) * h].copy_from_slice(&states[ep * h..(ep + 1) * h]);
        }
        let struct_proj = count_embed.forward(&fields, pred_count);
        let bp = struct_proj.len() / h; // count·n
        let scorer = Linear {
            w: struct_proj,
            b: vec![0.0; bp],
            n: bp,
            k: h,
            ..Default::default()
        };
        let mut flat = scorer.forward(&span_rep); // [L·K, bp]
        for v in flat.iter_mut() {
            *v = sigmoid(*v);
        }
        let kmax = self.max_width;

        // decode: only instance 0 (b=0 → bp index = type), per type.
        let mut out = Vec::new();
        for (ty, lab) in labels.iter().enumerate() {
            let mut spans: Vec<(String, f32, usize, usize)> = Vec::new();
            for li in 0..l {
                for ki in 0..kmax {
                    if li + ki >= l {
                        break;
                    }
                    let s = flat[(li * kmax + ki) * bp + ty]; // b=0
                    if s >= threshold {
                        let cs = words[li].char_start;
                        let ce = words[li + ki].char_end;
                        let txt: String = orig
                            .get(cs..ce)
                            .map(|c| c.iter().collect::<String>())
                            .unwrap_or_default();
                        spans.push((txt.trim().to_string(), s, cs, ce));
                    }
                }
            }
            spans.sort_by(|a, b| b.1.partial_cmp(&a.1).unwrap_or(std::cmp::Ordering::Equal));
            let mut selected: Vec<(String, f32, usize, usize)> = Vec::new();
            for sp in spans {
                let overlap = selected.iter().any(|s| !(sp.3 <= s.2 || sp.2 >= s.3));
                if !overlap {
                    selected.push(sp);
                }
            }
            for (txt, conf, cs, ce) in selected {
                out.push(Ner2Entity {
                    text: txt,
                    label: lab.clone(),
                    start: cs,
                    end: ce,
                    confidence: conf,
                });
            }
        }
        Ok(out)
    }

    /// Tokenize the CLASSIFICATION schema — per task `( [P] prompt ( [L] l1 [L] l2 … ) )`, tasks
    /// joined by `[SEP_STRUCT]` — then `[SEP_TEXT]` + lowercased text words. The `[P]` prompt is ONE
    /// element (`SchemaTransformer._transform_schema`): the task name, `": {prompt}"` if given, then
    /// ` [DESCRIPTION] label: desc` per (ordered) description and ` [EXAMPLE] in [OUTPUT] out` per
    /// example — descriptions/examples filtered to the task's labels, as inference mode "both" does.
    #[cfg(feature = "cli")]
    fn tokenize_classification(&self, text: &str, tasks: &[ClsTask]) -> (Vec<u32>, Vec<ClsBlock>) {
        let (mut ids, blocks) = self.classify_schema_ids(tasks);
        let low = text.to_lowercase();
        let push = |ids: &mut Vec<u32>, s: &str| {
            let enc = self
                .tokenizer
                .encode(s, false)
                .expect("gliner2 tokenize element");
            ids.extend_from_slice(enc.get_ids());
        };
        for m in self.splitter.find_iter(&low) {
            push(&mut ids, m.as_str());
        }
        (ids, blocks)
    }

    /// The classification schema prefix — all task blocks `( [P] prompt ( [L] … ) )` joined by
    /// `[SEP_STRUCT]`, then `[SEP_TEXT]` — the part BEFORE the text, fixed across chunk windows. The
    /// `[P]`/`[L]` positions live in this prefix, so they're identical in every window.
    #[cfg(feature = "cli")]
    fn classify_schema_ids(&self, tasks: &[ClsTask]) -> (Vec<u32>, Vec<ClsBlock>) {
        let mut ids: Vec<u32> = Vec::new();
        let mut blocks: Vec<ClsBlock> = Vec::new();
        let push = |ids: &mut Vec<u32>, s: &str| -> usize {
            let start = ids.len();
            let enc = self
                .tokenizer
                .encode(s, false)
                .expect("gliner2 tokenize element");
            ids.extend_from_slice(enc.get_ids());
            start
        };

        for (i, task) in tasks.iter().enumerate() {
            if i > 0 {
                push(&mut ids, "[SEP_STRUCT]");
            }
            let mut prompt_str = match &task.prompt {
                Some(p) => format!("{}: {p}", task.name),
                None => task.name.clone(),
            };
            for (label, desc) in &task.label_descriptions {
                if task.labels.iter().any(|l| l == label) {
                    prompt_str.push_str(&format!(" [DESCRIPTION] {label}: {desc}"));
                }
            }
            for (inp, out) in &task.examples {
                if task.labels.iter().any(|l| l == out) {
                    prompt_str.push_str(&format!(" [EXAMPLE] {inp} [OUTPUT] {out}"));
                }
            }
            push(&mut ids, "(");
            let p_pos = push(&mut ids, "[P]");
            push(&mut ids, &prompt_str);
            push(&mut ids, "(");
            let mut l_pos = Vec::with_capacity(task.labels.len());
            for label in &task.labels {
                l_pos.push(push(&mut ids, "[L]"));
                push(&mut ids, label);
            }
            push(&mut ids, ")");
            push(&mut ids, ")");
            blocks.push(ClsBlock { p_pos, l_pos });
        }
        push(&mut ids, "[SEP_TEXT]");
        (ids, blocks)
    }

    /// **Text classification via gliner2's schema path** (`classify_text`) — CHUNKED so ANY-length
    /// document works. Per task the `classifier` MLP scores each `[L]` marker (activation `class_act`,
    /// auto: sigmoid iff multi-label). For a document past the 512-token backbone window, the schema
    /// is held fixed and the text is split into windows; each window is classified and the per-label
    /// probabilities are aggregated by **max** across windows — the safety-conservative rule (a
    /// document is `unsafe` if ANY passage is), then decoded (single-label argmax / multi-label
    /// ≥ `cls_threshold` with argmax fallback). A single-window document is identical to the
    /// unchunked path (and to [`Self::classify_debug`], which the parity gate pins).
    #[cfg(feature = "cli")]
    pub fn classify_text(&mut self, text: &str, tasks: &[ClsTask]) -> Result<Vec<ClsResult>> {
        anyhow::ensure!(!tasks.is_empty(), "no classification tasks given");
        for task in tasks {
            anyhow::ensure!(
                !task.labels.is_empty(),
                "task {:?} has no labels",
                task.name
            );
        }
        let text = normalize_text(text);
        let (schema_ids, blocks) = self.classify_schema_ids(tasks);
        let words = self.word_toks(&text);

        let max_pos = self.backbone.config().max_pos;
        let fit = max_pos
            .saturating_sub(schema_ids.len())
            .saturating_sub(4)
            .max(1);
        // Cap the window well below the position limit: window classification is HOLISTIC, so a small
        // unsafe fragment (e.g. an injected instruction) diluted in a full 512-token window would be
        // out-voted by benign context. Smaller, overlapping windows keep local content salient so an
        // unsafe passage anywhere in a long document still flips its own window (then max-aggregated).
        let budget = fit.min(128);
        let overlap = (budget / 2).max(1);

        let mut windows: Vec<(usize, usize)> = Vec::new();
        let mut a = 0usize;
        while a < words.len() {
            let mut j = a;
            let mut toks = 0usize;
            while j < words.len() && (j == a || toks + words[j].ids.len() <= budget) {
                toks += words[j].ids.len();
                j += 1;
            }
            windows.push((a, j));
            if j >= words.len() {
                break;
            }
            a = j.saturating_sub(overlap).max(a + 1);
        }
        if windows.is_empty() {
            windows.push((0, 0));
        }

        // Decode EACH window, then aggregate: a label is chosen for the document iff some window
        // chose it (union), and its reported score is the max across windows. Aggregating the
        // *decisions* (not raw probabilities) is what makes it correct for single-label tasks — a
        // late "unsafe" window wins even though earlier benign windows have a higher "safe" score.
        let mut max_scores: Vec<Vec<f32>> =
            tasks.iter().map(|t| vec![0f32; t.labels.len()]).collect();
        let mut chosen_any: Vec<Vec<bool>> =
            tasks.iter().map(|t| vec![false; t.labels.len()]).collect();
        for (wa, wb) in windows {
            let mut ids = schema_ids.clone();
            for w in &words[wa..wb] {
                ids.extend_from_slice(&w.ids);
            }
            let per_task = self.classify_window(&ids, &blocks, tasks)?;
            for (ti, task) in tasks.iter().enumerate() {
                let probs = &per_task[ti];
                for (li, &p) in probs.iter().enumerate() {
                    max_scores[ti][li] = max_scores[ti][li].max(p);
                }
                // this window's decoded choice — single-label argmax, or multi-label ≥ threshold
                // with argmax fallback (identical to the unchunked decode on that window).
                if task.multi_label {
                    let mut any = false;
                    for (li, &p) in probs.iter().enumerate() {
                        if p >= task.cls_threshold {
                            chosen_any[ti][li] = true;
                            any = true;
                        }
                    }
                    if !any {
                        chosen_any[ti][argmax(probs)] = true;
                    }
                } else {
                    chosen_any[ti][argmax(probs)] = true;
                }
            }
        }

        // Build results: scores = max per label; chosen = the union set (each with its max score).
        let mut results = Vec::with_capacity(tasks.len());
        for (ti, task) in tasks.iter().enumerate() {
            let scores: Vec<ClsLabelScore> = task
                .labels
                .iter()
                .enumerate()
                .map(|(li, l)| ClsLabelScore {
                    label: l.clone(),
                    score: max_scores[ti][li],
                })
                .collect();
            let chosen: Vec<ClsLabelScore> = (0..task.labels.len())
                .filter(|&li| chosen_any[ti][li])
                .map(|li| scores[li].clone())
                .collect();
            results.push(ClsResult {
                task: task.name.clone(),
                scores,
                chosen,
            });
        }
        Ok(results)
    }

    /// Classify ONE assembled window (`input_ids` = schema ‖ window words) → per-task label
    /// probabilities. The `[L]` marker positions come from the fixed schema, so they're valid in
    /// every window.
    #[cfg(feature = "cli")]
    fn classify_window(
        &mut self,
        input_ids: &[u32],
        blocks: &[ClsBlock],
        tasks: &[ClsTask],
    ) -> Result<Vec<Vec<f32>>> {
        let states = self
            .backbone
            .forward_hidden(&EncBatch::from_seqs([input_ids.to_vec()]))
            .context("gliner2 backbone")?;
        let h = self.hidden;
        let mut out = Vec::with_capacity(tasks.len());
        for (task, blk) in tasks.iter().zip(blocks) {
            let n = blk.l_pos.len();
            let mut cls_embeds = vec![0f32; n * h];
            for (j, &lp) in blk.l_pos.iter().enumerate() {
                cls_embeds[j * h..(j + 1) * h].copy_from_slice(&states[lp * h..(lp + 1) * h]);
            }
            let logits = self.classifier.forward(&cls_embeds);
            let sigmoid_act = match task.class_act {
                ClsActivation::Sigmoid => true,
                ClsActivation::Softmax => false,
                ClsActivation::Auto => task.multi_label,
            };
            let probs: Vec<f32> = if sigmoid_act {
                logits.iter().map(|&v| sigmoid(v)).collect()
            } else {
                softmax(&logits)
            };
            out.push(probs);
        }
        Ok(out)
    }

    /// [`Self::classify_text`] that also returns every intermediate, for the parity gate.
    #[cfg(feature = "cli")]
    pub fn classify_debug(
        &mut self,
        text: &str,
        tasks: &[ClsTask],
    ) -> Result<(Vec<ClsResult>, ClsIntermediates)> {
        anyhow::ensure!(!tasks.is_empty(), "no classification tasks given");
        for task in tasks {
            anyhow::ensure!(
                !task.labels.is_empty(),
                "task {:?} has no labels",
                task.name
            );
        }
        let text = normalize_text(text);
        let (input_ids, blocks) = self.tokenize_classification(text.as_str(), tasks);

        let states = self
            .backbone
            .forward_hidden(&EncBatch::from_seqs([input_ids.clone()]))
            .context("gliner2 backbone")?;
        let h = self.hidden;

        let mut results = Vec::with_capacity(tasks.len());
        let mut dbg: Option<(Vec<f32>, Vec<f32>, Vec<f32>)> = None;
        for (task, blk) in tasks.iter().zip(&blocks) {
            let n = blk.l_pos.len();
            // cls_embeds: the [L] marker embeddings ([P] is skipped — engine.py `embs[1:]`).
            let mut cls_embeds = vec![0f32; n * h];
            for (j, &lp) in blk.l_pos.iter().enumerate() {
                cls_embeds[j * h..(j + 1) * h].copy_from_slice(&states[lp * h..(lp + 1) * h]);
            }
            let logits = self.classifier.forward(&cls_embeds); // [n, 1] → [n]
            let sigmoid_act = match task.class_act {
                ClsActivation::Sigmoid => true,
                ClsActivation::Softmax => false,
                ClsActivation::Auto => task.multi_label,
            };
            let probs: Vec<f32> = if sigmoid_act {
                logits.iter().map(|&v| sigmoid(v)).collect()
            } else {
                softmax(&logits)
            };

            let scores: Vec<ClsLabelScore> = task
                .labels
                .iter()
                .zip(&probs)
                .map(|(l, &p)| ClsLabelScore {
                    label: l.clone(),
                    score: p,
                })
                .collect();
            let chosen: Vec<ClsLabelScore> = if task.multi_label {
                let picked: Vec<ClsLabelScore> = scores
                    .iter()
                    .filter(|s| s.score >= task.cls_threshold)
                    .cloned()
                    .collect();
                if picked.is_empty() {
                    vec![scores[argmax(&probs)].clone()]
                } else {
                    picked
                }
            } else {
                vec![scores[argmax(&probs)].clone()]
            };

            if dbg.is_none() {
                let mut block0 = vec![0f32; (1 + n) * h];
                block0[0..h].copy_from_slice(&states[blk.p_pos * h..(blk.p_pos + 1) * h]);
                block0[h..].copy_from_slice(&cls_embeds);
                dbg = Some((block0, logits.clone(), probs.clone()));
            }
            results.push(ClsResult {
                task: task.name.clone(),
                scores,
                chosen,
            });
        }

        let (schema_block0, logits0, probs0) = dbg.unwrap_or_default();
        let inter = ClsIntermediates {
            input_ids,
            last_hidden_state: states,
            schema_block0,
            logits0,
            probs0,
        };
        Ok((results, inter))
    }
}

/// gliner2's text normalization (`inference/engine.py`): empty → ".", else append "." unless the
/// text already ends in `.`, `!` or `?`. Everything downstream — word split, offsets, span slices —
/// works on the NORMALIZED text, exactly as Python does. (Every early gate sentence already ended
/// in ".", which is how this hid until a bare single-word input.)
fn normalize_text(text: &str) -> String {
    if text.is_empty() {
        return ".".to_string();
    }
    if text.ends_with(['.', '!', '?']) {
        text.to_string()
    } else {
        format!("{text}.")
    }
}

fn argmax(v: &[f32]) -> usize {
    let mut best = 0usize;
    let mut bv = f32::NEG_INFINITY;
    for (i, &x) in v.iter().enumerate() {
        if x > bv {
            bv = x;
            best = i;
        }
    }
    best
}

/// `_find_spans` + `spans[0]`: the earliest-then-narrowest span ≥ threshold for field `p` of
/// instance `b`, or `None`. Spans running past the text (`l+k >= L`) are excluded. Recovers the
/// span text (from the ORIGINAL chars) and its character offsets.
#[allow(clippy::too_many_arguments)]
fn find_span(
    scores: &[f32],
    b: usize,
    p: usize,
    m: usize,
    l: usize,
    kmax: usize,
    threshold: f32,
    words: &[Word],
    orig: &[char],
) -> Option<RelEnd> {
    for li in 0..l {
        for ki in 0..kmax {
            if li + ki >= l {
                break; // wider spans from this start also run past the text
            }
            let s = scores[((b * m + p) * l + li) * kmax + ki];
            if s >= threshold {
                // Python: char_start = start_map[start], char_end = end_map[start+width]; the offsets
                // returned are these RAW word boundaries — only the text is `.strip()`ed, not the offsets.
                let cs = words[li].char_start;
                let ce = words[li + ki].char_end;
                let text: String = orig
                    .get(cs..ce)
                    .map(|c| c.iter().collect())
                    .unwrap_or_default();
                return Some(RelEnd {
                    text: text.trim().to_string(),
                    start: cs,
                    end: ce,
                    confidence: s,
                });
            }
        }
    }
    None
}