inferencelayer 0.2.1

//! Encoder-model configuration: `config.json` → the declarative fields the encoder executor
//! branches on.
//!
//! **The "any embedding model" contract**: the executor never matches on [`EncArch`] — only on
//! [`PosKind`], [`NormKind`], [`MlpKind`], [`MaskKind`], the per-layer `layer_window`, and the
//! [`Pooling`](crate::pooling::Pooling) mode. Supporting a new encoder architecture is therefore
//! (1) one `EncArch` variant + one `from_json` arm mapping its HF config keys onto these fields,
//! and (2) one tensor-name table in the weight loader. Anything a config asks for that the fields
//! cannot express is refused loudly at parse time — never silently mis-computed.

use crate::pooling::Pooling;
use anyhow::{Context, Result};

/// Which encoder architecture the checkpoint is (drives tensor names at load; NOT executor logic).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum EncArch {
    /// BERT family: learned absolute positions, post-LN, biased dense GELU MLP. Covers
    /// all-MiniLM-L6-v2 AND paraphrase-multilingual-MiniLM-L12-v2 (XLM-R vocab, BERT body).
    Bert,
    /// XLM-RoBERTa: BERT body with position ids offset by `pad_token_id + 1` (bge-m3 /
    /// multilingual-e5 class).
    XlmRoberta,
    /// ModernBERT: RoPE (dual thetas), alternating global/local attention, GeGLU, bias-free,
    /// pre-norm, layer-0 attention-norm skip.
    ModernBert,
    /// nomic-bert (nomic-embed-text): rotary, SwiGLU, Matryoshka-truncatable output.
    NomicBert,
    /// SigLIP text tower (dense dual-encoder text side).
    SiglipText,
    /// SigLIP vision tower (ViT patch encoder + MAP attention-pooling head).
    SiglipVision,
    /// Qwen3 checkpoint run as an embedder: causal mask + last-token pooling, f16 (never Q4).
    Qwen3Embed,
    /// LFM2 backbone run as a late-interaction embedder (LFM2-ColBERT): per-token projection.
    Lfm2Colbert,
    /// DeBERTa-v2/v3: BERT body with **disentangled attention** and NO absolute position
    /// embeddings (`position_biased_input: false`). The backbone of GLiNER. See [`RelAttn`].
    DebertaV2,
}

/// DeBERTa's disentangled (relative) attention. On top of the usual content→content term, the
/// score gains a content→position (`c2p`) and a position→content (`p2c`) term, both indexed by a
/// LOG-BUCKETED relative distance:
///
/// ```text
/// score[i][j] = ( q_i·k_j  +  q_i·Kr[m]  +  k_j·Qr[m] ) / sqrt(head_dim · scale_factor)
///     where m = clamp(bucket(i - j) + span, 0, 2·span - 1)
///           scale_factor = 1 + c2p + p2c   (= 3 when both are on, as in every real checkpoint)
/// ```
///
/// `Kr`/`Qr` are the (layer-norm'd) relative-embedding table projected through **that layer's own**
/// key/query projections — `share_att_key: true`, so DeBERTa-v3 checkpoints carry no separate
/// `pos_key_proj`/`pos_query_proj` weights.
///
/// The single index `m` serves BOTH terms because the bucket function is odd
/// (`bucket(-x) == -bucket(x)`), which collapses HF's separate `c2p_pos` / `-p2c_pos` gathers into
/// one — see `disentangled_attention_bias` in `modeling_deberta_v2.py`.
#[derive(Clone, Copy, Debug, PartialEq)]
pub struct RelAttn {
    /// `position_buckets` — the table holds `2 · span` rows and `m` is centered on `span`.
    pub span: usize,
    /// `max_relative_positions` (falls back to `max_position_embeddings` when `-1`): the distance
    /// at which the log-bucketing saturates.
    pub max_rel: usize,
    /// Content→position term present (`"c2p" in pos_att_type`).
    pub c2p: bool,
    /// Position→content term present (`"p2c" in pos_att_type`).
    pub p2c: bool,
}

impl RelAttn {
    /// `1 + c2p + p2c` — divides EVERY term of the score, the content one included.
    pub fn scale_factor(&self) -> usize {
        1 + usize::from(self.c2p) + usize::from(self.p2c)
    }

    /// HF `make_log_bucket_position`: identity within `±span/2`, logarithmic beyond, saturating at
    /// `max_rel`. Odd in `rel`.
    pub fn bucket(&self, rel: isize) -> isize {
        let mid = (self.span / 2) as isize;
        if rel.abs() <= mid {
            return rel;
        }
        let sign = rel.signum() as f32;
        // The `< mid` branch of HF's `abs_pos` only ever feeds `log_pos`, and it cannot be taken
        // here (we already returned for |rel| <= mid), so `abs_pos == |rel|`.
        let abs_pos = rel.abs() as f32;
        let mid_f = mid as f32;
        let log_pos = ((abs_pos / mid_f).ln() / ((self.max_rel as f32 - 1.0) / mid_f).ln()
            * (mid_f - 1.0))
            .ceil()
            + mid_f;
        (sign * log_pos) as isize
    }

    /// The row of the relative-embedding table that query `i` uses for key `j`.
    pub fn index(&self, i: usize, j: usize) -> usize {
        let span = self.span as isize;
        (self.bucket(i as isize - j as isize) + span).clamp(0, 2 * span - 1) as usize
    }
}

