fathomdb-embedder 0.8.20

//! `OrtBgeEmbedder` — cross-vendor ONNX-Runtime BGE-small embedder.
//!
//! A caller-supplied `impl fathomdb_embedder_api::Embedder`, sibling of
//! `candle_bge` / `nomic`, that produces `BAAI/bge-small-en-v1.5` vectors
//! (dim 384) through the `ort` ONNX-Runtime binding. Injected purely via
//! `EmbedderChoice::Caller(Arc::new(OrtBgeEmbedder::…))` — the engine never
//! names it, so there is ZERO engine change (ADR-0.8.16-onnx-embedder-backend
//! §2). The `Default` variant stays candle-only, preserving the footprint
//! invariant.
//!
//! Why ONNX at all: candle reaches only CPU / CUDA / Metal — no AMD ROCm,
//! Intel OpenVINO, or Windows DirectML. ONNX Runtime reaches all of those, so
//! this backend is the cross-vendor reach-hardware path (ADR §1). It is behind
//! the NON-default `onnx-embedder` Cargo feature so the thin default build
//! gains zero deps (EMB-3 wheel-size gate).
//!
//! Numeric equivalence to the candle reference is MEASURED (not enforced) at
//! Slice 15 (ADR §3 / design §5); the interim guard is same-backend
//! build-and-read, enforced here structurally by giving ONNX a DISTINCT
//! embedder identity name (`…-onnx`) so the engine's identity check never
//! silently reads candle-written vectors with the ONNX backend.

use std::path::Path;
use std::sync::Mutex;

use fathomdb_embedder_api::{Embedder, EmbedderError, EmbedderIdentity, Vector};
use ort::execution_providers::{
    CPUExecutionProvider, CUDAExecutionProvider, DirectMLExecutionProvider, ExecutionProvider,
    ExecutionProviderDispatch, OpenVINOExecutionProvider, ROCmExecutionProvider,
};
use ort::session::Session;
use ort::value::Tensor;
use tokenizers::{Tokenizer, TruncationParams};

use crate::device::{parse_device_request, DeviceRequest};

/// Engine-facing identity name. Deliberately DISTINCT from the candle default
/// (`fathomdb-bge-small-en-v1.5`) so the engine's identity check enforces the
/// R-ONNX-3 same-backend build-and-read discipline: candle-written vectors and
/// ONNX-read queries never silently mix until 0.8.18 #5 enforces a candle↔ONNX
/// tolerance (ADR §3).
pub const ORT_BGE_EMBEDDER_NAME: &str = "fathomdb-bge-small-en-v1.5-onnx";

/// Output dimension for `bge-small-en-v1.5` (matches the candle reference).
pub const ORT_BGE_EMBEDDER_DIM: u32 = 384;

/// Pinned HF revision of `BAAI/bge-small-en-v1.5` — same commit the candle
/// loader pins (`loader::HF_REVISION`), recorded so an ONNX build is traceable
/// to the same upstream weights the equivalence measurement compares against.
/// This is the BASE of the composed identity revision; the actually-loaded
/// asset digest is appended (`+onnx-<hex>`) so the identity self-describes the
/// bytes that were opened (see [`derive_asset_revision`]).
const HF_REVISION: &str = "5c38ec7c405ec4b44b94cc5a9bb96e735b38267a";

/// Tokenizer truncation ceiling — BGE-small's 512-slot learned position
/// embeddings, identical to the candle path (`candle_bge::MAX_SEQUENCE_TOKENS`).
const MAX_SEQUENCE_TOKENS: usize = 512;

/// Sentence-vector pooling. Mirrors `candle_bge::Pooling`. Default is
/// [`OrtPooling::Cls`] — the model-native, CLS-corrected mode the candle
/// reference is compared against at Slice 15 (design §5).
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub enum OrtPooling {
    /// Mean over the attention mask (candle's historical default).
    Mean,
    /// `[CLS]` token (position 0) — the mode BGE-small was trained for.
    Cls,
}

/// An ORT execution provider selection, resolved from the `FATHOMDB_EMBED_DEVICE`
/// grammar. Kept as a plain enum (no `ort` types) so the request→provider
/// mapping is pure + unit-testable without a model or a native runtime.
#[derive(Clone, Copy, Debug, PartialEq, Eq)]
pub(crate) enum OrtProvider {
    Cpu,
    Cuda(i32),
    Rocm(i32),
    DirectMl(i32),
    OpenVino,
}

impl OrtProvider {
    /// Stable string label for the provider (index dropped — the calibration
    /// leg records the requested device string separately). Used by the
    /// additive [`OrtBgeEmbedder::effective_provider`] accessor.
    fn label(self) -> &'static str {
        match self {
            OrtProvider::Cpu => "cpu",
            OrtProvider::Cuda(_) => "cuda",
            OrtProvider::Rocm(_) => "rocm",
            OrtProvider::DirectMl(_) => "directml",
            OrtProvider::OpenVino => "openvino",
        }
    }
}