/// Position-information scheme.
#[derive(Clone, Copy, Debug, PartialEq)]
pub enum PosKind {
    /// Learned absolute position embeddings, added at position `token_index + offset`
    /// (BERT: 0; XLM-R: `pad_token_id + 1`).
    Learned {
        /// First position id of a sequence.
        offset: usize,
    },
    /// Rotary embeddings applied per layer; `local_theta` serves windowed layers (ModernBERT).
    Rope {
        /// RoPE base frequency for global (full-attention) layers.
        theta: f32,
        /// RoPE base frequency for local (windowed) layers.
        local_theta: f32,
    },
}

/// Normalization operator.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum NormKind {
    /// Mean-subtracting LayerNorm, optionally with a learned bias.
    LayerNorm {
        /// Whether the norm carries a bias vector.
        bias: bool,
    },
    /// RMSNorm (decoder-based embedders).
    RmsNorm,
}

/// Activation function (shared by the GEMM epilogue and the GLU split kernel).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum Act {
    /// Exact erf-based GELU (BERT `"gelu"`).
    GeluErf,
    /// tanh-approximated GELU (`"gelu_new"` / `"gelu_pytorch_tanh"`).
    GeluTanh,
    /// SiLU / swish.
    Silu,
    /// tanh (SigLIP MAP-head MLP).
    Tanh,
    /// ReLU (GLiNER's `create_projection_layer`).
    Relu,
}

/// Feed-forward block shape.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MlpKind {
    /// Plain two-matmul MLP: `down(act(up(x)))`, optionally biased (BERT).
    Dense {
        /// Activation between the two matmuls.
        act: Act,
        /// Whether the linear layers carry biases.
        bias: bool,
    },
    /// Gated linear unit: `Wi` outputs `[T, 2·intermediate]`, `h = act(a) ⊙ b` (GeGLU/SwiGLU).
    Glu {
        /// Activation on the gate half.
        act: Act,
    },
}

/// Attention mask mode.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum MaskKind {
    /// Every token attends to every token of its sequence (encoders).
    Bidirectional,
    /// Token `t` attends to `0..=t` (decoder-based embedders).
    Causal,
}

/// Encoder hyper-parameters — the declarative surface the executor and loaders consume.
#[derive(Clone, Debug)]
pub struct EncoderConfig {
    /// Architecture tag (tensor-name table selection at load time only).
    pub arch: EncArch,
    /// Residual-stream width.
    pub hidden: usize,
    /// Transformer layer count.
    pub n_layers: usize,
    /// Attention heads.
    pub n_heads: usize,
    /// KV heads (== `n_heads` unless the architecture uses GQA).
    pub n_kv_heads: usize,
    /// Per-head dimension.
    pub head_dim: usize,
    /// Feed-forward width (for `Glu`, the width of ONE half — `Wi` is `2×` this).
    pub intermediate: usize,
    /// Vocabulary size.
    pub vocab: usize,
    /// Position-embedding table rows / maximum context (includes any learned-position offset).
    pub max_pos: usize,
    /// Token-type embedding rows (0 = none).
    pub type_vocab: usize,
    /// Position scheme.
    pub pos_kind: PosKind,
    /// Norm operator.
    pub norm_kind: NormKind,
    /// Feed-forward shape.
    pub mlp_kind: MlpKind,
    /// Attention mask mode.
    pub attn_mask: MaskKind,
    /// Per-layer attention window in tokens: `0` = full/global, else the total local span
    /// (ModernBERT: 128 = ±64 around the query).
    pub layer_window: Vec<u32>,
    /// Pre-norm residual placement (ModernBERT/decoders) vs post-LN (BERT).
    pub prenorm: bool,
    /// ModernBERT: layer 0 skips its attention norm (the embedding norm directly precedes it).
    pub skip_first_attn_norm: bool,
    /// QKV projections carry biases.
    pub qkv_bias: bool,
    /// Per-head RMS normalization of q and k before rotary (Qwen3 embedders).
    pub qk_norm: bool,
    /// LFM2-ColBERT short-conv kernel length (0 = no conv layers).
    pub conv_l: usize,
    /// Per-layer operator: `true` = attention, `false` = gated short-conv (LFM2-ColBERT only).
    pub layer_is_attn: Vec<bool>,
    /// Reduction from per-token states to the embedding output.
    pub pooling: Pooling,
    /// Norm epsilon.
    pub eps: f32,
    /// Disentangled relative attention (DeBERTa only; `None` everywhere else).
    pub rel_attn: Option<RelAttn>,
}

impl EncoderConfig {
    /// Parse encoder `config.json` bytes. Refuses (with the offending value named) anything the
    /// declarative fields cannot faithfully express.
    pub fn from_json(bytes: &[u8]) -> Result<Self> {
        let v: serde_json::Value = serde_json::from_slice(bytes)?;
        let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
        let f = |k: &str| v.get(k).and_then(|x| x.as_f64());
        let arch_name = v
            .get("architectures")
            .and_then(|a| a.get(0))
            .and_then(|x| x.as_str())
            .context("architectures[0] missing")?;
        let arch = encoder_arch(arch_name)
            .with_context(|| format!("`{arch_name}` is not a supported encoder architecture"))?;
        match arch {
            EncArch::Bert | EncArch::XlmRoberta => {
                // Absolute learned positions are the only implemented scheme for this family.
                let pos_type = v
                    .get("position_embedding_type")
                    .and_then(|x| x.as_str())
                    .unwrap_or("absolute");
                anyhow::ensure!(
                    pos_type == "absolute",
                    "BERT-family position_embedding_type `{pos_type}` not supported (absolute only)"
                );
                let act = act_from(
                    v.get("hidden_act")
                        .and_then(|x| x.as_str())
                        .unwrap_or("gelu"),
                )?;
                let hidden = g("hidden_size").context("hidden_size")?;
                let n_heads = g("num_attention_heads").context("num_attention_heads")?;
                let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
                // XLM-RoBERTa: position ids start at pad_token_id + 1 (HF
                // `create_position_ids_from_input_ids`); BERT starts at 0.
                let offset = if arch == EncArch::XlmRoberta {
                    g("pad_token_id").context("pad_token_id")? + 1
                } else {
                    0
                };
                // `*ForSequenceClassification` checkpoints (sentence-transformers CrossEncoders)
                // carry a scoring head, not a pooled embedding. `num_labels` (or the id2label map)
                // sizes the classifier output; a relevance scorer has 1.
                let pooling = if arch_name.ends_with("ForSequenceClassification") {
                    let n_labels = g("num_labels")
                        .or_else(|| {
                            v.get("id2label")
                                .and_then(|m| m.as_object())
                                .map(|m| m.len())
                        })
                        .unwrap_or(1)
                        .max(1);
                    Pooling::CrossEncoder { n_labels }
                } else {
                    Pooling::Mean
                };
                Ok(Self {
                    arch,
                    hidden,
                    n_layers,
                    n_heads,
                    n_kv_heads: n_heads,
                    head_dim: hidden / n_heads,
                    intermediate: g("intermediate_size").context("intermediate_size")?,
                    vocab: g("vocab_size").context("vocab_size")?,
                    max_pos: g("max_position_embeddings").context("max_position_embeddings")?,
                    type_vocab: g("type_vocab_size").unwrap_or(0),
                    pos_kind: PosKind::Learned { offset },
                    norm_kind: NormKind::LayerNorm { bias: true },
                    mlp_kind: MlpKind::Dense { act, bias: true },
                    attn_mask: MaskKind::Bidirectional,
                    layer_window: vec![0; n_layers],
                    prenorm: false,
                    skip_first_attn_norm: false,
                    qkv_bias: true,
                    qk_norm: false,
                    conv_l: 0,
                    layer_is_attn: vec![true; n_layers],
                    pooling,
                    eps: f("layer_norm_eps").unwrap_or(1e-12) as f32,
                    rel_attn: None,
                })
            }
            EncArch::ModernBert => {
                let hidden = g("hidden_size").context("hidden_size")?;
                let n_heads = g("num_attention_heads").context("num_attention_heads")?;
                let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
                let every = g("global_attn_every_n_layers").unwrap_or(3);
                anyhow::ensure!(every > 0, "global_attn_every_n_layers must be > 0");
                let window = g("local_attention").unwrap_or(128) as u32;
                let act = act_from(
                    v.get("hidden_activation")
                        .and_then(|x| x.as_str())
                        .unwrap_or("gelu"),
                )?;
                let norm_bias = v
                    .get("norm_bias")
                    .and_then(|x| x.as_bool())
                    .unwrap_or(false);
                Ok(Self {
                    arch,
                    hidden,
                    n_layers,
                    n_heads,
                    n_kv_heads: n_heads,
                    head_dim: hidden / n_heads,
                    intermediate: g("intermediate_size").context("intermediate_size")?,
                    vocab: g("vocab_size").context("vocab_size")?,
                    max_pos: g("max_position_embeddings").context("max_position_embeddings")?,
                    type_vocab: 0,
                    pos_kind: PosKind::Rope {
                        theta: f("global_rope_theta").unwrap_or(160_000.0) as f32,
                        local_theta: f("local_rope_theta").unwrap_or(10_000.0) as f32,
                    },
                    norm_kind: NormKind::LayerNorm { bias: norm_bias },
                    mlp_kind: MlpKind::Glu { act },
                    attn_mask: MaskKind::Bidirectional,
                    layer_window: (0..n_layers)
                        .map(|i| if i % every == 0 { 0 } else { window })
                        .collect(),
                    prenorm: true,
                    skip_first_attn_norm: true,
                    qkv_bias: v
                        .get("attention_bias")
                        .and_then(|x| x.as_bool())
                        .unwrap_or(false),
                    qk_norm: false,
                    conv_l: 0,
                    layer_is_attn: vec![true; n_layers],
                    pooling: Pooling::Mean,
                    eps: f("norm_eps").unwrap_or(1e-5) as f32,
                    rel_attn: None,
                })
            }
            EncArch::NomicBert => {
                // nomic-bert names its hyper-parameters `n_*`; the HF-standard duplicates
                // (`hidden_size`/`num_*`) are accepted as fallbacks so either export loads.
                let hidden = g("n_embd").or_else(|| g("hidden_size")).context("n_embd")?;
                let n_heads = g("n_head")
                    .or_else(|| g("num_attention_heads"))
                    .context("n_head")?;
                let n_layers = g("n_layer")
                    .or_else(|| g("num_hidden_layers"))
                    .context("n_layer")?;
                let intermediate = g("n_inner")
                    .or_else(|| g("intermediate_size"))
                    .context("n_inner")?;
                // Refuse — never silently mis-compute — every variant our forward cannot honor.
                let frac = f("rotary_emb_fraction").unwrap_or(1.0);
                anyhow::ensure!(
                    (frac - 1.0).abs() < 1e-6,
                    "nomic rotary_emb_fraction {frac} != 1.0 (partial rotary unsupported)"
                );
                anyhow::ensure!(
                    !v.get("rotary_emb_interleaved")
                        .and_then(|x| x.as_bool())
                        .unwrap_or(false),
                    "nomic interleaved rotary embeddings unsupported (rotate_half only)"
                );
                anyhow::ensure!(
                    v.get("rotary_emb_scale_base").map(|x| x.is_null()) != Some(false),
                    "nomic rotary_emb_scale_base scaling unsupported"
                );
                anyhow::ensure!(
                    v.get("rotary_scaling_factor").map(|x| x.is_null()) != Some(false),
                    "nomic rotary_scaling_factor unsupported"
                );
                anyhow::ensure!(
                    !v.get("use_rms_norm")
                        .and_then(|x| x.as_bool())
                        .unwrap_or(false),
                    "nomic use_rms_norm unsupported (biased LayerNorm only)"
                );
                anyhow::ensure!(
                    !v.get("causal").and_then(|x| x.as_bool()).unwrap_or(false),
                    "nomic causal attention unsupported (bidirectional embedder only)"
                );
                anyhow::ensure!(
                    !v.get("prenorm").and_then(|x| x.as_bool()).unwrap_or(false),
                    "nomic prenorm=true unsupported (nomic-embed is post-norm)"
                );
                // "swiglu" => SiLU gate; fall back to the explicit `hidden_act`.
                let act = match v.get("activation_function").and_then(|x| x.as_str()) {
                    Some("swiglu") | Some("silu") | Some("swish") => Act::Silu,
                    Some(other) => anyhow::bail!("unsupported nomic activation `{other}`"),
                    None => act_from(
                        v.get("hidden_act")
                            .and_then(|x| x.as_str())
                            .unwrap_or("silu"),
                    )?,
                };
                let theta = f("rotary_emb_base").unwrap_or(1000.0) as f32;
                Ok(Self {
                    arch,
                    hidden,
                    n_layers,
                    n_heads,
                    n_kv_heads: n_heads,
                    head_dim: hidden / n_heads,
                    intermediate,
                    vocab: g("vocab_size").context("vocab_size")?,
                    max_pos: g("max_position_embeddings")
                        .or_else(|| g("n_positions"))
                        .context("max_position_embeddings")?,
                    type_vocab: g("type_vocab_size").unwrap_or(0),
                    // Single-theta rotary over the full head_dim; no windowing (local == global).
                    pos_kind: PosKind::Rope {
                        theta,
                        local_theta: theta,
                    },
                    norm_kind: NormKind::LayerNorm { bias: true },
                    mlp_kind: MlpKind::Glu { act },
                    attn_mask: MaskKind::Bidirectional,
                    layer_window: vec![0; n_layers],
                    prenorm: false,
                    skip_first_attn_norm: false,
                    qkv_bias: v
                        .get("qkv_proj_bias")
                        .and_then(|x| x.as_bool())
                        .unwrap_or(false),
                    qk_norm: false,
                    conv_l: 0,
                    layer_is_attn: vec![true; n_layers],
                    pooling: Pooling::Mean,
                    eps: f("layer_norm_epsilon")
                        .or_else(|| f("layer_norm_eps"))
                        .unwrap_or(1e-12) as f32,
                    rel_attn: None,
                })
            }
            EncArch::DebertaV2 => {
                // DeBERTa is a BERT body with the position information moved OUT of the embedding
                // stage and INTO attention. Two consequences the declarative fields must carry:
                // `position_biased_input: false` → no absolute position table at all (the loader
                // leaves `pos` empty, and the embedding stage adds nothing), and the relative
                // scheme itself → `rel_attn`. Everything else is BERT: post-LN, biased dense GELU
                // MLP, bidirectional mask.
                let hidden = g("hidden_size").context("hidden_size")?;
                let n_heads = g("num_attention_heads").context("num_attention_heads")?;
                let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
                let max_pos = g("max_position_embeddings").context("max_position_embeddings")?;

                anyhow::ensure!(
                    v.get("relative_attention").and_then(|x| x.as_bool()) == Some(true),
                    "DeBERTa without relative_attention is not supported (nothing would position \
                     the tokens: position_biased_input is false)"
                );
                anyhow::ensure!(
                    !v.get("position_biased_input")
                        .and_then(|x| x.as_bool())
                        .unwrap_or(false),
                    "DeBERTa with position_biased_input=true (absolute position embeddings) is \
                     untested — refuse loudly"
                );
                // `norm_rel_ebd: layer_norm` means the relative table is LayerNorm'd before use.
                // Any other value would change the numbers, so refuse rather than ignore it.
                let norm_rel = v
                    .get("norm_rel_ebd")
                    .and_then(|x| x.as_str())
                    .unwrap_or("none");
                anyhow::ensure!(
                    norm_rel.contains("layer_norm"),
                    "DeBERTa norm_rel_ebd `{norm_rel}` not supported (layer_norm only)"
                );
                anyhow::ensure!(
                    v.get("share_att_key").and_then(|x| x.as_bool()) == Some(true),
                    "DeBERTa with share_att_key=false carries separate pos_key_proj/pos_query_proj \
                     weights this loader does not read — refuse loudly"
                );

                let pos_att: Vec<&str> = match v.get("pos_att_type") {
                    Some(serde_json::Value::Array(a)) => {
                        a.iter().filter_map(|x| x.as_str()).collect()
                    }
                    // Older configs spell it as a single "p2c|c2p" string.
                    Some(serde_json::Value::String(s)) => s.split('|').collect(),
                    _ => anyhow::bail!("pos_att_type missing"),
                };
                let rel = RelAttn {
                    span: g("position_buckets").context("position_buckets")?,
                    // `max_relative_positions: -1` (the usual value) means "use the context size".
                    max_rel: v
                        .get("max_relative_positions")
                        .and_then(|x| x.as_i64())
                        .filter(|x| *x > 0)
                        .map(|x| x as usize)
                        .unwrap_or(max_pos),
                    c2p: pos_att.contains(&"c2p"),
                    p2c: pos_att.contains(&"p2c"),
                };
                anyhow::ensure!(
                    rel.c2p || rel.p2c,
                    "DeBERTa pos_att_type has neither c2p nor p2c — nothing positions the tokens"
                );
                Ok(Self {
                    arch,
                    hidden,
                    n_layers,
                    n_heads,
                    n_kv_heads: n_heads,
                    head_dim: hidden / n_heads,
                    intermediate: g("intermediate_size").context("intermediate_size")?,
                    vocab: g("vocab_size").context("vocab_size")?,
                    max_pos,
                    type_vocab: g("type_vocab_size").unwrap_or(0),
                    // No absolute positions. `Learned { offset: 0 }` is inert here — the loader
                    // supplies no table, so the embedding stage adds nothing — and it keeps
                    // `max_pos` meaning "the tokenizer window", which is what callers read.
                    pos_kind: PosKind::Learned { offset: 0 },
                    norm_kind: NormKind::LayerNorm { bias: true },
                    mlp_kind: MlpKind::Dense {
                        act: act_from(
                            v.get("hidden_act")
                                .and_then(|x| x.as_str())
                                .unwrap_or("gelu"),
                        )?,
                        bias: true,
                    },
                    attn_mask: MaskKind::Bidirectional,
                    layer_window: vec![0; n_layers],
                    prenorm: false,
                    skip_first_attn_norm: false,
                    qkv_bias: true,
                    qk_norm: false,
                    conv_l: 0,
                    layer_is_attn: vec![true; n_layers],
                    pooling: Pooling::Mean,
                    // DeBERTa ships 1e-7, NOT BERT's 1e-12 — a wrong default here shows up as a
                    // small uniform drift in every hidden state.
                    eps: f("layer_norm_eps").unwrap_or(1e-7) as f32,
                    rel_attn: Some(rel),
                })
            }
            // Recognized encoder families whose config mapping lands with their loader phase.
            // Refusing here (rather than guessing a mapping) keeps "recognized" and "supported"
            // honest: `detect_model_kind` classifies them, loading them is a hard error until the
            // arm + golden gate exist.
            EncArch::SiglipText
            | EncArch::SiglipVision
            | EncArch::Qwen3Embed
            | EncArch::Lfm2Colbert => anyhow::bail!(
                "encoder architecture `{arch_name}` is recognized but not yet loadable by this build"
            ),
        }
    }
}