/// PURE map from the backend-agnostic [`DeviceRequest`] (parsed by the shared
/// `parse_device_request`, grammar parity with candle) to an [`OrtProvider`].
///
/// Returns the provider plus an optional LOUD-fallback message when the request
/// could not be honored and CPU was substituted (mirrors candle's loud CPU
/// fallback). The cross-vendor providers candle cannot reach — ROCm / DirectML
/// / OpenVINO — are requested through the base grammar's `Unknown` arm
/// (`FATHOMDB_EMBED_DEVICE=rocm|rocm:N|directml|openvino`), so the shared
/// parser stays unchanged and ONNX only extends the interpretation.
pub(crate) fn map_device_request(req: &DeviceRequest) -> (OrtProvider, Option<String>) {
    match req {
        DeviceRequest::Cpu => (OrtProvider::Cpu, None),
        DeviceRequest::Cuda(idx) => (OrtProvider::Cuda(*idx as i32), None),
        DeviceRequest::Metal => (
            OrtProvider::Cpu,
            Some(
                "FATHOMDB_EMBED_DEVICE=metal is a candle backend; the ONNX path has no Metal \
                 execution provider (use rocm|directml|openvino for cross-vendor GPUs, or \
                 candle's embed-metal build); using CPU"
                    .to_string(),
            ),
        ),
        DeviceRequest::Unknown(name) => map_extended_provider(name),
    }
}

/// Interpret an `Unknown` device token as a cross-vendor ORT provider.
/// Accepts `rocm`|`rocm:N`, `directml`|`dml`|`directml:N`, `openvino`|`ovep`.
/// Anything else is a LOUD CPU fallback.
fn map_extended_provider(raw: &str) -> (OrtProvider, Option<String>) {
    let (head, idx) = match raw.split_once(':') {
        Some((h, i)) => (h, i.parse::<i32>().unwrap_or(0)),
        None => (raw, 0),
    };
    match head {
        "rocm" => (OrtProvider::Rocm(idx), None),
        "directml" | "dml" => (OrtProvider::DirectMl(idx), None),
        "openvino" | "ovep" => (OrtProvider::OpenVino, None),
        other => (
            OrtProvider::Cpu,
            Some(format!(
                "FATHOMDB_EMBED_DEVICE={other} is not a recognized ONNX execution provider \
                 (expected cpu|cuda|cuda:N|rocm|rocm:N|directml|openvino); using CPU"
            )),
        ),
    }
}

/// Emit a LOUD construction-time fallback warning to stderr. Centralized so the
/// `clippy::print_stderr` allow is scoped to construction (never `embed()`), and
/// so the loud fallback is OUR OWN — `ort` is built `default-features = false`,
/// which compiles out its warn/error macros, so we cannot rely on it to surface
/// a silent CPU fallback (R-ONNX-2).
#[allow(clippy::print_stderr)] // construction-time only (not in `embed()`)
fn emit_onnx_warning(msg: &str) {
    eprintln!("fathomdb-embedder(onnx): {msg}");
}

/// RUNTIME resolution of the ORT provider from `FATHOMDB_EMBED_DEVICE`
/// (R-ONNX-2 — not a compile-time constant). Emits the LOUD stderr fallback
/// message at construction time, never inside `embed()`.
fn resolve_provider_from_env() -> OrtProvider {
    let raw = std::env::var("FATHOMDB_EMBED_DEVICE").unwrap_or_default();
    let (provider, warn) = map_device_request(&parse_device_request(&raw));
    if let Some(msg) = warn {
        emit_onnx_warning(&msg);
    }
    provider
}

/// PURE decision for the SESSION-BUILD stage of the loud fallback (R-ONNX-2):
/// given a requested [`OrtProvider`] and an availability probe, return the
/// EFFECTIVE provider plus an optional LOUD warning. A non-CPU provider the
/// probe reports unavailable is downgraded to CPU and the warning names the
/// requested provider; the caller emits it. This is distinct from the earlier
/// grammar-mapping warning ([`map_device_request`]): a request like `rocm` maps
/// cleanly to `Rocm`, but if this ONNX Runtime build lacks the ROCm EP, `ort`'s
/// own dispatch would fall back to CPU SILENTLY (its log macros are compiled out
/// under `default-features = false`), making a cross-vendor run look successful
/// while secretly on CPU. Kept pure (probe injected) so it is unit-testable
/// with no model and no ORT native lib.
fn resolve_effective_provider_with(
    requested: OrtProvider,
    is_available: impl Fn(OrtProvider) -> bool,
) -> (OrtProvider, Option<String>) {
    if matches!(requested, OrtProvider::Cpu) {
        return (OrtProvider::Cpu, None);
    }
    if is_available(requested) {
        (requested, None)
    } else {
        (
            OrtProvider::Cpu,
            Some(format!(
                "requested ONNX execution provider {requested:?} is unavailable in this ONNX \
                 Runtime build/runtime (ort is built default-features=false, so its own fallback \
                 log is compiled out); falling back to CPU"
            )),
        )
    }
}

/// Thin LIVE wrapper over [`resolve_effective_provider_with`] using `ort`'s real
/// `ExecutionProvider::is_available()` probe.
fn resolve_effective_provider(requested: OrtProvider) -> (OrtProvider, Option<String>) {
    resolve_effective_provider_with(requested, provider_is_available)
}

/// Probe whether this ONNX Runtime build was compiled with support for the
/// requested non-CPU provider (`ort`'s `ExecutionProvider::is_available()`).
/// Returns `false` when the probe errors — e.g. the ORT dylib is absent under
/// `load-dynamic` — which is precisely an unavailable provider, the silent-CPU
/// case we must surface.
fn provider_is_available(provider: OrtProvider) -> bool {
    match provider {
        OrtProvider::Cpu => true,
        OrtProvider::Cuda(_) => CUDAExecutionProvider::default().is_available().unwrap_or(false),
        OrtProvider::Rocm(_) => ROCmExecutionProvider::default().is_available().unwrap_or(false),
        OrtProvider::DirectMl(_) => {
            DirectMLExecutionProvider::default().is_available().unwrap_or(false)
        }
        OrtProvider::OpenVino => {
            OpenVINOExecutionProvider::default().is_available().unwrap_or(false)
        }
    }
}

/// Build the concrete `ort` execution-provider dispatch for a resolved
/// [`OrtProvider`].
///
/// ORT's DEFAULT dispatch (`fail_silently`) is non-fatal: if a compiled-in
/// non-CPU provider cannot REGISTER at runtime (missing CUDA/cuDNN/ROCm libs,
/// bad device id) ORT logs internally and silently falls back to CPU, so
/// `with_execution_providers` returns `Ok` and our loud-fallback machinery
/// never fires (codex §9 fix-3 root cause). To make a non-CPU registration
/// failure SURFACE as an `Err` — which then flows into
/// [`build_session_with_fallback`]'s CPU retry + LOUD warning — we mark every
/// non-CPU dispatch `.error_on_failure()`. The CPU dispatch keeps the default
/// (silent) behavior: it is the always-available floor and the retry target,
/// so it must be allowed to succeed rather than error.
fn provider_dispatch(provider: OrtProvider) -> ExecutionProviderDispatch {
    match provider {
        OrtProvider::Cpu => CPUExecutionProvider::default().build(),
        OrtProvider::Cuda(idx) => {
            CUDAExecutionProvider::default().with_device_id(idx).build().error_on_failure()
        }
        OrtProvider::Rocm(idx) => {
            ROCmExecutionProvider::default().with_device_id(idx).build().error_on_failure()
        }
        OrtProvider::DirectMl(idx) => {
            DirectMLExecutionProvider::default().with_device_id(idx).build().error_on_failure()
        }
        OrtProvider::OpenVino => OpenVINOExecutionProvider::default().build().error_on_failure(),
    }
}

/// SECOND, RUNTIME stage of the loud fallback (R-ONNX-2, codex §9 fix-2):
/// even when the availability probe ([`resolve_effective_provider`]) reports a
/// non-CPU EP as compiled-in, actually BUILDING/COMMITTING the session with it
/// can still fail at runtime — missing CUDA/cuDNN/ROCm runtime libs, a bad
/// device id, an incompatible driver. `ort`'s own dispatch would surface that
/// as a hard `Err`, so a user with a CUDA-enabled ORT but absent runtime deps
/// could not open the embedder AT ALL, violating the documented loud CPU
/// fallback. This helper attempts the build with `effective`; on failure of a
/// NON-CPU provider it produces a LOUD warning (naming the provider + the
/// error) and RETRIES with CPU, returning the CPU session. A CPU failure is a
/// real error (no retry). Generic over the session/error types with the build
/// step injected as a closure, so the retry control flow is unit-testable with
/// NO model and NO ORT native lib.
fn build_session_with_fallback<S, E, F>(
    effective: OrtProvider,
    build: F,
) -> Result<(S, Option<String>), E>
where
    F: Fn(OrtProvider) -> Result<S, E>,
    E: std::fmt::Display,
{
    match build(effective) {
        Ok(session) => Ok((session, None)),
        // A CPU build failure is a genuine error — there is nothing to fall
        // back to, so do not retry.
        Err(e) if matches!(effective, OrtProvider::Cpu) => Err(e),
        // The requested non-CPU EP was reported available but failed to build
        // at runtime: warn LOUDLY and retry on CPU. Only a CPU failure here is
        // fatal.
        Err(e) => {
            let warn = format!(
                "requested ONNX execution provider {effective:?} was reported available but \
                 FAILED during ONNX Runtime session creation ({e}); falling back to CPU"
            );
            let session = build(OrtProvider::Cpu)?;
            Ok((session, Some(warn)))
        }
    }
}

/// Stream a file through SHA-256, returning the 32-byte digest. Read in fixed
/// chunks rather than buffering the whole file so a ~130 MB `.onnx` export is
/// not held in memory (this runs once, at construction).
fn sha256_file(path: &Path) -> Result<[u8; 32], EmbedderError> {
    use sha2::{Digest, Sha256};
    let mut file = std::fs::File::open(path).map_err(|e| err("asset digest open", e))?;
    let mut hasher = Sha256::new();
    let mut buf = [0_u8; 64 * 1024];
    loop {
        let n =
            std::io::Read::read(&mut file, &mut buf).map_err(|e| err("asset digest read", e))?;
        if n == 0 {
            break;
        }
        hasher.update(&buf[..n]);
    }
    Ok(hasher.finalize().into())
}

/// PURE composition of the SELF-DESCRIBING identity revision from the two asset
/// content digests (codex §9 fix-5, remedy (b)). The revision is the pinned
/// base [`HF_REVISION`] plus a short hex of `SHA-256(model_digest ||
/// tokenizer_digest)`:
///
/// `"<HF_REVISION>+onnx-<12 hex>"`
///
/// Net effect: two different model/tokenizer assets → two DIFFERENT revisions
/// (hence different [`EmbedderIdentity`], so the engine rejects mixing vectors
/// written under a different asset), while the SAME asset yields a STABLE
/// revision across opens. Composing the two per-file digests into one final
/// hash is deterministic and unambiguous (both inputs are fixed 32-byte
/// digests, so there is no boundary ambiguity). Digest inputs make this a pure
/// function — no ORT session, unit-testable with arbitrary bytes.
fn compose_asset_revision(model_digest: &[u8; 32], tokenizer_digest: &[u8; 32]) -> String {
    use sha2::{Digest, Sha256};
    let mut hasher = Sha256::new();
    hasher.update(model_digest);
    hasher.update(tokenizer_digest);
    let combined = hasher.finalize();
    let mut hex = String::with_capacity(12);
    for byte in combined.iter().take(6) {
        use std::fmt::Write as _;
        let _ = write!(hex, "{byte:02x}");
    }
    format!("{HF_REVISION}+onnx-{hex}")
}