/// A ragged batch of token sequences: `tokens` is every sequence concatenated, `seq_starts`
/// (length `B+1`) delimits them. No padding exists anywhere — attention masking is exact by
/// construction because kernels iterate `seq_starts[s]..seq_starts[s+1]`.
#[derive(Clone, Debug)]
pub struct EncBatch {
    /// All sequences' token ids, concatenated.
    pub tokens: Vec<u32>,
    /// Sequence boundaries: sequence `s` spans `tokens[seq_starts[s]..seq_starts[s+1]]`.
    pub seq_starts: Vec<u32>,
    /// Optional per-token validity (1 = real, 0 = expansion/pad). Invalid tokens are excluded
    /// as attention keys and zeroed in conv-block inputs (the HF `apply_mask_to_padding_states`
    /// semantics) but STILL produce output vectors — ColBERT query expansion scores over them.
    /// `None` = all tokens valid.
    pub valid: Option<Vec<u32>>,
    /// Optional per-token segment id (BERT `token_type_ids`): 0 for the first text, 1 for the
    /// second in a sentence pair (the CrossEncoder `query [SEP] candidate` input). `None` = all 0
    /// (the single-sequence embedder case). Length, when present, equals `tokens`.
    pub type_ids: Option<Vec<u32>>,
}

impl EncBatch {
    /// A batch of one sequence.
    pub fn single(tokens: Vec<u32>) -> Self {
        let end = tokens.len() as u32;
        Self {
            tokens,
            seq_starts: vec![0, end],
            valid: None,
            type_ids: None,
        }
    }