/// Derive the self-describing identity revision from the ACTUALLY-loaded asset
/// files on disk: hash the model bytes and the tokenizer bytes, then compose
/// via [`compose_asset_revision`]. Takes paths (not an ORT session) so the
/// revision derivation is exercisable with arbitrary temp files in tests.
fn derive_asset_revision(
    model_path: &Path,
    tokenizer_path: &Path,
) -> Result<String, EmbedderError> {
    let model_digest = sha256_file(model_path)?;
    let tokenizer_digest = sha256_file(tokenizer_path)?;
    Ok(compose_asset_revision(&model_digest, &tokenizer_digest))
}

/// Cross-vendor ONNX-Runtime BGE-small embedder.
///
/// `Session::run` needs `&mut self` but the `Embedder` trait is `&self` +
/// `Send + Sync`, so the session lives behind a `Mutex`. Embedding is a short
/// forward pass, so lock contention is not a concern for the offline/eval use
/// this backend targets.
pub struct OrtBgeEmbedder {
    identity: EmbedderIdentity,
    tokenizer: Tokenizer,
    session: Mutex<Session>,
    pooling: OrtPooling,
    /// The EFFECTIVE ORT execution provider this session was built with, AFTER
    /// any availability-probe downgrade and any session-build CPU retry (R-D3-2
    /// / 0.8.18 U3). Captured at construction so a silent CPU fallback of a
    /// non-CPU request is recorded as durable calibration DATA, not merely a
    /// transient stderr warning. Exposed via [`OrtBgeEmbedder::effective_provider`].
    effective_provider: OrtProvider,
}

fn err(context: &str, e: impl std::fmt::Display) -> EmbedderError {
    EmbedderError::Failed { message: format!("ort_bge {context}: {e}") }
}

impl OrtBgeEmbedder {
    /// Construct from an on-disk `.onnx` model + `tokenizer.json`, selecting the
    /// ORT execution provider at RUNTIME from `FATHOMDB_EMBED_DEVICE` (R-ONNX-2).
    /// Paths are caller-supplied (no hardcoded absolute path) so the model is an
    /// offline-build/eval asset the caller provisions.
    pub fn from_files(model_path: &Path, tokenizer_path: &Path) -> Result<Self, EmbedderError> {
        Self::from_files_with_provider(model_path, tokenizer_path, resolve_provider_from_env())
    }

    /// Construct from `FATHOMDB_ONNX_MODEL_PATH` + `FATHOMDB_ONNX_TOKENIZER_PATH`
    /// (device from `FATHOMDB_EMBED_DEVICE`). The env-driven entry point an eval
    /// harness / caller uses to engage the ONNX backend without recompiling.
    pub fn from_env() -> Result<Self, EmbedderError> {
        let model = std::env::var("FATHOMDB_ONNX_MODEL_PATH")
            .map_err(|_| err("from_env", "FATHOMDB_ONNX_MODEL_PATH is unset"))?;
        let tok = std::env::var("FATHOMDB_ONNX_TOKENIZER_PATH")
            .map_err(|_| err("from_env", "FATHOMDB_ONNX_TOKENIZER_PATH is unset"))?;
        Self::from_files(Path::new(&model), Path::new(&tok))
    }

    fn from_files_with_provider(
        model_path: &Path,
        tokenizer_path: &Path,
        provider: OrtProvider,
    ) -> Result<Self, EmbedderError> {
        let mut tokenizer =
            Tokenizer::from_file(tokenizer_path).map_err(|e| err("tokenizer load", e))?;
        tokenizer
            .with_truncation(Some(TruncationParams {
                max_length: MAX_SEQUENCE_TOKENS,
                ..Default::default()
            }))
            .map_err(|e| err("tokenizer truncation", e))?;

        // SESSION-BUILD stage of the loud fallback (R-ONNX-2): if the requested
        // non-CPU provider is unavailable in this ORT build, downgrade to CPU
        // and warn LOUDLY ourselves rather than letting `ort`'s compiled-out
        // dispatch fall back silently. CPU functionality is preserved.
        let (effective, avail_warn) = resolve_effective_provider(provider);
        if let Some(msg) = avail_warn {
            emit_onnx_warning(&msg);
        }

        // SESSION-BUILD RUNTIME stage (codex §9 fix-2): the availability probe
        // above can pass yet the concrete build still fail (missing runtime
        // libs, bad device id). `build_session_with_fallback` warns LOUDLY and
        // retries on CPU for a non-CPU EP so an unavailable-at-runtime GPU never
        // blocks opening the embedder; a CPU failure stays a hard error.
        let (session, build_warn) = build_session_with_fallback(effective, |p| {
            Session::builder()?
                .with_execution_providers([provider_dispatch(p)])?
                .commit_from_file(model_path)
        })
        .map_err(|e| err("session build", e))?;
        // The session-build retry (`build_session_with_fallback`) downgrades a
        // non-CPU EP that was reported available but FAILED to build to CPU; a
        // `Some(build_warn)` therefore means the final session runs on CPU.
        let effective_provider = if build_warn.is_some() { OrtProvider::Cpu } else { effective };
        if let Some(msg) = build_warn {
            emit_onnx_warning(&msg);
        }

        // Self-describing identity (codex §9 fix-5): derive the revision from a
        // content digest of the ACTUALLY-loaded model + tokenizer assets rather
        // than advertising the pinned base revision unconditionally. A stale or
        // alternate FATHOMDB_ONNX_MODEL_PATH now yields a DIFFERENT identity, so
        // the engine's EmbedderIdentity check rejects reading vectors written
        // under a different embedding space (R-ONNX-3 same-asset discipline).
        let revision = derive_asset_revision(model_path, tokenizer_path)?;
        let identity = EmbedderIdentity::new(ORT_BGE_EMBEDDER_NAME, revision, ORT_BGE_EMBEDDER_DIM);

        Ok(Self {
            identity,
            tokenizer,
            session: Mutex::new(session),
            pooling: OrtPooling::Cls,
            effective_provider,
        })
    }