    /// A batch of one sequence with an explicit validity mask (ColBERT query expansion).
    pub fn single_with_validity(tokens: Vec<u32>, valid: Vec<u32>) -> Self {
        let end = tokens.len() as u32;
        Self {
            tokens,
            seq_starts: vec![0, end],
            valid: Some(valid),
            type_ids: None,
        }
    }

    /// A single sentence-PAIR sequence with segment ids (CrossEncoder scoring input): `tokens` is
    /// `[CLS] a [SEP] b [SEP]`, `type_ids` the matching `0…0 1…1`.
    pub fn single_pair(tokens: Vec<u32>, type_ids: Vec<u32>) -> Self {
        let end = tokens.len() as u32;
        Self {
            tokens,
            seq_starts: vec![0, end],
            valid: None,
            type_ids: Some(type_ids),
        }
    }

    /// Build from per-sequence token lists.
    pub fn from_seqs<I: IntoIterator<Item = Vec<u32>>>(seqs: I) -> Self {
        let mut tokens = Vec::new();
        let mut seq_starts = vec![0u32];
        for s in seqs {
            tokens.extend_from_slice(&s);
            seq_starts.push(tokens.len() as u32);
        }
        Self {
            tokens,
            seq_starts,
            valid: None,
            type_ids: None,
        }
    }

    /// Build from per-sequence `(tokens, type_ids)` pairs (batched CrossEncoder scoring).
    pub fn from_pairs<I: IntoIterator<Item = (Vec<u32>, Vec<u32>)>>(seqs: I) -> Self {
        let mut tokens = Vec::new();
        let mut type_ids = Vec::new();
        let mut seq_starts = vec![0u32];
        for (s, t) in seqs {
            tokens.extend_from_slice(&s);
            type_ids.extend_from_slice(&t);
            seq_starts.push(tokens.len() as u32);
        }
        Self {
            tokens,
            seq_starts,
            valid: None,
            type_ids: Some(type_ids),
        }
    }