    /// Select the pooling strategy (default [`OrtPooling::Cls`]). Does NOT change
    /// the identity — use only on a fresh workspace / in the equivalence harness.
    #[must_use]
    pub fn with_pooling(mut self, pooling: OrtPooling) -> Self {
        self.pooling = pooling;
        self
    }

    /// The EFFECTIVE ORT execution provider label (`"cpu"` / `"cuda"` / `"rocm"`
    /// / `"directml"` / `"openvino"`) this session was built with, AFTER any
    /// availability-probe downgrade and session-build CPU retry (R-D3-2).
    ///
    /// Additive read-only accessor (0.8.18 U3 calibration). A request for a
    /// non-CPU EP that is unavailable in this ONNX Runtime build/runtime is
    /// downgraded to CPU at construction; this returns the provider actually in
    /// force, so the calibration harness records a silent CPU fallback as DATA
    /// (never a `cuda`-requested leg mislabeled as GPU). Does not change identity.
    #[must_use]
    pub fn effective_provider(&self) -> &'static str {
        self.effective_provider.label()
    }
}

impl Embedder for OrtBgeEmbedder {
    fn identity(&self) -> EmbedderIdentity {
        self.identity.clone()
    }

    fn embed(&self, input: &str) -> Result<Vector, EmbedderError> {
        let encoding = self.tokenizer.encode(input, true).map_err(|e| err("tokenize", e))?;
        let ids: Vec<i64> = encoding.get_ids().iter().map(|&x| i64::from(x)).collect();
        let mask: Vec<i64> = encoding.get_attention_mask().iter().map(|&x| i64::from(x)).collect();
        let len = ids.len();
        let token_type: Vec<i64> = vec![0; len];
        let shape = vec![1_i64, len as i64];

        let ids_t = Tensor::from_array((shape.clone(), ids)).map_err(|e| err("input_ids", e))?;
        let mask_t =
            Tensor::from_array((shape.clone(), mask)).map_err(|e| err("attention_mask", e))?;
        let tt_t = Tensor::from_array((shape, token_type)).map_err(|e| err("token_type_ids", e))?;

        let mut session = self.session.lock().map_err(|_| err("session lock", "poisoned"))?;
        let outputs = session
            .run(ort::inputs![
                "input_ids" => ids_t,
                "attention_mask" => mask_t,
                "token_type_ids" => tt_t,
            ])
            .map_err(|e| err("forward", e))?;

        // last_hidden_state — first output, shape (1, L, H).
        let (out_shape, data) =
            outputs[0].try_extract_tensor::<f32>().map_err(|e| err("extract", e))?;
        let dims: Vec<usize> = out_shape.iter().map(|&d| d as usize).collect();
        if dims.len() != 3 {
            return Err(err("output shape", format!("expected rank-3 (1,L,H), got {dims:?}")));
        }
        let (seq_len, hidden) = (dims[1], dims[2]);
        if hidden != ORT_BGE_EMBEDDER_DIM as usize {
            return Err(err(
                "output dim",
                format!("expected hidden {ORT_BGE_EMBEDDER_DIM}, got {hidden}"),
            ));
        }

        // Pool. `embed()` is single-input (no padding), so the attention mask is
        // all-ones and mean-pool reduces to a plain mean over `seq_len`.
        let mut pooled = vec![0.0_f32; hidden];
        match self.pooling {
            OrtPooling::Cls => {
                pooled.copy_from_slice(&data[0..hidden]);
            }
            OrtPooling::Mean => {
                for pos in 0..seq_len {
                    let base = pos * hidden;
                    for (j, slot) in pooled.iter_mut().enumerate() {
                        *slot += data[base + j];
                    }
                }
                let denom = seq_len.max(1) as f32;
                for slot in &mut pooled {
                    *slot /= denom;
                }
            }
        }

        // L2-normalize (matches candle's `l2_normalize`).
        let norm = pooled.iter().map(|v| v * v).sum::<f32>().sqrt().max(1e-12);
        for slot in &mut pooled {
            *slot /= norm;
        }
        Ok(pooled)
    }
}

#[cfg(test)]
mod tests {
    //! R-ONNX-2 device-mapping unit tests: the `FATHOMDB_EMBED_DEVICE` grammar
    //! (parsed by the shared `parse_device_request`) → the correct ORT execution
    //! provider. Pure — no model, no ONNX Runtime native lib, no GPU required.
    use super::{map_device_request, OrtProvider};
    use crate::device::parse_device_request;

    fn resolve(raw: &str) -> (OrtProvider, Option<String>) {
        map_device_request(&parse_device_request(raw))
    }

    #[test]
    fn cpu_and_unset_map_to_cpu_no_warning() {
        for raw in ["", "cpu", "  CPU  "] {
            let (p, warn) = resolve(raw);
            assert_eq!(p, OrtProvider::Cpu, "{raw:?}");
            assert!(warn.is_none(), "{raw:?} should not warn");
        }
    }