    /// Number of sequences.
    pub fn n_seqs(&self) -> usize {
        self.seq_starts.len().saturating_sub(1)
    }
}

/// Map an HF `architectures[0]` name to its encoder family, if it is one.
pub fn encoder_arch(name: &str) -> Option<EncArch> {
    if name.starts_with("Bert") {
        Some(EncArch::Bert)
    } else if name.starts_with("XLMRoberta") {
        Some(EncArch::XlmRoberta)
    } else if name.starts_with("ModernBert") {
        Some(EncArch::ModernBert)
    } else if name.starts_with("NomicBert") {
        Some(EncArch::NomicBert)
    } else if name.starts_with("Siglip") {
        Some(EncArch::SiglipText)
    } else if name.starts_with("DebertaV2") || name.starts_with("DebertaV3") {
        // HF names deberta-v3 checkpoints `DebertaV2*` (v3 is a training recipe, not a new
        // architecture); accept both spellings.
        Some(EncArch::DebertaV2)
    } else {
        None
    }
}

/// Map an HF activation name onto [`Act`]; unknown names are refused loudly.
fn act_from(name: &str) -> Result<Act> {
    Ok(match name {
        "relu" => Act::Relu,
        "gelu" => Act::GeluErf,
        "gelu_new" | "gelu_pytorch_tanh" => Act::GeluTanh,
        "silu" | "swish" => Act::Silu,
        "tanh" => Act::Tanh,
        other => anyhow::bail!("unsupported activation `{other}`"),
    })
}

/// Read the sentence-transformers pooling marker (`1_Pooling/config.json`) from a checkpoint
/// directory, if present. This is the data-driven signal that turns a decoder-architecture
/// checkpoint (e.g. Qwen3-Embedding, whose `config.json` says `Qwen3ForCausalLM`) into an
/// EMBEDDER — no per-model special-casing.
pub fn pooling_marker(dir: &std::path::Path) -> Option<Pooling> {
    let bytes = std::fs::read(dir.join("1_Pooling").join("config.json")).ok()?;
    let v: serde_json::Value = serde_json::from_slice(&bytes).ok()?;
    let flag = |k: &str| v.get(k).and_then(|x| x.as_bool()).unwrap_or(false);
    if flag("pooling_mode_lasttoken") {
        Some(Pooling::LastToken)
    } else if flag("pooling_mode_cls_token") {
        Some(Pooling::Cls)
    } else if flag("pooling_mode_mean_tokens") {
        Some(Pooling::Mean)
    } else {
        None
    }
}

/// Resolve a checkpoint DIRECTORY to an encoder configuration: encoder architectures parse via
/// [`EncoderConfig::from_json`] (the pooling marker, when present, overrides the per-arch
/// default); decoder architectures are accepted only with an explicit embedder marker —
/// `Qwen3ForCausalLM` + `pooling_mode_lasttoken` is Qwen3-Embedding. Everything else refuses
/// loudly: a generation checkpoint must never be silently embedded.
pub fn encoder_config_from_dir(dir: &std::path::Path) -> Result<EncoderConfig> {
    let bytes = std::fs::read(dir.join("config.json"))
        .with_context(|| format!("read {}/config.json", dir.display()))?;
    let v: serde_json::Value = serde_json::from_slice(&bytes)?;
    let arch_name = v
        .get("architectures")
        .and_then(|a| a.get(0))
        .and_then(|x| x.as_str())
        .context("architectures[0] missing")?;
    let marker = pooling_marker(dir);
    if encoder_arch(arch_name).is_some() {
        let mut cfg = EncoderConfig::from_json(&bytes)?;
        if let Some(p) = marker {
            cfg.pooling = p;
        }
        return Ok(cfg);
    }
    if arch_name == "Lfm2Model" {
        let dense_path = dir.join("1_Dense").join("config.json");
        anyhow::ensure!(
            dense_path.exists(),
            "`Lfm2Model` is a bare backbone; embedding it requires the pylate Dense module \
             (1_Dense/config.json) — LFM2-ColBERT checkpoints ship it, plain backbones are \
             refused rather than silently mis-embedded"
        );
        let dense: serde_json::Value = serde_json::from_slice(&std::fs::read(&dense_path)?)?;
        return lfm2_colbert_from_json(&v, &dense);
    }
    if arch_name == "Qwen3ForCausalLM" {
        anyhow::ensure!(
            marker == Some(Pooling::LastToken),
            "`{arch_name}` is a decoder; embedding it requires the sentence-transformers \
             last-token marker (1_Pooling/config.json with pooling_mode_lasttoken) — a plain \
             generation checkpoint is refused rather than silently mis-embedded"
        );
        return qwen3_embedder_from_json(&v);
    }
    anyhow::bail!("`{arch_name}` is not a supported encoder architecture")
}