    #[test]
    fn cuda_maps_to_cuda_provider_with_index() {
        assert_eq!(resolve("cuda").0, OrtProvider::Cuda(0));
        assert_eq!(resolve("cuda:1").0, OrtProvider::Cuda(1));
        assert_eq!(resolve("cuda:2").0, OrtProvider::Cuda(2));
        assert!(resolve("cuda:1").1.is_none());
    }

    #[test]
    fn rocm_maps_to_rocm_provider() {
        // ROCm is unreachable through candle — the cross-vendor payoff.
        assert_eq!(resolve("rocm").0, OrtProvider::Rocm(0));
        assert_eq!(resolve("rocm:1").0, OrtProvider::Rocm(1));
        assert!(resolve("rocm").1.is_none());
        assert!(resolve("ROCm").1.is_none()); // grammar lower-cases
    }

    #[test]
    fn directml_maps_to_directml_provider() {
        assert_eq!(resolve("directml").0, OrtProvider::DirectMl(0));
        assert_eq!(resolve("dml").0, OrtProvider::DirectMl(0));
        assert_eq!(resolve("directml:1").0, OrtProvider::DirectMl(1));
        assert!(resolve("directml").1.is_none());
    }

    #[test]
    fn openvino_maps_to_openvino_provider() {
        assert_eq!(resolve("openvino").0, OrtProvider::OpenVino);
        assert_eq!(resolve("ovep").0, OrtProvider::OpenVino);
        assert!(resolve("openvino").1.is_none());
    }

    #[test]
    fn metal_falls_back_to_cpu_loudly() {
        // ORT has no Metal EP; candle owns that lane. Loud, never silent.
        let (p, warn) = resolve("metal");
        assert_eq!(p, OrtProvider::Cpu);
        assert!(warn.is_some(), "metal must warn on CPU fallback");
    }

    #[test]
    fn unrecognized_device_falls_back_to_cpu_loudly() {
        for raw in ["vulkan", "tpu", "gpu"] {
            let (p, warn) = resolve(raw);
            assert_eq!(p, OrtProvider::Cpu, "{raw:?}");
            assert!(warn.is_some(), "{raw:?} must warn on CPU fallback");
        }
    }

    /// SESSION-BUILD loud fallback (codex §9 fix-1): a requested GPU provider
    /// that the ORT build does not support must downgrade to CPU with a LOUD
    /// warning that names the requested provider — never a silent CPU fallback.
    /// Probe is injected (`|_| false`), so this runs with no ORT lib / no GPU.
    #[test]
    fn requested_gpu_unavailable_falls_back_to_cpu_loudly() {
        for requested in [
            OrtProvider::Cuda(0),
            OrtProvider::Rocm(1),
            OrtProvider::DirectMl(0),
            OrtProvider::OpenVino,
        ] {
            let (eff, warn) = super::resolve_effective_provider_with(requested, |_| false);
            assert_eq!(eff, OrtProvider::Cpu, "{requested:?} must downgrade to CPU");
            let msg = warn.expect("unavailable GPU provider must emit a warning");
            assert!(
                msg.contains(&format!("{requested:?}")),
                "warning must name the requested provider {requested:?}, got {msg:?}"
            );
            assert!(
                msg.to_lowercase().contains("cpu"),
                "warning must state the CPU fallback, got {msg:?}"
            );
        }
    }

    /// When the probe reports the requested provider available, it is honored
    /// and there is NO warning (no spurious loud fallback).
    #[test]
    fn requested_available_provider_is_honored_without_warning() {
        let (eff, warn) = super::resolve_effective_provider_with(OrtProvider::Cuda(2), |_| true);
        assert_eq!(eff, OrtProvider::Cuda(2));
        assert!(warn.is_none(), "an available provider must not warn");
    }

    /// A CPU request is always honored, never warns, and never probes (CPU is
    /// unconditionally available) — the probe closure must not run.
    #[test]
    fn cpu_request_never_probes_and_never_warns() {
        let (eff, warn) = super::resolve_effective_provider_with(OrtProvider::Cpu, |_| {
            panic!("CPU request must not probe provider availability")
        });
        assert_eq!(eff, OrtProvider::Cpu);
        assert!(warn.is_none());
    }