/// Map a Qwen3 decoder config onto the declarative encoder fields: causal mask, last-token
/// pooling, RMSNorm with per-head qk-norm, single-theta full rotary, SiLU GLU, GQA geometry.
fn qwen3_embedder_from_json(v: &serde_json::Value) -> Result<EncoderConfig> {
    let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
    let f = |k: &str| v.get(k).and_then(|x| x.as_f64());
    let hidden = g("hidden_size").context("hidden_size")?;
    let n_heads = g("num_attention_heads").context("num_attention_heads")?;
    let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
    let head_dim = g("head_dim").unwrap_or(hidden / n_heads);
    let n_kv_heads = g("num_key_value_heads").unwrap_or(n_heads);
    anyhow::ensure!(
        n_heads % n_kv_heads == 0,
        "GQA requires n_heads ({n_heads}) divisible by n_kv_heads ({n_kv_heads})"
    );
    let theta = f("rope_theta").unwrap_or(1_000_000.0) as f32;
    Ok(EncoderConfig {
        arch: EncArch::Qwen3Embed,
        hidden,
        n_layers,
        n_heads,
        n_kv_heads,
        head_dim,
        intermediate: g("intermediate_size").context("intermediate_size")?,
        vocab: g("vocab_size").context("vocab_size")?,
        max_pos: g("max_position_embeddings").context("max_position_embeddings")?,
        type_vocab: 0,
        pos_kind: PosKind::Rope {
            theta,
            local_theta: theta,
        },
        norm_kind: NormKind::RmsNorm,
        mlp_kind: MlpKind::Glu { act: Act::Silu },
        attn_mask: MaskKind::Causal,
        layer_window: vec![0; n_layers],
        prenorm: true,
        skip_first_attn_norm: false,
        qkv_bias: false,
        qk_norm: true,
        conv_l: 0,
        layer_is_attn: vec![true; n_layers],
        pooling: Pooling::LastToken,
        eps: f("rms_norm_eps").unwrap_or(1e-6) as f32,
        rel_attn: None,
    })
}

/// Map an LFM2 backbone config + pylate Dense config onto the declarative encoder fields
/// (verified against LiquidAI/LFM2-ColBERT-350M): conv/attention hybrid per `layer_types`,
/// GQA with per-head qk-norm, RMSNorm prenorm, SwiGLU (`w1`/`w3` fused), causal mask (the
/// backbone's trained masking), per-token projection to the Dense `out_features` (ColBERT
/// late interaction).
fn lfm2_colbert_from_json(
    v: &serde_json::Value,
    dense: &serde_json::Value,
) -> Result<EncoderConfig> {
    let g = |k: &str| v.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
    let f = |k: &str| v.get(k).and_then(|x| x.as_f64());
    let hidden = g("hidden_size").context("hidden_size")?;
    let n_heads = g("num_attention_heads").context("num_attention_heads")?;
    let n_layers = g("num_hidden_layers").context("num_hidden_layers")?;
    let n_kv_heads = g("num_key_value_heads").unwrap_or(n_heads);
    anyhow::ensure!(n_heads % n_kv_heads == 0, "GQA head divisibility");
    // The pylate Dense head: identity activation, no bias — anything else is unverified.
    anyhow::ensure!(
        !dense.get("bias").and_then(|x| x.as_bool()).unwrap_or(false),
        "ColBERT Dense with bias is unverified — refused"
    );
    if let Some(act) = dense.get("activation_function").and_then(|x| x.as_str()) {
        anyhow::ensure!(
            act.ends_with("Identity"),
            "ColBERT Dense activation `{act}` unverified (Identity only)"
        );
    }
    let dim = dense
        .get("out_features")
        .and_then(|x| x.as_u64())
        .context("1_Dense out_features")? as usize;
    anyhow::ensure!(
        dense.get("in_features").and_then(|x| x.as_u64()) == Some(hidden as u64),
        "Dense in_features must equal hidden"
    );
    // block_auto_adjust_ff_dim: same recompute the decoder loader applies.
    let raw_ff = g("intermediate_size").context("intermediate_size")?;
    let intermediate = if v
        .get("block_auto_adjust_ff_dim")
        .and_then(|x| x.as_bool())
        .unwrap_or(false)
    {
        let m = g("block_multiple_of").unwrap_or(256);
        let ff = (2 * raw_ff) / 3;
        m * ff.div_ceil(m)
    } else {
        raw_ff
    };
    let layer_is_attn: Vec<bool> = v
        .get("layer_types")
        .and_then(|x| x.as_array())
        .context("layer_types")?
        .iter()
        .map(|x| x.as_str() == Some("full_attention"))
        .collect();
    anyhow::ensure!(layer_is_attn.len() == n_layers, "layer_types length");
    let theta = f("rope_theta").unwrap_or(1_000_000.0) as f32;
    Ok(EncoderConfig {
        arch: EncArch::Lfm2Colbert,
        hidden,
        n_layers,
        n_heads,
        n_kv_heads,
        head_dim: g("head_dim").unwrap_or(hidden / n_heads),
        intermediate,
        vocab: g("vocab_size").context("vocab_size")?,
        max_pos: g("max_position_embeddings").context("max_position_embeddings")?,
        type_vocab: 0,
        pos_kind: PosKind::Rope {
            theta,
            local_theta: theta,
        },
        norm_kind: NormKind::RmsNorm,
        mlp_kind: MlpKind::Glu { act: Act::Silu },
        attn_mask: MaskKind::Causal,
        layer_window: vec![0; n_layers],
        prenorm: true,
        skip_first_attn_norm: false,
        qkv_bias: false,
        qk_norm: true,
        conv_l: g("conv_L_cache").unwrap_or(3),
        layer_is_attn,
        pooling: Pooling::PerToken { dim },
        eps: f("norm_eps").unwrap_or(1e-5) as f32,
        rel_attn: None,
    })
}

/// Build the SigLIP TEXT-tower config from the joint `config.json` (HF `SiglipConfig` defaults
/// apply when the sub-config omits keys — verified against google/siglip2-base-patch16-224).
/// Text tower: learned positions (offset 0, no embedding-stage norm), biased pre-LN layers,
/// `final_layer_norm`, last-token pooling followed by the `head` projection. NOTE the canonical
/// SigLIP text tower attends WITHOUT a padding mask — the adapter pads to `max_pos` and the
/// final (padded) position is the pooled token, exactly as trained.
pub fn siglip_text_config(v: &serde_json::Value) -> Result<EncoderConfig> {
    let t = v.get("text_config").context("text_config")?;
    let g = |k: &str| t.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
    let f = |k: &str| t.get(k).and_then(|x| x.as_f64());
    let hidden = g("hidden_size").unwrap_or(768);
    let n_heads = g("num_attention_heads").unwrap_or(12);
    let n_layers = g("num_hidden_layers").unwrap_or(12);
    let act = act_from(
        t.get("hidden_act")
            .and_then(|x| x.as_str())
            .unwrap_or("gelu_pytorch_tanh"),
    )?;
    Ok(EncoderConfig {
        arch: EncArch::SiglipText,
        hidden,
        n_layers,
        n_heads,
        n_kv_heads: n_heads,
        head_dim: hidden / n_heads,
        intermediate: g("intermediate_size").unwrap_or(3072),
        vocab: g("vocab_size").unwrap_or(256_000),
        max_pos: g("max_position_embeddings").unwrap_or(64),
        type_vocab: 0,
        pos_kind: PosKind::Learned { offset: 0 },
        norm_kind: NormKind::LayerNorm { bias: true },
        mlp_kind: MlpKind::Dense { act, bias: true },
        attn_mask: MaskKind::Bidirectional,
        layer_window: vec![0; n_layers],
        prenorm: true,
        skip_first_attn_norm: false,
        qkv_bias: true,
        qk_norm: false,
        conv_l: 0,
        layer_is_attn: vec![true; n_layers],
        pooling: Pooling::LastToken,
        eps: f("layer_norm_eps").unwrap_or(1e-6) as f32,
        rel_attn: None,
    })
}

/// Build the SigLIP VISION-tower config: 2D patch embedding (`image_size`/`patch_size` ⇒
/// `(image/patch)²` positions), learned per-patch positions, biased pre-LN layers,
/// `post_layernorm`, MAP attention-pooling head.
pub fn siglip_vision_config(v: &serde_json::Value) -> Result<SiglipVisionSpec> {
    let c = v.get("vision_config").context("vision_config")?;
    let g = |k: &str| c.get(k).and_then(|x| x.as_u64()).map(|x| x as usize);
    let f = |k: &str| c.get(k).and_then(|x| x.as_f64());
    let hidden = g("hidden_size").unwrap_or(768);
    let n_heads = g("num_attention_heads").unwrap_or(12);
    let n_layers = g("num_hidden_layers").unwrap_or(12);
    let image_size = g("image_size").unwrap_or(224);
    let patch_size = g("patch_size").unwrap_or(16);
    anyhow::ensure!(
        image_size % patch_size == 0,
        "image_size {image_size} not divisible by patch_size {patch_size}"
    );
    let n_patches = (image_size / patch_size) * (image_size / patch_size);
    let act = act_from(
        c.get("hidden_act")
            .and_then(|x| x.as_str())
            .unwrap_or("gelu_pytorch_tanh"),
    )?;
    Ok(SiglipVisionSpec {
        config: EncoderConfig {
            arch: EncArch::SiglipVision,
            hidden,
            n_layers,
            n_heads,
            n_kv_heads: n_heads,
            head_dim: hidden / n_heads,
            intermediate: g("intermediate_size").unwrap_or(3072),
            vocab: 0,
            max_pos: n_patches,
            type_vocab: 0,
            pos_kind: PosKind::Learned { offset: 0 },
            norm_kind: NormKind::LayerNorm { bias: true },
            mlp_kind: MlpKind::Dense { act, bias: true },
            attn_mask: MaskKind::Bidirectional,
            layer_window: vec![0; n_layers],
            prenorm: true,
            skip_first_attn_norm: false,
            qkv_bias: true,
            qk_norm: false,
            conv_l: 0,
            layer_is_attn: vec![true; n_layers],
            pooling: Pooling::MapHead,
            eps: f("layer_norm_eps").unwrap_or(1e-6) as f32,
            rel_attn: None,
        },
        image_size,
        patch_size,
    })
}

/// Vision-tower spec: the declarative config plus the patch geometry the preprocessing needs.
#[derive(Clone, Debug)]
pub struct SiglipVisionSpec {
    /// The tower's declarative config (`max_pos` = number of patches).
    pub config: EncoderConfig,
    /// Input image side length in pixels.
    pub image_size: usize,
    /// Square patch side length in pixels.
    pub patch_size: usize,
}

/// Detect a SigLIP checkpoint directory (joint text+vision `config.json`; `architectures` may
/// be absent — `model_type: "siglip"`/`"siglip2"` is the signal) and return both tower configs.
pub fn siglip_configs_from_dir(dir: &std::path::Path) -> Result<(EncoderConfig, SiglipVisionSpec)> {
    let bytes = std::fs::read(dir.join("config.json"))
        .with_context(|| format!("read {}/config.json", dir.display()))?;
    let v: serde_json::Value = serde_json::from_slice(&bytes)?;
    let is_siglip = v
        .get("model_type")
        .and_then(|x| x.as_str())
        .is_some_and(|t| t.starts_with("siglip"))
        || v.get("architectures")
            .and_then(|a| a.get(0))
            .and_then(|x| x.as_str())
            .is_some_and(|a| a.starts_with("Siglip"));
    anyhow::ensure!(is_siglip, "{} is not a SigLIP checkpoint", dir.display());
    Ok((siglip_text_config(&v)?, siglip_vision_config(&v)?))
}