    /// SESSION-BUILD RUNTIME loud fallback (codex §9 fix-2): a NON-CPU provider
    /// that the availability probe reports as compiled-in can STILL fail when
    /// the session is actually built (missing runtime libs / bad device id).
    /// The helper must warn LOUDLY (naming the provider) and RETRY on CPU,
    /// returning the CPU session. Build step injected, so no model / ORT lib.
    #[test]
    fn session_build_gpu_fails_retries_cpu_with_warning() {
        use std::cell::RefCell;
        for requested in [OrtProvider::Cuda(0), OrtProvider::Rocm(1), OrtProvider::OpenVino] {
            let attempts = RefCell::new(Vec::new());
            let build = |p: OrtProvider| -> Result<&'static str, String> {
                attempts.borrow_mut().push(p);
                match p {
                    OrtProvider::Cpu => Ok("cpu-session"),
                    other => Err(format!("no runtime lib for {other:?}")),
                }
            };
            let (session, warn) = super::build_session_with_fallback(requested, build)
                .expect("must retry CPU and succeed");
            assert_eq!(session, "cpu-session", "{requested:?} must yield the CPU session");
            let msg = warn.expect("a runtime GPU-build failure must emit a warning");
            assert!(
                msg.contains(&format!("{requested:?}")),
                "warning must name the requested provider {requested:?}, got {msg:?}"
            );
            assert!(
                msg.to_lowercase().contains("cpu"),
                "warning must state the CPU fallback, got {msg:?}"
            );
            assert_eq!(
                *attempts.borrow(),
                vec![requested, OrtProvider::Cpu],
                "must attempt the GPU provider first, then CPU"
            );
        }
    }

    /// When BOTH the requested non-CPU build and the CPU retry fail, the helper
    /// must surface an `Err` (no session can be opened).
    #[test]
    fn session_build_gpu_and_cpu_both_fail_errors() {
        let build =
            |p: OrtProvider| -> Result<&'static str, String> { Err(format!("hard fail {p:?}")) };
        let res = super::build_session_with_fallback(OrtProvider::Cuda(0), build);
        assert!(res.is_err(), "both GPU and CPU failing must be an error");
    }

    /// A CPU-effective build failure is a genuine error — there is nothing to
    /// fall back to, so the build closure must run EXACTLY once (no retry).
    #[test]
    fn session_build_cpu_failure_is_error_without_retry() {
        use std::cell::Cell;
        let calls = Cell::new(0);
        let build = |_p: OrtProvider| -> Result<&'static str, String> {
            calls.set(calls.get() + 1);
            Err("cpu build failed".to_string())
        };
        let res = super::build_session_with_fallback(OrtProvider::Cpu, build);
        assert!(res.is_err(), "a CPU build failure must be an error");
        assert_eq!(calls.get(), 1, "a CPU failure must NOT retry");
    }

    /// The happy path: an effective provider whose build succeeds returns the
    /// session with NO warning and attempts the build exactly once.
    #[test]
    fn session_build_success_no_warning_single_attempt() {
        use std::cell::Cell;
        let calls = Cell::new(0);
        let build = |_p: OrtProvider| -> Result<&'static str, String> {
            calls.set(calls.get() + 1);
            Ok("session")
        };
        let (session, warn) = super::build_session_with_fallback(OrtProvider::Cuda(2), build)
            .expect("build succeeds");
        assert_eq!(session, "session");
        assert!(warn.is_none(), "a successful build must not warn");
        assert_eq!(calls.get(), 1, "a successful build must attempt exactly once");
    }

    /// codex §9 fix-3 ROOT: a non-CPU dispatch must be built with
    /// `error_on_failure` so a runtime EP-registration failure surfaces as an
    /// `Err` from `with_execution_providers` (feeding the CPU-retry + loud
    /// warning). The CPU dispatch keeps the default silent behavior so the
    /// retry target can succeed. Inspected via the dispatch `Debug` output,
    /// which renders the private `error_on_failure` flag. No native ORT lib is
    /// loaded — `.build()` only constructs the dispatch config struct.
    #[test]
    fn non_cpu_dispatch_errors_on_failure_cpu_does_not() {
        for p in [
            OrtProvider::Cuda(0),
            OrtProvider::Rocm(1),
            OrtProvider::DirectMl(0),
            OrtProvider::OpenVino,
        ] {
            let dbg = format!("{:?}", super::provider_dispatch(p));
            assert!(
                dbg.contains("error_on_failure: true"),
                "non-CPU provider {p:?} must set error_on_failure, got {dbg:?}"
            );
        }
        let cpu = format!("{:?}", super::provider_dispatch(OrtProvider::Cpu));
        assert!(
            cpu.contains("error_on_failure: false"),
            "CPU provider must keep the silent (fallback-target) default, got {cpu:?}"
        );
    }

    /// codex §9 fix-5 — the identity revision self-describes the loaded assets.
    /// Written a set of arbitrary temp "model"/"tokenizer" byte contents and
    /// assert the derived revision distinguishes distinct assets, is stable for
    /// identical bytes, is sensitive to EITHER file changing, keeps the pinned
    /// base + `+onnx-<12 hex>` shape, and does not depend on file NAMES (only
    /// bytes). Pure — arbitrary temp files, no ORT session / native lib.
    mod asset_revision {
        use std::io::Write as _;
        use std::path::Path;

        use super::super::{compose_asset_revision, derive_asset_revision, HF_REVISION};

        fn write_temp(name: &str, bytes: &[u8]) -> std::path::PathBuf {
            // Honors TMPDIR — the sandboxed scratch dir set for this build.
            let dir = std::env::temp_dir();
            let path = dir.join(format!("fathomdb-ort-rev-{}-{name}", std::process::id()));
            let mut f = std::fs::File::create(&path).expect("create temp asset");
            f.write_all(bytes).expect("write temp asset");
            f.flush().expect("flush temp asset");
            path
        }

        fn revision_of(model_bytes: &[u8], tok_bytes: &[u8], tag: &str) -> String {
            let m = write_temp(&format!("model-{tag}.onnx"), model_bytes);
            let t = write_temp(&format!("tok-{tag}.json"), tok_bytes);
            let rev = derive_asset_revision(Path::new(&m), Path::new(&t)).expect("derive revision");
            let _ = std::fs::remove_file(&m);
            let _ = std::fs::remove_file(&t);
            rev
        }

        #[test]
        fn distinct_assets_yield_distinct_revisions() {
            let a = revision_of(b"model-bytes-AAAA", b"tokenizer-bytes-AAAA", "a");
            let b = revision_of(b"model-bytes-BBBB", b"tokenizer-bytes-AAAA", "b");
            assert_ne!(a, b, "a different MODEL must change the identity revision");

            let c = revision_of(b"model-bytes-AAAA", b"tokenizer-bytes-CCCC", "c");
            assert_ne!(a, c, "a different TOKENIZER must change the identity revision");
            assert_ne!(b, c, "distinct model+tokenizer pairs must differ");
        }

        #[test]
        fn identical_bytes_yield_stable_revision() {
            // Same bytes, DIFFERENT file names/paths + separate opens → identical
            // revision (stable/deterministic; depends only on content).
            let first = revision_of(b"same-model-bytes", b"same-tokenizer-bytes", "stable1");
            let second = revision_of(b"same-model-bytes", b"same-tokenizer-bytes", "stable2");
            assert_eq!(first, second, "identical asset bytes must yield a stable revision");
        }

        #[test]
        fn revision_has_pinned_base_and_onnx_digest_shape() {
            let rev = revision_of(b"m", b"t", "shape");
            let prefix = format!("{HF_REVISION}+onnx-");
            let hex = rev.strip_prefix(&prefix).expect("revision must start with pinned base");
            assert_eq!(hex.len(), 12, "digest suffix must be 12 hex chars, got {hex:?}");
            assert!(
                hex.chars().all(|c| c.is_ascii_hexdigit()),
                "digest suffix must be lower-hex, got {hex:?}"
            );
        }

        #[test]
        fn compose_is_pure_and_order_sensitive() {
            let m = [1_u8; 32];
            let t = [2_u8; 32];
            // Pure: same inputs → same output.
            assert_eq!(compose_asset_revision(&m, &t), compose_asset_revision(&m, &t));
            // Model/tokenizer roles are distinct positions in the composed hash,
            // so swapping the two digests generally changes the revision.
            assert_ne!(
                compose_asset_revision(&m, &t),
                compose_asset_revision(&t, &m),
                "the two asset roles must not be interchangeable"
            );
        }
    }

    /// R-ONNX-1 real-vector test — runs FOR REAL against the offline ONNX asset
    /// on the ONNX-Runtime CPU EP (the same-backend fidelity baseline per policy
    /// `649a8d45`; GPU is a Slice-15 re-embed speed concern, not needed here).
    ///
    /// ENV-GATED (not `#[ignore]`) so CI without the provisioned asset skips it
    /// cleanly while a provisioned host runs it: set `ORT_DYLIB_PATH` (the
    /// on-host `libonnxruntime.so`, load-dynamic), `FATHOMDB_ONNX_MODEL_PATH`
    /// (the offline export from `dev/tools/onnx/export_bge_small_onnx.py`), and
    /// `FATHOMDB_ONNX_TOKENIZER_PATH` (the pinned `tokenizer.json`). When those
    /// are unset the test returns early (recorded skip); when set it asserts a
    /// fixture text → a 384-dim, finite, L2-normalized, deterministic vector via
    /// the `Embedder` trait, and that the identity revision self-describes the
    /// loaded asset digest (fix-5). See `dev/tools/onnx/README.md` for the exact
    /// invocation and the chosen ORT lib.
    #[test]
    fn ort_bge_embeds_384_dim_finite_deterministic_vector() {
        use std::path::Path;

        use fathomdb_embedder_api::Embedder;

        use super::{OrtBgeEmbedder, OrtProvider};

        let (Ok(_dylib), Ok(model), Ok(tok)) = (
            std::env::var("ORT_DYLIB_PATH"),
            std::env::var("FATHOMDB_ONNX_MODEL_PATH"),
            std::env::var("FATHOMDB_ONNX_TOKENIZER_PATH"),
        ) else {
            eprintln!(
                "SKIP ort_bge_embeds_384_dim_finite_deterministic_vector: set ORT_DYLIB_PATH + \
                 FATHOMDB_ONNX_MODEL_PATH + FATHOMDB_ONNX_TOKENIZER_PATH to run the real-vector \
                 R-ONNX-1 test (see dev/tools/onnx/README.md)"
            );
            return;
        };

        // CPU same-backend baseline (policy 649a8d45): force the ONNX CPU EP
        // explicitly (not via the ambient FATHOMDB_EMBED_DEVICE) so this
        // fidelity assertion is deterministic and does not depend on env order.
        let embedder = OrtBgeEmbedder::from_files_with_provider(
            Path::new(&model),
            Path::new(&tok),
            OrtProvider::Cpu,
        )
        .expect("open OrtBgeEmbedder on the CPU EP from the provisioned asset");

        // Identity self-describes the loaded asset (fix-5): distinct ONNX name,
        // dim 384, revision = pinned base + "+onnx-<12 hex>" asset digest.
        let id = embedder.identity();
        assert_eq!(id.name, super::ORT_BGE_EMBEDDER_NAME, "distinct ONNX identity name");
        assert_eq!(id.dimension, super::ORT_BGE_EMBEDDER_DIM, "384-dim identity");
        assert!(
            id.revision.contains("+onnx-"),
            "revision must carry the asset digest, got {:?}",
            id.revision
        );

        let v1 = embedder.embed("the quick brown fox").expect("embed");
        let v2 = embedder.embed("the quick brown fox").expect("embed");
        assert_eq!(v1.len(), super::ORT_BGE_EMBEDDER_DIM as usize);
        assert!(v1.iter().all(|x| x.is_finite()), "all components finite");
        assert_eq!(v1, v2, "deterministic for identical input");
        let norm = v1.iter().map(|x| x * x).sum::<f32>().sqrt();
        assert!((norm - 1.0).abs() < 1e-3, "L2-normalized (got {norm})");
    }
}